ixa 2.0.0-beta2.3

/*!

A [`PropertyStore`] implements the registry pattern for property value stores: A [`PropertyStore`]
wraps a vector of `PropertyValueStore`s, one for each concrete property type. The implementor
of [`crate::entity::property::Property`] is the value type. Since there's a 1-1 correspondence between property types
and their value stores, we implement the `index` method for each property type to make
property lookup fast. The [`PropertyStore`] stores a list of all properties in the form of
boxed `PropertyValueStore` instances, which provide a type-erased interface to the backing
storage (including index) of the property. Storage is only allocated as-needed, so the
instantiation of a `PropertyValueStore` for a property that is never used is negligible.
There's no need, then, for lazy initialization of the `PropertyValueStore`s themselves.

This module also implements the initialization of "static" data associated with a property,
that is, data that is the same across all [`crate::context::Context`] instances, which is computed before `main()`
using `ctor` magic. (Each property implements a ctor that calls [`add_to_property_registry()`].)
For simplicity, a property's ctor implementation, supplied by a macro, just calls
`add_to_property_registry<E: Entity, P: Property<E>>()`, which does all the work. The
`add_to_property_registry` function adds the following metadata to global metadata stores:

Metadata stored on `PROPERTY_METADATA`, which for each property stores:
- a list of dependent (derived) properties, and
- a constructor function to create a new `PropertyValueStore` instance for the property.

Metadata stored on `ENTITY_METADATA`, which for each entity stores:
- a list of properties associated with the entity, and
- a list of _required_ properties for the entity. These are properties for
  which values must be supplied to `add_entity` when creating a new entity.

*/

use std::any::{Any, TypeId};
use std::collections::HashMap;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{LazyLock, Mutex, OnceLock};

use crate::entity::entity::Entity;
use crate::entity::entity_store::register_property_with_entity;
use crate::entity::events::PartialPropertyChangeEventBox;
use crate::entity::index::{IndexCountResult, IndexSetResult};
use crate::entity::property::Property;
use crate::entity::property_list::PropertyList;
use crate::entity::property_value_store::PropertyValueStore;
use crate::entity::property_value_store_core::PropertyValueStoreCore;
use crate::entity::value_change_counter::StratifiedValueChangeCounter;
use crate::entity::{EntityId, PropertyIndexType};
use crate::Context;

/// A map from Entity ID to a count of the properties already associated with the entity. The value for the key is
/// equivalent to the next property ID that will be assigned to the next property that requests an ID. Each `Entity`
/// type has its own series of increasing property IDs.
///
/// Note: The mechanism to assign property IDs needs to be distinct from the rest of property registration, because
/// properties often need to have an ID assigned _before_ its registration proper so that it can be recorded as a
/// dependency of some other property.
static NEXT_PROPERTY_ID: LazyLock<Mutex<HashMap<usize, usize>>> =
    LazyLock::new(|| Mutex::new(HashMap::default()));

/// A container struct to hold the (global) metadata for a single property.
///
/// At program startup (before `main()`, using ctors) we compute metadata for all properties
/// that are linked into the binary, and this data remains unchanged for the life of the program.
#[derive(Default)]
pub(super) struct PropertyMetadata<E: Entity> {
    /// The (derived) properties that depend on this property, as represented by their
    /// `Property::index` value. This list is used to update the index (if applicable)
    /// and emit change events for these properties when this property changes.
    pub dependents: Vec<usize>,
    /// A function that constructs a new `PropertyValueStoreCore<E, P>` instance in a type-erased
    /// way, used in the constructor of `PropertyStore`. This is an `Option` because this
    /// function pointer is recorded possibly out-of-order from when the `PropertyMetadata`
    /// instance for this property needs to exist (when its dependents are recorded).
    #[allow(clippy::type_complexity)]
    pub value_store_constructor: Option<fn() -> Box<dyn PropertyValueStore<E>>>,
}

