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 struct View<'a, Q: AsQuery> {
    world: &'a World,
    query: Q::Query,
    archetypes: Vec<&'a Archetype>,
}

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

impl<'a, Q> IntoIterator for View<'a, Q>
where
    Q: AsQuery,
{
    type Item = <Q::Query as Query>::Item<'a>;

    type IntoIter = ViewIter<'a, Q::Query>;

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

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

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

impl<'a, Q> Iterator for ViewIter<'a, Q>
where
    Q: 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 entity_index = ArchetypeEntityId(entity_index);
                if !fetcher.can_visit_item(entity_index) {
                    continue;
                } else {
                    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.entities.is_empty() {
                    continue;
                }

                if !self.query.can_visit_archetype(arch) {
                    continue;
                }

                self.fetcher = unsafe { Some(self.query.fetch(self.world, arch, self.tick)) };
                self.component_indices = 0..arch.entities.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.entity_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
    }
}