cranpose_core/

retention.rs

use std::cmp::Ordering;

#[cfg(any(test, debug_assertions))]
use crate::slot::{AnchorState, PayloadAnchorLifecycle, SlotInvariantError};
#[cfg(any(test, debug_assertions))]
use crate::{AnchorId, SlotTable};
use crate::{
    ScopeId,
    collections::map::HashMap,
    slot::{DetachedSubtree, GroupKey, NodeLifecycle},
};

/// What a group's content becomes when the group leaves the composition.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum RetentionMode {
    /// Disposes the group's content as soon as the group leaves the composition.
    #[default]
    DisposeWhenInactive,
    /// Detaches the group and keeps its content, so the group returns with its nodes and state when
    /// its key is composed again. The [`RetentionPolicy`] budget can still evict it.
    RetainWhenInactive,
}

/// Limits on how much content a composition keeps for groups that left it. Content over a limit is
/// evicted in the order of the [`RetentionEvictionPolicy`].
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct RetentionBudget {
    /// Most detached subtrees to keep at once. `None` sets no limit.
    pub max_retained_subtrees: Option<usize>,
    /// Most heap bytes the retained subtrees may take together. `None` sets no limit.
    pub max_retained_bytes: Option<usize>,
    /// Most composition passes a subtree may stay retained before it is evicted. `None` sets no age
    /// limit.
    pub max_age_passes: Option<u64>,
}

impl RetentionBudget {
    /// A budget with no limit on any measure.
    pub const UNBOUNDED: Self = Self {
        max_retained_subtrees: None,
        max_retained_bytes: None,
        max_age_passes: None,
    };
}

impl Default for RetentionBudget {
    fn default() -> Self {
        Self::UNBOUNDED
    }
}

/// Which retained subtree is evicted first when a [`RetentionBudget`] is exceeded.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum RetentionEvictionPolicy {
    /// The subtree that was detached the longest ago.
    #[default]
    LeastRecentlyDetached,
    /// The subtree whose last restore is the oldest. Subtrees never restored go first.
    LeastRecentlyRestored,
    /// The subtree that takes the most heap memory.
    LargestFirst,
}

/// How a composition limits and evicts retained content: the budget it may keep and the order in
/// which it drops subtrees that exceed the budget.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct RetentionPolicy {
    /// The limits on how much retained content the composition keeps.
    pub budget: RetentionBudget,
    /// Which retained subtree is evicted first when the budget is exceeded.
    pub eviction: RetentionEvictionPolicy,
}

impl RetentionPolicy {
    /// A policy with an unbounded budget that evicts the least recently detached subtree first.
    pub const UNBOUNDED: Self = Self {
        budget: RetentionBudget::UNBOUNDED,
        eviction: RetentionEvictionPolicy::LeastRecentlyDetached,
    };
}

