ixa/entity/entity_set/

entity_set_iterator.rs

//! Iterator implementation for [`EntitySet`].
//!
//! `EntitySetIterator` mirrors an `EntitySet` expression tree and evaluates nodes lazily.
//! - `Source` iterates a concrete backing source (`Population`, contiguous
//!   `PopulationRange`, index set, or property-backed source).
//! - `Intersection` and `Difference` drive iteration from one branch and filter candidates using
//!   membership checks.
//! - `Union` yields the left branch first, then lazily activates the right branch (`Pending` to
//!   `Active`) and filters out any IDs already present in the left branch via membership checks.
//!
//! The iterator is created through `EntitySet::into_iter()`.

use crate::entity::entity_set::entity_set::{EntitySet, EntitySetInner};
use crate::entity::entity_set::source_iterator::{PopulationRangeIterator, SourceIterator};
use crate::entity::entity_set::source_set::SourceSet;
use crate::entity::{Entity, EntityId, PopulationIterator};
use crate::hashing::IndexSet;

enum EntitySetIteratorInner<'a, E: Entity> {
    Source(SourceIterator<'a, E>),
    // The `IntersectionSources` variant is a micro-optimization to avoid recursive
    // `EntitySet` membership checks in the most common case, improving tight-loop
    // benchmark performance by 5%-15%.
    IntersectionSources {
        driver: SourceIterator<'a, E>,
        filters: Vec<SourceSet<'a, E>>,
    },
    Intersection {
        driver: Box<EntitySetIteratorInner<'a, E>>,
        filters: Vec<EntitySet<'a, E>>,
    },
    Difference {
        left: Box<EntitySetIteratorInner<'a, E>>,
        right: EntitySet<'a, E>,
    },
    Union {
        left: Box<EntitySetIteratorInner<'a, E>>,
        right: UnionRightState<'a, E>,
    },
}

enum UnionRightState<'a, E: Entity> {
    Pending(Option<EntitySet<'a, E>>),
    Active(Box<EntitySetIteratorInner<'a, E>>),
}

impl<'a, E: Entity> EntitySetIteratorInner<'a, E> {
    fn empty_source() -> Self {
        Self::Source(SourceIterator::PopulationRange(
            PopulationRangeIterator::new(0..0),
        ))
    }

    fn from_entity_set(set: EntitySet<'a, E>) -> Self {
        match set.into_inner() {
            EntitySetInner::Source(source) => {
                if source.try_len() == Some(0) {
                    Self::empty_source()
                } else {
                    Self::Source(source.into_iter())
                }
            }
            EntitySetInner::Intersection(mut sets) => {
                if sets.is_empty() {
                    return Self::empty_source();
                }
                if sets.len() == 1 {
                    return Self::from_entity_set(sets.pop().unwrap());
                }

                if sets.iter().all(EntitySet::is_source_leaf) {
                    // `EntitySet::Intersection` stores operands sorted ascending by `cost_hint`.
                    // Keep that order so membership checks short-circuit early on small filters.
                    let mut set_iter = sets.into_iter();

                    let first = set_iter.next().unwrap().into_source_leaf().unwrap();
                    let driver = first.into_iter();

                    let filters = set_iter
                        .map(|set| set.into_source_leaf().unwrap())
                        .collect();

                    return Self::IntersectionSources { driver, filters };
                }

                // `EntitySet::Intersection` stores operands sorted ascending by `cost_hint`.
                // Use the first set as the iteration driver and keep the remaining filters in
                // ascending order for short-circuit-friendly `contains` checks.
                let mut set_iter = sets.into_iter();
                let driver = Box::new(Self::from_entity_set(set_iter.next().unwrap()));
                Self::Intersection {
                    driver,
                    filters: set_iter.collect(),
                }
            }
            EntitySetInner::Difference(left, right) => Self::Difference {
                left: Box::new(Self::from_entity_set(*left)),
                right: *right,
            },
            EntitySetInner::Union(left, right) => Self::Union {
                left: Box::new(Self::from_entity_set(*left)),
                right: UnionRightState::Pending(Some(*right)),
            },
        }
    }

    fn contains(&self, entity_id: EntityId<E>) -> bool {
        match self {
            Self::Source(iter) => iter.contains(entity_id),
            Self::IntersectionSources { driver, filters } => {
                driver.contains(entity_id) && filters.iter().all(|set| set.contains(entity_id))
            }
            Self::Intersection { driver, filters } => {
                driver.contains(entity_id) && filters.iter().all(|set| set.contains(entity_id))
            }
            Self::Difference { left, right } => {
                left.contains(entity_id) && !right.contains(entity_id)
            }
            Self::Union { left, right } => left.contains(entity_id) || right.contains(entity_id),
        }
    }
}

impl<'a, E: Entity> UnionRightState<'a, E> {
    fn contains(&self, entity_id: EntityId<E>) -> bool {
        match self {
            Self::Pending(Some(set)) => set.contains(entity_id),
            Self::Pending(None) => false,
            Self::Active(iter) => iter.contains(entity_id),
        }
    }
}

impl<'a, E: Entity> EntitySetIteratorInner<'a, E> {
    #[inline]
    fn next_inner(&mut self) -> Option<EntityId<E>> {
        match self {
            Self::Source(source) => source.next(),
            Self::IntersectionSources { driver, filters } => driver
                .by_ref()
                .find(|&entity_id| filters.iter().all(|filter| filter.contains(entity_id))),
            Self::Intersection { driver, filters } => {
                while let Some(entity_id) = driver.next_inner() {
                    if filters.iter().all(|filter| filter.contains(entity_id)) {
                        return Some(entity_id);
                    }
                }
                None
            }
            Self::Difference { left, right } => {
                while let Some(entity_id) = left.next_inner() {
                    if !right.contains(entity_id) {
                        return Some(entity_id);
                    }
                }
                None
            }
            Self::Union { left, right } => loop {
                if let Some(entity_id) = left.next_inner() {
                    return Some(entity_id);
                }

                match right {
                    UnionRightState::Pending(maybe_set) => {
                        if let Some(set) = maybe_set.take() {
                            *right = UnionRightState::Active(Box::new(Self::from_entity_set(set)));
                        }
                        continue;
                    }
                    UnionRightState::Active(right_iter) => {
                        while let Some(entity_id) = right_iter.next_inner() {
                            if !left.contains(entity_id) {
                                return Some(entity_id);
                            }
                        }
                        return None;
                    }
                }
            },
        }
    }

    #[inline]
    fn size_hint_inner(&self) -> (usize, Option<usize>) {
        match self {
            Self::Source(source) => source.size_hint(),
            Self::IntersectionSources { driver, .. } => {
                let (_, upper) = driver.size_hint();
                (0, upper)
            }
            Self::Intersection { driver, .. } => {
                let (_, upper) = driver.size_hint_inner();
                (0, upper)
            }
            Self::Difference { left, .. } => {
                let (_, upper) = left.size_hint_inner();
                (0, upper)
            }
            Self::Union { left, right } => {
                let (_, left_upper) = left.size_hint_inner();
                let right_upper = match right {
                    UnionRightState::Pending(_) => None,
                    UnionRightState::Active(right_iter) => right_iter.size_hint_inner().1,
                };
                let upper = match (left_upper, right_upper) {
                    (Some(a), Some(b)) => Some(a.saturating_add(b)),
                    _ => None,
                };
                (0, upper)
            }
        }
    }
}

/// An iterator over the IDs in an entity set, producing `EntityId<E>`s until exhausted.
pub struct EntitySetIterator<'c, E: Entity> {
    inner: EntitySetIteratorInner<'c, E>,
}

