mtb_entity_slab/container/

order_cached_list.rs

//! Order-cached linked-list container for entity allocation. Maintains an optional
//! “order cache” on each node.
//!
//! ## Order representation 顺序表示
//!
//! The cache can be treated as either an exact position index (`OrderRepr::Exact`)
//!  or a relative ordering hint (`OrderRepr::Relative`). The cache is only trusted
//!  while the list considers it valid; once invalid, it will be rebuilt on demand.
//!
//! - **Exact mode:** the cache stays valid only when inserting at the end or
//!   removing the last element; any other insertion/removal invalidates it.
//!
//! - **Relative mode:** insertions/removals do not invalidate the cache, but the
//!   cached values are not required to be a contiguous index.
//!
//! The representation is bound to the node type via `IOrderCachedListNodeID::ORDER_REPR`.
//! It is a compile-time constant: you cannot change it at runtime, and you cannot create
//! two `OrderCachedList`s with different `OrderRepr` for the same node type.

use crate::{
    EntityListError, EntityListIter, EntityListNodeHead, EntityListRange, EntityListRangeDebug,
    EntityListRes, IBasicEntityListID, IDBoundAlloc, container::list_base::NodeOps,
};
use std::cell::Cell;

/// Representation of order cache in an order-cached list. This will
/// determine when the order cache is invalidated.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OrderRepr {
    /// The order exactly represents the position in the list. e.g.
    ///
    /// ```text
    /// List: (Head = 0) <-> A(1) <-> B(2) <-> C(3) <-> (Tail = usize::MAX)
    /// ```
    ///
    /// Invalidates the order cache on any insertion or removal except
    /// when inserting at the end or removing the last element.
    Exact,

    /// The order represents a relative position in the list, but not
    /// necessarily the exact index. e.g.
    ///
    /// ```text
    /// List: (Head = 0) <-> A(1) <-> B(3) <-> C(4) <-> (Tail = usize::MAX)
    /// ```
    Relative,
}

/// An entity pointer list node that supports order-cached lists.
pub trait IOrderCachedListNodeID: IBasicEntityListID {
    /// Representation type of the order cache.
    const ORDER_REPR: OrderRepr;

    /// Load the order from the entity object.
    fn obj_load_order(obj: &Self::ObjectT) -> usize;

    /// Store the order into the entity object.
    fn obj_store_order(obj: &Self::ObjectT, order: usize);
}

/// Order-cached doubly-linked list.
pub struct OrderCachedList<I: IOrderCachedListNodeID> {
    /// Head sentinel node ID. Order is always `0` and do not represent a valid item.
    pub head: I,
    /// Tail sentinel node ID. Order is always `usize::MAX` and does not represent a valid item.
    pub tail: I,
    /// Combined length and order validity flag.
    len_flag: Cell<usize>,
}

impl<I: IOrderCachedListNodeID> OrderCachedList<I> {
    const VALID_BITMASK: usize = 1 << (usize::BITS - 1);
    const LEN_MASK: usize = !Self::VALID_BITMASK;

    /// Order value of the head sentinel.
    pub const HEAD_ORDER: usize = 0;
    /// Order value of the tail sentinel.
    pub const TAIL_ORDER: usize = usize::MAX;
    /// Maximum length supported by the list: `2^(usize::BITS - 1) - 1`.
    pub const MAX_LEN: usize = Self::LEN_MASK;

    /// Create a new empty order-cached list with given head and tail sentinels.
    pub fn new(alloc: &IDBoundAlloc<I>) -> Self {
        let head = I::new_sentinel(alloc);
        let tail = I::new_sentinel(alloc);
        head.set_next_id(alloc, Some(tail));
        tail.set_prev_id(alloc, Some(head));
        I::obj_store_order(head.deref_alloc(alloc), Self::HEAD_ORDER);
        I::obj_store_order(tail.deref_alloc(alloc), Self::TAIL_ORDER);
        Self {
            head,
            tail,
            len_flag: Cell::new(0),
        }
    }

