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

use std::ops::Range;

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

use super::{Query, Fetch};

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

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

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

    type IntoIter = ViewIter<'a, Q>;

    fn into_iter(self) -> Self::IntoIter {
        ViewIter {
            world: self.world,
            query: self.query,
            fetcher: None,
            archetypes: self.archetypes,
            next_archetype: 0,
            component_indices: 0..0,
        }
    }
}

pub struct ViewIter<'a, Q: Query> {
    world: &'a World,
    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.len() == 0 {
                    continue;
                }

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

                self.fetcher = unsafe { Some(self.query.fetch(self.world, ArchetypeId(arch_id as u64), arch)) };
                self.component_indices = 0..arch.entities.len() as u64;
            }
        }
    }
}