impl<'c, E: Entity> EntitySetIterator<'c, E> {
    pub(crate) fn empty() -> EntitySetIterator<'c, E> {
        EntitySetIterator {
            inner: EntitySetIteratorInner::empty_source(),
        }
    }

    pub(crate) fn from_population_iterator(iter: PopulationIterator<E>) -> Self {
        EntitySetIterator {
            inner: EntitySetIteratorInner::Source(SourceIterator::PopulationRange(
                PopulationRangeIterator::from_population_iterator(iter),
            )),
        }
    }

    pub(crate) fn from_sources(mut sources: Vec<SourceSet<'c, E>>) -> Self {
        if sources.is_empty() {
            return Self::empty();
        }
        if sources.len() == 1 {
            return EntitySetIterator {
                inner: EntitySetIteratorInner::Source(sources.pop().unwrap().into_iter()),
            };
        }

        // This path constructs intersections from raw source vectors, so we sort here.
        // We keep ascending order so filters checked by `all()` are smallest-first.
        sources.sort_unstable_by_key(SourceSet::sort_key);
        let mut source_iter = sources.into_iter();
        let driver = source_iter.next().unwrap().into_iter();
        EntitySetIterator {
            inner: EntitySetIteratorInner::IntersectionSources {
                driver,
                filters: source_iter.collect(),
            },
        }
    }

    pub(crate) fn from_index_set(set: &'c IndexSet<EntityId<E>>) -> EntitySetIterator<'c, E> {
        EntitySetIterator {
            inner: EntitySetIteratorInner::Source(SourceSet::IndexSet(set).into_iter()),
        }
    }

    pub(super) fn new(set: EntitySet<'c, E>) -> Self {
        EntitySetIterator {
            inner: EntitySetIteratorInner::from_entity_set(set),
        }
    }
}