/// This maps `(entity_type_id, property_type_index)` to `PropertyMetadata<E>`, which holds a vector of dependents (as IDs)
/// and a function pointer to the constructor that constucts a `PropertyValueStoreCore<E, P>` type erased as
/// 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`()].
#[allow(clippy::type_complexity)]
static PROPERTY_METADATA_BUILDER: LazyLock<
    Mutex<HashMap<(usize, usize), Box<dyn Any + Send + Sync>>>,
> = LazyLock::new(|| Mutex::new(HashMap::default()));

/// The frozen property metadata registry, created exactly once on first read.
///
/// This is derived from `PROPERTY_METADATA_BUILDER` by moving the builder `HashMap` out. After this point,
/// registration is no longer allowed.
#[allow(clippy::type_complexity)]
static PROPERTY_METADATA: OnceLock<HashMap<(usize, usize), Box<dyn Any + Send + Sync>>> =
    OnceLock::new();

/// Private helper to fetch or initialize the frozen metadata.
fn property_metadata() -> &'static HashMap<(usize, usize), Box<dyn Any + Send + Sync>> {
    PROPERTY_METADATA.get_or_init(|| {
        let mut builder = PROPERTY_METADATA_BUILDER.lock().unwrap();
        std::mem::take(&mut *builder)
    })
}

/// The public getter for the dependents of a property with index `property_index` (as stored in
/// `PROPERTY_METADATA`). The `Property<E: Entity>::dependents()` method defers to this.
///
/// This function should only be called once `main()` starts, that is, not in `ctors` constructors,
/// as it assumes `PROPERTY_METADATA` has been correctly initialized. Hence, the "static" suffix.
#[must_use]
pub(super) fn get_property_dependents_static<E: Entity>(property_index: usize) -> &'static [usize] {
    let map = property_metadata();
    let property_metadata = map
        .get(&(E::id(), property_index))
                               .unwrap_or_else(|| panic!("No registered property found with index = {property_index:?}. You must use the `define_property!` macro to create a registered property."));
    let property_metadata: &PropertyMetadata<E> = property_metadata.downcast_ref().unwrap_or_else(
        || panic!(
            "Property type at index {:?} does not match registered property type. You must use the `define_property!` macro to create a registered property.",
            property_index
        )
    );

    property_metadata.dependents.as_slice()
}

/// Adds a new item to the registry. The job of this method is to create whatever "singleton"
/// data/metadata is associated with the [`crate::entity::property::Property`] if it doesn't already exist. In
/// our use case, this method is called in the `ctor` function of each `Property<E>` type.
pub fn add_to_property_registry<E: Entity, P: Property<E>>() {
    // Ensure the ID of the property type is initialized.
    let property_index = P::id();

    // Registers the property with the entity type.
    register_property_with_entity(
        <E as Entity>::type_id(),
        <P as Property<E>>::type_id(),
        P::is_required(),
    );

    let mut property_metadata = PROPERTY_METADATA_BUILDER.lock().unwrap();
    if PROPERTY_METADATA.get().is_some() {
        panic!(
            "`add_to_property_registry()` called after property metadata was frozen; registration must occur during startup/ctors."
        );
    }

    // Register the `PropertyValueStoreCore<E, P>` constructor.
    {
        let metadata = property_metadata
            .entry((E::id(), property_index))
            .or_insert_with(|| Box::new(PropertyMetadata::<E>::default()));
        let metadata: &mut PropertyMetadata<E> = metadata.downcast_mut().unwrap();
        metadata
            .value_store_constructor
            .get_or_insert(PropertyValueStoreCore::<E, P>::new_boxed);
    }

    // Construct the dependency graph
    for dependency in P::non_derived_dependencies() {
        // Add `property_index` as a dependent of the dependency
        let dependency_meta = property_metadata
            .entry((E::id(), dependency))
            .or_insert_with(|| Box::new(PropertyMetadata::<E>::default()));
        let dependency_meta: &mut PropertyMetadata<E> = dependency_meta.downcast_mut().unwrap();
        dependency_meta.dependents.push(property_index);
    }
}