impl Default for RetentionPolicy {
    fn default() -> Self {
        Self::UNBOUNDED
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub(crate) struct RetainKey {
    pub(crate) parent_scope: Option<ScopeId>,
    pub(crate) key: GroupKey,
}

impl RetainKey {
    pub(crate) fn for_group(parent_scope: Option<ScopeId>, key: GroupKey) -> Self {
        Self {
            parent_scope: if key.is_movable() { None } else { parent_scope },
            key,
        }
    }
}

pub(crate) struct RetainedGroup {
    pub(crate) subtree: DetachedSubtree,
    detached_pass: u64,
    detached_order: u64,
    last_restored_order: u64,
    pinned: bool,
}

impl RetainedGroup {
    fn node_count(&self) -> usize {
        self.subtree.node_count()
    }

    fn payload_count(&self) -> usize {
        self.subtree.payload_count()
    }

    fn scope_count(&self) -> usize {
        self.subtree.scope_count()
    }

    fn anchor_count(&self) -> usize {
        self.subtree.anchor_count()
    }

    fn heap_bytes(&self) -> usize {
        self.subtree.heap_bytes()
    }
}

#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub(crate) struct RetentionDebugStats {
    pub(crate) subtree_count: usize,
    pub(crate) group_count: usize,
    pub(crate) payload_count: usize,
    pub(crate) node_count: usize,
    pub(crate) scope_count: usize,
    pub(crate) anchor_count: usize,
    pub(crate) heap_bytes: usize,
    pub(crate) evictions_total: usize,
}

pub(crate) struct RetentionManager {
    groups: HashMap<RetainKey, RetainedGroup>,
    restored_at_by_key: HashMap<RetainKey, u64>,
    policy: RetentionPolicy,
    pass_clock: u64,
    operation_clock: u64,
    evictions_total: usize,
}

impl Default for RetentionManager {
    fn default() -> Self {
        Self::new(RetentionPolicy::default())
    }
}

impl RetentionManager {
    pub(crate) fn new(policy: RetentionPolicy) -> Self {
        Self {
            groups: HashMap::default(),
            restored_at_by_key: HashMap::default(),
            policy,
            pass_clock: 0,
            operation_clock: 0,
            evictions_total: 0,
        }
    }

    pub(crate) fn set_policy(&mut self, policy: RetentionPolicy) {
        self.policy = policy;
    }

    pub(crate) fn take(&mut self, key: RetainKey) -> Option<DetachedSubtree> {
        let retained = self.groups.get(&key)?;
        if !retained_subtree_matches_key(&retained.subtree, key, "restore") {
            return None;
        }
        if !retained_subtree_root_is_detached(&retained.subtree, key, "restore") {
            return None;
        }
        if !retained_subtree_nodes_have_lifecycle(
            &retained.subtree,
            key,
            NodeLifecycle::RetainedDetached,
            "restore",
        ) {
            return None;
        }
        let restored_order = self.tick_operation();
        let retained = self.groups.remove(&key)?;
        self.restored_at_by_key.insert(key, restored_order);
        Some(retained.subtree)
    }

    pub(crate) fn take_after_restore_preflight(
        &mut self,
        key: RetainKey,
        preflight: impl FnOnce(&mut DetachedSubtree) -> bool,
    ) -> Option<DetachedSubtree> {
        if !preflight(&mut self.groups.get_mut(&key)?.subtree) {
            log::error!(
                "retention restore preflight rejected subtree for parent_scope={:?} key={:?}",
                key.parent_scope,
                key.key
            );
            return None;
        }
        self.take(key)
    }

    pub(crate) fn contains(&self, key: RetainKey) -> bool {
        self.groups.contains_key(&key)
    }

    pub(crate) fn insert(
        &mut self,
        key: RetainKey,
        subtree: DetachedSubtree,
    ) -> Vec<DetachedSubtree> {
        self.insert_with_pin(key, subtree, false)
    }

    pub(crate) fn insert_pinned(
        &mut self,
        key: RetainKey,
        subtree: DetachedSubtree,
    ) -> Vec<DetachedSubtree> {
        self.insert_with_pin(key, subtree, true)
    }

    fn insert_with_pin(
        &mut self,
        key: RetainKey,
        mut subtree: DetachedSubtree,
        pinned: bool,
    ) -> Vec<DetachedSubtree> {
        if self.groups.contains_key(&key) {
            log::error!(
                "retention insert rejected duplicate key for parent_scope={:?} key={:?}",
                key.parent_scope,
                key.key
            );
            return vec![subtree];
        }
        if !retained_subtree_matches_key(&subtree, key, "insert") {
            return vec![subtree];
        }
        if !retained_subtree_root_is_detached(&subtree, key, "insert") {
            return vec![subtree];
        }
        let detached_order = self.tick_operation();
        let last_restored_order = self.restored_at_by_key.remove(&key).unwrap_or_default();
        subtree.mark_nodes_retained_detached();
        self.groups.insert(
            key,
            RetainedGroup {
                subtree,
                detached_pass: self.pass_clock,
                detached_order,
                last_restored_order,
                pinned,
            },
        );
        self.evict_to_budget()
    }