impl<'a, E: Entity> Iterator for EntitySetIterator<'a, E> {
    type Item = EntityId<E>;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        self.inner.next_inner()
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        self.inner.size_hint_inner()
    }

    fn count(self) -> usize {
        let EntitySetIterator { inner } = self;
        match inner {
            EntitySetIteratorInner::Source(source) => source.count(),
            EntitySetIteratorInner::IntersectionSources {
                mut driver,
                filters,
            } => driver
                .by_ref()
                .filter(|&entity_id| filters.iter().all(|filter| filter.contains(entity_id)))
                .count(),
            other => {
                let mut it = EntitySetIterator { inner: other };
                let mut n = 0usize;
                while it.next().is_some() {
                    n += 1;
                }
                n
            }
        }
    }

    #[inline]
    fn nth(&mut self, n: usize) -> Option<Self::Item> {
        match &mut self.inner {
            EntitySetIteratorInner::Source(source) => source.nth(n),
            EntitySetIteratorInner::IntersectionSources { driver, filters } => driver
                .by_ref()
                .filter(|&entity_id| filters.iter().all(|filter| filter.contains(entity_id)))
                .nth(n),
            _ => {
                for _ in 0..n {
                    self.next()?;
                }
                self.next()
            }
        }
    }

    fn for_each<F>(self, mut f: F)
    where
        F: FnMut(Self::Item),
    {
        let EntitySetIterator { inner } = self;
        match inner {
            EntitySetIteratorInner::Source(source) => source.for_each(f),
            other => {
                let it = EntitySetIterator { inner: other };
                for item in it {
                    f(item);
                }
            }
        }
    }

    fn fold<B, F>(self, init: B, mut f: F) -> B
    where
        F: FnMut(B, Self::Item) -> B,
    {
        let EntitySetIterator { inner } = self;
        match inner {
            EntitySetIteratorInner::Source(source) => source.fold(init, f),
            other => {
                let it = EntitySetIterator { inner: other };
                let mut acc = init;
                for item in it {
                    acc = f(acc, item);
                }
                acc
            }
        }
    }
}

impl<'c, E: Entity> std::iter::FusedIterator for EntitySetIterator<'c, E> {}

#[cfg(test)]
mod tests {
    /*!
    ## Test Matrix

    | Test # | PropertyInitializationKind | is_default_value | Indexed | Initial Source Position |
    | ------ | -------------------------- | ---------------- | ------- | ----------------------- |
    | 1      | Explicit                   | N/A              | No      | Yes                     |
    | 2      | Explicit                   | N/A              | Yes     | Yes                     |
    | 3      | Constant                   | true             | No      | Yes                     |
    | 4      | Constant                   | true             | Yes     | Yes                     |
    | 5      | Constant                   | false            | No      | Yes                     |
    | 6      | Constant                   | false            | Yes     | Yes                     |
    | 7      | Derived                    | N/A              | No      | Yes                     |
    | 8      | Derived                    | N/A              | Yes     | Yes                     |
    | 9      | Explicit                   | N/A              | No      | No                      |
    | 10     | Explicit                   | N/A              | Yes     | No                      |
    | 11     | Constant                   | true             | No      | No                      |
    | 12     | Constant                   | false            | Yes     | No                      |
    | 13     | Derived                    | N/A              | No      | No                      |
    | 14     | Derived                    | N/A              | Yes     | No                      |

    The tests use multi-property queries (tests 9-14) where one property
    is indexed with fewer results to ensure it becomes the "smallest source
    set", making the tested property NOT the initial source position.
    */

    use indexmap::IndexSet;

    use crate::entity::entity_set::{EntitySet, SourceSet};
    use crate::hashing::IndexSet as FxIndexSet;
    use crate::prelude::*;
    use crate::{define_derived_property, define_property, with};

    define_entity!(Person);

    // Test properties covering different initialization kinds

    // Explicit (Normal) property - no default, requires explicit initialization
    define_property!(struct ExplicitProp(u8), Person);

    // Constant property - has a constant default value
    define_property!(struct ConstantProp(u8), Person, default_const = ConstantProp(42));

    // Derived property - computed from other properties
    define_derived_property!(struct DerivedProp(bool), Person, [ExplicitProp], |explicit| {
        DerivedProp(explicit.0 % 2 == 0)
    });

    // Additional properties for multi-property queries
    define_property!(struct ConstantProp2(u16), Person, default_const = ConstantProp2(100));
    define_property!(struct ExplicitProp2(bool), Person);

    define_property!(struct Age(u8), Person, default_const = Age(0));
    define_property!(struct Alive(bool), Person, default_const = Alive(true));

    define_derived_property!(
        enum AgeGroupRisk {
            NewBorn,
            General,
            OldAdult,
        },
        Person,
        [Age],
        [],
        |age| {
            if age.0 <= 1 {
                AgeGroupRisk::NewBorn
            } else if age.0 <= 65 {
                AgeGroupRisk::General
            } else {
                AgeGroupRisk::OldAdult
            }
        }
    );

