lyra-engine/lyra-ecs/src/query/optional.rs

use crate::{Archetype, World};

use super::{AsQuery, Fetch, Query};

#[derive(Default)]
pub struct OptionalFetcher<'a, Q: AsQuery> {
    fetcher: Option<<Q::Query as Query>::Fetch<'a>>,
}

impl<'a, Q: AsQuery> Fetch<'a> for OptionalFetcher<'a, Q> {
    type Item = Option<<Q::Query as Query>::Item<'a>>;

    fn dangling() -> Self {
        unreachable!()
    }

    unsafe fn get_item(&mut self, entity: crate::ArchetypeEntityId) -> Self::Item {
        self.fetcher.as_mut()
            .map(|f| f.get_item(entity))
    }
    
    fn can_visit_item(&mut self, entity: crate::ArchetypeEntityId) -> bool {
        self.fetcher.as_mut()
            .map(|f| f.can_visit_item(entity))
            .unwrap_or(true)
    }
}

#[derive(Default)]
pub struct Optional<Q: AsQuery> {
    query: Q::Query,
}

impl<Q: AsQuery> Copy for Optional<Q> { }

impl<Q: AsQuery> Clone for Optional<Q> {
    fn clone(&self) -> Self {
        Self { query: self.query.clone() }
    }
}

impl<Q: AsQuery> Query for Optional<Q> {
    type Item<'a> = Option<<Q::Query as Query>::Item<'a>>;
    
    type Fetch<'a> = OptionalFetcher<'a, Q>;

    fn new() -> Self {
        Optional {
            query: Q::Query::new(),
        }
    }

    fn can_visit_archetype(&self, _: &Archetype) -> bool {
        true
    }

    unsafe fn fetch<'a>(&self, world: &'a World, arch: &'a Archetype, tick: crate::Tick) -> Self::Fetch<'a> {
        let fetcher = if self.query.can_visit_archetype(arch) {
            Some(self.query.fetch(world, arch, tick))
        } else {
            None
        };

        OptionalFetcher {
            fetcher,
        }
    }
}

impl<Q: AsQuery> AsQuery for Optional<Q> {
    type Query = Self;
}

impl<Q: AsQuery> AsQuery for Option<Q> {
    type Query = Optional<Q>;
}