    pub(crate) fn is_empty(&self) -> bool {
        self.groups.is_empty()
    }

    pub(crate) fn debug_stats(&self) -> RetentionDebugStats {
        RetentionDebugStats {
            subtree_count: self.groups.len(),
            group_count: self
                .groups
                .values()
                .map(|retained| retained.subtree.group_count())
                .sum(),
            payload_count: self.groups.values().map(RetainedGroup::payload_count).sum(),
            node_count: self.groups.values().map(RetainedGroup::node_count).sum(),
            scope_count: self.groups.values().map(RetainedGroup::scope_count).sum(),
            anchor_count: self.groups.values().map(RetainedGroup::anchor_count).sum(),
            heap_bytes: self.groups.values().map(RetainedGroup::heap_bytes).sum(),
            evictions_total: self.evictions_total,
        }
    }

    pub(crate) fn into_subtrees(self) -> Vec<DetachedSubtree> {
        self.groups
            .into_values()
            .map(|retained| retained.subtree)
            .collect()
    }

    pub(crate) fn subtrees(&self) -> impl Iterator<Item = &DetachedSubtree> + '_ {
        self.groups.values().map(|retained| &retained.subtree)
    }

    #[cfg(test)]
    pub(crate) fn subtrees_mut(&mut self) -> impl Iterator<Item = &mut DetachedSubtree> + '_ {
        self.groups
            .values_mut()
            .map(|retained| &mut retained.subtree)
    }

    #[cfg(any(test, debug_assertions))]
    pub(crate) fn validate(&self, table: &SlotTable) -> Result<(), SlotInvariantError> {
        for (key, retained) in &self.groups {
            let subtree = &retained.subtree;
            subtree.validate_detached()?;
            let Some(root_key) = subtree.root_key_checked() else {
                return Err(SlotInvariantError::DetachedSubtreeEmpty);
            };
            if root_key != key.key {
                return Err(SlotInvariantError::RetainedRootKeyMismatch {
                    parent_scope: key.parent_scope,
                    expected: key.key,
                    actual: root_key,
                });
            }

            let root_parent_anchor = subtree
                .root_parent_anchor_checked()
                .unwrap_or(AnchorId::INVALID);
            if root_parent_anchor.is_valid() {
                return Err(SlotInvariantError::RetainedRootHasActiveParent {
                    root_key,
                    parent_anchor: root_parent_anchor,
                });
            }

            for anchor in subtree.group_anchors() {
                match table.anchor_state(anchor) {
                    Some(AnchorState::Detached) => {}
                    Some(AnchorState::Active(active_index)) => {
                        return Err(SlotInvariantError::RetainedSubtreeAnchorStillActive {
                            root_key,
                            anchor,
                            active_index: active_index as usize,
                        });
                    }
                    actual => {
                        return Err(SlotInvariantError::RetainedAnchorStateMismatch {
                            root_key,
                            anchor,
                            actual,
                        });
                    }
                }
            }

            for scope in subtree.scopes() {
                if table.active_group_for_scope(scope).is_some() {
                    return Err(SlotInvariantError::RetainedScopeStillActive {
                        root_key,
                        scope_id: scope.id(),
                        active_anchor: scope.group_anchor(),
                    });
                }
            }

            for payload_anchor in subtree.payload_anchors() {
                match table.payload_anchor_lifecycle(payload_anchor) {
                    Some(PayloadAnchorLifecycle::Detached) => {}
                    Some(PayloadAnchorLifecycle::Active) => {
                        let Some((active_owner, active_index)) =
                            table.payload_anchor_active_location(payload_anchor)
                        else {
                            return Err(SlotInvariantError::RetainedPayloadAnchorStateMismatch {
                                root_key,
                                payload_anchor,
                                actual: Some(PayloadAnchorLifecycle::Active),
                            });
                        };
                        return Err(SlotInvariantError::RetainedPayloadAnchorStillActive {
                            root_key,
                            payload_anchor,
                            active_owner,
                            active_index,
                        });
                    }
                    actual => {
                        return Err(SlotInvariantError::RetainedPayloadAnchorStateMismatch {
                            root_key,
                            payload_anchor,
                            actual,
                        });
                    }
                }
            }

            for (node_id, lifecycle) in subtree.node_states() {
                if lifecycle != NodeLifecycle::RetainedDetached {
                    return Err(SlotInvariantError::RetainedNodeLifecycleMismatch {
                        root_key,
                        node_id,
                        actual: lifecycle,
                    });
                }
            }
        }
        Ok(())
    }

    #[cfg(any(test, debug_assertions))]
    pub(crate) fn debug_verify(&self, table: &SlotTable) {
        if crate::slot_validation_diagnostics_enabled()
            && let Err(err) = self.validate(table)
        {
            panic!("retention invariant violation: {err:?}");
        }
    }

    pub(crate) fn evictions_total(&self) -> usize {
        self.evictions_total
    }

    pub(crate) fn advance_pass(&mut self) -> Vec<DetachedSubtree> {
        self.pass_clock = self.pass_clock.saturating_add(1);
        self.evict_to_budget()
    }

    fn tick_operation(&mut self) -> u64 {
        self.operation_clock = self.operation_clock.saturating_add(1);
        self.operation_clock
    }

    fn evict_to_budget(&mut self) -> Vec<DetachedSubtree> {
        let mut evicted = Vec::new();
        while let Some(key) = self.budget_eviction_key() {
            let Some(retained) = self.groups.remove(&key) else {
                break;
            };
            self.restored_at_by_key.remove(&key);
            self.evictions_total = self.evictions_total.saturating_add(1);
            if !retained_subtree_matches_key(&retained.subtree, key, "eviction")
                || !retained_subtree_root_is_detached(&retained.subtree, key, "eviction")
                || !retained_subtree_nodes_have_lifecycle(
                    &retained.subtree,
                    key,
                    NodeLifecycle::RetainedDetached,
                    "eviction",
                )
            {
                log::error!(
                    "retention eviction returned malformed retained subtree for caller-owned disposal"
                );
                evicted.push(retained.subtree);
                continue;
            }
            evicted.push(retained.subtree);
        }
        evicted
    }

    fn budget_eviction_key(&self) -> Option<RetainKey> {
        if self.groups.is_empty() {
            return None;
        }

        if let Some(max_age_passes) = self.policy.budget.max_age_passes
            && let Some(key) = self.age_eviction_key(max_age_passes)
        {
            return Some(key);
        }

        let over_count = self
            .policy
            .budget
            .max_retained_subtrees
            .is_some_and(|max| self.groups.len() > max);
        let over_bytes = self
            .policy
            .budget
            .max_retained_bytes
            .is_some_and(|max| self.retained_heap_bytes() > max);

        (over_count || over_bytes)
            .then(|| self.policy_eviction_key())
            .flatten()
    }

    fn evictable(&self) -> impl Iterator<Item = (&RetainKey, &RetainedGroup)> + '_ {
        self.groups.iter().filter(|(_, retained)| !retained.pinned)
    }

    fn age_eviction_key(&self, max_age_passes: u64) -> Option<RetainKey> {
        self.evictable()
            .filter(|(_, retained)| {
                self.pass_clock.saturating_sub(retained.detached_pass) > max_age_passes
            })
            .min_by(|(left_key, left), (right_key, right)| {
                left.detached_pass
                    .cmp(&right.detached_pass)
                    .then_with(|| left.detached_order.cmp(&right.detached_order))
                    .then_with(|| retain_key_cmp(left_key, right_key))
            })
            .map(|(key, _)| *key)
    }

    fn policy_eviction_key(&self) -> Option<RetainKey> {
        match self.policy.eviction {
            RetentionEvictionPolicy::LeastRecentlyDetached => self.least_recently_detached_key(),
            RetentionEvictionPolicy::LeastRecentlyRestored => self.least_recently_restored_key(),
            RetentionEvictionPolicy::LargestFirst => self.largest_first_key(),
        }
    }

    fn least_recently_detached_key(&self) -> Option<RetainKey> {
        self.evictable()
            .min_by(|(left_key, left), (right_key, right)| {
                left.detached_order
                    .cmp(&right.detached_order)
                    .then_with(|| retain_key_cmp(left_key, right_key))
            })
            .map(|(key, _)| *key)
    }

    fn least_recently_restored_key(&self) -> Option<RetainKey> {
        self.evictable()
            .min_by(|(left_key, left), (right_key, right)| {
                left.last_restored_order
                    .cmp(&right.last_restored_order)
                    .then_with(|| left.detached_order.cmp(&right.detached_order))
                    .then_with(|| retain_key_cmp(left_key, right_key))
            })
            .map(|(key, _)| *key)
    }

    fn largest_first_key(&self) -> Option<RetainKey> {
        self.evictable()
            .max_by(|(left_key, left), (right_key, right)| {
                left.heap_bytes()
                    .cmp(&right.heap_bytes())
                    .then_with(|| retain_key_cmp(left_key, right_key))
            })
            .map(|(key, _)| *key)
    }

    fn retained_heap_bytes(&self) -> usize {
        self.groups.values().map(RetainedGroup::heap_bytes).sum()
    }
}