    // Helper function to create a population for testing
    fn setup_test_population(context: &mut Context, size: usize) -> Vec<EntityId<Person>> {
        let mut people = Vec::new();
        for i in 0..size {
            let person = context
                .add_entity(with!(
                    Person,
                    ExplicitProp((i % 20) as u8),
                    ExplicitProp2(i % 2 == 0)
                ))
                .unwrap();
            people.push(person);
        }
        people
    }

    // region: Test Matrix Tests

    // Test 1: Explicit property, non-default value, not indexed, initial source position = Yes
    #[test]
    fn test_explicit_non_default_not_indexed_initial_source_yes() {
        let mut context = Context::new();
        setup_test_population(&mut context, 100);

        let results = context
            .query_result_iterator(with!(Person, ExplicitProp(5)))
            .collect::<Vec<_>>();

        assert_eq!(results.len(), 5); // 5, 25, 45, 65, 85
        for person in results {
            assert_eq!(
                context.get_property::<_, ExplicitProp>(person),
                ExplicitProp(5)
            );
        }
    }

    // Test 2: Explicit property, non-default value, indexed, initial source position = Yes
    #[test]
    fn test_explicit_non_default_indexed_initial_source_yes() {
        let mut context = Context::new();
        context.index_property::<Person, ExplicitProp>();
        setup_test_population(&mut context, 100);

        let results = context
            .query_result_iterator(with!(Person, ExplicitProp(7)))
            .collect::<Vec<_>>();

        assert_eq!(results.len(), 5); // 7, 27, 47, 67, 87
        for person in results {
            assert_eq!(
                context.get_property::<_, ExplicitProp>(person),
                ExplicitProp(7)
            );
        }
    }

    // Test 3: Constant property, default value, not indexed, initial source position = Yes
    #[test]
    fn test_constant_default_not_indexed_initial_source_yes() {
        let mut context = Context::new();
        // Create people without setting ConstantProp - they'll use default value 42
        for _ in 0..50 {
            context
                .add_entity(with!(Person, ExplicitProp(1), ExplicitProp2(false)))
                .unwrap();
        }

        let results = context
            .query_result_iterator(with!(Person, ConstantProp(42), ExplicitProp2(false)))
            .collect::<Vec<_>>();

        assert_eq!(results.len(), 50);
        for person in results {
            assert_eq!(
                context.get_property::<_, ConstantProp>(person),
                ConstantProp(42)
            );
        }
    }

    // Test 4: Constant property, default value, indexed, initial source position = Yes
    #[test]
    fn test_constant_default_indexed_initial_source_yes() {
        let mut context = Context::new();
        context.index_property::<Person, ConstantProp>();

        for _ in 0..50 {
            context
                .add_entity(with!(Person, ExplicitProp(1), ExplicitProp2(false)))
                .unwrap();
        }

        let results = context
            .query_result_iterator(with!(Person, ConstantProp(42)))
            .collect::<Vec<_>>();
        assert_eq!(results.len(), 50);
    }

    // Test 5: Constant property, non-default value, not indexed, initial source position = Yes
    #[test]
    fn test_constant_non_default_not_indexed_initial_source_yes() {
        let mut context = Context::new();

        for i in 0..50 {
            if i < 10 {
                context
                    .add_entity(with!(
                        Person,
                        ExplicitProp(1),
                        ExplicitProp2(false),
                        ConstantProp(99)
                    ))
                    .unwrap();
            } else {
                context
                    .add_entity(with!(Person, ExplicitProp(1), ExplicitProp2(false)))
                    .unwrap();
            }
        }

        let results = context
            .query_result_iterator(with!(Person, ConstantProp(99)))
            .collect::<Vec<_>>();

        assert_eq!(results.len(), 10);
        for person in results {
            assert_eq!(
                context.get_property::<_, ConstantProp>(person),
                ConstantProp(99)
            );
        }
    }

    // Test 6: Constant property, non-default value, indexed, initial source position = Yes
    #[test]
    fn test_constant_non_default_indexed_initial_source_yes() {
        let mut context = Context::new();
        context.index_property::<Person, ConstantProp>();

        for i in 0..50 {
            if i < 10 {
                context
                    .add_entity(with!(
                        Person,
                        ExplicitProp(1),
                        ExplicitProp2(false),
                        ConstantProp(99)
                    ))
                    .unwrap();
            } else {
                context
                    .add_entity(with!(Person, ExplicitProp(1), ExplicitProp2(false)))
                    .unwrap();
            }
        }

        let results = context
            .query_result_iterator(with!(Person, ConstantProp(99)))
            .collect::<Vec<_>>();

        assert_eq!(results.len(), 10);
    }

    // Test 7: Derived property, not indexed, initial source position = Yes
    #[test]
    fn test_derived_not_indexed_initial_source_yes() {
        let mut context = Context::new();

        for i in 0..100 {
            context
                .add_entity(with!(Person, ExplicitProp(i as u8), ExplicitProp2(false)))
                .unwrap();
        }

        let results = context
            .query_result_iterator(with!(Person, DerivedProp(true)))
            .collect::<Vec<_>>();

        // DerivedProp is true when ExplicitProp is even
        assert_eq!(results.len(), 50);
        for person in results {
            assert_eq!(
                context.get_property::<Person, DerivedProp>(person),
                DerivedProp(true)
            );
        }
    }

