ratio_metadata/

category.rs

//! # Metadata category store module
//!
//! A category store is a store for a certain category of metadata, such as names, kinds, labels,
//! weights, or annotations. It contains a field name interner that does the heavy lifting of
//! reducing the need to clone field names (except once for reverse indexing), and exposes different
//! methods for different types of values such as names and kinds for "single value" options, sets
//! such as labels, and maps such as weights and annotations.
//!
//! ## License
//!
//! This Source Code Form is subject to the terms of the Mozilla Public License, v. 2.0.
//! If a copy of the MPL was not distributed with this file,
//! You can obtain one at <https://mozilla.org/MPL/2.0/>.
//!
//! **Code examples both in the docstrings and rendered documentation are free to use.**

use std::borrow::{Borrow, Cow};
use std::collections::{BTreeMap, BTreeSet};
use std::fmt::Debug;
use std::iter::once;

use slotmap::{Key, SecondaryMap};

use crate::interner::{DenseSlottedBTreeInterner, Interner};

/// Storage of a certain categories of metadata.
#[derive(Clone, Debug)]
#[cfg_attr(
    feature = "serde",
    derive(serde::Serialize, serde::Deserialize),
    serde(default, rename_all = "camelCase")
)]
#[cfg_attr(feature = "reactive", derive(reactive_stores::Store))]
pub struct Category<ObjKey, Value, FieldKey, Field>
where
    ObjKey: Key,
    FieldKey: Key,
    Field: Clone + Debug + Ord,
{
    /// External keys to values expressed as collections or values using inner field keys.
    values: SecondaryMap<ObjKey, Value>,

    /// Internal keys to external keys that use them.
    objects_by_field: SecondaryMap<FieldKey, BTreeSet<ObjKey>>,

    /// Field name interner of the category store.
    field_interner: DenseSlottedBTreeInterner<FieldKey, Field>,
}

impl<ObjKey, Value, FieldKey, Field> Default for Category<ObjKey, Value, FieldKey, Field>
where
    ObjKey: Key,
    FieldKey: Key,
    Field: Clone + Debug + Ord,
{
    fn default() -> Self {
        Self {
            values: SecondaryMap::default(),
            objects_by_field: SecondaryMap::default(),
            field_interner: DenseSlottedBTreeInterner::default(),
        }
    }
}

/// Behavior for any Category store.
impl<ObjKey, Value, FieldKey, Field> Category<ObjKey, Value, FieldKey, Field>
where
    ObjKey: Key,
    FieldKey: Key,
    Field: Clone + Debug + Ord,
{
    /// Create a new category store.
    pub fn new() -> Self {
        Self::default()
    }

    /// Interned key for a field name.
    /// Returns `None` if the field name is unknown (not interned).
    pub fn field_key_interned(&self, name: &Field) -> Option<FieldKey> {
        self.field_interner.get_key(name)
    }

    /// Intern a field name and return the key.
    pub fn field_key(&mut self, name: Cow<'_, Field>) -> FieldKey {
        self.field_interner.intern(name)
    }

    /// Obtain a reference to an interned field name via its key.
    pub fn field_name(&self, key: FieldKey) -> Option<&Field> {
        self.field_interner.get_value(key)
    }

    /// Number of tracked objects.
    pub fn num_objects(&self) -> usize {
        self.values.len()
    }

    /// All object keys in store.
    pub fn object_keys(&self) -> impl Iterator<Item = ObjKey> {
        self.values.keys()
    }

    /// Number of tracked fields.
    pub fn num_fields(&self) -> usize {
        self.field_interner.len()
    }

    /// All field names in store.
    pub fn field_names(&self) -> impl Iterator<Item = &Field> {
        self.field_interner.values()
    }

    /// All field keys by field name.
    pub fn field_keys(&self) -> impl Iterator<Item = FieldKey> {
        self.field_interner.keys()
    }

    /// All field names and keys in store.
    pub fn field_items(&self) -> impl Iterator<Item = (FieldKey, &Field)> {
        self.field_interner.items()
    }

    /// Get the current value expressed in internal keys.
    pub fn object_value(&self, obj: ObjKey) -> Option<&Value> {
        self.values.get(obj)
    }

    /// Get all objects which have a certain field.
    pub fn objects_with_field(&self, field: FieldKey) -> Option<&BTreeSet<ObjKey>> {
        self.objects_by_field.get(field)
    }

    /// Get all objects which have a certain field name.
    pub fn objects_with_field_name(&self, name: &Field) -> Option<&BTreeSet<ObjKey>> {
        let key = self.field_key_interned(name)?;
        self.objects_with_field(key)
    }

    /// Whether this manager tracks this object.
    pub fn contains_object(&self, obj: ObjKey) -> bool {
        self.values.contains_key(obj)
    }

    /// Whether this field key is being tracked by this category.
    /// Should hold as long as the key is valid.
    pub fn contains_field(&self, field: FieldKey) -> bool {
        self.objects_by_field.contains_key(field)
    }

    /// Whether this field name is being tracked by this category.
    pub fn contains_field_name<Name: Borrow<Field>>(&self, name: Name) -> bool {
        self.field_interner.contains_value(name.borrow())
    }

    /// Whether an object has this field set.
    pub fn object_contains_field(&self, obj: ObjKey, field: FieldKey) -> bool {
        self.objects_with_field(field)
            .map(|objs| objs.contains(&obj))
            .unwrap_or_default()
    }
}