/// A convenience getter for `NEXT_ENTITY_INDEX`.
pub fn get_registered_property_count<E: Entity>() -> usize {
    let map = NEXT_PROPERTY_ID.lock().unwrap();
    *map.get(&E::id()).unwrap_or(&0)
}

/// Encapsulates the synchronization logic for initializing an item's index.
///
/// Acquires a global lock on the next available property ID, but only increments
/// it if we successfully initialize the provided ID. The ID of a property is
/// assigned at runtime but only once per type. It's possible for a single
/// type to attempt to initialize its index multiple times from different threads,
/// which is why all this synchronization is required. However, the overhead
/// is negligible, as this initialization only happens once upon first access.
///
/// In fact, for our use case we know we are calling this function
/// once for each type in each `Property`'s `ctor` function, which
/// should be the only time this method is ever called for the type.
pub fn initialize_property_id<E: Entity>(property_id: &AtomicUsize) -> usize {
    // Acquire a global lock.
    let mut guard = NEXT_PROPERTY_ID.lock().unwrap();
    let candidate = guard.entry(E::id()).or_insert_with(|| 0);

    // Try to claim the candidate index. Here we guard against the potential race condition that
    // another instance of this plugin in another thread just initialized the index prior to us
    // obtaining the lock. If the index has been initialized beneath us, we do not update
    // NEXT_PROPERTY_INDEX, we just return the value `index` was initialized to.
    // For a justification of the data ordering, see:
    //     https://github.com/CDCgov/ixa/pull/477#discussion_r2244302872
    match property_id.compare_exchange(usize::MAX, *candidate, Ordering::AcqRel, Ordering::Acquire)
    {
        Ok(_) => {
            // We won the race — increment the global next plugin index and return the new index
            *candidate += 1;
            *candidate - 1
        }
        Err(existing) => {
            // Another thread beat us — don’t increment the global next plugin index,
            // just return existing
            existing
        }
    }
}

/// A wrapper around a vector of property value stores.
pub struct PropertyStore<E: Entity> {
    /// A vector of `Box<PropertyValueStoreCore<E, P>>`, type-erased to `Box<dyn PropertyValueStore<E>>`
    items: Vec<Box<dyn PropertyValueStore<E>>>,
}

impl<E: Entity> Default for PropertyStore<E> {
    fn default() -> Self {
        PropertyStore::new()
    }
}

impl<E: Entity> PropertyStore<E> {
    /// Creates a new [`PropertyStore`].
    pub fn new() -> Self {
        let num_items = get_registered_property_count::<E>();
        // The constructors for each `PropertyValueStoreCore<E, P>` are stored in the `PROPERTY_METADATA` global.
        let property_metadata = property_metadata();

        // We construct the correct concrete `PropertyValueStoreCore<E, P>` value for each index (=`P::index()`).
        let items = (0..num_items)
            .map(|idx| {
                let metadata = property_metadata
                    .get(&(E::id(), idx))
                    .unwrap_or_else(|| panic!("No property metadata entry for index {idx}"))
                    .downcast_ref::<PropertyMetadata<E>>()
                    .unwrap_or_else(|| {
                        panic!(
                            "Property metadata entry for index {idx} does not match expected type"
                        )
                    });
                let constructor = metadata
                    .value_store_constructor
                    .unwrap_or_else(|| panic!("No PropertyValueStore constructor for index {idx}"));
                constructor()
            })
            .collect();

        Self { items }
    }

    /// Fetches an immutable reference to the `PropertyValueStoreCore<E, P>`.
    #[must_use]
    pub fn get<P: Property<E>>(&self) -> &PropertyValueStoreCore<E, P> {
        let index = P::id();
        let property_value_store =
            self.items
                .get(index)
                .unwrap_or_else(||
                    panic!(
                        "No registered property found with index = {:?} while trying to get property {}. You must use the `define_property!` macro to create a registered property.",
                        index,
                        P::name()
                    )
                );
        let property_value_store: &PropertyValueStoreCore<E, P> = property_value_store
            .as_any()
            .downcast_ref::<PropertyValueStoreCore<E, P>>()
            .unwrap_or_else(||
                {
                    panic!(
                        "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.",
                        index,
                        (**property_value_store).type_id(),
                        TypeId::of::<PropertyValueStoreCore<E, P>>()
                    )
                }
            );
        property_value_store
    }