    // Test 8: Derived property, indexed, initial source position = Yes
    #[test]
    fn test_derived_indexed_initial_source_yes() {
        let mut context = Context::new();
        context.index_property::<Person, DerivedProp>();

        for i in 0..100 {
            context
                .add_entity(with!(Person, ExplicitProp(i as u8), ExplicitProp2(false)))
                .unwrap();
        }

        let results = context
            .query_result_iterator(with!(Person, DerivedProp(false)))
            .collect::<Vec<_>>();

        // DerivedProp is false when ExplicitProp is odd
        assert_eq!(results.len(), 50);
        for person in results {
            assert_eq!(
                context.get_property::<Person, DerivedProp>(person),
                DerivedProp(false)
            );
        }
    }

    // Test 9-14: Initial source position = No (multi-property queries where property is NOT the smallest source)

    // Test 9: Explicit property, non-default, not indexed, initial source = No
    #[test]
    fn test_explicit_non_default_not_indexed_initial_source_no() {
        let mut context = Context::new();
        context.index_property::<Person, ExplicitProp2>(); // Index the other property so it's the smallest

        for i in 0..100 {
            context
                .add_entity(with!(
                    Person,
                    ExplicitProp((i % 20) as u8),
                    ExplicitProp2(i % 2 == 0)
                ))
                .unwrap();
        }

        // ExplicitProp2 has only 2 values, so it will be the smaller source
        let results = context
            .query_result_iterator(with!(Person, ExplicitProp(5), ExplicitProp2(false)))
            .collect::<Vec<_>>();

        // Looking for ExplicitProp=5 AND ExplicitProp2=true
        // ExplicitProp cycles 0-19, ExplicitProp2 alternates
        // Matches: 4 (since 5,25,45,65,85 but only even indices)
        let expected = results.len();
        assert!(expected > 0);
        for person in results {
            assert_eq!(
                context.get_property::<Person, ExplicitProp>(person),
                ExplicitProp(5)
            );
            assert_eq!(
                context.get_property::<Person, ExplicitProp2>(person),
                ExplicitProp2(false)
            );
        }
    }

    // Test 10: Explicit property, non-default, indexed, initial source = No
    #[test]
    fn test_explicit_non_default_indexed_initial_source_no() {
        let mut context = Context::new();
        context.index_property::<Person, ExplicitProp>();
        context.index_property::<Person, ConstantProp2>(); // ConstantProp2 will likely be the smaller source

        for i in 0..100 {
            if i < 10 {
                context
                    .add_entity(with!(
                        Person,
                        ExplicitProp(7),
                        ExplicitProp2(false),
                        ConstantProp2(200), // Non-default for smaller source
                    ))
                    .unwrap();
            } else {
                context
                    .add_entity(with!(
                        Person,
                        ExplicitProp((i % 20) as u8),
                        ExplicitProp2(false)
                    ))
                    .unwrap();
            }
        }

        let results = context
            .query_result_iterator(with!(Person, ExplicitProp(7), ConstantProp2(200)))
            .collect::<Vec<_>>();

        assert_eq!(results.len(), 10);
    }

    // Test 11: Constant property, default value, not indexed, initial source = No
    #[test]
    fn test_constant_default_not_indexed_initial_source_no() {
        let mut context = Context::new();
        context.index_property::<Person, ExplicitProp>(); // Make ExplicitProp the smaller source

        for i in 0..100 {
            if i < 5 {
                context
                    .add_entity(with!(Person, ExplicitProp(99), ExplicitProp2(false)))
                    .unwrap(); // ConstantProp uses default
            } else {
                context
                    .add_entity(with!(
                        Person,
                        ExplicitProp((i % 20) as u8),
                        ExplicitProp2(false)
                    ))
                    .unwrap();
            }
        }

        let results = context
            .query_result_iterator(with!(Person, ExplicitProp(99), ConstantProp(42)))
            .collect::<Vec<_>>();

        assert_eq!(results.len(), 5);
        for person in results {
            assert_eq!(
                context.get_property::<Person, ConstantProp>(person),
                ConstantProp(42)
            );
        }
    }

    // Test 12: Constant property, non-default, indexed, initial source = No
    #[test]
    fn test_constant_non_default_indexed_initial_source_no() {
        let mut context = Context::new();
        context.index_property::<Person, ConstantProp>();
        context.index_property::<Person, ExplicitProp2>();

        for i in 0..100 {
            if i < 10 {
                context
                    .add_entity(with!(
                        Person,
                        ConstantProp(99),
                        ExplicitProp(0),
                        ExplicitProp2(true)
                    ))
                    .unwrap();
            } else {
                context
                    .add_entity(with!(Person, ExplicitProp(0), ExplicitProp2(false)))
                    .unwrap();
            }
        }

        let results = context
            .query_result_iterator(with!(Person, ConstantProp(99), ExplicitProp2(true)))
            .collect::<Vec<_>>();

        assert_eq!(results.len(), 10);
    }

