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

ixa/entity/entity_set/
entity_set_iterator.rs
