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

use std::ops::Range;

use crate::{archetype::Archetype, world::{ArchetypeEntityId, World}, EntityId, Tick};

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

pub type View<'a, Q, F = ()> = ViewState<'a, <Q as AsQuery>::Query, <F as AsQuery>::Query>;

pub struct ViewState<'a, Q: Query, F: Query> {
    world: &'a World,
    query: Q,
    filter: F,
    archetypes: Vec<&'a Archetype>,
}

impl<'a, Q, F> ViewState<'a, Q, F>
where
    Q: Query,
    F: Query,
{
    pub fn new(world: &'a World, query: Q, filter: F, archetypes: Vec<&'a Archetype>) -> Self {
        Self {
            world,
            query,
            filter,
            archetypes,
        }
    }

    /// Converts self into an iterator
    pub fn iter(self) -> ViewIter<'a, Q, F> {
        self.into_iter()
    }

    /// Consumes `self`, adding a query to the view.
    pub fn expand<U: AsQuery>(self, query: U::Query) -> ViewState<'a, (Q, U::Query), F> {
        ViewState::new(self.world, (self.query, query), self.filter, self.archetypes)
    }

    /// Consumes `self`, adding a filter to the view.
    pub fn with<U: AsQuery>(self, filter: U::Query) -> ViewState<'a, Q, (F, U::Query)> {
        ViewState::new(self.world, self.query, (self.filter, filter), self.archetypes)
    }
}

impl<'a, Q, F> IntoIterator for ViewState<'a, Q, F>
where
    Q: Query,
    F: Query,
{
    type Item = Q::Item<'a>;

    type IntoIter = ViewIter<'a, Q, F>;

    fn into_iter(self) -> Self::IntoIter {
        let tick = self.world.tick_tracker().tick_when(Q::MUTATES);

        ViewIter {
            world: self.world,
            tick,
            query: self.query,
            filter: self.filter,
            fetcher: None,
            filter_fetcher: None,
            archetypes: self.archetypes,
            next_archetype: 0,
            component_indices: 0..0,
        }
    }
}

pub struct ViewIter<'a, Q: Query, F: Query> {
    world: &'a World,
    tick: Tick,
    query: Q,
    filter: F,
    fetcher: Option<Q::Fetch<'a>>,
    filter_fetcher: Option<F::Fetch<'a>>,
    archetypes: Vec<&'a Archetype>,
    next_archetype: usize,
    component_indices: Range<u64>,
}

impl<'a, Q, F> Iterator for ViewIter<'a, Q, F>
where
    Q: Query,
    F: Query,
{
    type Item = Q::Item<'a>;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            if Q::ALWAYS_FETCHES {
                // only fetch this query once.
                // fetcher gets set to Some after this `next` call.
                if self.fetcher.is_none() {
                    if let Some(mut fetch) = unsafe { self.query.fetch_world(self.world) } {
                        let res =  unsafe { Some(fetch.get_item(ArchetypeEntityId(0))) };
                        self.fetcher = Some(fetch);
                        return res;
                    }
                } else {
                    return None;
                }
            }

            if let Some(entity_index) = self.component_indices.next() {
                let fetcher = self.fetcher.as_mut().unwrap();
                let filter_fetcher = self.filter_fetcher.as_mut().unwrap();
                let entity_index = ArchetypeEntityId(entity_index);

                if fetcher.can_visit_item(entity_index) && filter_fetcher.can_visit_item(entity_index) {
                    let i = unsafe { fetcher.get_item(entity_index) };
                    return Some(i);
                }
            } else {
                if self.next_archetype >= self.archetypes.len() {
                    return None; // ran out of archetypes to go through
                }

                let arch_id = self.next_archetype;
                self.next_archetype += 1;
                let arch = unsafe { self.archetypes.get_unchecked(arch_id) };

                if arch.entity_ids.is_empty() {
                    continue;
                }

                if self.query.can_visit_archetype(arch) && self.filter.can_visit_archetype(arch) {
                    unsafe {
                        self.fetcher = Some(self.query.fetch(self.world, arch, self.tick));
                        self.filter_fetcher = Some(self.filter.fetch(self.world, arch, self.tick));
                    }
                    self.component_indices = 0..arch.entity_ids.len() as u64;
                }
            }
        }
    }
}

pub struct ViewOne<'a, Q: Query> {
    world: &'a World,
    tick: Tick,
    entity: EntityId,
    query: Q,
}

impl<'a, Q: Query> ViewOne<'a, Q> {
    pub fn new(world: &'a World, entity: EntityId, query: Q) -> Self {
        let tick = world.tick_tracker().tick_when(Q::MUTATES);

        Self {
            world,
            tick,
            entity,
            query
        }
    }

    pub fn get(&self) -> Option<Q::Item<'a>> {
        if let Some(record) = self.world.entities.arch_index.get(&self.entity) {
            let arch = self.world.archetypes.get(&record.id)
                .expect("An invalid record was specified for an entity");

            if self.query.can_visit_archetype(arch) {
                let mut fetch = unsafe { self.query.fetch(self.world, arch, self.tick) };
                if fetch.can_visit_item(record.index) {
                    return Some(unsafe { fetch.get_item(record.index) });
                }
            }
        }
        
        None
    }
}