Skip to main content

ixa/entity/
property_store.rs

1/*!
2
3A [`PropertyStore`] implements the registry pattern for property value stores: A [`PropertyStore`]
4wraps a vector of `PropertyValueStore`s, one for each concrete property type. The implementor
5of [`crate::entity::property::Property`] is the value type. Since there's a 1-1 correspondence between property types
6and their value stores, we assign an ID to each property type to make
7property lookup fast. The [`PropertyStore`] stores a list of all properties in the form of
8boxed `PropertyValueStore` instances, which provide a type-erased interface to the backing
9storage (including index) of the property. Storage is only allocated as-needed, so the
10instantiation of a `PropertyValueStore` for a property that is never used is negligible.
11There's no need, then, for lazy initialization of the `PropertyValueStore`s themselves.
12
13This module also implements the initialization of "static" data associated with a property,
14that is, data that is the same across all [`crate::context::Context`] instances, which is computed before `main()`
15using `ctor` magic. (Each property implements a ctor that calls [`add_to_property_registry()`].)
16For simplicity, a property's ctor implementation, supplied by a macro, just calls
17`add_to_property_registry<E: Entity, P: Property<E>>()`, which does all the work. The
18`add_to_property_registry` function adds the following metadata to global metadata stores:
19
20Metadata stored on `PROPERTY_METADATA`, which for each property stores:
21- a list of dependent (derived) properties, and
22- a constructor function to create a new `PropertyValueStore` instance for the property.
23
24Metadata stored on `ENTITY_METADATA`, which for each entity stores:
25- a list of properties associated with the entity, and
26- a list of _required_ properties for the entity. These are properties for
27  which values must be supplied to `add_entity` when creating a new entity.
28
29*/
30
31use std::any::{Any, TypeId};
32use std::collections::HashMap;
33use std::sync::atomic::{AtomicUsize, Ordering};
34use std::sync::{LazyLock, Mutex, OnceLock};
35
36use crate::entity::entity::Entity;
37use crate::entity::entity_store::register_property_with_entity;
38use crate::entity::events::PartialPropertyChangeEventBox;
39use crate::entity::index::{IndexCountResult, IndexSetResult, PropertyIndex};
40use crate::entity::property::{IndexableProperty, Property};
41use crate::entity::property_list::PropertyList;
42use crate::entity::property_value_store::PropertyValueStore;
43use crate::entity::property_value_store_core::PropertyValueStoreCore;
44use crate::entity::value_change_counter::StratifiedValueChangeCounter;
45use crate::entity::EntityId;
46use crate::{Context, ContextEntitiesExt};
47
48pub(in crate::entity) type IndexNewEntityFn<E> = fn(&mut Context, EntityId<E>);
49
50fn index_new_entity<E, P>(context: &mut Context, entity_id: EntityId<E>)
51where
52    E: Entity,
53    P: IndexableProperty<E>,
54{
55    // This may compute a derived or multi-property. `P` is copied, so no reference into Context
56    // survives into the subsequent mutable borrow.
57    let value: P = context.get_property(entity_id);
58
59    let property_value_store = context.get_property_value_store_mut::<E, P>();
60    let index = property_value_store
61        .index
62        .as_mut()
63        .expect("index_new_entity dispatch invoked for an unindexed property");
64
65    index.add_entity(&value, entity_id);
66}
67
68/// A map from Entity ID to a count of the properties already associated with the entity. The value for the key is
69/// equivalent to the next property ID that will be assigned to the next property that requests an ID. Each `Entity`
70/// type has its own series of increasing property IDs.
71///
72/// Note: The mechanism to assign property IDs needs to be distinct from the rest of property registration, because
73/// properties often need to have an ID assigned _before_ its registration proper so that it can be recorded as a
74/// dependency of some other property.
75static NEXT_PROPERTY_ID: LazyLock<Mutex<HashMap<usize, usize>>> =
76    LazyLock::new(|| Mutex::new(HashMap::default()));
77
78/// A container struct to hold the (global) metadata for a single property.
79///
80/// At program startup (before `main()`, using ctors) we compute metadata for all properties
81/// that are linked into the binary, and this data remains unchanged for the life of the program.
82#[derive(Default)]
83pub(super) struct PropertyMetadata<E: Entity> {
84    /// The (derived) properties that depend on this property, as represented by their
85    /// `Property::id` value. This list is used to update the index (if applicable)
86    /// and emit change events for these properties when this property changes.
87    pub dependents: Vec<usize>,
88    /// A function that constructs a new `PropertyValueStoreCore<E, P>` instance in a type-erased
89    /// way, used in the constructor of `PropertyStore`. This is an `Option` because this
90    /// function pointer is recorded possibly out-of-order from when the `PropertyMetadata`
91    /// instance for this property needs to exist (when its dependents are recorded).
92    #[allow(clippy::type_complexity)]
93    pub value_store_constructor: Option<fn() -> Box<dyn PropertyValueStore<E>>>,
94}
95
96/// This maps `(entity_type_id, property_type_index)` to `PropertyMetadata<E>`, which holds a vector of dependents (as IDs)
97/// and a function pointer to the constructor that constucts a `PropertyValueStoreCore<E, P>` type erased as
98/// a `Box<dyn PropertyValueStore<E>>`. This data is actually written by the property `ctor`s with a call to [`crate::entity::entity_store::register_property_with_entity`()].
99#[allow(clippy::type_complexity)]
100static PROPERTY_METADATA_BUILDER: LazyLock<
101    Mutex<HashMap<(usize, usize), Box<dyn Any + Send + Sync>>>,
102> = LazyLock::new(|| Mutex::new(HashMap::default()));
103
104/// The frozen property metadata registry, created exactly once on first read.
105///
106/// This is derived from `PROPERTY_METADATA_BUILDER` by moving the builder `HashMap` out. After this point,
107/// registration is no longer allowed.
108static PROPERTY_METADATA: OnceLock<HashMap<(usize, usize), Box<dyn Any + Send + Sync>>> =
109    OnceLock::new();
110
111/// Private helper to fetch or initialize the frozen metadata.
112fn property_metadata() -> &'static HashMap<(usize, usize), Box<dyn Any + Send + Sync>> {
113    PROPERTY_METADATA.get_or_init(|| {
114        let mut builder = PROPERTY_METADATA_BUILDER.lock().unwrap();
115        std::mem::take(&mut *builder)
116    })
117}
118
119/// The public getter for the dependents of a property with index `property_index` (as stored in
120/// `PROPERTY_METADATA`). The `Property<E: Entity>::dependents()` method defers to this.
121///
122/// This function should only be called once `main()` starts, that is, not in `ctors` constructors,
123/// as it assumes `PROPERTY_METADATA` has been correctly initialized. Hence, the "static" suffix.
124#[must_use]
125pub(super) fn get_property_dependents_static<E: Entity>(property_index: usize) -> &'static [usize] {
126    let map = property_metadata();
127    let property_metadata = map
128        .get(&(E::id(), property_index))
129                               .unwrap_or_else(|| panic!("No registered property found with index = {property_index:?}. You must use the `define_property!` macro to create a registered property."));
130    let property_metadata: &PropertyMetadata<E> = property_metadata.downcast_ref().unwrap_or_else(
131        || panic!(
132            "Property type at index {:?} does not match registered property type. You must use the `define_property!` macro to create a registered property.",
133            property_index
134        )
135    );
136
137    property_metadata.dependents.as_slice()
138}
139
140/// Adds a new item to the registry. The job of this method is to create whatever "singleton"
141/// data/metadata is associated with the [`crate::entity::property::Property`] if it doesn't already exist. In
142/// our use case, this method is called in the `ctor` function of each `Property<E>` type.
143pub fn add_to_property_registry<E: Entity, P: Property<E>>() {
144    // Ensure the ID of the property type is initialized.
145    let property_index = P::id();
146
147    // Registers the property with the entity type.
148    register_property_with_entity(
149        <E as Entity>::type_id(),
150        <P as Property<E>>::type_id(),
151        P::is_required(),
152    );
153
154    let mut property_metadata = PROPERTY_METADATA_BUILDER.lock().unwrap();
155    if PROPERTY_METADATA.get().is_some() {
156        panic!(
157            "`add_to_property_registry()` called after property metadata was frozen; registration must occur during startup/ctors."
158        );
159    }
160
161    // Register the `PropertyValueStoreCore<E, P>` constructor.
162    {
163        let metadata = property_metadata
164            .entry((E::id(), property_index))
165            .or_insert_with(|| Box::new(PropertyMetadata::<E>::default()));
166        let metadata: &mut PropertyMetadata<E> = metadata.downcast_mut().unwrap();
167        metadata
168            .value_store_constructor
169            .get_or_insert(PropertyValueStoreCore::<E, P>::new_boxed);
170    }
171
172    // Construct the dependency graph
173    for dependency in P::non_derived_dependencies() {
174        // Add `property_index` as a dependent of the dependency
175        let dependency_meta = property_metadata
176            .entry((E::id(), dependency))
177            .or_insert_with(|| Box::new(PropertyMetadata::<E>::default()));
178        let dependency_meta: &mut PropertyMetadata<E> = dependency_meta.downcast_mut().unwrap();
179        dependency_meta.dependents.push(property_index);
180    }
181}
182
183/// A convenience getter for `NEXT_ENTITY_INDEX`.
184pub fn get_registered_property_count<E: Entity>() -> usize {
185    let map = NEXT_PROPERTY_ID.lock().unwrap();
186    *map.get(&E::id()).unwrap_or(&0)
187}
188
189/// Encapsulates the synchronization logic for initializing an item's index.
190///
191/// Acquires a global lock on the next available property ID, but only increments
192/// it if we successfully initialize the provided ID. The ID of a property is
193/// assigned at runtime but only once per type. It's possible for a single
194/// type to attempt to initialize its index multiple times from different threads,
195/// which is why all this synchronization is required. However, the overhead
196/// is negligible, as this initialization only happens once upon first access.
197///
198/// In fact, for our use case we know we are calling this function
199/// once for each type in each `Property`'s `ctor` function, which
200/// should be the only time this method is ever called for the type.
201pub fn initialize_property_id<E: Entity>(property_id: &AtomicUsize) -> usize {
202    // Acquire a global lock.
203    let mut guard = NEXT_PROPERTY_ID.lock().unwrap();
204    let candidate = guard.entry(E::id()).or_insert_with(|| 0);
205
206    // Try to claim the candidate index. Here we guard against the potential race condition that
207    // another instance of this plugin in another thread just initialized the index prior to us
208    // obtaining the lock. If the index has been initialized beneath us, we do not update
209    // NEXT_PROPERTY_INDEX, we just return the value `index` was initialized to.
210    // For a justification of the data ordering, see:
211    //     https://github.com/CDCgov/ixa/pull/477#discussion_r2244302872
212    match property_id.compare_exchange(usize::MAX, *candidate, Ordering::AcqRel, Ordering::Acquire)
213    {
214        Ok(_) => {
215            // We won the race — increment the global next plugin index and return the new index
216            *candidate += 1;
217            *candidate - 1
218        }
219        Err(existing) => {
220            // Another thread beat us — don’t increment the global next plugin index,
221            // just return existing
222            existing
223        }
224    }
225}
226
227/// A wrapper around a vector of property value stores.
228pub struct PropertyStore<E: Entity> {
229    /// A vector of `Box<PropertyValueStoreCore<E, P>>`, type-erased to `Box<dyn PropertyValueStore<E>>`
230    items: Vec<Box<dyn PropertyValueStore<E>>>,
231
232    /// One entry for every property whose `PropertyValueStoreCore` currently has an index.
233    /// The property ID supports removal without relying on deduplicable function addresses.
234    pub(in crate::entity) index_new_entity_fns: Vec<(usize, IndexNewEntityFn<E>)>,
235}
236
237impl<E: Entity> Default for PropertyStore<E> {
238    fn default() -> Self {
239        PropertyStore::new()
240    }
241}
242
243impl<E: Entity> PropertyStore<E> {
244    /// Creates a new [`PropertyStore`].
245    #[must_use]
246    pub fn new() -> Self {
247        let num_items = get_registered_property_count::<E>();
248        // The constructors for each `PropertyValueStoreCore<E, P>` are stored in the `PROPERTY_METADATA` global.
249        let property_metadata = property_metadata();
250
251        // We construct the correct concrete `PropertyValueStoreCore<E, P>` value for each ID.
252        let items = (0..num_items)
253            .map(|idx| {
254                let metadata = property_metadata
255                    .get(&(E::id(), idx))
256                    .unwrap_or_else(|| panic!("No property metadata entry for index {idx}"))
257                    .downcast_ref::<PropertyMetadata<E>>()
258                    .unwrap_or_else(|| {
259                        panic!(
260                            "Property metadata entry for index {idx} does not match expected type"
261                        )
262                    });
263                let constructor = metadata
264                    .value_store_constructor
265                    .unwrap_or_else(|| panic!("No PropertyValueStore constructor for index {idx}"));
266                constructor()
267            })
268            .collect();
269
270        Self {
271            items,
272            index_new_entity_fns: Vec::new(),
273        }
274    }
275
276    /// Fetches an immutable reference to the `PropertyValueStoreCore<E, P>`.
277    #[must_use]
278    pub fn get<P: Property<E>>(&self) -> &PropertyValueStoreCore<E, P> {
279        let index = P::id();
280        let property_value_store =
281            self.items
282                .get(index)
283                .unwrap_or_else(||
284                    panic!(
285                        "No registered property found with index = {:?} while trying to get property {}. You must use the `define_property!` macro to create a registered property.",
286                        index,
287                        P::name()
288                    )
289                );
290        let property_value_store: &PropertyValueStoreCore<E, P> = property_value_store
291            .as_any()
292            .downcast_ref::<PropertyValueStoreCore<E, P>>()
293            .unwrap_or_else(||
294                {
295                    panic!(
296                        "Property type at index {:?} does not match registered property type. Found type_id {:?} while getting type_id {:?}. You must use the `define_property!` macro to create a registered property.",
297                        index,
298                        (**property_value_store).type_id(),
299                        TypeId::of::<PropertyValueStoreCore<E, P>>()
300                    )
301                }
302            );
303        property_value_store
304    }
305
306    /// Fetches a mutable reference to the `PropertyValueStoreCore<E, P>`.
307    #[must_use]
308    pub fn get_mut<P: Property<E>>(&mut self) -> &mut PropertyValueStoreCore<E, P> {
309        let index = P::id();
310        let property_value_store =
311            self.items
312                .get_mut(index)
313                .unwrap_or_else(||
314                    panic!(
315                        "No registered property found with index = {:?} while trying to get property {}. You must use the `define_property!` macro to create a registered property.",
316                        index,
317                        P::name()
318                    )
319                );
320        let type_id = (**property_value_store).type_id(); // Only used for error message if error occurs.
321        let property_value_store: &mut PropertyValueStoreCore<E, P> = property_value_store
322            .as_any_mut()
323            .downcast_mut::<PropertyValueStoreCore<E, P>>()
324            .unwrap_or_else(||
325                {
326                    panic!(
327                        "Property type at index {:?} does not match registered property type. Found type_id {:?} while getting type_id {:?}. You must use the `define_property!` macro to create a registered property.",
328                        index,
329                        type_id,
330                        TypeId::of::<PropertyValueStoreCore<E, P>>()
331                    )
332                }
333            );
334        property_value_store
335    }
336
337    /// Creates a `PartialPropertyChangeEvent` instance for the `entity_id` and `property_index`. This method is only
338    /// called for derived dependents of some property that has changed (one of `P`'s non-derived dependencies).
339    #[must_use]
340    pub(crate) fn create_partial_property_change(
341        &self,
342        property_index: usize,
343        entity_id: EntityId<E>,
344        context: &Context,
345    ) -> PartialPropertyChangeEventBox {
346        let property_value_store = self.items
347                                       .get(property_index)
348            .unwrap_or_else(|| panic!("No registered property found with index = {property_index:?}. You must use the `define_property!` macro to create a registered property."));
349
350        property_value_store.create_partial_property_change(entity_id, context)
351    }
352
353    /// Returns whether the property with `property_index` needs partial change-event processing.
354    #[must_use]
355    pub(crate) fn should_create_partial_property_change(
356        &self,
357        property_index: usize,
358        context: &Context,
359    ) -> bool {
360        let property_value_store = self.items
361                                       .get(property_index)
362            .unwrap_or_else(|| panic!("No registered property found with index = {property_index:?}. You must use the `define_property!` macro to create a registered property."));
363
364        property_value_store.should_create_partial_change(context)
365    }
366
367    /// Returns whether or not the property `P` is indexed.
368    ///
369    #[cfg(test)]
370    #[must_use]
371    pub fn is_property_indexed<P: Property<E>>(&self) -> bool {
372        self.get::<P>().index.is_some()
373    }
374
375    pub(in crate::entity) fn install_property_index<P>(
376        &mut self,
377        new_index: Option<Box<dyn PropertyIndex<E, P>>>,
378    ) where
379        P: IndexableProperty<E>,
380    {
381        let property_id = P::id();
382        let was_indexed = self.get::<P>().index.is_some();
383        let will_be_indexed = new_index.is_some();
384
385        // Property IDs are stable dispatcher identities; function-pointer addresses are not,
386        // because the linker may deduplicate distinct monomorphizations.
387        let dispatcher_position = self
388            .index_new_entity_fns
389            .iter()
390            .position(|(id, _)| *id == property_id);
391
392        if was_indexed {
393            assert!(
394                dispatcher_position.is_some(),
395                "indexed property missing its index_new_entity dispatcher",
396            );
397        } else {
398            debug_assert!(dispatcher_position.is_none());
399        }
400
401        // Reserve before replacing a valid index so allocation failure cannot leave the index
402        // installed without its dispatcher.
403        if !was_indexed && will_be_indexed {
404            self.index_new_entity_fns.reserve(1);
405        }
406
407        // All invariant checks and potentially allocating preparation are complete. Replacing the
408        // installed index is the commit point.
409        self.get_mut::<P>().index = new_index;
410
411        match (was_indexed, will_be_indexed) {
412            (false, true) => {
413                self.index_new_entity_fns
414                    .push((property_id, index_new_entity::<E, P> as IndexNewEntityFn<E>));
415            }
416            (true, false) => {
417                self.index_new_entity_fns
418                    .swap_remove(dispatcher_position.unwrap());
419            }
420            _ => {}
421        }
422    }
423
424    /// Creates a stratified value change counter for tracked property `P` with strata `PL`.
425    ///
426    /// Returns the counter ID.
427    #[must_use]
428    pub fn create_value_change_counter<PL, P>(&mut self) -> usize
429    where
430        PL: PropertyList<E> + Eq + std::hash::Hash,
431        P: Property<E> + Eq + std::hash::Hash,
432    {
433        let property_value_store = self.get_mut::<P>();
434        property_value_store.add_value_change_counter(Box::new(StratifiedValueChangeCounter::<
435            E,
436            PL,
437            P,
438        >::new()))
439    }
440
441    #[must_use]
442    pub fn get_index_set_for_query_parts(
443        &self,
444        property_id: usize,
445        query_parts: &[&dyn Any],
446    ) -> IndexSetResult<'_, E> {
447        self.items[property_id].get_index_set_for_query_parts(query_parts)
448    }
449
450    #[must_use]
451    pub fn get_index_count_for_query_parts(
452        &self,
453        property_id: usize,
454        query_parts: &[&dyn Any],
455    ) -> IndexCountResult {
456        self.items[property_id].get_index_count_for_query_parts(query_parts)
457    }
458}
459
460#[cfg(test)]
461mod tests {
462    #![allow(dead_code)]
463    use std::any::Any;
464    use std::panic::{catch_unwind, AssertUnwindSafe};
465
466    use super::*;
467    use crate::entity::index::{FullIndex, IndexCountResult, IndexSetResult, ValueCountIndex};
468    use crate::entity::PropertyIndexType;
469    use crate::prelude::*;
470    use crate::{define_derived_property, define_entity, define_property, with, Context};
471
472    define_entity!(Person);
473
474    define_property!(struct Age(u8), Person);
475    define_property!(
476        enum InfectionStatus {
477            Susceptible,
478            Infected,
479            Recovered,
480        },
481        Person,
482        default_const = InfectionStatus::Susceptible
483    );
484    define_property!(struct Vaccinated(bool), Person, default_const = Vaccinated(false));
485    define_property!(struct PanicDependency(u8), Person, default_const = PanicDependency(0));
486
487    define_derived_property!(
488        struct PanickingDerived(u8),
489        Person,
490        [PanicDependency],
491        [],
492        |dependency| {
493            let dependency: PanicDependency = dependency;
494            assert_ne!(dependency, PanicDependency(255), "sentinel property value");
495            PanickingDerived(dependency.0)
496        }
497    );
498
499    #[test]
500    fn property_store_default_matches_new() {
501        let property_store = PropertyStore::<Person>::default();
502        assert_eq!(
503            property_store.items.len(),
504            get_registered_property_count::<Person>()
505        );
506    }
507
508    #[test]
509    fn install_property_index_maintains_active_dispatchers() {
510        let mut context = Context::new();
511
512        {
513            let property_store = context.entity_store.get_property_store_mut::<Person>();
514            assert_eq!(property_store.index_new_entity_fns.len(), 0);
515
516            property_store.install_property_index::<Age>(Some(Box::new(ValueCountIndex::new())));
517            assert_eq!(property_store.index_new_entity_fns.len(), 1);
518
519            property_store.install_property_index::<Age>(Some(Box::new(ValueCountIndex::new())));
520            assert_eq!(property_store.index_new_entity_fns.len(), 1);
521
522            property_store.install_property_index::<Age>(Some(Box::new(FullIndex::new())));
523            assert_eq!(property_store.index_new_entity_fns.len(), 1);
524
525            property_store.install_property_index::<Age>(None);
526            assert_eq!(property_store.index_new_entity_fns.len(), 0);
527
528            property_store.install_property_index::<Age>(Some(Box::new(ValueCountIndex::new())));
529            property_store.install_property_index::<Vaccinated>(Some(Box::new(FullIndex::new())));
530            property_store.install_property_index::<Age>(None);
531
532            assert_eq!(property_store.index_new_entity_fns.len(), 1);
533            assert_eq!(property_store.index_new_entity_fns[0].0, Vaccinated::id());
534        }
535
536        context.add_entity(with!(Person, Age(10))).unwrap();
537        let property_store = context.entity_store.get_property_store::<Person>();
538        assert_eq!(
539            property_store.get::<Age>().index_type(),
540            PropertyIndexType::Unindexed
541        );
542        assert_eq!(
543            property_store.get_index_count_for_query_parts(
544                Vaccinated::id(),
545                &[&Vaccinated(false) as &dyn Any],
546            ),
547            IndexCountResult::Count(1),
548        );
549    }
550
551    #[test]
552    fn failed_replacement_keeps_old_index_and_dispatcher() {
553        let mut context = Context::new();
554        let ordinary = context
555            .add_entity(with!(Person, Age(10), PanicDependency(10)))
556            .unwrap();
557        let sentinel = context
558            .add_entity(with!(Person, Age(20), PanicDependency(255)))
559            .unwrap();
560
561        let mut old_index = ValueCountIndex::<Person, PanickingDerived>::new();
562        old_index.add_entity(&PanickingDerived(10), ordinary);
563        old_index.add_entity(&PanickingDerived(255), sentinel);
564        context
565            .entity_store
566            .get_property_store_mut::<Person>()
567            .install_property_index::<PanickingDerived>(Some(Box::new(old_index)));
568
569        let (old_type, old_dispatcher_count, old_dispatcher_property_id) = {
570            let property_store = context.entity_store.get_property_store::<Person>();
571            (
572                property_store.get::<PanickingDerived>().index_type(),
573                property_store.index_new_entity_fns.len(),
574                property_store.index_new_entity_fns[0].0,
575            )
576        };
577        assert_eq!(old_type, PropertyIndexType::ValueCountIndex);
578        assert_eq!(
579            context.query_entity_count(with!(Person, PanickingDerived(10))),
580            1
581        );
582        assert_eq!(
583            context.query_entity_count(with!(Person, PanickingDerived(255))),
584            1
585        );
586
587        let result = catch_unwind(AssertUnwindSafe(|| {
588            context.index_property::<Person, PanickingDerived>();
589        }));
590        assert!(result.is_err());
591
592        let property_store = context.entity_store.get_property_store::<Person>();
593        assert_eq!(
594            property_store.get::<PanickingDerived>().index_type(),
595            old_type
596        );
597        assert_eq!(
598            property_store.index_new_entity_fns.len(),
599            old_dispatcher_count
600        );
601        assert_eq!(
602            property_store.index_new_entity_fns[0].0,
603            old_dispatcher_property_id
604        );
605        assert_eq!(
606            context.query_entity_count(with!(Person, PanickingDerived(10))),
607            1
608        );
609        assert_eq!(
610            context.query_entity_count(with!(Person, PanickingDerived(255))),
611            1
612        );
613
614        context
615            .add_entity(with!(Person, Age(30), PanicDependency(10)))
616            .unwrap();
617        assert_eq!(
618            context.query_entity_count(with!(Person, PanickingDerived(10))),
619            2
620        );
621    }
622
623    #[test]
624    fn test_get_property_store() {
625        let mut property_store = PropertyStore::new();
626
627        {
628            let ages: &mut PropertyValueStoreCore<_, Age> = property_store.get_mut();
629            ages.set(EntityId::<Person>::new(0), Age(12));
630            ages.set(EntityId::<Person>::new(1), Age(33));
631            ages.set(EntityId::<Person>::new(2), Age(44));
632
633            let infection_statuses: &mut PropertyValueStoreCore<_, InfectionStatus> =
634                property_store.get_mut();
635            infection_statuses.set(EntityId::<Person>::new(0), InfectionStatus::Susceptible);
636            infection_statuses.set(EntityId::<Person>::new(1), InfectionStatus::Susceptible);
637            infection_statuses.set(EntityId::<Person>::new(2), InfectionStatus::Infected);
638
639            let vaccine_status: &mut PropertyValueStoreCore<_, Vaccinated> =
640                property_store.get_mut();
641            vaccine_status.set(EntityId::<Person>::new(0), Vaccinated(true));
642            vaccine_status.set(EntityId::<Person>::new(1), Vaccinated(false));
643            vaccine_status.set(EntityId::<Person>::new(2), Vaccinated(true));
644        }
645
646        // Verify that `get` returns the expected values
647        {
648            let ages: &PropertyValueStoreCore<_, Age> = property_store.get();
649            assert_eq!(ages.get(EntityId::<Person>::new(0)), Age(12));
650            assert_eq!(ages.get(EntityId::<Person>::new(1)), Age(33));
651            assert_eq!(ages.get(EntityId::<Person>::new(2)), Age(44));
652
653            let infection_statuses: &PropertyValueStoreCore<_, InfectionStatus> =
654                property_store.get();
655            assert_eq!(
656                infection_statuses.get(EntityId::<Person>::new(0)),
657                InfectionStatus::Susceptible
658            );
659            assert_eq!(
660                infection_statuses.get(EntityId::<Person>::new(1)),
661                InfectionStatus::Susceptible
662            );
663            assert_eq!(
664                infection_statuses.get(EntityId::<Person>::new(2)),
665                InfectionStatus::Infected
666            );
667
668            let vaccine_status: &PropertyValueStoreCore<_, Vaccinated> = property_store.get();
669            assert_eq!(
670                vaccine_status.get(EntityId::<Person>::new(0)),
671                Vaccinated(true)
672            );
673            assert_eq!(
674                vaccine_status.get(EntityId::<Person>::new(1)),
675                Vaccinated(false)
676            );
677            assert_eq!(
678                vaccine_status.get(EntityId::<Person>::new(2)),
679                Vaccinated(true)
680            );
681        }
682    }
683
684    #[test]
685    fn test_index_query_results_for_property_store() {
686        let mut context = Context::new();
687        context.index_property::<Person, Age>();
688
689        let existing_value = Age(12);
690        let missing_value = Age(99);
691        let existing_query_parts = [&existing_value as &dyn Any];
692        let missing_query_parts = [&missing_value as &dyn Any];
693
694        let _ = context.add_entity(with!(Person, existing_value)).unwrap();
695        let _ = context.add_entity(with!(Person, existing_value)).unwrap();
696
697        let property_store = context.entity_store.get_property_store::<Person>();
698
699        // FullIndex + count
700        assert_eq!(
701            property_store.get_index_count_for_query_parts(Age::id(), &missing_query_parts,),
702            IndexCountResult::Count(0)
703        );
704        assert_eq!(
705            property_store.get_index_count_for_query_parts(Age::id(), &existing_query_parts,),
706            IndexCountResult::Count(2)
707        );
708
709        // FullIndex + set
710        assert!(matches!(
711            property_store.get_index_set_for_query_parts(Age::id(), &missing_query_parts,),
712            IndexSetResult::Empty
713        ));
714        assert!(matches!(
715            property_store.get_index_set_for_query_parts(
716                Age::id(),
717                &existing_query_parts,
718            ),
719            IndexSetResult::Set(set) if set.len() == 2
720        ));
721    }
722
723    #[test]
724    fn test_index_query_results_for_property_store_value_count_index() {
725        let mut context = Context::new();
726        context.index_property_counts::<Person, Age>();
727
728        let existing_value = Age(12);
729        let missing_value = Age(99);
730        let existing_query_parts = [&existing_value as &dyn Any];
731        let missing_query_parts = [&missing_value as &dyn Any];
732
733        let _ = context.add_entity(with!(Person, existing_value)).unwrap();
734        let _ = context.add_entity(with!(Person, existing_value)).unwrap();
735
736        let property_store = context.entity_store.get_property_store::<Person>();
737
738        // ValueCountIndex + count
739        assert_eq!(
740            property_store.get_index_count_for_query_parts(Age::id(), &missing_query_parts,),
741            IndexCountResult::Count(0)
742        );
743        assert_eq!(
744            property_store.get_index_count_for_query_parts(Age::id(), &existing_query_parts,),
745            IndexCountResult::Count(2)
746        );
747
748        // ValueCountIndex + set (unsupported)
749        assert!(matches!(
750            property_store.get_index_set_for_query_parts(Age::id(), &missing_query_parts,),
751            IndexSetResult::Unsupported
752        ));
753        assert!(matches!(
754            property_store.get_index_set_for_query_parts(Age::id(), &existing_query_parts,),
755            IndexSetResult::Unsupported
756        ));
757    }
758
759    #[test]
760    fn test_index_query_results_for_property_store_unindexed() {
761        let mut context = Context::new();
762        let existing_value = Age(12);
763        let missing_value = Age(99);
764        let existing_query_parts = [&existing_value as &dyn Any];
765        let missing_query_parts = [&missing_value as &dyn Any];
766
767        let _ = context.add_entity(with!(Person, existing_value)).unwrap();
768        let _ = context.add_entity(with!(Person, existing_value)).unwrap();
769
770        let property_store = context.entity_store.get_property_store::<Person>();
771
772        // Unindexed + count
773        assert_eq!(
774            property_store.get_index_count_for_query_parts(Age::id(), &missing_query_parts,),
775            IndexCountResult::Unsupported
776        );
777        assert_eq!(
778            property_store.get_index_count_for_query_parts(Age::id(), &existing_query_parts,),
779            IndexCountResult::Unsupported
780        );
781
782        // Unindexed + set
783        assert!(matches!(
784            property_store.get_index_set_for_query_parts(Age::id(), &missing_query_parts,),
785            IndexSetResult::Unsupported
786        ));
787        assert!(matches!(
788            property_store.get_index_set_for_query_parts(Age::id(), &existing_query_parts,),
789            IndexSetResult::Unsupported
790        ));
791    }
792}