/// Behavior for any category store with field-key encoded values.
impl<ObjKey, Value, FieldKey, Field> Category<ObjKey, Value, FieldKey, Field>
where
    ObjKey: Key,
    FieldKey: Key,
    Field: Clone + Debug + Ord,
    Value: FieldValue<FieldKey>,
{
    /// Clean the innards of this store, clearing the fields which are not assigned to any
    /// objects from the reverse index and field name interner.
    pub fn clean(&mut self) {
        self.values.retain(|_, fields| fields.should_retain());

        let unused = self
            .field_interner
            .keys()
            .filter(|&field| match self.objects_by_field.get(field) {
                Some(objects) => objects.is_empty(),
                None => true,
            })
            .collect::<Vec<_>>();

        for field in unused {
            self.objects_by_field.remove(field);
            self.field_interner.remove_key(field);
        }
    }

    /// Get a value entry by its object key.
    fn object_value_entry(&mut self, obj: ObjKey) -> Option<&mut Value> {
        self.values.entry(obj).map(|entry| entry.or_default())
    }

    /// Get a objects entry by field.
    fn objects_by_field_entry(&mut self, field: FieldKey) -> Option<&mut BTreeSet<ObjKey>> {
        self.objects_by_field
            .entry(field)
            .map(|entry| entry.or_default())
    }

    /// Forget an object's entire value, removing all fields from it and itself from all fields.
    /// Returns whether anything was removed.
    pub fn remove_object(&mut self, obj: ObjKey) -> Option<Value> {
        self.values.remove(obj).inspect(|value| {
            value.fields().for_each(|field| {
                self.remove_object_from_field_objects(field, obj);
            });
        })
    }

    /// Add an object to the field reverse search.
    fn add_object_to_field_objects(&mut self, field: FieldKey, obj: ObjKey) -> bool {
        self.objects_by_field_entry(field)
            .map(|objects| objects.insert(obj))
            .unwrap_or_default()
    }

    /// Remove an object from the reverse field search.
    /// Removes the reverse search entry if it was the last object.
    /// Does not remove it from the interned value, as to not invalidate any interned key.
    /// Returns whether anything was removed.
    /// Private method, as users should turn to `remove_field_from_object`.
    fn remove_object_from_field_objects(&mut self, field: FieldKey, obj: ObjKey) -> bool {
        let (is_empty, removed) = self
            .objects_by_field
            .get_mut(field)
            .map(|objects| {
                let removed = objects.remove(&obj);
                (objects.is_empty(), removed)
            })
            .unwrap_or_default();
        if is_empty {
            self.objects_by_field.remove(field);
        }
        removed
    }

    /// Set an object value, replacing any pre-existing value.
    pub fn insert_object_value(&mut self, obj: ObjKey, value: Value) {
        self.remove_object(obj);
        value.fields().for_each(|field| {
            self.objects_by_field
                .entry(field)
                .map(|entry| entry.or_default().insert(obj));
        });
        self.values.insert(obj, value);
    }

    /// All fields attached to an object.
    pub fn object_fields(&self, obj: ObjKey) -> impl Iterator<Item = FieldKey> {
        self.values
            .get(obj)
            .into_iter()
            .flat_map(|value| value.fields())
    }

    /// All field names attached to an object.
    pub fn object_field_names(&self, obj: ObjKey) -> impl Iterator<Item = &Field> {
        self.object_fields(obj).flat_map(|key| self.field_name(key))
    }
}