    /// Get the length of the list.
    pub fn len(&self) -> usize {
        self.len_flag.get() & Self::LEN_MASK
    }
    /// Check if the list is empty.
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }
    /// try to get the front node ID, or None if the list is empty
    #[inline]
    pub fn get_front_id(&self, alloc: &IDBoundAlloc<I>) -> Option<I> {
        let next = self.head.get_next_id(alloc)?;
        if next == self.tail { None } else { Some(next) }
    }
    /// try to get the back node ID, or None if the list is empty
    #[inline]
    pub fn get_back_id(&self, alloc: &IDBoundAlloc<I>) -> Option<I> {
        let prev = self.tail.get_prev_id(alloc)?;
        if prev == self.head { None } else { Some(prev) }
    }

    /// Check whether the order cache is valid.
    pub fn order_valid(&self) -> bool {
        (self.len_flag.get() & Self::VALID_BITMASK) != 0
    }

    /// Invalidate the order cache.
    pub fn invalidate_order(&self) {
        self.len_flag.set(self.len_flag.get() & Self::LEN_MASK);
    }

    /// Rebuild the order cache regardless of its current validity.
    pub fn rebuild_order(&self, alloc: &IDBoundAlloc<I>) {
        let mut order = 1;
        let mut curr_opt = self.head.get_next_id(alloc);
        while let Some(curr) = curr_opt {
            if curr == self.tail {
                break;
            }
            I::obj_store_order(curr.deref_alloc(alloc), order);
            order += 1;
            curr_opt = curr.get_next_id(alloc);
        }
        // Now enable the valid flag
        self.len_flag.set(self.len_flag.get() | Self::VALID_BITMASK);
    }

    /// Get the order of a node. If the order cache is invalid, it will be rebuilt first.
    pub fn get_node_order(&self, alloc: &IDBoundAlloc<I>, node: I) -> usize {
        if !self.order_valid() {
            self.rebuild_order(alloc);
        }
        I::obj_load_order(node.deref_alloc(alloc))
    }

    /// Check whether node `a` comes before node `b` in the list.
    pub fn node_comes_before(&self, alloc: &IDBoundAlloc<I>, a: I, b: I) -> bool {
        let order_a = self.get_node_order(alloc, a);
        let order_b = self.get_node_order(alloc, b);
        order_a < order_b
    }
    /// Check whether node `a` comes after node `b` in the list.
    pub fn node_comes_after(&self, alloc: &IDBoundAlloc<I>, a: I, b: I) -> bool {
        let order_a = self.get_node_order(alloc, a);
        let order_b = self.get_node_order(alloc, b);
        order_a > order_b
    }

    /// Get the range covering all nodes in the order-cached list (excluding sentinels).
    pub fn get_range(&self, alloc: &IDBoundAlloc<I>) -> EntityListRange<I> {
        let front = self
            .head
            .get_next_id(alloc)
            .expect("Broken order-cached list detected in get_range()");
        EntityListRange {
            start: front,
            end: Some(self.tail),
        }
    }
    /// Get the range covering all nodes including sentinels.
    pub fn get_range_with_sentinels(&self) -> EntityListRange<I> {
        EntityListRange {
            start: self.head,
            end: None,
        }
    }
    /// Create a debug view for this range.
    pub fn debug<'a>(&self, alloc: &'a IDBoundAlloc<I>) -> EntityListRangeDebug<'a, I> {
        self.get_range(alloc).debug(alloc)
    }
    /// Create an iterator over the nodes in the order-cached list.
    pub fn iter<'a>(&self, alloc: &'a IDBoundAlloc<I>) -> EntityListIter<'a, I> {
        self.get_range(alloc).iter(alloc)
    }

    /// Create an iterator over all nodes including sentinels.
    pub fn iter_with_sentinels<'a>(&self, alloc: &'a IDBoundAlloc<I>) -> EntityListIter<'a, I> {
        self.get_range_with_sentinels().iter(alloc)
    }

    /// Apply a function to all nodes in the list (including sentinels).
    pub fn forall_with_sentinel(
        &self,
        alloc: &IDBoundAlloc<I>,
        mut f: impl FnMut(I, &I::ObjectT) -> EntityListRes<I>,
    ) -> EntityListRes<I, ()> {
        let mut curr = self.head;
        loop {
            f(curr, curr.deref_alloc(alloc))?;
            if curr == self.tail {
                break;
            }
            let next = curr.get_next_id(alloc).ok_or(EntityListError::ListBroken)?;
            curr = next;
        }
        Ok(())
    }

    /// Unplug a node from the list.
    ///
    /// ### Signal invocation
    ///
    /// `node_unplug` will invoke `on_unplug` on `node` before modifying any links.
    /// If the signal returns an error, the operation is aborted without modifying any links.
    ///
    /// ### Order maintaining
    ///
    /// If the order representation is `Exact`, the order cache will be invalidated unless the
    /// unplugged node is the last node in the list. Otherwise, the order cache validity is left
    /// unchanged.
    pub fn node_unplug(&self, node: I, alloc: &IDBoundAlloc<I>) -> EntityListRes<I> {
        if node == self.head || node == self.tail {
            return Err(EntityListError::NodeIsSentinel);
        }
        node.on_unplug(alloc)?;
        let node_obj = node.deref_alloc(alloc);

        let Some(prev) = I::obj_get_prev_id(node_obj) else {
            return Err(EntityListError::ListBroken);
        };
        let Some(next) = I::obj_get_next_id(node_obj) else {
            return Err(EntityListError::ListBroken);
        };

        prev.set_next_id(alloc, Some(next));
        next.set_prev_id(alloc, Some(prev));
        I::obj_store_head(node_obj, EntityListNodeHead::none());
        I::obj_store_order(node_obj, 0usize);

        // Update order cache.
        //
        // When `Exact`: removing the last node keeps the cache valid; any other removal clears it.
        // When `Relative`: just decrease the length part. `-1` does not affect the valid bit.
        if I::ORDER_REPR == OrderRepr::Exact && next != self.tail {
            self.len_flag.set(self.len() - 1);
        } else {
            self.len_flag.set(self.len_flag.get() - 1);
        }
        Ok(())
    }

    /// Add a new node after an existing node linked to `this` list.
    ///
    /// If the current node is linked to another list, then the operation is undefined.
    /// You'll get a broken list without reporting an error.
    ///
    /// ### Signal invocation
    ///
    /// `node_add_next` will invoke `on_push_next` on `node` before modifying any links.
    /// If the signal returns an error, the operation is aborted without modifying any links.
    ///
    /// ### Order maintaining
    ///
    /// The order cache will be invalidated unless the new node is added at the end of the list.
    pub fn node_add_next(&self, node: I, new_node: I, alloc: &IDBoundAlloc<I>) -> EntityListRes<I> {
        if node == new_node {
            return Err(EntityListError::RepeatedNode);
        }
        if node == self.tail {
            return Err(EntityListError::NodeIsSentinel);
        }
        node.on_push_next(new_node, alloc)?;

        let node_obj = node.deref_alloc(alloc);
        let new_node_obj = new_node.deref_alloc(alloc);

        debug_assert!(
            I::obj_is_attached(node_obj),
            "Cannot add next to a detached node"
        );
        debug_assert!(
            !I::obj_is_attached(new_node_obj),
            "Cannot add a node that is already attached"
        );

        let Some(old_next) = I::obj_get_next_id(node_obj) else {
            return Err(EntityListError::ListBroken);
        };
        I::obj_set_next_id(node_obj, Some(new_node));
        I::obj_store_head(new_node_obj, EntityListNodeHead::from_id(node, old_next));
        old_next.set_prev_id(alloc, Some(new_node));

        debug_assert!(self.len() < Self::MAX_LEN, "List length overflow");
        if old_next == self.tail && self.order_valid() {
            I::obj_store_order(new_node_obj, I::obj_load_order(node_obj) + 1);
            self.len_flag.set(self.len_flag.get() + 1);
        } else {
            // Invalidate order cache (use `len()` then `set` to clear valid bit)
            self.len_flag.set(self.len() + 1);
        }
        Ok(())
    }

    /// Add a new node before an existing node linked to `this` list.
    ///
    /// If the current node is linked to another list, then the operation is undefined.
    /// You'll get a broken list without reporting an error.
    ///
    /// ### Signal invocation
    ///
    /// `node_add_prev` will invoke `on_push_prev` on `node` before modifying any links.
    /// If the signal returns an error, the operation is aborted without modifying any links.
    ///
    /// ### Order maintaining
    ///
    /// The order cache will be invalidated unless the `node` is the tail sentinel.
    pub fn node_add_prev(&self, node: I, new_node: I, alloc: &IDBoundAlloc<I>) -> EntityListRes<I> {
        if node == new_node {
            return Err(EntityListError::RepeatedNode);
        }
        if node == self.head {
            return Err(EntityListError::NodeIsSentinel);
        }
        node.on_push_prev(new_node, alloc)?;

        let node_obj = node.deref_alloc(alloc);
        let new_node_obj = new_node.deref_alloc(alloc);

        debug_assert!(
            I::obj_is_attached(node_obj),
            "Cannot add prev to a detached node"
        );
        debug_assert!(
            !I::obj_is_attached(new_node_obj),
            "Cannot add a node that is already attached"
        );
        debug_assert!(self.len() < Self::MAX_LEN, "List length overflow");

        let Some(old_prev) = I::obj_get_prev_id(node_obj) else {
            return Err(EntityListError::ListBroken);
        };
        let old_prev_obj = old_prev.deref_alloc(alloc);

        NodeOps::obj_set_prev_id(node_obj, Some(new_node));
        I::obj_store_head(new_node_obj, EntityListNodeHead::from_id(old_prev, node));
        NodeOps::obj_set_next_id(old_prev_obj, Some(new_node));

        if node == self.tail && self.order_valid() {
            let old_order = I::obj_load_order(old_prev_obj);
            I::obj_store_order(new_node_obj, old_order + 1);
            self.len_flag.set(self.len_flag.get() + 1);
        } else {
            // Invalidate order cache (use `len()` then `set` to clear valid bit)
            self.len_flag.set(self.len() + 1);
        }
        Ok(())
    }
    /// Push a new node to the back of the list. Will not invalidate order cache if possible.
    #[inline]
    pub fn push_back_id(&self, new_node: I, alloc: &IDBoundAlloc<I>) -> EntityListRes<I> {
        self.node_add_prev(self.tail, new_node, alloc)
    }

    /// Push a new node to the front of the list. Will invalidate order cache.
    #[inline]
    pub fn push_front_id(&self, new_node: I, alloc: &IDBoundAlloc<I>) -> EntityListRes<I> {
        self.node_add_next(self.head, new_node, alloc)
    }

    /// Pop a node from the back of the list. Will not invalidate order cache if possible.
    pub fn pop_back(&self, alloc: &IDBoundAlloc<I>) -> EntityListRes<I, I> {
        let back_id = self.get_back_id(alloc).ok_or(EntityListError::EmptyList)?;
        self.node_unplug(back_id, alloc)?;
        Ok(back_id)
    }

    /// Pop a node from the front of the list. Will invalidate order cache if repr is `Exact`.
    pub fn pop_front(&self, alloc: &IDBoundAlloc<I>) -> EntityListRes<I, I> {
        let front_id = self.get_front_id(alloc).ok_or(EntityListError::EmptyList)?;
        self.node_unplug(front_id, alloc)?;
        Ok(front_id)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{IEntityAllocID, IPoliciedID, IndexedID, entity_id};

    #[derive(Debug, Clone)]
    #[entity_id(ExactInstID, policy = 256, backend = index)]
    struct ExactInst {
        head: Cell<EntityListNodeHead<ExactInstID>>,
        order: Cell<usize>,
        value: usize,
    }

    type ExactInstAlloc = IDBoundAlloc<ExactInstID>;

    impl ExactInst {
        fn new(value: usize) -> Self {
            Self {
                head: Cell::new(EntityListNodeHead::none()),
                order: Cell::new(0),
                value,
            }
        }

        fn new_id(alloc: &IDBoundAlloc<ExactInstID>, value: usize) -> ExactInstID {
            ExactInstID::from_backend(IndexedID::allocate_from(alloc, Self::new(value)))
        }
    }

    impl IBasicEntityListID for ExactInstID {
        fn obj_load_head(obj: &Self::ObjectT) -> EntityListNodeHead<Self> {
            obj.head.get()
        }
        fn obj_store_head(obj: &Self::ObjectT, head: EntityListNodeHead<Self>) {
            obj.head.set(head);
        }

        fn obj_is_sentinel(obj: &Self::ObjectT) -> bool {
            obj.value == usize::MAX
        }
        fn new_sentinel_obj() -> Self::ObjectT {
            ExactInst {
                head: Cell::new(EntityListNodeHead::none()),
                order: Cell::new(0),
                value: usize::MAX,
            }
        }

        fn on_push_prev(self, _: Self, _: &IDBoundAlloc<ExactInstID>) -> EntityListRes<Self> {
            Ok(())
        }
        fn on_push_next(self, _: Self, _: &IDBoundAlloc<ExactInstID>) -> EntityListRes<Self> {
            Ok(())
        }
        fn on_unplug(self, _: &IDBoundAlloc<ExactInstID>) -> EntityListRes<Self> {
            Ok(())
        }
    }

    impl IOrderCachedListNodeID for ExactInstID {
        const ORDER_REPR: OrderRepr = OrderRepr::Exact;

        fn obj_load_order(obj: &Self::ObjectT) -> usize {
            obj.order.get()
        }
        fn obj_store_order(obj: &Self::ObjectT, order: usize) {
            obj.order.set(order);
        }
    }

    #[derive(Debug, Clone)]
    #[entity_id(RelativeInstID, policy = 256, backend = index)]
    struct RelativeInst {
        head: Cell<EntityListNodeHead<RelativeInstID>>,
        order: Cell<usize>,
        value: usize,
    }

    type RelativeInstAlloc = IDBoundAlloc<RelativeInstID>;

    impl RelativeInst {
        fn new(value: usize) -> Self {
            Self {
                head: Cell::new(EntityListNodeHead::none()),
                order: Cell::new(0),
                value,
            }
        }

        fn new_id(alloc: &IDBoundAlloc<RelativeInstID>, value: usize) -> RelativeInstID {
            type RelativeBackID = <RelativeInstID as IPoliciedID>::BackID;
            RelativeInstID::from_backend(RelativeBackID::allocate_from(alloc, Self::new(value)))
        }
    }

    impl IBasicEntityListID for RelativeInstID {
        fn obj_load_head(obj: &Self::ObjectT) -> EntityListNodeHead<Self> {
            obj.head.get()
        }
        fn obj_store_head(obj: &Self::ObjectT, head: EntityListNodeHead<Self>) {
            obj.head.set(head);
        }

        fn obj_is_sentinel(obj: &Self::ObjectT) -> bool {
            obj.value == usize::MAX
        }
        fn new_sentinel_obj() -> Self::ObjectT {
            RelativeInst {
                head: Cell::new(EntityListNodeHead::none()),
                order: Cell::new(0),
                value: usize::MAX,
            }
        }

        fn on_push_prev(self, _: Self, _: &RelativeInstAlloc) -> EntityListRes<Self> {
            Ok(())
        }
        fn on_push_next(self, _: Self, _: &RelativeInstAlloc) -> EntityListRes<Self> {
            Ok(())
        }
        fn on_unplug(self, _: &RelativeInstAlloc) -> EntityListRes<Self> {
            Ok(())
        }
    }

    impl IOrderCachedListNodeID for RelativeInstID {
        const ORDER_REPR: OrderRepr = OrderRepr::Relative;

        fn obj_load_order(obj: &Self::ObjectT) -> usize {
            obj.order.get()
        }
        fn obj_store_order(obj: &Self::ObjectT, order: usize) {
            obj.order.set(order);
        }
    }

    fn assert_contiguous_orders<I: IOrderCachedListNodeID>(
        list: &OrderCachedList<I>,
        alloc: &IDBoundAlloc<I>,
        expected_len: usize,
    ) {
        let mut expected = 1;
        let mut count = 0;
        for (id, _) in list.iter(alloc) {
            let order = list.get_node_order(alloc, id);
            assert_eq!(order, expected, "Node order mismatch");
            expected += 1;
            count += 1;
        }
        assert_eq!(count, expected_len, "List length mismatch");
    }

    fn assert_strictly_increasing_orders<I: IOrderCachedListNodeID>(
        list: &OrderCachedList<I>,
        alloc: &IDBoundAlloc<I>,
        expected_len: usize,
    ) {
        let mut last: Option<usize> = None;
        let mut count = 0;
        for (id, _) in list.iter(alloc) {
            let order = list.get_node_order(alloc, id);
            if let Some(prev) = last {
                assert!(order > prev, "Node order is not strictly increasing");
            }
            last = Some(order);
            count += 1;
        }
        assert_eq!(count, expected_len, "List length mismatch");
    }

    #[test]
    fn test_order_cached_list_exact_basic() {
        let alloc = ExactInstAlloc::new();
        let list = OrderCachedList::<ExactInstID>::new(&alloc);

        // Empty list behavior and error on popping from empty.
        assert!(list.is_empty());
        assert!(matches!(
            list.pop_back(&alloc),
            Err(EntityListError::EmptyList)
        ));

        // Push a few nodes and verify length is tracked.
        for i in 0..5 {
            let inst = ExactInst::new_id(&alloc, i);
            list.push_back_id(inst, &alloc).unwrap();
        }
        assert_eq!(list.len(), 5);

        // Rebuild cache and verify contiguous orders.
        list.rebuild_order(&alloc);
        assert!(list.order_valid());
        assert_contiguous_orders(&list, &alloc, 5);

        // Removing last keeps cache valid in Exact mode.
        list.pop_back(&alloc).unwrap();
        assert!(list.order_valid());
        assert_eq!(list.len(), 4);
        assert_contiguous_orders(&list, &alloc, 4);

        // Removing a middle node invalidates cache in Exact mode.
        let ids: Vec<_> = list.iter(&alloc).map(|(id, _)| id).collect();
        let mid = ids[1];
        list.node_unplug(mid, &alloc).unwrap();
        assert!(!list.order_valid());
        list.rebuild_order(&alloc);
        assert_contiguous_orders(&list, &alloc, 3);

        // Inserting not at end invalidates cache in Exact mode.
        let front = list.get_front_id(&alloc).unwrap();
        let new_id = ExactInst::new_id(&alloc, 99);
        list.node_add_next(front, new_id, &alloc).unwrap();
        assert!(!list.order_valid());
        list.rebuild_order(&alloc);
        assert_contiguous_orders(&list, &alloc, 4);

        // Sentinel misuse should return an error.
        let bad_new = ExactInst::new_id(&alloc, 100);
        assert!(matches!(
            list.node_add_next(list.tail, bad_new, &alloc),
            Err(EntityListError::NodeIsSentinel)
        ));
    }

    #[test]
    fn test_order_cached_list_relative_semantics() {
        let alloc = RelativeInstAlloc::new();
        let list = OrderCachedList::<RelativeInstID>::new(&alloc);

        // Initial push and rebuild gives contiguous orders.
        for i in 0..4 {
            let inst = RelativeInst::new_id(&alloc, i);
            list.push_back_id(inst, &alloc).unwrap();
        }
        list.rebuild_order(&alloc);
        assert!(list.order_valid());
        assert_contiguous_orders(&list, &alloc, 4);

        // Append keeps cache valid; orders remain strictly increasing.
        let new_id = RelativeInst::new_id(&alloc, 10);
        list.push_back_id(new_id, &alloc).unwrap();
        assert!(list.order_valid());
        assert_contiguous_orders(&list, &alloc, 5);

        // Middle removal keeps cache valid in Relative mode.
        let ids: Vec<_> = list.iter(&alloc).map(|(id, _)| id).collect();
        let mid = ids[2];
        list.node_unplug(mid, &alloc).unwrap();
        assert!(list.order_valid());
        assert_strictly_increasing_orders(&list, &alloc, 4);

        // Inserting not at end invalidates cache; rebuild restores contiguous orders.
        let front = list.get_front_id(&alloc).unwrap();
        let insert_id = RelativeInst::new_id(&alloc, 11);
        list.node_add_next(front, insert_id, &alloc).unwrap();
        assert!(!list.order_valid());
        list.rebuild_order(&alloc);
        assert_contiguous_orders(&list, &alloc, 5);
    }
}