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

ixa/entity/entity_set/
entity_set_iterator.rs