/// Behavior for any category store with field-key encoded values *without* sub-values.
impl<ObjKey, Value, FieldKey, Field> Category<ObjKey, Value, FieldKey, Field>
where
    ObjKey: Key,
    FieldKey: Key,
    Field: Clone + Debug + Ord,
    Value: FieldValue<FieldKey> + FieldWithoutSubValue<FieldKey>,
{
    /// Insert a field name for an object.
    pub fn insert_field_name_for_object(&mut self, obj: ObjKey, name: Cow<'_, Field>) -> bool {
        let field = self.field_key(name);
        self.insert_field_for_object(obj, field)
    }

    /// Insert a field key for an object.
    pub fn insert_field_for_object(&mut self, obj: ObjKey, field: FieldKey) -> bool {
        self.add_object_to_field_objects(field, obj);
        self.object_value_entry(obj)
            .map(|value| value.insert_field(field))
            .unwrap_or_default()
    }

    /// Remove a field name from the store from all objects by its name.
    /// Returns whether it had to be removed.
    pub fn remove_field_name(&mut self, name: &Field) -> Option<FieldKey> {
        self.field_key_interned(name).map(|field| {
            self.remove_field(field);
            field
        })
    }

    /// Remove a field entirely from the store from all objects.
    /// Returns whether it had to be removed.
    pub fn remove_field(&mut self, field: FieldKey) -> Option<Field> {
        self.objects_by_field
            .remove(field)
            .into_iter()
            .for_each(|objects| {
                objects.into_iter().for_each(|obj| {
                    self.remove_field_from_object(obj, field);
                });
            });
        self.field_interner.remove_key(field)
    }

    /// Remove a field by its name from an object.
    /// Returns whether anything was actually removed.
    pub fn remove_field_name_from_object(&mut self, obj: ObjKey, field: &Field) -> bool {
        self.field_key_interned(field)
            .map(|field| self.remove_field_from_object(obj, field))
            .unwrap_or_default()
    }

    /// Remove a field from an object.
    /// Returns whether anything was actually removed.
    pub fn remove_field_from_object(&mut self, obj: ObjKey, field: FieldKey) -> bool {
        let retain = self
            .values
            .get_mut(obj)
            .map(|value| {
                value.remove_field(field);
                value.should_retain()
            })
            .unwrap_or_default();

        if !retain {
            self.values.remove(obj);
        }

        self.remove_object_from_field_objects(field, obj)
    }
}

/// Behavior for any category store with field-key encoded values *with* sub-values.
impl<ObjKey, Value, SubValue, FieldKey, Field> Category<ObjKey, Value, FieldKey, Field>
where
    ObjKey: Key,
    FieldKey: Key,
    Field: Clone + Debug + Ord,
    Value: FieldValue<FieldKey> + FieldWithSubValue<FieldKey, SubValue = SubValue>,
{
    /// Insert a field sub-value by its name for an object.
    /// Returns the previous value if any.
    pub fn insert_field_name_value_for_object(
        &mut self,
        obj: ObjKey,
        field: Cow<Field>,
        sub_value: SubValue,
    ) -> Option<SubValue> {
        let field = self.field_key(field);

        self.insert_field_value_for_object(obj, field, sub_value)
    }

    /// Insert a field sub-value for an object by its field key.
    /// Returns the previous value if any.
    pub fn insert_field_value_for_object(
        &mut self,
        obj: ObjKey,
        field: FieldKey,
        sub_value: SubValue,
    ) -> Option<SubValue> {
        self.add_object_to_field_objects(field, obj);
        self.object_value_entry(obj)
            .and_then(|value| value.insert_field_value(field, sub_value))
    }

    /// Remove a field from all objects.
    pub fn remove_field_name_values(&mut self, name: &Field) -> Option<FieldKey> {
        self.field_key_interned(name).map(move |field| {
            self.remove_field_values(field);
            field
        })
    }

    /// Remove field values from all objects.
    pub fn remove_field_values(&mut self, field: FieldKey) -> Option<Field> {
        self.objects_by_field
            .remove(field)
            .into_iter()
            .flat_map(|objects| objects.into_iter())
            .for_each(|obj| {
                self.remove_field_value_from_object(obj, field);
            });

        self.field_interner.remove_key(field)
    }

    /// Remove a field sub-value by its name from an object.
    /// Returns whether anything was actually removed.
    pub fn remove_field_name_value_from_object(
        &mut self,
        obj: ObjKey,
        name: &Field,
    ) -> Option<SubValue> {
        self.field_key_interned(name)
            .and_then(|field| self.remove_field_value_from_object(obj, field))
    }

    /// Remove a field from an object.
    /// Returns whether anything was actually removed.
    pub fn remove_field_value_from_object(
        &mut self,
        obj: ObjKey,
        field: FieldKey,
    ) -> Option<SubValue> {
        self.remove_object_from_field_objects(field, obj);

        let (retain, sub_value) = self
            .values
            .get_mut(obj)
            .map(|value| {
                let sub_value = value.remove_field_value(field);
                (value.should_retain(), sub_value)
            })
            .unwrap_or_default();

        if !retain {
            self.values.remove(obj);
        }

        sub_value
    }

    /// Get a field sub-value for an object by its field key.
    pub fn field_value_for_object(&self, obj: ObjKey, field: FieldKey) -> Option<&SubValue> {
        self.values
            .get(obj)
            .and_then(|value| value.field_value(field))
    }

    /// Get a field sub-value by its name for an object.
    pub fn field_name_value_for_object(&self, obj: ObjKey, field: &Field) -> Option<&SubValue> {
        self.field_key_interned(field)
            .and_then(|field| self.field_value_for_object(obj, field))
    }
}

