ixa/entity/entity_set/
entity_set_iterator.rs

1//! Iterator implementation for [`EntitySet`].
2//!
3//! `EntitySetIterator` mirrors an `EntitySet` expression tree and evaluates nodes lazily.
4//! - `Source` iterates a concrete backing source (`Population`, contiguous
5//!   `PopulationRange`, index set, or property-backed source).
6//! - `Intersection` and `Difference` drive iteration from one branch and filter candidates using
7//!   membership checks.
8//! - `Union` yields the left branch first, then lazily activates the right branch (`Pending` to
9//!   `Active`) and filters out any IDs already present in the left branch via membership checks.
10//!
11//! The iterator is created through `EntitySet::into_iter()`.
12
13use crate::entity::entity_set::entity_set::{EntitySet, EntitySetInner};
14use crate::entity::entity_set::source_iterator::{PopulationRangeIterator, SourceIterator};
15use crate::entity::entity_set::source_set::SourceSet;
16use crate::entity::{Entity, EntityId, PopulationIterator};
17use crate::hashing::IndexSet;
18
19enum EntitySetIteratorInner<'a, E: Entity> {
20    Source(SourceIterator<'a, E>),
21    // The `IntersectionSources` variant is a micro-optimization to avoid recursive
22    // `EntitySet` membership checks in the most common case, improving tight-loop
23    // benchmark performance by 5%-15%.
24    IntersectionSources {
25        driver: SourceIterator<'a, E>,
26        filters: Vec<SourceSet<'a, E>>,
27    },
28    Intersection {
29        driver: Box<EntitySetIteratorInner<'a, E>>,
30        filters: Vec<EntitySet<'a, E>>,
31    },
32    Difference {
33        left: Box<EntitySetIteratorInner<'a, E>>,
34        right: EntitySet<'a, E>,
35    },
36    Union {
37        left: Box<EntitySetIteratorInner<'a, E>>,
38        right: UnionRightState<'a, E>,
39    },
40}
41
42enum UnionRightState<'a, E: Entity> {
43    Pending(Option<EntitySet<'a, E>>),
44    Active(Box<EntitySetIteratorInner<'a, E>>),
45}
46
47impl<'a, E: Entity> EntitySetIteratorInner<'a, E> {
48    fn empty_source() -> Self {
49        Self::Source(SourceIterator::PopulationRange(
50            PopulationRangeIterator::new(0..0),
51        ))
52    }
53
54    fn from_entity_set(set: EntitySet<'a, E>) -> Self {
55        match set.into_inner() {
56            EntitySetInner::Source(source) => {
57                if source.try_len() == Some(0) {
58                    Self::empty_source()
59                } else {
60                    Self::Source(source.into_iter())
61                }
62            }
63            EntitySetInner::Intersection(mut sets) => {
64                if sets.is_empty() {
65                    return Self::empty_source();
66                }
67                if sets.len() == 1 {
68                    return Self::from_entity_set(sets.pop().unwrap());
69                }
70
71                if sets.iter().all(EntitySet::is_source_leaf) {
72                    // `EntitySet::Intersection` stores operands sorted ascending by `cost_hint`.
73                    // Keep that order so membership checks short-circuit early on small filters.
74                    let mut set_iter = sets.into_iter();
75
76                    let first = set_iter.next().unwrap().into_source_leaf().unwrap();
77                    let driver = first.into_iter();
78
79                    let filters = set_iter
80                        .map(|set| set.into_source_leaf().unwrap())
81                        .collect();
82
83                    return Self::IntersectionSources { driver, filters };
84                }
85
86                // `EntitySet::Intersection` stores operands sorted ascending by `cost_hint`.
87                // Use the first set as the iteration driver and keep the remaining filters in
88                // ascending order for short-circuit-friendly `contains` checks.
89                let mut set_iter = sets.into_iter();
90                let driver = Box::new(Self::from_entity_set(set_iter.next().unwrap()));
91                Self::Intersection {
92                    driver,
93                    filters: set_iter.collect(),
94                }
95            }
96            EntitySetInner::Difference(left, right) => Self::Difference {
97                left: Box::new(Self::from_entity_set(*left)),
98                right: *right,
99            },
100            EntitySetInner::Union(left, right) => Self::Union {
101                left: Box::new(Self::from_entity_set(*left)),
102                right: UnionRightState::Pending(Some(*right)),
103            },
104        }
105    }
106
107    fn contains(&self, entity_id: EntityId<E>) -> bool {
108        match self {
109            Self::Source(iter) => iter.contains(entity_id),
110            Self::IntersectionSources { driver, filters } => {
111                driver.contains(entity_id) && filters.iter().all(|set| set.contains(entity_id))
112            }
113            Self::Intersection { driver, filters } => {
114                driver.contains(entity_id) && filters.iter().all(|set| set.contains(entity_id))
115            }
116            Self::Difference { left, right } => {
117                left.contains(entity_id) && !right.contains(entity_id)
118            }
119            Self::Union { left, right } => left.contains(entity_id) || right.contains(entity_id),
120        }
121    }
122}
123
124impl<'a, E: Entity> UnionRightState<'a, E> {
125    fn contains(&self, entity_id: EntityId<E>) -> bool {
126        match self {
127            Self::Pending(Some(set)) => set.contains(entity_id),
128            Self::Pending(None) => false,
129            Self::Active(iter) => iter.contains(entity_id),
130        }
131    }
132}
133
134impl<'a, E: Entity> EntitySetIteratorInner<'a, E> {
135    #[inline]
136    fn next_inner(&mut self) -> Option<EntityId<E>> {
137        match self {
138            Self::Source(source) => source.next(),
139            Self::IntersectionSources { driver, filters } => driver
140                .by_ref()
141                .find(|&entity_id| filters.iter().all(|filter| filter.contains(entity_id))),
142            Self::Intersection { driver, filters } => {
143                while let Some(entity_id) = driver.next_inner() {
144                    if filters.iter().all(|filter| filter.contains(entity_id)) {
145                        return Some(entity_id);
146                    }
147                }
148                None
149            }
150            Self::Difference { left, right } => {
151                while let Some(entity_id) = left.next_inner() {
152                    if !right.contains(entity_id) {
153                        return Some(entity_id);
154                    }
155                }
156                None
157            }
158            Self::Union { left, right } => loop {
159                if let Some(entity_id) = left.next_inner() {
160                    return Some(entity_id);
161                }
162
163                match right {
164                    UnionRightState::Pending(maybe_set) => {
165                        if let Some(set) = maybe_set.take() {
166                            *right = UnionRightState::Active(Box::new(Self::from_entity_set(set)));
167                        }
168                        continue;
169                    }
170                    UnionRightState::Active(right_iter) => {
171                        while let Some(entity_id) = right_iter.next_inner() {
172                            if !left.contains(entity_id) {
173                                return Some(entity_id);
174                            }
175                        }
176                        return None;
177                    }
178                }
179            },
180        }
181    }
182
183    #[inline]
184    fn size_hint_inner(&self) -> (usize, Option<usize>) {
185        match self {
186            Self::Source(source) => source.size_hint(),
187            Self::IntersectionSources { driver, .. } => {
188                let (_, upper) = driver.size_hint();
189                (0, upper)
190            }
191            Self::Intersection { driver, .. } => {
192                let (_, upper) = driver.size_hint_inner();
193                (0, upper)
194            }
195            Self::Difference { left, .. } => {
196                let (_, upper) = left.size_hint_inner();
197                (0, upper)
198            }
199            Self::Union { left, right } => {
200                let (_, left_upper) = left.size_hint_inner();
201                let right_upper = match right {
202                    UnionRightState::Pending(_) => None,
203                    UnionRightState::Active(right_iter) => right_iter.size_hint_inner().1,
204                };
205                let upper = match (left_upper, right_upper) {
206                    (Some(a), Some(b)) => Some(a.saturating_add(b)),
207                    _ => None,
208                };
209                (0, upper)
210            }
211        }
212    }
213}
214
215/// An iterator over the IDs in an entity set, producing `EntityId<E>`s until exhausted.
216pub struct EntitySetIterator<'c, E: Entity> {
217    inner: EntitySetIteratorInner<'c, E>,
218}
219
220impl<'c, E: Entity> EntitySetIterator<'c, E> {
221    pub(crate) fn empty() -> EntitySetIterator<'c, E> {
222        EntitySetIterator {
223            inner: EntitySetIteratorInner::empty_source(),
224        }
225    }
226
227    pub(crate) fn from_population_iterator(iter: PopulationIterator<E>) -> Self {
228        EntitySetIterator {
229            inner: EntitySetIteratorInner::Source(SourceIterator::PopulationRange(
230                PopulationRangeIterator::from_population_iterator(iter),
231            )),
232        }
233    }
234
235    pub(crate) fn from_sources(mut sources: Vec<SourceSet<'c, E>>) -> Self {
236        if sources.is_empty() {
237            return Self::empty();
238        }
239        if sources.len() == 1 {
240            return EntitySetIterator {
241                inner: EntitySetIteratorInner::Source(sources.pop().unwrap().into_iter()),
242            };
243        }
244
245        // This path constructs intersections from raw source vectors, so we sort here.
246        // We keep ascending order so filters checked by `all()` are smallest-first.
247        sources.sort_unstable_by_key(SourceSet::sort_key);
248        let mut source_iter = sources.into_iter();
249        let driver = source_iter.next().unwrap().into_iter();
250        EntitySetIterator {
251            inner: EntitySetIteratorInner::IntersectionSources {
252                driver,
253                filters: source_iter.collect(),
254            },
255        }
256    }
257
258    pub(crate) fn from_index_set(set: &'c IndexSet<EntityId<E>>) -> EntitySetIterator<'c, E> {
259        EntitySetIterator {
260            inner: EntitySetIteratorInner::Source(SourceSet::IndexSet(set).into_iter()),
261        }
262    }
263
264    pub(super) fn new(set: EntitySet<'c, E>) -> Self {
265        EntitySetIterator {
266            inner: EntitySetIteratorInner::from_entity_set(set),
267        }
268    }
269}
270
271impl<'a, E: Entity> Iterator for EntitySetIterator<'a, E> {
272    type Item = EntityId<E>;
273
274    #[inline]
275    fn next(&mut self) -> Option<Self::Item> {
276        self.inner.next_inner()
277    }
278
279    #[inline]
280    fn size_hint(&self) -> (usize, Option<usize>) {
281        self.inner.size_hint_inner()
282    }
283
284    fn count(self) -> usize {
285        let EntitySetIterator { inner } = self;
286        match inner {
287            EntitySetIteratorInner::Source(source) => source.count(),
288            EntitySetIteratorInner::IntersectionSources {
289                mut driver,
290                filters,
291            } => driver
292                .by_ref()
293                .filter(|&entity_id| filters.iter().all(|filter| filter.contains(entity_id)))
294                .count(),
295            other => {
296                let mut it = EntitySetIterator { inner: other };
297                let mut n = 0usize;
298                while it.next().is_some() {
299                    n += 1;
300                }
301                n
302            }
303        }
304    }
305
306    #[inline]
307    fn nth(&mut self, n: usize) -> Option<Self::Item> {
308        match &mut self.inner {
309            EntitySetIteratorInner::Source(source) => source.nth(n),
310            EntitySetIteratorInner::IntersectionSources { driver, filters } => driver
311                .by_ref()
312                .filter(|&entity_id| filters.iter().all(|filter| filter.contains(entity_id)))
313                .nth(n),
314            _ => {
315                for _ in 0..n {
316                    self.next()?;
317                }
318                self.next()
319            }
320        }
321    }
322
323    fn for_each<F>(self, mut f: F)
324    where
325        F: FnMut(Self::Item),
326    {
327        let EntitySetIterator { inner } = self;
328        match inner {
329            EntitySetIteratorInner::Source(source) => source.for_each(f),
330            other => {
331                let it = EntitySetIterator { inner: other };
332                for item in it {
333                    f(item);
334                }
335            }
336        }
337    }
338
339    fn fold<B, F>(self, init: B, mut f: F) -> B
340    where
341        F: FnMut(B, Self::Item) -> B,
342    {
343        let EntitySetIterator { inner } = self;
344        match inner {
345            EntitySetIteratorInner::Source(source) => source.fold(init, f),
346            other => {
347                let it = EntitySetIterator { inner: other };
348                let mut acc = init;
349                for item in it {
350                    acc = f(acc, item);
351                }
352                acc
353            }
354        }
355    }
356}
357
358impl<'c, E: Entity> std::iter::FusedIterator for EntitySetIterator<'c, E> {}
359
360#[cfg(test)]
361mod tests {
362    /*!
363    ## Test Matrix
364
365    | Test # | PropertyInitializationKind | is_default_value | Indexed | Initial Source Position |
366    | ------ | -------------------------- | ---------------- | ------- | ----------------------- |
367    | 1      | Explicit                   | N/A              | No      | Yes                     |
368    | 2      | Explicit                   | N/A              | Yes     | Yes                     |
369    | 3      | Constant                   | true             | No      | Yes                     |
370    | 4      | Constant                   | true             | Yes     | Yes                     |
371    | 5      | Constant                   | false            | No      | Yes                     |
372    | 6      | Constant                   | false            | Yes     | Yes                     |
373    | 7      | Derived                    | N/A              | No      | Yes                     |
374    | 8      | Derived                    | N/A              | Yes     | Yes                     |
375    | 9      | Explicit                   | N/A              | No      | No                      |
376    | 10     | Explicit                   | N/A              | Yes     | No                      |
377    | 11     | Constant                   | true             | No      | No                      |
378    | 12     | Constant                   | false            | Yes     | No                      |
379    | 13     | Derived                    | N/A              | No      | No                      |
380    | 14     | Derived                    | N/A              | Yes     | No                      |
381
382    The tests use multi-property queries (tests 9-14) where one property
383    is indexed with fewer results to ensure it becomes the "smallest source
384    set", making the tested property NOT the initial source position.
385    */
386
387    use indexmap::IndexSet;
388
389    use crate::entity::entity_set::{EntitySet, SourceSet};
390    use crate::hashing::IndexSet as FxIndexSet;
391    use crate::prelude::*;
392    use crate::{define_derived_property, define_property, with};
393
394    define_entity!(Person);
395
396    // Test properties covering different initialization kinds
397
398    // Explicit (Normal) property - no default, requires explicit initialization
399    define_property!(struct ExplicitProp(u8), Person);
400
401    // Constant property - has a constant default value
402    define_property!(struct ConstantProp(u8), Person, default_const = ConstantProp(42));
403
404    // Derived property - computed from other properties
405    define_derived_property!(struct DerivedProp(bool), Person, [ExplicitProp], |explicit| {
406        DerivedProp(explicit.0 % 2 == 0)
407    });
408
409    // Additional properties for multi-property queries
410    define_property!(struct ConstantProp2(u16), Person, default_const = ConstantProp2(100));
411    define_property!(struct ExplicitProp2(bool), Person);
412
413    define_property!(struct Age(u8), Person, default_const = Age(0));
414    define_property!(struct Alive(bool), Person, default_const = Alive(true));
415
416    define_derived_property!(
417        enum AgeGroupRisk {
418            NewBorn,
419            General,
420            OldAdult,
421        },
422        Person,
423        [Age],
424        [],
425        |age| {
426            if age.0 <= 1 {
427                AgeGroupRisk::NewBorn
428            } else if age.0 <= 65 {
429                AgeGroupRisk::General
430            } else {
431                AgeGroupRisk::OldAdult
432            }
433        }
434    );
435
436    // Helper function to create a population for testing
437    fn setup_test_population(context: &mut Context, size: usize) -> Vec<EntityId<Person>> {
438        let mut people = Vec::new();
439        for i in 0..size {
440            let person = context
441                .add_entity(with!(
442                    Person,
443                    ExplicitProp((i % 20) as u8),
444                    ExplicitProp2(i % 2 == 0)
445                ))
446                .unwrap();
447            people.push(person);
448        }
449        people
450    }
451
452    // region: Test Matrix Tests
453
454    // Test 1: Explicit property, non-default value, not indexed, initial source position = Yes
455    #[test]
456    fn test_explicit_non_default_not_indexed_initial_source_yes() {
457        let mut context = Context::new();
458        setup_test_population(&mut context, 100);
459
460        let results = context
461            .query_result_iterator(with!(Person, ExplicitProp(5)))
462            .collect::<Vec<_>>();
463
464        assert_eq!(results.len(), 5); // 5, 25, 45, 65, 85
465        for person in results {
466            assert_eq!(
467                context.get_property::<_, ExplicitProp>(person),
468                ExplicitProp(5)
469            );
470        }
471    }
472
473    // Test 2: Explicit property, non-default value, indexed, initial source position = Yes
474    #[test]
475    fn test_explicit_non_default_indexed_initial_source_yes() {
476        let mut context = Context::new();
477        context.index_property::<Person, ExplicitProp>();
478        setup_test_population(&mut context, 100);
479
480        let results = context
481            .query_result_iterator(with!(Person, ExplicitProp(7)))
482            .collect::<Vec<_>>();
483
484        assert_eq!(results.len(), 5); // 7, 27, 47, 67, 87
485        for person in results {
486            assert_eq!(
487                context.get_property::<_, ExplicitProp>(person),
488                ExplicitProp(7)
489            );
490        }
491    }
492
493    // Test 3: Constant property, default value, not indexed, initial source position = Yes
494    #[test]
495    fn test_constant_default_not_indexed_initial_source_yes() {
496        let mut context = Context::new();
497        // Create people without setting ConstantProp - they'll use default value 42
498        for _ in 0..50 {
499            context
500                .add_entity(with!(Person, ExplicitProp(1), ExplicitProp2(false)))
501                .unwrap();
502        }
503
504        let results = context
505            .query_result_iterator(with!(Person, ConstantProp(42), ExplicitProp2(false)))
506            .collect::<Vec<_>>();
507
508        assert_eq!(results.len(), 50);
509        for person in results {
510            assert_eq!(
511                context.get_property::<_, ConstantProp>(person),
512                ConstantProp(42)
513            );
514        }
515    }
516
517    // Test 4: Constant property, default value, indexed, initial source position = Yes
518    #[test]
519    fn test_constant_default_indexed_initial_source_yes() {
520        let mut context = Context::new();
521        context.index_property::<Person, ConstantProp>();
522
523        for _ in 0..50 {
524            context
525                .add_entity(with!(Person, ExplicitProp(1), ExplicitProp2(false)))
526                .unwrap();
527        }
528
529        let results = context
530            .query_result_iterator(with!(Person, ConstantProp(42)))
531            .collect::<Vec<_>>();
532        assert_eq!(results.len(), 50);
533    }
534
535    // Test 5: Constant property, non-default value, not indexed, initial source position = Yes
536    #[test]
537    fn test_constant_non_default_not_indexed_initial_source_yes() {
538        let mut context = Context::new();
539
540        for i in 0..50 {
541            if i < 10 {
542                context
543                    .add_entity(with!(
544                        Person,
545                        ExplicitProp(1),
546                        ExplicitProp2(false),
547                        ConstantProp(99)
548                    ))
549                    .unwrap();
550            } else {
551                context
552                    .add_entity(with!(Person, ExplicitProp(1), ExplicitProp2(false)))
553                    .unwrap();
554            }
555        }
556
557        let results = context
558            .query_result_iterator(with!(Person, ConstantProp(99)))
559            .collect::<Vec<_>>();
560
561        assert_eq!(results.len(), 10);
562        for person in results {
563            assert_eq!(
564                context.get_property::<_, ConstantProp>(person),
565                ConstantProp(99)
566            );
567        }
568    }
569
570    // Test 6: Constant property, non-default value, indexed, initial source position = Yes
571    #[test]
572    fn test_constant_non_default_indexed_initial_source_yes() {
573        let mut context = Context::new();
574        context.index_property::<Person, ConstantProp>();
575
576        for i in 0..50 {
577            if i < 10 {
578                context
579                    .add_entity(with!(
580                        Person,
581                        ExplicitProp(1),
582                        ExplicitProp2(false),
583                        ConstantProp(99)
584                    ))
585                    .unwrap();
586            } else {
587                context
588                    .add_entity(with!(Person, ExplicitProp(1), ExplicitProp2(false)))
589                    .unwrap();
590            }
591        }
592
593        let results = context
594            .query_result_iterator(with!(Person, ConstantProp(99)))
595            .collect::<Vec<_>>();
596
597        assert_eq!(results.len(), 10);
598    }
599
600    // Test 7: Derived property, not indexed, initial source position = Yes
601    #[test]
602    fn test_derived_not_indexed_initial_source_yes() {
603        let mut context = Context::new();
604
605        for i in 0..100 {
606            context
607                .add_entity(with!(Person, ExplicitProp(i as u8), ExplicitProp2(false)))
608                .unwrap();
609        }
610
611        let results = context
612            .query_result_iterator(with!(Person, DerivedProp(true)))
613            .collect::<Vec<_>>();
614
615        // DerivedProp is true when ExplicitProp is even
616        assert_eq!(results.len(), 50);
617        for person in results {
618            assert_eq!(
619                context.get_property::<Person, DerivedProp>(person),
620                DerivedProp(true)
621            );
622        }
623    }
624
625    // Test 8: Derived property, indexed, initial source position = Yes
626    #[test]
627    fn test_derived_indexed_initial_source_yes() {
628        let mut context = Context::new();
629        context.index_property::<Person, DerivedProp>();
630
631        for i in 0..100 {
632            context
633                .add_entity(with!(Person, ExplicitProp(i as u8), ExplicitProp2(false)))
634                .unwrap();
635        }
636
637        let results = context
638            .query_result_iterator(with!(Person, DerivedProp(false)))
639            .collect::<Vec<_>>();
640
641        // DerivedProp is false when ExplicitProp is odd
642        assert_eq!(results.len(), 50);
643        for person in results {
644            assert_eq!(
645                context.get_property::<Person, DerivedProp>(person),
646                DerivedProp(false)
647            );
648        }
649    }
650
651    // Test 9-14: Initial source position = No (multi-property queries where property is NOT the smallest source)
652
653    // Test 9: Explicit property, non-default, not indexed, initial source = No
654    #[test]
655    fn test_explicit_non_default_not_indexed_initial_source_no() {
656        let mut context = Context::new();
657        context.index_property::<Person, ExplicitProp2>(); // Index the other property so it's the smallest
658
659        for i in 0..100 {
660            context
661                .add_entity(with!(
662                    Person,
663                    ExplicitProp((i % 20) as u8),
664                    ExplicitProp2(i % 2 == 0)
665                ))
666                .unwrap();
667        }
668
669        let results = context.query_result_iterator(Person).collect::<Vec<_>>();
670        for person in results {
671            let explicit_prop = context.get_property::<Person, ExplicitProp>(person);
672            let explicit_prop2 = context.get_property::<Person, ExplicitProp2>(person);
673            println!("({:?} {:?} {:?})", person, explicit_prop, explicit_prop2);
674        }
675
676        // ExplicitProp2 has only 2 values, so it will be the smaller source
677        let results = context
678            .query_result_iterator(with!(Person, ExplicitProp(5), ExplicitProp2(false)))
679            .collect::<Vec<_>>();
680
681        // Looking for ExplicitProp=5 AND ExplicitProp2=true
682        // ExplicitProp cycles 0-19, ExplicitProp2 alternates
683        // Matches: 4 (since 5,25,45,65,85 but only even indices)
684        let expected = results.len();
685        assert!(expected > 0);
686        for person in results {
687            assert_eq!(
688                context.get_property::<Person, ExplicitProp>(person),
689                ExplicitProp(5)
690            );
691            assert_eq!(
692                context.get_property::<Person, ExplicitProp2>(person),
693                ExplicitProp2(false)
694            );
695        }
696    }
697
698    // Test 10: Explicit property, non-default, indexed, initial source = No
699    #[test]
700    fn test_explicit_non_default_indexed_initial_source_no() {
701        let mut context = Context::new();
702        context.index_property::<Person, ExplicitProp>();
703        context.index_property::<Person, ConstantProp2>(); // ConstantProp2 will likely be the smaller source
704
705        for i in 0..100 {
706            if i < 10 {
707                context
708                    .add_entity(with!(
709                        Person,
710                        ExplicitProp(7),
711                        ExplicitProp2(false),
712                        ConstantProp2(200), // Non-default for smaller source
713                    ))
714                    .unwrap();
715            } else {
716                context
717                    .add_entity(with!(
718                        Person,
719                        ExplicitProp((i % 20) as u8),
720                        ExplicitProp2(false)
721                    ))
722                    .unwrap();
723            }
724        }
725
726        let results = context
727            .query_result_iterator(with!(Person, ExplicitProp(7), ConstantProp2(200)))
728            .collect::<Vec<_>>();
729
730        assert_eq!(results.len(), 10);
731    }
732
733    // Test 11: Constant property, default value, not indexed, initial source = No
734    #[test]
735    fn test_constant_default_not_indexed_initial_source_no() {
736        let mut context = Context::new();
737        context.index_property::<Person, ExplicitProp>(); // Make ExplicitProp the smaller source
738
739        for i in 0..100 {
740            if i < 5 {
741                context
742                    .add_entity(with!(Person, ExplicitProp(99), ExplicitProp2(false)))
743                    .unwrap(); // ConstantProp uses default
744            } else {
745                context
746                    .add_entity(with!(
747                        Person,
748                        ExplicitProp((i % 20) as u8),
749                        ExplicitProp2(false)
750                    ))
751                    .unwrap();
752            }
753        }
754
755        let results = context
756            .query_result_iterator(with!(Person, ExplicitProp(99), ConstantProp(42)))
757            .collect::<Vec<_>>();
758
759        assert_eq!(results.len(), 5);
760        for person in results {
761            assert_eq!(
762                context.get_property::<Person, ConstantProp>(person),
763                ConstantProp(42)
764            );
765        }
766    }
767
768    // Test 12: Constant property, non-default, indexed, initial source = No
769    #[test]
770    fn test_constant_non_default_indexed_initial_source_no() {
771        let mut context = Context::new();
772        context.index_property::<Person, ConstantProp>();
773        context.index_property::<Person, ExplicitProp2>();
774
775        for i in 0..100 {
776            if i < 10 {
777                context
778                    .add_entity(with!(
779                        Person,
780                        ConstantProp(99),
781                        ExplicitProp(0),
782                        ExplicitProp2(true)
783                    ))
784                    .unwrap();
785            } else {
786                context
787                    .add_entity(with!(Person, ExplicitProp(0), ExplicitProp2(false)))
788                    .unwrap();
789            }
790        }
791
792        let results = context
793            .query_result_iterator(with!(Person, ConstantProp(99), ExplicitProp2(true)))
794            .collect::<Vec<_>>();
795
796        assert_eq!(results.len(), 10);
797    }
798
799    // Test 13: Derived property, not indexed, initial source = No
800    #[test]
801    fn test_derived_not_indexed_initial_source_no() {
802        let mut context = Context::new();
803        context.index_property::<Person, ExplicitProp2>();
804
805        for i in 0..100 {
806            context
807                .add_entity(with!(Person, ExplicitProp(i as u8), ExplicitProp2(i < 50)))
808                .unwrap();
809        }
810
811        let results = context
812            .query_result_iterator(with!(Person, ExplicitProp2(true), DerivedProp(true)))
813            .collect::<Vec<_>>();
814
815        // ExplicitProp2=true for i<50, DerivedProp=true when ExplicitProp is even
816        // So we want i<50 AND i is even: 0,2,4,...,48 = 25 people
817        assert_eq!(results.len(), 25);
818    }
819
820    // Test 14: Derived property, indexed, initial source = No
821    #[test]
822    fn test_derived_indexed_initial_source_no() {
823        let mut context = Context::new();
824        context.index_property::<Person, DerivedProp>();
825        context.index_property::<Person, ExplicitProp2>();
826
827        for i in 0..100 {
828            context
829                .add_entity(with!(Person, ExplicitProp(i as u8), ExplicitProp2(i < 30)))
830                .unwrap();
831        }
832
833        let results = context
834            .query_result_iterator(with!(Person, ExplicitProp2(true), DerivedProp(false)))
835            .collect::<Vec<_>>();
836
837        // ExplicitProp2=true for i<30, DerivedProp=false when ExplicitProp is odd
838        // So we want i < 30 AND i is odd: 1,3,5,...,29 = 15 people
839        assert_eq!(results.len(), 15);
840    }
841
842    // endregion Test Matrix Tests
843
844    #[test]
845    fn test_multiple_query_result_iterators() {
846        let mut context = Context::new();
847        context.index_property::<Person, Age>();
848
849        for age in 0..100 {
850            context
851                .add_entity(with!(
852                    Person,
853                    Age(age),
854                    ExplicitProp(age.wrapping_mul(7) % 100),
855                    ExplicitProp2(false),
856                ))
857                .unwrap();
858        }
859        for age in 0..100 {
860            context
861                .add_entity(with!(
862                    Person,
863                    Age(age),
864                    ExplicitProp(age.wrapping_mul(14) % 100),
865                    ExplicitProp2(false),
866                ))
867                .unwrap();
868        }
869
870        // Since both queries include `Age`, both will attempt to index unindexed entities. This tests that there is
871        // no double borrow error.
872        let results = context.query_result_iterator(with!(Person, Age(25)));
873        let more_results = context.query_result_iterator(with!(Person, Age(25), ExplicitProp(75)));
874
875        let collected_results = results.collect::<IndexSet<_>>();
876        let other_collected_results = more_results.collect::<IndexSet<_>>();
877        let intersection_count = collected_results
878            .intersection(&other_collected_results)
879            .count();
880        assert_eq!(intersection_count, 1);
881    }
882
883    #[test]
884    fn test_expression_intersection_iteration() {
885        let set = EntitySet::from_source(SourceSet::<Person>::PopulationRange(0..10)).intersection(
886            EntitySet::from_source(SourceSet::<Person>::PopulationRange(0..6)),
887        );
888
889        let ids = set.into_iter().collect::<Vec<_>>();
890        let expected = (0..6).map(EntityId::new).collect::<Vec<_>>();
891        assert_eq!(ids, expected);
892    }
893
894    #[test]
895    fn test_expression_difference_iteration() {
896        let set = EntitySet::from_source(SourceSet::<Person>::PopulationRange(0..5)).difference(
897            EntitySet::from_source(SourceSet::<Person>::PopulationRange(2..3)),
898        );
899
900        let ids = set.into_iter().collect::<Vec<_>>();
901        let expected = vec![
902            EntityId::new(0),
903            EntityId::new(1),
904            EntityId::new(3),
905            EntityId::new(4),
906        ];
907        assert_eq!(ids, expected);
908    }
909
910    #[test]
911    fn test_expression_union_deduplicates() {
912        let left = EntitySet::from_source(SourceSet::<Person>::PopulationRange(0..3)).difference(
913            EntitySet::from_source(SourceSet::<Person>::PopulationRange(99..100)),
914        );
915        let right = EntitySet::from_source(SourceSet::<Person>::PopulationRange(0..5)).difference(
916            EntitySet::from_source(SourceSet::<Person>::PopulationRange(99..100)),
917        );
918        let set = left.union(right);
919
920        let ids = set.into_iter().collect::<Vec<_>>();
921        let expected = (0..5).map(EntityId::new).collect::<Vec<_>>();
922        assert_eq!(ids, expected);
923    }
924
925    #[test]
926    fn test_expression_union_overlap_no_duplicates() {
927        let left = EntitySet::from_source(SourceSet::<Person>::PopulationRange(0..5)).difference(
928            EntitySet::from_source(SourceSet::<Person>::PopulationRange(4..5)),
929        );
930        let right = EntitySet::from_source(SourceSet::<Person>::PopulationRange(0..7)).difference(
931            EntitySet::from_source(SourceSet::<Person>::PopulationRange(0..1)),
932        );
933
934        let ids = left.union(right).into_iter().collect::<IndexSet<_>>();
935        let expected = (0..7).map(EntityId::new).collect::<IndexSet<_>>();
936        assert_eq!(ids, expected);
937    }
938
939    #[test]
940    fn test_expression_intersection_of_unions() {
941        let ab = EntitySet::from_source(SourceSet::<Person>::PopulationRange(1..2))
942            .union(EntitySet::from_source(
943                SourceSet::<Person>::PopulationRange(2..3),
944            ))
945            .union(EntitySet::from_source(
946                SourceSet::<Person>::PopulationRange(3..4),
947            ));
948        let cd = EntitySet::from_source(SourceSet::<Person>::PopulationRange(2..3))
949            .union(EntitySet::from_source(
950                SourceSet::<Person>::PopulationRange(3..4),
951            ))
952            .union(EntitySet::from_source(
953                SourceSet::<Person>::PopulationRange(4..5),
954            ));
955
956        let ids = ab.intersection(cd).into_iter().collect::<Vec<_>>();
957        assert_eq!(ids, vec![EntityId::new(2), EntityId::new(3)]);
958    }
959
960    #[test]
961    fn test_expression_difference_not_commutative() {
962        let left = EntitySet::from_source(SourceSet::<Person>::PopulationRange(1..2))
963            .union(EntitySet::from_source(
964                SourceSet::<Person>::PopulationRange(2..3),
965            ))
966            .union(EntitySet::from_source(
967                SourceSet::<Person>::PopulationRange(3..4),
968            ));
969        let right = EntitySet::from_source(SourceSet::<Person>::PopulationRange(2..3))
970            .union(EntitySet::from_source(
971                SourceSet::<Person>::PopulationRange(3..4),
972            ))
973            .union(EntitySet::from_source(
974                SourceSet::<Person>::PopulationRange(4..5),
975            ));
976
977        let left_minus_right = left.difference(right).into_iter().collect::<Vec<_>>();
978        let right_minus_left = EntitySet::from_source(SourceSet::<Person>::PopulationRange(2..3))
979            .union(EntitySet::from_source(
980                SourceSet::<Person>::PopulationRange(3..4),
981            ))
982            .union(EntitySet::from_source(
983                SourceSet::<Person>::PopulationRange(4..5),
984            ))
985            .difference(
986                EntitySet::from_source(SourceSet::<Person>::PopulationRange(1..2))
987                    .union(EntitySet::from_source(
988                        SourceSet::<Person>::PopulationRange(2..3),
989                    ))
990                    .union(EntitySet::from_source(
991                        SourceSet::<Person>::PopulationRange(3..4),
992                    )),
993            )
994            .into_iter()
995            .collect::<Vec<_>>();
996
997        assert_eq!(left_minus_right, vec![EntityId::new(1)]);
998        assert_eq!(right_minus_left, vec![EntityId::new(4)]);
999    }
1000
1001    #[test]
1002    fn test_union_size_hint_pending_right_is_unknown() {
1003        let left = EntitySet::from_source(SourceSet::<Person>::PopulationRange(0..2)).difference(
1004            EntitySet::from_source(SourceSet::<Person>::PopulationRange(99..100)),
1005        );
1006        let right = EntitySet::from_source(SourceSet::<Person>::PopulationRange(0..4)).difference(
1007            EntitySet::from_source(SourceSet::<Person>::PopulationRange(98..99)),
1008        );
1009        let iter = left.union(right).into_iter();
1010
1011        assert_eq!(iter.size_hint(), (4, Some(4)));
1012    }
1013
1014    #[test]
1015    fn test_nth_and_count_on_source() {
1016        let mut iter =
1017            EntitySet::from_source(SourceSet::<Person>::PopulationRange(0..5)).into_iter();
1018        assert_eq!(iter.nth(2), Some(EntityId::new(2)));
1019
1020        let remaining = iter.count();
1021        assert_eq!(remaining, 2);
1022    }
1023
1024    #[test]
1025    fn population_range_source_iteration_and_size_hint() {
1026        let mut iter =
1027            EntitySet::from_source(SourceSet::<Person>::PopulationRange(3..7)).into_iter();
1028        assert_eq!(iter.size_hint(), (4, Some(4)));
1029        assert_eq!(iter.next(), Some(EntityId::new(3)));
1030        assert_eq!(iter.size_hint(), (3, Some(3)));
1031        assert_eq!(
1032            iter.collect::<Vec<_>>(),
1033            vec![EntityId::new(4), EntityId::new(5), EntityId::new(6)]
1034        );
1035    }
1036
1037    #[test]
1038    fn optimized_range_results_iterate_as_single_source() {
1039        let union = EntitySet::from_source(SourceSet::<Person>::PopulationRange(0..3)).union(
1040            EntitySet::from_source(SourceSet::<Person>::PopulationRange(3..5)),
1041        );
1042        assert_eq!(
1043            union.into_iter().collect::<Vec<_>>(),
1044            vec![
1045                EntityId::new(0),
1046                EntityId::new(1),
1047                EntityId::new(2),
1048                EntityId::new(3),
1049                EntityId::new(4),
1050            ]
1051        );
1052
1053        let intersection = EntitySet::from_source(SourceSet::<Person>::PopulationRange(2..8))
1054            .intersection(EntitySet::from_source(
1055                SourceSet::<Person>::PopulationRange(4..6),
1056            ));
1057        assert_eq!(
1058            intersection.into_iter().collect::<Vec<_>>(),
1059            vec![EntityId::new(4), EntityId::new(5)]
1060        );
1061    }
1062
1063    #[test]
1064    fn split_difference_still_uses_generic_iteration_path() {
1065        let split = EntitySet::from_source(SourceSet::<Person>::PopulationRange(2..8)).difference(
1066            EntitySet::from_source(SourceSet::<Person>::PopulationRange(4..6)),
1067        );
1068        assert_eq!(
1069            split.into_iter().collect::<Vec<_>>(),
1070            vec![
1071                EntityId::new(2),
1072                EntityId::new(3),
1073                EntityId::new(6),
1074                EntityId::new(7)
1075            ]
1076        );
1077    }
1078
1079    fn finite_set(ids: &[usize]) -> FxIndexSet<EntityId<Person>> {
1080        ids.iter()
1081            .copied()
1082            .map(EntityId::new)
1083            .collect::<FxIndexSet<_>>()
1084    }
1085
1086    fn as_entity_set(set: &FxIndexSet<EntityId<Person>>) -> EntitySet<Person> {
1087        EntitySet::from_source(SourceSet::IndexSet(set))
1088    }
1089
1090    #[test]
1091    fn prototype_empty_entity_set_yields_nothing() {
1092        let ids = EntitySet::<Person>::empty().into_iter().collect::<Vec<_>>();
1093        assert!(ids.is_empty());
1094    }
1095
1096    #[test]
1097    fn prototype_iter_union_disjoint() {
1098        let a = finite_set(&[1, 2]);
1099        let b = finite_set(&[3, 4]);
1100        let ids = as_entity_set(&a)
1101            .union(as_entity_set(&b))
1102            .into_iter()
1103            .collect::<IndexSet<_>>();
1104        let expected = [1usize, 2, 3, 4]
1105            .into_iter()
1106            .map(EntityId::new)
1107            .collect::<IndexSet<_>>();
1108        assert_eq!(ids, expected);
1109    }
1110
1111    #[test]
1112    fn prototype_iter_union_overlapping() {
1113        let a = finite_set(&[1, 2, 3]);
1114        let b = finite_set(&[2, 3, 4]);
1115        let ids = as_entity_set(&a)
1116            .union(as_entity_set(&b))
1117            .into_iter()
1118            .collect::<Vec<_>>();
1119        let unique = ids.iter().copied().collect::<IndexSet<_>>();
1120        assert_eq!(ids.len(), unique.len());
1121        assert!(unique.contains(&EntityId::new(1)));
1122        assert!(unique.contains(&EntityId::new(4)));
1123    }
1124
1125    #[test]
1126    fn prototype_iter_intersection_disjoint() {
1127        let a = finite_set(&[1, 2]);
1128        let b = finite_set(&[3, 4]);
1129        let ids = as_entity_set(&a)
1130            .intersection(as_entity_set(&b))
1131            .into_iter()
1132            .collect::<Vec<_>>();
1133        assert!(ids.is_empty());
1134    }
1135
1136    #[test]
1137    fn prototype_iter_difference_basic() {
1138        let a = finite_set(&[1, 2, 3, 4]);
1139        let b = finite_set(&[3, 4, 5, 6]);
1140        let ids = as_entity_set(&a)
1141            .difference(as_entity_set(&b))
1142            .into_iter()
1143            .collect::<IndexSet<_>>();
1144        assert_eq!(ids.len(), 2);
1145        assert!(ids.contains(&EntityId::new(1)));
1146        assert!(ids.contains(&EntityId::new(2)));
1147    }
1148
1149    #[test]
1150    fn prototype_iter_matches_contains_compound() {
1151        let a = finite_set(&[1, 2, 3, 4]);
1152        let b = finite_set(&[3, 4, 5]);
1153        let c = finite_set(&[7, 8, 9, 10]);
1154        let d = finite_set(&[9, 10, 11]);
1155        let left = as_entity_set(&a).intersection(as_entity_set(&b));
1156        let right = as_entity_set(&c).difference(as_entity_set(&d));
1157        let iterated = left.union(right).into_iter().collect::<IndexSet<_>>();
1158
1159        let a2 = finite_set(&[1, 2, 3, 4]);
1160        let b2 = finite_set(&[3, 4, 5]);
1161        let c2 = finite_set(&[7, 8, 9, 10]);
1162        let d2 = finite_set(&[9, 10, 11]);
1163        let check = as_entity_set(&a2)
1164            .intersection(as_entity_set(&b2))
1165            .union(as_entity_set(&c2).difference(as_entity_set(&d2)));
1166
1167        for value in 0..15 {
1168            let entity = EntityId::new(value);
1169            assert_eq!(iterated.contains(&entity), check.contains(entity));
1170        }
1171    }
1172
1173    #[test]
1174    fn prototype_size_hint_single_source_and_partial_consume() {
1175        let mut iter =
1176            EntitySet::from_source(SourceSet::<Person>::PopulationRange(0..5)).into_iter();
1177        assert_eq!(iter.size_hint(), (5, Some(5)));
1178        iter.next();
1179        assert_eq!(iter.size_hint(), (4, Some(4)));
1180    }
1181}