    /// Fetches a mutable reference to the `PropertyValueStoreCore<E, P>`.
    #[must_use]
    pub fn get_mut<P: Property<E>>(&mut self) -> &mut PropertyValueStoreCore<E, P> {
        let index = P::id();
        let property_value_store =
            self.items
                .get_mut(index)
                .unwrap_or_else(||
                    panic!(
                        "No registered property found with index = {:?} while trying to get property {}. You must use the `define_property!` macro to create a registered property.",
                        index,
                        P::name()
                    )
                );
        let type_id = (**property_value_store).type_id(); // Only used for error message if error occurs.
        let property_value_store: &mut PropertyValueStoreCore<E, P> = property_value_store
            .as_any_mut()
            .downcast_mut::<PropertyValueStoreCore<E, P>>()
            .unwrap_or_else(||
                {
                    panic!(
                        "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.",
                        index,
                        type_id,
                        TypeId::of::<PropertyValueStoreCore<E, P>>()
                    )
                }
            );
        property_value_store
    }

    /// Creates a `PartialPropertyChangeEvent` instance for the `entity_id` and `property_index`. This method is only
    /// called for derived dependents of some property that has changed (one of `P`'s non-derived dependencies).
    pub(crate) fn create_partial_property_change(
        &self,
        property_index: usize,
        entity_id: EntityId<E>,
        context: &Context,
    ) -> PartialPropertyChangeEventBox {
        let property_value_store = self.items
                                       .get(property_index)
            .unwrap_or_else(|| panic!("No registered property found with index = {property_index:?}. You must use the `define_property!` macro to create a registered property."));

        property_value_store.create_partial_property_change(entity_id, context)
    }

    /// Returns whether the property with `property_index` needs partial change-event processing.
    pub(crate) fn should_create_partial_property_change(
        &self,
        property_index: usize,
        context: &Context,
    ) -> bool {
        let property_value_store = self.items
                                       .get(property_index)
            .unwrap_or_else(|| panic!("No registered property found with index = {property_index:?}. You must use the `define_property!` macro to create a registered property."));

        property_value_store.should_create_partial_change(context)
    }

    /// Returns whether or not the property `P` is indexed.
    ///
    /// This method can return `true` even if `context.index_property::<P>()` has never been called. For example,
    /// if a multi-property is indexed, all equivalent multi-properties are automatically also indexed, as they
    /// share a single index.
    #[cfg(test)]
    pub fn is_property_indexed<P: Property<E>>(&self) -> bool {
        self.items
            .get(P::index_id())
            .unwrap_or_else(|| panic!("No registered property {} found with index = {:?}. You must use the `define_property!` macro to create a registered property.", P::name(), P::index_id()))
            .index_type()
            != PropertyIndexType::Unindexed
    }

    /// Sets the index type for `P`. Passing `PropertyIndexType::Unindexed` removes any existing index for `P`.
    ///
    /// Note that the index might not live in the `PropertyValueStore` associated with `P` itself, as in the case
    /// of multi-properties which share a single index among all equivalent multi-properties.
    pub fn set_property_indexed<P: Property<E>>(&mut self, index_type: PropertyIndexType) {
        let property_value_store = self.items
            .get_mut(P::index_id())
            .unwrap_or_else(|| panic!("No registered property {} found with index = {:?}. You must use the `define_property!` macro to create a registered property.", P::name(), P::index_id()));
        property_value_store.set_indexed(index_type);
    }

    /// Creates a stratified value change counter for tracked property `P` with strata `PL`.
    ///
    /// Returns the counter ID.
    pub fn create_value_change_counter<PL, P>(&mut self) -> usize
    where
        PL: PropertyList<E> + Eq + std::hash::Hash,
        P: Property<E> + Eq + std::hash::Hash,
    {
        let property_value_store = self.get_mut::<P>();
        property_value_store.add_value_change_counter(Box::new(StratifiedValueChangeCounter::<
            E,
            PL,
            P,
        >::new()))
    }