/// Behavior for a value based on interned field keys.
pub trait FieldValue<FieldKey>: Default {
    /// Whether to retain this value.
    fn should_retain(&self) -> bool;

    /// All involved fields for this value.
    fn fields(&self) -> impl Iterator<Item = FieldKey>;
}
impl<FieldKey: Key> FieldValue<FieldKey> for FieldKey {
    fn should_retain(&self) -> bool {
        self.is_null()
    }
    fn fields(&self) -> impl Iterator<Item = FieldKey> {
        once(*self)
    }
}
impl<FieldKey: Key> FieldValue<FieldKey> for Option<FieldKey> {
    fn should_retain(&self) -> bool {
        self.is_some()
    }

    fn fields(&self) -> impl Iterator<Item = FieldKey> {
        self.iter().copied()
    }
}
impl<FieldKey: Key> FieldValue<FieldKey> for BTreeSet<FieldKey> {
    fn should_retain(&self) -> bool {
        !self.is_empty()
    }

    fn fields(&self) -> impl Iterator<Item = FieldKey> {
        self.iter().copied()
    }
}
impl<FieldKey: Key, SubValue> FieldValue<FieldKey> for BTreeMap<FieldKey, SubValue> {
    fn should_retain(&self) -> bool {
        !self.is_empty()
    }

    fn fields(&self) -> impl Iterator<Item = FieldKey> {
        self.keys().copied()
    }
}

/// Behavior for a value without sub-values per interned field key such as `Option<FieldKey>` and
/// `BTreeSet<FieldKey>`.
pub trait FieldWithoutSubValue<FieldKey: Key> {
    /// Set or insert this field in the value. Returns whether it was newly inserted.
    fn insert_field(&mut self, field: FieldKey) -> bool;

    /// Remove field. Returns if anything was removed.
    fn remove_field(&mut self, field: FieldKey) -> bool;
}

impl<FieldKey: Key> FieldWithoutSubValue<FieldKey> for FieldKey {
    fn insert_field(&mut self, mut field: FieldKey) -> bool {
        if self == &field {
            false
        } else {
            std::mem::swap(self, &mut field);
            true
        }
    }
    fn remove_field(&mut self, field: FieldKey) -> bool {
        if self == &field {
            std::mem::take(self);
            true
        } else {
            false
        }
    }
}
impl<FieldKey: Key> FieldWithoutSubValue<FieldKey> for Option<FieldKey> {
    fn insert_field(&mut self, field: FieldKey) -> bool {
        if self == &Some(field) {
            false
        } else {
            self.replace(field);
            true
        }
    }

    fn remove_field(&mut self, field: FieldKey) -> bool {
        if self.as_ref() == Some(&field) {
            self.take().is_some()
        } else {
            false
        }
    }
}
impl<FieldKey: Key> FieldWithoutSubValue<FieldKey> for BTreeSet<FieldKey> {
    fn insert_field(&mut self, field: FieldKey) -> bool {
        self.insert(field)
    }