    // Test 13: Derived property, not indexed, initial source = No
    #[test]
    fn test_derived_not_indexed_initial_source_no() {
        let mut context = Context::new();
        context.index_property::<Person, ExplicitProp2>();

        for i in 0..100 {
            context
                .add_entity(with!(Person, ExplicitProp(i as u8), ExplicitProp2(i < 50)))
                .unwrap();
        }

        let results = context
            .query_result_iterator(with!(Person, ExplicitProp2(true), DerivedProp(true)))
            .collect::<Vec<_>>();

        // ExplicitProp2=true for i<50, DerivedProp=true when ExplicitProp is even
        // So we want i<50 AND i is even: 0,2,4,...,48 = 25 people
        assert_eq!(results.len(), 25);
    }

    // Test 14: Derived property, indexed, initial source = No
    #[test]
    fn test_derived_indexed_initial_source_no() {
        let mut context = Context::new();
        context.index_property::<Person, DerivedProp>();
        context.index_property::<Person, ExplicitProp2>();

        for i in 0..100 {
            context
                .add_entity(with!(Person, ExplicitProp(i as u8), ExplicitProp2(i < 30)))
                .unwrap();
        }

        let results = context
            .query_result_iterator(with!(Person, ExplicitProp2(true), DerivedProp(false)))
            .collect::<Vec<_>>();

        // ExplicitProp2=true for i<30, DerivedProp=false when ExplicitProp is odd
        // So we want i < 30 AND i is odd: 1,3,5,...,29 = 15 people
        assert_eq!(results.len(), 15);
    }

    // endregion Test Matrix Tests

    #[test]
    fn test_multiple_query_result_iterators() {
        let mut context = Context::new();
        context.index_property::<Person, Age>();

        for age in 0..100 {
            context
                .add_entity(with!(
                    Person,
                    Age(age),
                    ExplicitProp(age.wrapping_mul(7) % 100),
                    ExplicitProp2(false),
                ))
                .unwrap();
        }
        for age in 0..100 {
            context
                .add_entity(with!(
                    Person,
                    Age(age),
                    ExplicitProp(age.wrapping_mul(14) % 100),
                    ExplicitProp2(false),
                ))
                .unwrap();
        }

        // Since both queries include `Age`, both will attempt to index unindexed entities. This tests that there is
        // no double borrow error.
        let results = context.query_result_iterator(with!(Person, Age(25)));
        let more_results = context.query_result_iterator(with!(Person, Age(25), ExplicitProp(75)));

        let collected_results = results.collect::<IndexSet<_>>();
        let other_collected_results = more_results.collect::<IndexSet<_>>();
        let intersection_count = collected_results
            .intersection(&other_collected_results)
            .count();
        assert_eq!(intersection_count, 1);
    }

    #[test]
    fn test_expression_intersection_iteration() {
        let set = EntitySet::from_source(SourceSet::<Person>::PopulationRange(0..10)).intersection(
            EntitySet::from_source(SourceSet::<Person>::PopulationRange(0..6)),
        );

        let ids = set.into_iter().collect::<Vec<_>>();
        let expected = (0..6).map(EntityId::new).collect::<Vec<_>>();
        assert_eq!(ids, expected);
    }

    #[test]
    fn test_expression_difference_iteration() {
        let set = EntitySet::from_source(SourceSet::<Person>::PopulationRange(0..5)).difference(
            EntitySet::from_source(SourceSet::<Person>::PopulationRange(2..3)),
        );

        let ids = set.into_iter().collect::<Vec<_>>();
        let expected = vec![
            EntityId::new(0),
            EntityId::new(1),
            EntityId::new(3),
            EntityId::new(4),
        ];
        assert_eq!(ids, expected);
    }

    #[test]
    fn test_expression_union_deduplicates() {
        let left = EntitySet::from_source(SourceSet::<Person>::PopulationRange(0..3)).difference(
            EntitySet::from_source(SourceSet::<Person>::PopulationRange(99..100)),
        );
        let right = EntitySet::from_source(SourceSet::<Person>::PopulationRange(0..5)).difference(
            EntitySet::from_source(SourceSet::<Person>::PopulationRange(99..100)),
        );
        let set = left.union(right);

        let ids = set.into_iter().collect::<Vec<_>>();
        let expected = (0..5).map(EntityId::new).collect::<Vec<_>>();
        assert_eq!(ids, expected);
    }