    /// Updates the index of the property having the given ID for any entities that have been added to the context
    /// since the last time the index was updated.
    pub fn index_unindexed_entities_for_property_id(
        &mut self,
        context: &Context,
        property_id: usize,
    ) {
        self.items[property_id].index_unindexed_entities(context)
    }

    /// Updates all indexed properties for any entities that have been added since the last update.
    pub fn index_unindexed_entities_for_all_properties(&mut self, context: &Context) {
        for store in &mut self.items {
            store.index_unindexed_entities(context);
        }
    }

    pub fn get_index_set_for_query_parts(
        &self,
        property_id: usize,
        query_parts: &[&dyn Any],
    ) -> IndexSetResult<'_, E> {
        self.items[property_id].get_index_set_for_query_parts(query_parts)
    }

    pub fn get_index_count_for_query_parts(
        &self,
        property_id: usize,
        query_parts: &[&dyn Any],
    ) -> IndexCountResult {
        self.items[property_id].get_index_count_for_query_parts(query_parts)
    }
}

#[cfg(test)]
mod tests {
    #![allow(dead_code)]
    use std::any::Any;

    use super::*;
    use crate::entity::index::{IndexCountResult, IndexSetResult};
    use crate::prelude::*;
    use crate::{define_entity, define_property, with, Context};

    define_entity!(Person);

    define_property!(struct Age(u8), Person);
    define_property!(
        enum InfectionStatus {
            Susceptible,
            Infected,
            Recovered,
        },
        Person,
        default_const = InfectionStatus::Susceptible
    );
    define_property!(struct Vaccinated(bool), Person, default_const = Vaccinated(false));

    #[test]
    fn test_get_property_store() {
        let mut property_store = PropertyStore::new();

        {
            let ages: &mut PropertyValueStoreCore<_, Age> = property_store.get_mut();
            ages.set(EntityId::<Person>::new(0), Age(12));
            ages.set(EntityId::<Person>::new(1), Age(33));
            ages.set(EntityId::<Person>::new(2), Age(44));

            let infection_statuses: &mut PropertyValueStoreCore<_, InfectionStatus> =
                property_store.get_mut();
            infection_statuses.set(EntityId::<Person>::new(0), InfectionStatus::Susceptible);
            infection_statuses.set(EntityId::<Person>::new(1), InfectionStatus::Susceptible);
            infection_statuses.set(EntityId::<Person>::new(2), InfectionStatus::Infected);

            let vaccine_status: &mut PropertyValueStoreCore<_, Vaccinated> =
                property_store.get_mut();
            vaccine_status.set(EntityId::<Person>::new(0), Vaccinated(true));
            vaccine_status.set(EntityId::<Person>::new(1), Vaccinated(false));
            vaccine_status.set(EntityId::<Person>::new(2), Vaccinated(true));
        }

        // Verify that `get` returns the expected values
        {
            let ages: &PropertyValueStoreCore<_, Age> = property_store.get();
            assert_eq!(ages.get(EntityId::<Person>::new(0)), Age(12));
            assert_eq!(ages.get(EntityId::<Person>::new(1)), Age(33));
            assert_eq!(ages.get(EntityId::<Person>::new(2)), Age(44));

            let infection_statuses: &PropertyValueStoreCore<_, InfectionStatus> =
                property_store.get();
            assert_eq!(
                infection_statuses.get(EntityId::<Person>::new(0)),
                InfectionStatus::Susceptible
            );
            assert_eq!(
                infection_statuses.get(EntityId::<Person>::new(1)),
                InfectionStatus::Susceptible
            );
            assert_eq!(
                infection_statuses.get(EntityId::<Person>::new(2)),
                InfectionStatus::Infected
            );

            let vaccine_status: &PropertyValueStoreCore<_, Vaccinated> = property_store.get();
            assert_eq!(
                vaccine_status.get(EntityId::<Person>::new(0)),
                Vaccinated(true)
            );
            assert_eq!(
                vaccine_status.get(EntityId::<Person>::new(1)),
                Vaccinated(false)
            );
            assert_eq!(
                vaccine_status.get(EntityId::<Person>::new(2)),
                Vaccinated(true)
            );
        }
    }