    fn remove_field(&mut self, field: FieldKey) -> bool {
        self.remove(&field)
    }
}
impl<FieldKey: Key, SubValue: Default> FieldWithoutSubValue<FieldKey>
    for BTreeMap<FieldKey, SubValue>
{
    fn insert_field(&mut self, field: FieldKey) -> bool {
        self.insert(field, SubValue::default()).is_none()
    }

    fn remove_field(&mut self, field: FieldKey) -> bool {
        self.remove(&field).is_some()
    }
}

/// Behavior for a value with sub-values for interned field keys.
pub trait FieldWithSubValue<FieldKey: Key> {
    type SubValue;

    /// Access a field's sub-value for it's interned field key.
    fn field_value(&self, field: FieldKey) -> Option<&Self::SubValue>;

    /// Insert a field's sub-value for it's interned field key.
    fn insert_field_value(
        &mut self,
        field: FieldKey,
        value: Self::SubValue,
    ) -> Option<Self::SubValue>;

    /// Remove a sub-value for an interned field key.
    fn remove_field_value(&mut self, field: FieldKey) -> Option<Self::SubValue>;
}
impl<FieldKey: Key, SubValue> FieldWithSubValue<FieldKey> for BTreeMap<FieldKey, SubValue> {
    type SubValue = SubValue;

    fn field_value(&self, field: FieldKey) -> Option<&SubValue> {
        self.get(&field)
    }

    fn insert_field_value(&mut self, field: FieldKey, value: SubValue) -> Option<SubValue> {
        self.insert(field, value)
    }

    fn remove_field_value(&mut self, field: FieldKey) -> Option<SubValue> {
        self.remove(&field)
    }
}

/// Behavior for a category store whose values are just a single optional field key.
impl<ObjKey, FieldKey, Field> Category<ObjKey, FieldKey, FieldKey, Field>
where
    ObjKey: Key,
    FieldKey: Key,
    Field: Clone + Debug + Ord,
{
    /// Get the current value expressed in terms of the external name.
    pub fn value_name_of_object(&self, obj: ObjKey) -> Option<&Field> {
        let key = self.object_value(obj)?;
        self.field_name(*key)
    }

    /// Set the value expressed in terms of the external name.
    pub fn set_value_name_of_object(&mut self, obj: ObjKey, value: Option<Cow<Field>>) {
        if let Some(field) = value {
            let field = self.field_key(field);
            self.insert_object_value(obj, field);
        } else {
            self.remove_object(obj);
        }
    }
}

/// Behavior for a category store whose values are just a single optional field key.
impl<ObjKey, FieldKey, Field> Category<ObjKey, Option<FieldKey>, FieldKey, Field>
where
    ObjKey: Key,
    FieldKey: Key,
    Field: Clone + Debug + Ord,
{
    /// Get the current value expressed in terms of the external name.
    pub fn value_name_of_object(&self, obj: ObjKey) -> Option<&Field> {
        let key = self.object_value(obj)?.as_ref()?;
        self.field_name(*key)
    }

    /// Set the value expressed in terms of the external name.
    pub fn set_value_name_of_object(&mut self, obj: ObjKey, value: Option<Cow<Field>>) {
        if let Some(field) = value {
            let field = self.field_key(field);
            self.insert_object_value(obj, Some(field));
        } else {
            self.remove_object(obj);
        }
    }
}

/// Behavior for a category store whose values are a BTreeSet.
impl<ObjKey, FieldKey, Field> Category<ObjKey, BTreeSet<FieldKey>, FieldKey, Field>
where
    ObjKey: Key,
    FieldKey: Key,
    Field: Clone + Debug + Ord,
{
    /// Get the current value expressed in references to field names.
    pub fn named_value_of_object(&self, obj: ObjKey) -> BTreeSet<&Field> {
        self.object_value(obj)
            .into_iter()
            .flat_map(|set| set.iter())
            .flat_map(|key| self.field_name(*key))
            .collect()
    }
}

/// Behavior for a category store whose values are a BTreeMap with SubValues.
impl<ObjKey, FieldKey, Field, SubValue>
    Category<ObjKey, BTreeMap<FieldKey, SubValue>, FieldKey, Field>
where
    ObjKey: Key,
    FieldKey: Key,
    Field: Clone + Debug + Ord,
{
    /// Get the current value expressed in references to field names and values.
    pub fn named_value_of_object(&self, obj: ObjKey) -> BTreeMap<&Field, &SubValue> {
        self.object_value(obj)
            .into_iter()
            .flat_map(|mapping| mapping.iter())
            .flat_map(|(key, value)| self.field_name(*key).map(|field| (field, value)))
            .collect()
    }
}