    #[test]
    fn test_expression_union_overlap_no_duplicates() {
        let left = EntitySet::from_source(SourceSet::<Person>::PopulationRange(0..5)).difference(
            EntitySet::from_source(SourceSet::<Person>::PopulationRange(4..5)),
        );
        let right = EntitySet::from_source(SourceSet::<Person>::PopulationRange(0..7)).difference(
            EntitySet::from_source(SourceSet::<Person>::PopulationRange(0..1)),
        );

        let ids = left.union(right).into_iter().collect::<IndexSet<_>>();
        let expected = (0..7).map(EntityId::new).collect::<IndexSet<_>>();
        assert_eq!(ids, expected);
    }

    #[test]
    fn test_expression_intersection_of_unions() {
        let ab = EntitySet::from_source(SourceSet::<Person>::PopulationRange(1..2))
            .union(EntitySet::from_source(
                SourceSet::<Person>::PopulationRange(2..3),
            ))
            .union(EntitySet::from_source(
                SourceSet::<Person>::PopulationRange(3..4),
            ));
        let cd = EntitySet::from_source(SourceSet::<Person>::PopulationRange(2..3))
            .union(EntitySet::from_source(
                SourceSet::<Person>::PopulationRange(3..4),
            ))
            .union(EntitySet::from_source(
                SourceSet::<Person>::PopulationRange(4..5),
            ));

        let ids = ab.intersection(cd).into_iter().collect::<Vec<_>>();
        assert_eq!(ids, vec![EntityId::new(2), EntityId::new(3)]);
    }

    #[test]
    fn test_expression_difference_not_commutative() {
        let left = EntitySet::from_source(SourceSet::<Person>::PopulationRange(1..2))
            .union(EntitySet::from_source(
                SourceSet::<Person>::PopulationRange(2..3),
            ))
            .union(EntitySet::from_source(
                SourceSet::<Person>::PopulationRange(3..4),
            ));
        let right = EntitySet::from_source(SourceSet::<Person>::PopulationRange(2..3))
            .union(EntitySet::from_source(
                SourceSet::<Person>::PopulationRange(3..4),
            ))
            .union(EntitySet::from_source(
                SourceSet::<Person>::PopulationRange(4..5),
            ));

        let left_minus_right = left.difference(right).into_iter().collect::<Vec<_>>();
        let right_minus_left = EntitySet::from_source(SourceSet::<Person>::PopulationRange(2..3))
            .union(EntitySet::from_source(
                SourceSet::<Person>::PopulationRange(3..4),
            ))
            .union(EntitySet::from_source(
                SourceSet::<Person>::PopulationRange(4..5),
            ))
            .difference(
                EntitySet::from_source(SourceSet::<Person>::PopulationRange(1..2))
                    .union(EntitySet::from_source(
                        SourceSet::<Person>::PopulationRange(2..3),
                    ))
                    .union(EntitySet::from_source(
                        SourceSet::<Person>::PopulationRange(3..4),
                    )),
            )
            .into_iter()
            .collect::<Vec<_>>();

        assert_eq!(left_minus_right, vec![EntityId::new(1)]);
        assert_eq!(right_minus_left, vec![EntityId::new(4)]);
    }

    #[test]
    fn test_union_size_hint_pending_right_is_unknown() {
        let left = EntitySet::from_source(SourceSet::<Person>::PopulationRange(0..2)).difference(
            EntitySet::from_source(SourceSet::<Person>::PopulationRange(99..100)),
        );
        let right = EntitySet::from_source(SourceSet::<Person>::PopulationRange(0..4)).difference(
            EntitySet::from_source(SourceSet::<Person>::PopulationRange(98..99)),
        );
        let iter = left.union(right).into_iter();

        assert_eq!(iter.size_hint(), (4, Some(4)));
    }

    #[test]
    fn test_nth_and_count_on_source() {
        let mut iter =
            EntitySet::from_source(SourceSet::<Person>::PopulationRange(0..5)).into_iter();
        assert_eq!(iter.nth(2), Some(EntityId::new(2)));

        let remaining = iter.count();
        assert_eq!(remaining, 2);
    }

    #[test]
    fn population_range_source_iteration_and_size_hint() {
        let mut iter =
            EntitySet::from_source(SourceSet::<Person>::PopulationRange(3..7)).into_iter();
        assert_eq!(iter.size_hint(), (4, Some(4)));
        assert_eq!(iter.next(), Some(EntityId::new(3)));
        assert_eq!(iter.size_hint(), (3, Some(3)));
        assert_eq!(
            iter.collect::<Vec<_>>(),
            vec![EntityId::new(4), EntityId::new(5), EntityId::new(6)]
        );
    }