fn retained_subtree_matches_key(
    subtree: &DetachedSubtree,
    key: RetainKey,
    context: &'static str,
) -> bool {
    let Some(root_key) = subtree.root_key_checked() else {
        log::error!(
            "retention {context} rejected empty subtree for parent_scope={:?} key={:?}",
            key.parent_scope,
            key.key
        );
        return false;
    };
    if root_key != key.key {
        log::error!(
            "retention {context} rejected root key mismatch for parent_scope={:?}: expected {:?}, actual {:?}",
            key.parent_scope,
            key.key,
            root_key
        );
        return false;
    }
    true
}

fn retained_subtree_root_is_detached(
    subtree: &DetachedSubtree,
    key: RetainKey,
    context: &'static str,
) -> bool {
    let Some(parent_anchor) = subtree.root_parent_anchor_checked() else {
        log::error!(
            "retention {context} rejected empty subtree for parent_scope={:?} key={:?}",
            key.parent_scope,
            key.key
        );
        return false;
    };
    if parent_anchor.is_valid() {
        log::error!(
            "retention {context} rejected attached root parent {:?} for parent_scope={:?} key={:?}",
            parent_anchor,
            key.parent_scope,
            key.key
        );
        return false;
    }
    true
}

fn retained_subtree_nodes_have_lifecycle(
    subtree: &DetachedSubtree,
    key: RetainKey,
    expected: NodeLifecycle,
    context: &'static str,
) -> bool {
    if let Some((node_id, actual)) = subtree.first_node_lifecycle_mismatch(expected) {
        log::error!(
            "retention {context} rejected node lifecycle mismatch for node {node_id} parent_scope={:?} key={:?}: expected {:?}, actual {:?}",
            key.parent_scope,
            key.key,
            expected,
            actual
        );
        return false;
    }
    true
}

fn retain_key_cmp(left: &RetainKey, right: &RetainKey) -> Ordering {
    (
        left.parent_scope,
        left.key.static_key,
        left.key.explicit_key,
        left.key.ordinal,
    )
        .cmp(&(
            right.parent_scope,
            right.key.static_key,
            right.key.explicit_key,
            right.key.ordinal,
        ))
}

#[cfg(test)]
#[path = "tests/retention_tests.rs"]
mod tests;