ixa/entity/
context_extension.rs

1use std::any::{Any, TypeId};
2
3use crate::entity::entity_set::EntitySetIterator;
4use crate::entity::events::{EntityCreatedEvent, PartialPropertyChangeEvent};
5use crate::entity::index::{IndexCountResult, IndexSetResult, PropertyIndexType};
6use crate::entity::property::Property;
7use crate::entity::property_list::PropertyList;
8use crate::entity::query::Query;
9use crate::entity::{Entity, EntityId, PopulationIterator};
10use crate::hashing::IndexSet;
11use crate::rand::Rng;
12use crate::{warn, Context, ContextRandomExt, IxaError, RngId};
13
14/// A trait extension for [`Context`] that exposes entity-related
15/// functionality.
16pub trait ContextEntitiesExt {
17    fn add_entity<E: Entity, PL: PropertyList<E>>(
18        &mut self,
19        property_list: PL,
20    ) -> Result<EntityId<E>, IxaError>;
21
22    /// Fetches the property value set for the given `entity_id`.
23    ///
24    /// The easiest way to call this method is by assigning it to a variable with an explicit type:
25    /// ```rust, ignore
26    /// let vaccine_status: VaccineStatus = context.get_property(entity_id);
27    /// ```
28    fn get_property<E: Entity, P: Property<E>>(&self, entity_id: EntityId<E>) -> P;
29
30    /// Sets the value of the given property. This method unconditionally emits a `PropertyChangeEvent`.
31    fn set_property<E: Entity, P: Property<E>>(
32        &mut self,
33        entity_id: EntityId<E>,
34        property_value: P,
35    );
36
37    /// Enables indexing of property values for the property `P`.
38    ///
39    /// This method is called with the turbo-fish syntax:
40    ///     `context.index_property::<Person, Age>()`
41    /// The actual computation of the index is done lazily as needed upon execution of queries,
42    /// not when this method is called.
43    fn index_property<E: Entity, P: Property<E>>(&mut self);
44
45    /// Enables value-count indexing of property values for the property `P`.
46    ///
47    /// If the property already has a full index, that index is left unchanged, as it
48    /// already supports value-count queries.
49    fn index_property_counts<E: Entity, P: Property<E>>(&mut self);
50
51    /// Checks if a property `P` is indexed.
52    ///
53    /// This method is called with the turbo-fish syntax:
54    ///     `context.index_property::<Person, Age>()`
55    ///
56    /// This method can return `true` even if `context.index_property::<P>()` has never been called. For example,
57    /// if a multi-property is indexed, all equivalent multi-properties are automatically also indexed, as they
58    /// share a single index.
59    #[cfg(test)]
60    fn is_property_indexed<E: Entity, P: Property<E>>(&self) -> bool;
61
62    /// This method gives client code direct immutable access to the fully realized set of
63    /// entity IDs. This is especially efficient for indexed queries, as this method reduces
64    /// to a simple lookup of a hash bucket. Otherwise, the set is allocated and computed.
65    fn with_query_results<E: Entity, Q: Query<E>>(
66        &self,
67        query: Q,
68        callback: &mut dyn FnMut(&IndexSet<EntityId<E>>),
69    );
70
71    /// Gives the count of distinct entity IDs satisfying the query. This is especially
72    /// efficient for indexed queries.
73    ///
74    /// Supplying an empty query `()` is equivalent to calling `get_entity_count::<E>()`.
75    fn query_entity_count<E: Entity, Q: Query<E>>(&self, query: Q) -> usize;
76
77    /// Sample a single entity uniformly from the query results. Returns `None` if the
78    /// query's result set is empty.
79    ///
80    /// To sample from the entire population, pass in the empty query `()`.
81    fn sample_entity<E, Q, R>(&self, rng_id: R, query: Q) -> Option<EntityId<E>>
82    where
83        E: Entity,
84        Q: Query<E>,
85        R: RngId + 'static,
86        R::RngType: Rng;
87
88    /// Sample up to `requested` entities uniformly from the query results. If the
89    /// query's result set has fewer than `requested` entities, the entire result
90    /// set is returned.
91    ///
92    /// To sample from the entire population, pass in the empty query `()`.
93    fn sample_entities<E, Q, R>(&self, rng_id: R, query: Q, n: usize) -> Vec<EntityId<E>>
94    where
95        E: Entity,
96        Q: Query<E>,
97        R: RngId + 'static,
98        R::RngType: Rng;
99
100    /// Returns a total count of all created entities of type `E`.
101    fn get_entity_count<E: Entity>(&self) -> usize;
102
103    /// Returns an iterator over all created entities of type `E`.
104    fn get_entity_iterator<E: Entity>(&self) -> PopulationIterator<E>;
105
106    /// Generates an iterator over the results of the query.
107    fn query_result_iterator<E: Entity, Q: Query<E>>(&self, query: Q) -> EntitySetIterator<E>;
108
109    /// Determines if the given person matches this query.
110    fn match_entity<E: Entity, Q: Query<E>>(&self, entity_id: EntityId<E>, query: Q) -> bool;
111
112    /// Removes all `EntityId`s from the given vector that do not match the given query.
113    fn filter_entities<E: Entity, Q: Query<E>>(&self, entities: &mut Vec<EntityId<E>>, query: Q);
114}
115
116impl ContextEntitiesExt for Context {
117    fn add_entity<E: Entity, PL: PropertyList<E>>(
118        &mut self,
119        property_list: PL,
120    ) -> Result<EntityId<E>, IxaError> {
121        // Check that the properties in the list are distinct.
122        PL::validate()?;
123
124        // Check that all required properties are present.
125        if !PL::contains_required_properties() {
126            return Err("initialization list is missing required properties".into());
127        }
128
129        // Now that we know we will succeed, we create the entity.
130        let new_entity_id = self.entity_store.new_entity_id::<E>();
131
132        // Assign the properties in the list to the new entity.
133        // This does not generate a property change event.
134        property_list
135            .set_values_for_entity(new_entity_id, self.entity_store.get_property_store::<E>());
136
137        // Emit an `EntityCreatedEvent<Entity>`.
138        self.emit_event(EntityCreatedEvent::<E>::new(new_entity_id));
139
140        Ok(new_entity_id)
141    }
142
143    fn get_property<E: Entity, P: Property<E>>(&self, entity_id: EntityId<E>) -> P {
144        if P::is_derived() {
145            P::compute_derived(self, entity_id)
146        } else {
147            let property_store = self.get_property_value_store::<E, P>();
148            property_store.get(entity_id)
149        }
150    }
151
152    fn set_property<E: Entity, P: Property<E>>(
153        &mut self,
154        entity_id: EntityId<E>,
155        property_value: P,
156    ) {
157        debug_assert!(!P::is_derived(), "cannot set a derived property");
158
159        // The algorithm is as follows
160        // 1. Get the previous value of the property.
161        //    1.1 If it's the same as `property_value`, exit.
162        //    1.2 Otherwise, create a `PartialPropertyChangeEvent<E, P>`.
163        // 2. Remove the `entity_id` from the index bucket corresponding to its old value.
164        // 3. For each dependent of the property, do the analog of steps 1 & 2:
165        //    3.1 Compute the previous value of the dependent property `Q`, creating a
166        //        `PartialPropertyChangeEvent<E, Q>` instance if necessary.
167        //    3.2 Remove the `entity_id` from the index bucket corresponding to the old value of `Q`.
168        // 4. Set the new value of the (main) property in the property store.
169        // 5. Update the property index: Insert the `entity_id` into the index bucket corresponding to the new value.
170        // 6. Emit the property change event: convert the `PartialPropertyChangeEvent<E, P>` into a
171        //    `event: PropertyChangeEvent<E, P>` and call `Context::emit_event(event)`.
172        // 7. For each dependent of the property, do the analog of steps 4-6:
173        //    7.1 Compute the new value of the dependent property
174        //    7.2 Add `entity_id` to the index bucket corresponding to the new value.
175        //    7.3 convert the `PartialPropertyChangeEvent<E, Q>` into a
176        //        `event: PropertyChangeEvent<E, Q>` and call `Context::emit_event(event)`.
177
178        // We need two passes over the dependents: one pass to compute all the old values and
179        // another to compute all the new values. We group the steps for each dependent (and, it
180        // turns out, for the main property `P` as well) into two parts:
181        //  1. Before setting the main property `P`, factored out into
182        //     `self.property_store.create_partial_property_change`
183        //  2. After setting the main property `P`, factored out into
184        //     `PartialPropertyChangeEvent::emit_in_context`
185
186        // We decided not to do the following check:
187        // ```rust
188        // let previous_value = { self.get_property_value_store::<E, P>().get(entity_id) };
189        // if property_value == previous_value {
190        //     return;
191        // }
192        // ```
193        // The reasoning is:
194        // - It should be rare that we ever set a property to its present value.
195        // - It's not a significant burden on client code to check `property_value == previous_value` on
196        //   their own if they need to.
197        // - There may be use cases for listening to "writes" that don't actually change values.
198
199        let mut dependents: Vec<Box<dyn PartialPropertyChangeEvent>> = vec![];
200        let property_store = self.entity_store.get_property_store::<E>();
201
202        // Create the partial property change for this value.
203        dependents.push(property_store.create_partial_property_change(P::id(), entity_id, self));
204        // Now create partial property change events for each dependent.
205        for dependent_idx in P::dependents() {
206            dependents.push(property_store.create_partial_property_change(
207                *dependent_idx,
208                entity_id,
209                self,
210            ));
211        }
212
213        // Update the value
214        let property_value_store = self.get_property_value_store::<E, P>();
215        property_value_store.set(entity_id, property_value);
216
217        // After updating the value
218        for dependent in dependents.into_iter() {
219            dependent.emit_in_context(self)
220        }
221    }
222
223    fn index_property<E: Entity, P: Property<E>>(&mut self) {
224        let property_store = self.entity_store.get_property_store_mut::<E>();
225        property_store.set_property_indexed::<P>(PropertyIndexType::FullIndex);
226    }
227
228    fn index_property_counts<E: Entity, P: Property<E>>(&mut self) {
229        let property_store = self.entity_store.get_property_store_mut::<E>();
230        let current_index_type = property_store.get::<P>().index_type();
231        if current_index_type != PropertyIndexType::FullIndex {
232            property_store.set_property_indexed::<P>(PropertyIndexType::ValueCountIndex);
233        }
234    }
235
236    #[cfg(test)]
237    fn is_property_indexed<E: Entity, P: Property<E>>(&self) -> bool {
238        let property_store = self.entity_store.get_property_store::<E>();
239        property_store.is_property_indexed::<P>()
240    }
241
242    fn with_query_results<E: Entity, Q: Query<E>>(
243        &self,
244        query: Q,
245        callback: &mut dyn FnMut(&IndexSet<EntityId<E>>),
246    ) {
247        // The fast path for indexed queries.
248
249        // This mirrors the indexed case in `SourceSet<'a, E>::new()` and `Query:: new_query_result_iterator`.
250        // The difference is, we access the index set if we find it.
251        if let Some(multi_property_id) = query.multi_property_id() {
252            let property_store = self.entity_store.get_property_store::<E>();
253            match property_store.get_index_set_with_hash_for_property_id(
254                self,
255                multi_property_id,
256                query.multi_property_value_hash(),
257            ) {
258                IndexSetResult::Set(people_set) => {
259                    callback(&people_set);
260                    return;
261                }
262                IndexSetResult::Empty => {
263                    let people_set = IndexSet::default();
264                    callback(&people_set);
265                    return;
266                }
267                IndexSetResult::Unsupported => {}
268            }
269            // If the property is not indexed, we fall through.
270        }
271
272        // Special case the empty query, which creates a set containing the entire population.
273        if query.type_id() == TypeId::of::<()>() {
274            warn!("Called Context::with_query_results() with an empty query. Prefer Context::get_entity_iterator::<E>() for working with the entire population.");
275            let entity_set = self.get_entity_iterator::<E>().collect::<IndexSet<_>>();
276            callback(&entity_set);
277            return;
278        }
279
280        // The slow path of computing the full query set.
281        warn!("Called Context::with_query_results() with an unindexed query. It's almost always better to use Context::query_result_iterator() for unindexed queries.");
282
283        // Fall back to `EntitySetIterator`.
284        let people_set = query
285            .new_query_result_iterator(self)
286            .collect::<IndexSet<_>>();
287        callback(&people_set);
288    }
289
290    fn query_entity_count<E: Entity, Q: Query<E>>(&self, query: Q) -> usize {
291        // The fast path for indexed queries.
292        //
293        // This mirrors the indexed case in `SourceSet<'a, E>::new()` and `Query:: new_query_result_iterator`.
294        if let Some(multi_property_id) = query.multi_property_id() {
295            let property_store = self.entity_store.get_property_store::<E>();
296            match property_store.get_index_count_with_hash_for_property_id(
297                self,
298                multi_property_id,
299                query.multi_property_value_hash(),
300            ) {
301                IndexCountResult::Count(count) => return count,
302                IndexCountResult::Unsupported => {}
303            }
304            // If the property is not indexed, we fall through.
305        }
306
307        self.query_result_iterator(query).count()
308    }
309    fn sample_entity<E, Q, R>(&self, rng_id: R, query: Q) -> Option<EntityId<E>>
310    where
311        E: Entity,
312        Q: Query<E>,
313        R: RngId + 'static,
314        R::RngType: Rng,
315    {
316        let query_result = self.query_result_iterator(query);
317        self.sample(rng_id, move |rng| query_result.sample_entity(rng))
318    }
319
320    fn sample_entities<E, Q, R>(&self, rng_id: R, query: Q, n: usize) -> Vec<EntityId<E>>
321    where
322        E: Entity,
323        Q: Query<E>,
324        R: RngId + 'static,
325        R::RngType: Rng,
326    {
327        let query_result = self.query_result_iterator(query);
328        self.sample(rng_id, move |rng| query_result.sample_entities(rng, n))
329    }
330
331    fn get_entity_count<E: Entity>(&self) -> usize {
332        self.entity_store.get_entity_count::<E>()
333    }
334
335    fn get_entity_iterator<E: Entity>(&self) -> PopulationIterator<E> {
336        self.entity_store.get_entity_iterator::<E>()
337    }
338
339    fn query_result_iterator<E: Entity, Q: Query<E>>(&self, query: Q) -> EntitySetIterator<E> {
340        query.new_query_result_iterator(self)
341    }
342
343    fn match_entity<E: Entity, Q: Query<E>>(&self, entity_id: EntityId<E>, query: Q) -> bool {
344        query.match_entity(entity_id, self)
345    }
346
347    fn filter_entities<E: Entity, Q: Query<E>>(&self, entities: &mut Vec<EntityId<E>>, query: Q) {
348        query.filter_entities(entities, self);
349    }
350}
351
352#[cfg(test)]
353mod tests {
354    use std::cell::{Ref, RefCell};
355    use std::rc::Rc;
356
357    use super::*;
358    use crate::hashing::IndexSet;
359    use crate::prelude::PropertyChangeEvent;
360    use crate::{define_derived_property, define_entity, define_multi_property, define_property};
361
362    define_entity!(Animal);
363    define_property!(struct Legs(u8), Animal, default_const = Legs(4));
364
365    define_entity!(Person);
366
367    define_property!(struct Age(u8), Person);
368
369    define_property!(
370        enum InfectionStatus {
371            Susceptible,
372            Infected,
373            Recovered,
374        },
375        Person,
376        default_const = InfectionStatus::Susceptible
377    );
378
379    define_property!(
380        struct Vaccinated(bool),
381        Person,
382        default_const = Vaccinated(false)
383    );
384
385    define_derived_property!(
386        enum AgeGroup {
387            Child,
388            Adult,
389            Senior,
390        },
391        Person,
392        [Age],
393        |age| {
394            if age.0 <= 18 {
395                AgeGroup::Child
396            } else if age.0 <= 65 {
397                AgeGroup::Adult
398            } else {
399                AgeGroup::Senior
400            }
401        }
402    );
403
404    define_derived_property!(
405        enum RiskLevel {
406            Low,
407            Medium,
408            High,
409        },
410        Person,
411        [AgeGroup, Vaccinated, InfectionStatus],
412        |age_group, vaccinated, infection_status| {
413            match (age_group, vaccinated, infection_status) {
414                (AgeGroup::Senior, Vaccinated(false), InfectionStatus::Susceptible) => {
415                    RiskLevel::High
416                }
417                (_, Vaccinated(false), InfectionStatus::Susceptible) => RiskLevel::Medium,
418                _ => RiskLevel::Low,
419            }
420        }
421    );
422
423    // ToDo(RobertJacobsonCDC): Enable this once #691 is resolved, https://github.com/CDCgov/ixa/issues/691.
424    // define_global_property!(GlobalDummy, u8);
425    // define_derived_property!(
426    //     struct MyDerivedProperty(u8),
427    //     Person,
428    //     [Age],
429    //     [GlobalDummy],
430    //     |age, global_dummy| {
431    //         MyDerivedProperty(age.0 + global_dummy)
432    //     }
433    // );
434
435    // Derived properties in a diamond dependency relationship
436    define_property!(struct IsRunner(bool), Person, default_const = IsRunner(false));
437    define_property!(struct IsSwimmer(bool), Person, default_const = IsSwimmer(false));
438    define_derived_property!(
439        struct AdultRunner(bool),
440        Person,
441        [AgeGroup, IsRunner],
442        | age_group, is_runner | {
443            AdultRunner(
444                age_group == AgeGroup::Adult
445                && is_runner.0
446            )
447        }
448    );
449    define_derived_property!(
450        struct AdultSwimmer(bool),
451        Person,
452        [AgeGroup, IsSwimmer],
453        | age_group, is_swimmer | {
454            AdultSwimmer(
455                age_group == AgeGroup::Adult
456                && is_swimmer.0
457            )
458        }
459    );
460    define_derived_property!(
461        struct AdultAthlete(bool),
462        Person,
463        [AdultSwimmer, AdultRunner],
464        | adult_swimmer, adult_runner | {
465            AdultAthlete(
466                adult_swimmer.0 || adult_runner.0
467            )
468        }
469    );
470
471    #[test]
472    fn add_and_count_entities() {
473        let mut context = Context::new();
474
475        let _person1 = context
476            .add_entity((Age(12), InfectionStatus::Susceptible, Vaccinated(true)))
477            .unwrap();
478        assert_eq!(context.get_entity_count::<Person>(), 1);
479
480        let _person2 = context.add_entity((Age(34), Vaccinated(true))).unwrap();
481        assert_eq!(context.get_entity_count::<Person>(), 2);
482
483        // Age is the only required property
484        let _person3 = context.add_entity((Age(120),)).unwrap();
485        assert_eq!(context.get_entity_count::<Person>(), 3);
486    }
487
488    // Helper for index tests
489    #[derive(Copy, Clone)]
490    enum IndexMode {
491        Unindexed,
492        FullIndex,
493        ValueCountIndex,
494    }
495
496    // Returns `(context, existing_value, missing_value)`
497    fn setup_context_for_index_tests(index_mode: IndexMode) -> (Context, Age, Age) {
498        let mut context = Context::new();
499        match index_mode {
500            IndexMode::Unindexed => {}
501            IndexMode::FullIndex => context.index_property::<Person, Age>(),
502            IndexMode::ValueCountIndex => context.index_property_counts::<Person, Age>(),
503        }
504
505        let existing_value = Age(12);
506        let missing_value = Age(99);
507
508        let _ = context.add_entity((existing_value,)).unwrap();
509        let _ = context.add_entity((existing_value,)).unwrap();
510
511        (context, existing_value, missing_value)
512    }
513
514    #[test]
515    fn query_results_respect_index_modes() {
516        let modes = [
517            IndexMode::Unindexed,
518            IndexMode::FullIndex,
519            IndexMode::ValueCountIndex,
520        ];
521
522        for mode in modes {
523            let (context, existing_value, missing_value) = setup_context_for_index_tests(mode);
524
525            let mut existing_len = 0;
526            context.with_query_results((existing_value,), &mut |people_set| {
527                existing_len = people_set.len();
528            });
529            assert_eq!(existing_len, 2);
530
531            let mut missing_len = 0;
532            context.with_query_results((missing_value,), &mut |people_set| {
533                missing_len = people_set.len();
534            });
535            assert_eq!(missing_len, 0);
536
537            let existing_count = context.query_result_iterator((existing_value,)).count();
538            assert_eq!(existing_count, 2);
539
540            let missing_count = context.query_result_iterator((missing_value,)).count();
541            assert_eq!(missing_count, 0);
542
543            assert_eq!(context.query_entity_count((existing_value,)), 2);
544            assert_eq!(context.query_entity_count((missing_value,)), 0);
545        }
546    }
547
548    #[test]
549    fn add_an_entity_without_required_properties() {
550        let mut context = Context::new();
551        let result = context.add_entity((InfectionStatus::Susceptible, Vaccinated(true)));
552
553        assert!(matches!(
554            result,
555            Err(crate::IxaError::IxaError(ref msg)) if msg == "initialization list is missing required properties"
556        ));
557    }
558
559    #[test]
560    fn new_entities_have_default_values() {
561        let mut context = Context::new();
562
563        // Create a person with required Age property
564        let person = context.add_entity((Age(25),)).unwrap();
565
566        // Retrieve and check their values
567        let age: Age = context.get_property(person);
568        assert_eq!(age, Age(25));
569        let infection_status: InfectionStatus = context.get_property(person);
570        assert_eq!(infection_status, InfectionStatus::Susceptible);
571        let vaccinated: Vaccinated = context.get_property(person);
572        assert_eq!(vaccinated, Vaccinated(false));
573
574        // Change them
575        context.set_property(person, Age(26));
576        context.set_property(person, InfectionStatus::Infected);
577        context.set_property(person, Vaccinated(true));
578
579        // Retrieve and check their values
580        let age: Age = context.get_property(person);
581        assert_eq!(age, Age(26));
582        let infection_status: InfectionStatus = context.get_property(person);
583        assert_eq!(infection_status, InfectionStatus::Infected);
584        let vaccinated: Vaccinated = context.get_property(person);
585        assert_eq!(vaccinated, Vaccinated(true));
586    }
587
588    #[test]
589    fn get_and_set_property_explicit() {
590        let mut context = Context::new();
591
592        // Create a person with explicit property values
593        let person = context
594            .add_entity((Age(25), InfectionStatus::Recovered, Vaccinated(true)))
595            .unwrap();
596
597        // Retrieve and check their values
598        let age: Age = context.get_property(person);
599        assert_eq!(age, Age(25));
600        let infection_status: InfectionStatus = context.get_property(person);
601        assert_eq!(infection_status, InfectionStatus::Recovered);
602        let vaccinated: Vaccinated = context.get_property(person);
603        assert_eq!(vaccinated, Vaccinated(true));
604
605        // Change them
606        context.set_property(person, Age(26));
607        context.set_property(person, InfectionStatus::Infected);
608        context.set_property(person, Vaccinated(false));
609
610        // Retrieve and check their values
611        let age: Age = context.get_property(person);
612        assert_eq!(age, Age(26));
613        let infection_status: InfectionStatus = context.get_property(person);
614        assert_eq!(infection_status, InfectionStatus::Infected);
615        let vaccinated: Vaccinated = context.get_property(person);
616        assert_eq!(vaccinated, Vaccinated(false));
617    }
618
619    #[test]
620    fn count_entities() {
621        let mut context = Context::new();
622
623        assert_eq!(context.get_entity_count::<Animal>(), 0);
624        assert_eq!(context.get_entity_count::<Person>(), 0);
625
626        // Create entities of different kinds
627        for _ in 0..7 {
628            let _: PersonId = context.add_entity((Age(25),)).unwrap();
629        }
630        for _ in 0..5 {
631            let _: AnimalId = context.add_entity((Legs(2),)).unwrap();
632        }
633
634        assert_eq!(context.get_entity_count::<Animal>(), 5);
635        assert_eq!(context.get_entity_count::<Person>(), 7);
636
637        let _: PersonId = context.add_entity((Age(30),)).unwrap();
638        let _: AnimalId = context.add_entity((Legs(8),)).unwrap();
639
640        assert_eq!(context.get_entity_count::<Animal>(), 6);
641        assert_eq!(context.get_entity_count::<Person>(), 8);
642    }
643
644    #[test]
645    fn get_derived_property_multiple_deps() {
646        let mut context = Context::new();
647
648        let expected_high_id: PersonId = context
649            .add_entity((Age(77), Vaccinated(false), InfectionStatus::Susceptible))
650            .unwrap();
651        let expected_med_id: PersonId = context
652            .add_entity((Age(30), Vaccinated(false), InfectionStatus::Susceptible))
653            .unwrap();
654        let expected_low_id: PersonId = context
655            .add_entity((Age(3), Vaccinated(true), InfectionStatus::Recovered))
656            .unwrap();
657
658        let actual_high: RiskLevel = context.get_property(expected_high_id);
659        assert_eq!(actual_high, RiskLevel::High);
660        let actual_med: RiskLevel = context.get_property(expected_med_id);
661        assert_eq!(actual_med, RiskLevel::Medium);
662        let actual_low: RiskLevel = context.get_property(expected_low_id);
663        assert_eq!(actual_low, RiskLevel::Low);
664    }
665
666    #[test]
667    fn listen_to_derived_property_change_event() {
668        let mut context = Context::new();
669
670        let expected_high_id = PersonId::new(0);
671
672        // Listen for derived property change events and record how many times it fires
673        // For `RiskLevel`
674        let risk_flag = Rc::new(RefCell::new(0));
675        let risk_flag_clone = risk_flag.clone();
676        context.subscribe_to_event(
677            move |_context, event: PropertyChangeEvent<Person, RiskLevel>| {
678                assert_eq!(event.entity_id, expected_high_id);
679                assert_eq!(event.previous, RiskLevel::High);
680                assert_eq!(event.current, RiskLevel::Medium);
681                *risk_flag_clone.borrow_mut() += 1;
682            },
683        );
684        // For `AgeGroup`
685        let age_group_flag = Rc::new(RefCell::new(0));
686        let age_group_flag_clone = age_group_flag.clone();
687        context.subscribe_to_event(
688            move |_context, event: PropertyChangeEvent<Person, AgeGroup>| {
689                assert_eq!(event.entity_id, expected_high_id);
690                assert_eq!(event.previous, AgeGroup::Senior);
691                assert_eq!(event.current, AgeGroup::Adult);
692                *age_group_flag_clone.borrow_mut() += 1;
693            },
694        );
695
696        // Should not emit change events
697        let expected_high_id: PersonId = context
698            .add_entity((Age(77), Vaccinated(false), InfectionStatus::Susceptible))
699            .unwrap();
700
701        // Should emit change events
702        context.set_property(expected_high_id, Age(20));
703
704        // Execute queued event handlers
705        context.execute();
706        // Should have exactly one event recorded
707        assert_eq!(*risk_flag.borrow(), 1);
708        assert_eq!(*age_group_flag.borrow(), 1);
709    }
710
711    /*
712    ToDo(RobertJacobsonCDC): Enable this once #691 is resolved, https://github.com/CDCgov/ixa/issues/691.
713
714    #[test]
715    fn get_derived_property_with_globals() {
716        let mut context = Context::new();
717
718        context.set_global_property_value(GlobalDummy, 18).unwrap();
719        let child = context.add_entity((Age(17),)).unwrap();
720        let adult = context.add_entity((Age(19),)).unwrap();
721
722        let child_computed: MyDerivedProperty = context.get_property(child);
723        assert_eq!(child_computed, MyDerivedProperty(17+18));
724
725        let adult_computed: MyDerivedProperty = context.get_property(adult);
726        assert_eq!(adult_computed, MyDerivedProperty(19+18));
727    }
728    */
729
730    #[test]
731    fn observe_diamond_property_change() {
732        let mut context = Context::new();
733        let person = context.add_entity((Age(17), IsSwimmer(true))).unwrap();
734
735        let is_adult_athlete: AdultAthlete = context.get_property(person);
736        assert!(!is_adult_athlete.0);
737
738        let flag = Rc::new(RefCell::new(0));
739        let flag_clone = flag.clone();
740        context.subscribe_to_event(
741            move |_context, event: PropertyChangeEvent<Person, AdultAthlete>| {
742                assert_eq!(event.entity_id, person);
743                assert_eq!(event.previous, AdultAthlete(false));
744                assert_eq!(event.current, AdultAthlete(true));
745                *flag_clone.borrow_mut() += 1;
746            },
747        );
748
749        context.set_property(person, Age(20));
750        // Make sure the derived property is what we expect.
751        let is_adult_athlete: AdultAthlete = context.get_property(person);
752        assert!(is_adult_athlete.0);
753
754        // Execute queued event handlers
755        context.execute();
756        // Should have exactly one event recorded
757        assert_eq!(*flag.borrow(), 1);
758    }
759
760    // Tests related to queries and indexing
761
762    define_multi_property!((InfectionStatus, Vaccinated), Person);
763    define_multi_property!((Vaccinated, InfectionStatus), Person);
764
765    #[test]
766    fn with_query_results_finds_multi_index() {
767        use crate::rand::seq::IndexedRandom;
768        let mut rng = crate::rand::rng();
769        let mut context = Context::new();
770
771        for _ in 0..10_000usize {
772            let infection_status = *[
773                InfectionStatus::Susceptible,
774                InfectionStatus::Infected,
775                InfectionStatus::Recovered,
776            ]
777            .choose(&mut rng)
778            .unwrap();
779            let vaccination_status: bool = rng.random_bool(0.5);
780            let age: u8 = rng.random_range(0..100);
781            context
782                .add_entity((Age(age), infection_status, Vaccinated(vaccination_status)))
783                .unwrap();
784        }
785        context.index_property::<Person, InfectionStatusVaccinated>();
786        // Force an index build by running a query.
787        let _ = context.query_result_iterator((InfectionStatus::Susceptible, Vaccinated(true)));
788
789        // Capture the address of the has set given by `with_query_result`
790        let mut address: *const IndexSet<EntityId<Person>> = std::ptr::null();
791        context.with_query_results(
792            (InfectionStatus::Susceptible, Vaccinated(true)),
793            &mut |result_set| {
794                address = result_set as *const _;
795            },
796        );
797
798        // Check that the order doesn't matter.
799        assert_eq!(
800            InfectionStatusVaccinated::index_id(),
801            VaccinatedInfectionStatus::index_id()
802        );
803        assert_eq!(
804            InfectionStatusVaccinated::index_id(),
805            (InfectionStatus::Susceptible, Vaccinated(true))
806                .multi_property_id()
807                .unwrap()
808        );
809
810        // Check if it matches the expected bucket.
811        let index_id = InfectionStatusVaccinated::index_id();
812
813        let property_store = context.entity_store.get_property_store::<Person>();
814        let property_value_store = property_store.get_with_id(index_id);
815        let bucket: Ref<IndexSet<EntityId<Person>>> = property_value_store
816            .get_index_set_with_hash(
817                (InfectionStatus::Susceptible, Vaccinated(true)).multi_property_value_hash(),
818            )
819            .unwrap();
820
821        let address2 = &*bucket as *const _;
822        assert_eq!(address2, address);
823    }
824
825    #[test]
826    fn set_property_correctly_maintains_index() {
827        let mut context = Context::new();
828        context.index_property::<Person, InfectionStatus>();
829        context.index_property::<Person, AgeGroup>();
830
831        let person1 = context.add_entity((Age(22),)).unwrap();
832        let person2 = context.add_entity((Age(22),)).unwrap();
833        for _ in 0..4 {
834            let _: PersonId = context.add_entity((Age(22),)).unwrap();
835        }
836
837        // Check non-derived property index is correctly maintained
838        assert_eq!(
839            context.query_entity_count((InfectionStatus::Susceptible,)),
840            6
841        );
842        assert_eq!(context.query_entity_count((InfectionStatus::Infected,)), 0);
843        assert_eq!(context.query_entity_count((InfectionStatus::Recovered,)), 0);
844
845        context.set_property(person1, InfectionStatus::Infected);
846
847        assert_eq!(
848            context.query_entity_count((InfectionStatus::Susceptible,)),
849            5
850        );
851        assert_eq!(context.query_entity_count((InfectionStatus::Infected,)), 1);
852        assert_eq!(context.query_entity_count((InfectionStatus::Recovered,)), 0);
853
854        context.set_property(person1, InfectionStatus::Recovered);
855
856        assert_eq!(
857            context.query_entity_count((InfectionStatus::Susceptible,)),
858            5
859        );
860        assert_eq!(context.query_entity_count((InfectionStatus::Infected,)), 0);
861        assert_eq!(context.query_entity_count((InfectionStatus::Recovered,)), 1);
862
863        // Check derived property index is correctly maintained.
864        assert_eq!(context.query_entity_count((AgeGroup::Child,)), 0);
865        assert_eq!(context.query_entity_count((AgeGroup::Adult,)), 6);
866        assert_eq!(context.query_entity_count((AgeGroup::Senior,)), 0);
867
868        context.set_property(person2, Age(12));
869
870        assert_eq!(context.query_entity_count((AgeGroup::Child,)), 1);
871        assert_eq!(context.query_entity_count((AgeGroup::Adult,)), 5);
872        assert_eq!(context.query_entity_count((AgeGroup::Senior,)), 0);
873
874        context.set_property(person1, Age(75));
875
876        assert_eq!(context.query_entity_count((AgeGroup::Child,)), 1);
877        assert_eq!(context.query_entity_count((AgeGroup::Adult,)), 4);
878        assert_eq!(context.query_entity_count((AgeGroup::Senior,)), 1);
879
880        context.set_property(person2, Age(77));
881
882        assert_eq!(context.query_entity_count((AgeGroup::Child,)), 0);
883        assert_eq!(context.query_entity_count((AgeGroup::Adult,)), 4);
884        assert_eq!(context.query_entity_count((AgeGroup::Senior,)), 2);
885    }
886
887    #[test]
888    fn query_unindexed_default_properties() {
889        let mut context = Context::new();
890
891        // Half will have the default value.
892        for idx in 0..10 {
893            if idx % 2 == 0 {
894                context.add_entity((Age(22),)).unwrap();
895            } else {
896                context
897                    .add_entity((Age(22), InfectionStatus::Recovered))
898                    .unwrap();
899            }
900        }
901        // The tail also has the default value
902        for _ in 0..10 {
903            let _: PersonId = context.add_entity((Age(22),)).unwrap();
904        }
905
906        assert_eq!(context.query_entity_count((InfectionStatus::Recovered,)), 5);
907        assert_eq!(
908            context.query_entity_count((InfectionStatus::Susceptible,)),
909            15
910        );
911    }
912
913    #[test]
914    fn query_unindexed_derived_properties() {
915        let mut context = Context::new();
916
917        for _ in 0..10 {
918            let _: PersonId = context.add_entity((Age(22),)).unwrap();
919        }
920
921        assert_eq!(context.query_entity_count((AdultAthlete(false),)), 10);
922    }
923}
ixa/entity/context_extension.rs

ixa/entity/
context_extension.rs
