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: Q::Fetch::dangling(),
            archetypes: self.archetypes,
            next_archetype: 0,
            component_indices: 0..0,
        }
    }
}

pub struct ViewIter<'a, Q: Query> {
    world: &'a World,
    query: Q,
    fetcher: 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 let Some(entity_index) = self.component_indices.next() {
                let entity_index = ArchetypeEntityId(entity_index);
                if !self.fetcher.can_visit_item(entity_index) {
                    continue;
                } else {
                    let i = unsafe { self.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 { self.query.fetch(self.world, ArchetypeId(arch_id as u64), arch) };
                self.component_indices = 0..arch.entities.len() as u64;
            }
        }
    }
}