    #[test]
    fn optimized_range_results_iterate_as_single_source() {
        let union = EntitySet::from_source(SourceSet::<Person>::PopulationRange(0..3)).union(
            EntitySet::from_source(SourceSet::<Person>::PopulationRange(3..5)),
        );
        assert_eq!(
            union.into_iter().collect::<Vec<_>>(),
            vec![
                EntityId::new(0),
                EntityId::new(1),
                EntityId::new(2),
                EntityId::new(3),
                EntityId::new(4),
            ]
        );

        let intersection = EntitySet::from_source(SourceSet::<Person>::PopulationRange(2..8))
            .intersection(EntitySet::from_source(
                SourceSet::<Person>::PopulationRange(4..6),
            ));
        assert_eq!(
            intersection.into_iter().collect::<Vec<_>>(),
            vec![EntityId::new(4), EntityId::new(5)]
        );
    }

    #[test]
    fn split_difference_still_uses_generic_iteration_path() {
        let split = EntitySet::from_source(SourceSet::<Person>::PopulationRange(2..8)).difference(
            EntitySet::from_source(SourceSet::<Person>::PopulationRange(4..6)),
        );
        assert_eq!(
            split.into_iter().collect::<Vec<_>>(),
            vec![
                EntityId::new(2),
                EntityId::new(3),
                EntityId::new(6),
                EntityId::new(7)
            ]
        );
    }

    fn finite_set(ids: &[usize]) -> FxIndexSet<EntityId<Person>> {
        ids.iter()
            .copied()
            .map(EntityId::new)
            .collect::<FxIndexSet<_>>()
    }

    fn as_entity_set(set: &FxIndexSet<EntityId<Person>>) -> EntitySet<Person> {
        EntitySet::from_source(SourceSet::IndexSet(set))
    }

    #[test]
    fn prototype_empty_entity_set_yields_nothing() {
        let ids = EntitySet::<Person>::empty().into_iter().collect::<Vec<_>>();
        assert!(ids.is_empty());
    }

    #[test]
    fn prototype_iter_union_disjoint() {
        let a = finite_set(&[1, 2]);
        let b = finite_set(&[3, 4]);
        let ids = as_entity_set(&a)
            .union(as_entity_set(&b))
            .into_iter()
            .collect::<IndexSet<_>>();
        let expected = [1usize, 2, 3, 4]
            .into_iter()
            .map(EntityId::new)
            .collect::<IndexSet<_>>();
        assert_eq!(ids, expected);
    }

    #[test]
    fn prototype_iter_union_overlapping() {
        let a = finite_set(&[1, 2, 3]);
        let b = finite_set(&[2, 3, 4]);
        let ids = as_entity_set(&a)
            .union(as_entity_set(&b))
            .into_iter()
            .collect::<Vec<_>>();
        let unique = ids.iter().copied().collect::<IndexSet<_>>();
        assert_eq!(ids.len(), unique.len());
        assert!(unique.contains(&EntityId::new(1)));
        assert!(unique.contains(&EntityId::new(4)));
    }

    #[test]
    fn prototype_iter_intersection_disjoint() {
        let a = finite_set(&[1, 2]);
        let b = finite_set(&[3, 4]);
        let ids = as_entity_set(&a)
            .intersection(as_entity_set(&b))
            .into_iter()
            .collect::<Vec<_>>();
        assert!(ids.is_empty());
    }

    #[test]
    fn prototype_iter_difference_basic() {
        let a = finite_set(&[1, 2, 3, 4]);
        let b = finite_set(&[3, 4, 5, 6]);
        let ids = as_entity_set(&a)
            .difference(as_entity_set(&b))
            .into_iter()
            .collect::<IndexSet<_>>();
        assert_eq!(ids.len(), 2);
        assert!(ids.contains(&EntityId::new(1)));
        assert!(ids.contains(&EntityId::new(2)));
    }

    #[test]
    fn prototype_iter_matches_contains_compound() {
        let a = finite_set(&[1, 2, 3, 4]);
        let b = finite_set(&[3, 4, 5]);
        let c = finite_set(&[7, 8, 9, 10]);
        let d = finite_set(&[9, 10, 11]);
        let left = as_entity_set(&a).intersection(as_entity_set(&b));
        let right = as_entity_set(&c).difference(as_entity_set(&d));
        let iterated = left.union(right).into_iter().collect::<IndexSet<_>>();

        let a2 = finite_set(&[1, 2, 3, 4]);
        let b2 = finite_set(&[3, 4, 5]);
        let c2 = finite_set(&[7, 8, 9, 10]);
        let d2 = finite_set(&[9, 10, 11]);
        let check = as_entity_set(&a2)
            .intersection(as_entity_set(&b2))
            .union(as_entity_set(&c2).difference(as_entity_set(&d2)));

        for value in 0..15 {
            let entity = EntityId::new(value);
            assert_eq!(iterated.contains(&entity), check.contains(entity));
        }
    }

    #[test]
    fn prototype_size_hint_single_source_and_partial_consume() {
        let mut iter =
            EntitySet::from_source(SourceSet::<Person>::PopulationRange(0..5)).into_iter();
        assert_eq!(iter.size_hint(), (5, Some(5)));
        iter.next();
        assert_eq!(iter.size_hint(), (4, Some(4)));
    }
}