    #[test]
    fn test_index_query_results_for_property_store() {
        let mut context = Context::new();
        context.index_property::<Person, Age>();

        let existing_value = Age(12);
        let missing_value = Age(99);
        let existing_query_parts = [&existing_value as &dyn Any];
        let missing_query_parts = [&missing_value as &dyn Any];

        let _ = context.add_entity(with!(Person, existing_value)).unwrap();
        let _ = context.add_entity(with!(Person, existing_value)).unwrap();

        let property_store = context.entity_store.get_property_store::<Person>();

        // FullIndex + count
        assert_eq!(
            property_store.get_index_count_for_query_parts(Age::index_id(), &missing_query_parts,),
            IndexCountResult::Count(0)
        );
        assert_eq!(
            property_store.get_index_count_for_query_parts(Age::index_id(), &existing_query_parts,),
            IndexCountResult::Count(2)
        );

        // FullIndex + set
        assert!(matches!(
            property_store.get_index_set_for_query_parts(Age::index_id(), &missing_query_parts,),
            IndexSetResult::Empty
        ));
        assert!(matches!(
            property_store.get_index_set_for_query_parts(
                Age::index_id(),
                &existing_query_parts,
            ),
            IndexSetResult::Set(set) if set.len() == 2
        ));
    }

    #[test]
    fn test_index_query_results_for_property_store_value_count_index() {
        let mut context = Context::new();
        context.index_property_counts::<Person, Age>();

        let existing_value = Age(12);
        let missing_value = Age(99);
        let existing_query_parts = [&existing_value as &dyn Any];
        let missing_query_parts = [&missing_value as &dyn Any];

        let _ = context.add_entity(with!(Person, existing_value)).unwrap();
        let _ = context.add_entity(with!(Person, existing_value)).unwrap();

        let property_store = context.entity_store.get_property_store::<Person>();

        // ValueCountIndex + count
        assert_eq!(
            property_store.get_index_count_for_query_parts(Age::index_id(), &missing_query_parts,),
            IndexCountResult::Count(0)
        );
        assert_eq!(
            property_store.get_index_count_for_query_parts(Age::index_id(), &existing_query_parts,),
            IndexCountResult::Count(2)
        );

        // ValueCountIndex + set (unsupported)
        assert!(matches!(
            property_store.get_index_set_for_query_parts(Age::index_id(), &missing_query_parts,),
            IndexSetResult::Unsupported
        ));
        assert!(matches!(
            property_store.get_index_set_for_query_parts(Age::index_id(), &existing_query_parts,),
            IndexSetResult::Unsupported
        ));
    }

    #[test]
    fn test_index_query_results_for_property_store_unindexed() {
        let mut context = Context::new();
        let existing_value = Age(12);
        let missing_value = Age(99);
        let existing_query_parts = [&existing_value as &dyn Any];
        let missing_query_parts = [&missing_value as &dyn Any];

        let _ = context.add_entity(with!(Person, existing_value)).unwrap();
        let _ = context.add_entity(with!(Person, existing_value)).unwrap();

        let property_store = context.entity_store.get_property_store::<Person>();

        // Unindexed + count
        assert_eq!(
            property_store.get_index_count_for_query_parts(Age::index_id(), &missing_query_parts,),
            IndexCountResult::Unsupported
        );
        assert_eq!(
            property_store.get_index_count_for_query_parts(Age::index_id(), &existing_query_parts,),
            IndexCountResult::Unsupported
        );

        // Unindexed + set
        assert!(matches!(
            property_store.get_index_set_for_query_parts(Age::index_id(), &missing_query_parts,),
            IndexSetResult::Unsupported
        ));
        assert!(matches!(
            property_store.get_index_set_for_query_parts(Age::index_id(), &existing_query_parts,),
            IndexSetResult::Unsupported
        ));
    }
}