mtb-entity-slab 0.2.4

//! Entity slab allocator implementation.
//!
//! 实体 Slab 分块分配器实现.

use crate::{
    IndexedID, PtrID, SlicePtrErr, SlicePtrRes,
    bitalloc::IBitAlloc,
    chunk::{Chunk, Unit},
    gen_index::GenIndex,
    iter::{EntityAllocConsumeIter, EntityAllocEditIter, EntityAllocReadIter},
    policy::{AllocPolicy256, IAllocPolicy},
};
use std::{
    cell::{BorrowMutError, Cell, RefCell},
    ptr::NonNull,
};

/// Possible errors when accessing allocated entities.
///
/// 分配实体访问时可能出现的错误.
#[derive(Debug)]
pub enum EntityAccessErr {
    /// The entity has been freed and cannot be accessed.
    ///
    /// 实体已被释放，无法访问。
    UseAfterFree,

    /// The index is invalid and does not reference any entity.
    ///
    /// 索引无效，未引用任何实体。
    InvalidReference,
}
impl std::fmt::Display for EntityAccessErr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let msg = match self {
            EntityAccessErr::UseAfterFree => "Entity has been freed (use-after-free).",
            EntityAccessErr::InvalidReference => "Index is invalid (no entity referenced).",
        };
        f.write_str(msg)
    }
}
impl std::error::Error for EntityAccessErr {}

/// Result type for entity access operations.
///
/// 实体访问操作的结果类型.
pub type EntityAccessRes<T = ()> = Result<T, EntityAccessErr>;

/// Entity slab allocator for entities of type `E` using allocation policy `P`.
///
/// This allocator maintains a growable vector of fixed-size `Chunk<E, P>` slabs.
/// Each `Chunk` stores up to `P::CHUNK_SIZE` units and keeps per-unit generation
/// indices plus a bitset of allocated slots. `EntityAlloc` provides fast
/// allocation/deallocation inside chunks, keeps a free-chain of chunks that have
/// available slots (`free_head`), and tracks the total number of allocated
/// entities (`num_allocated`).
///
/// Primary responsibilities and guarantees:
/// - Allocate and deallocate entities with O(1) operations inside a chunk;
///   the allocator grows by creating new `Chunk`s when needed.
/// - Maintain a `free_head` pointing to a chunk with available slots to avoid
///   scanning all chunks on each allocation.
/// - Expose helper methods that return either raw unit pointers or generation-
///   aware indexes (`GenIndex`). Callers must respect generation checks to
///   avoid use-after-free; the allocator provides checks in `unit_ptr_of_indexed`
///   and related helpers.
///
/// Concurrency and safety notes:
/// - `EntityAlloc` uses interior mutability (`RefCell` and `Cell`) and is not
///   `Sync`; external synchronization is required for concurrent access.
/// - Several `pub(crate)` functions operate on raw pointers and rely on
///   invariants (generation matching, non-invalid indexed values). Use the public
///   crate APIs where possible; internal functions are optimized for
///   performance and assume correct usage.
///
/// Field overview:
/// - `chunks`: `RefCell<Vec<Chunk<E, P>>>` — growable storage of slab chunks.
/// - `free_head`: `Cell<u32>` — index of the first chunk with free slots (or a
///   sentinel when none available).
/// - `num_allocated`: `Cell<usize>` — current count of allocated entities.
///
/// Usage tips:
/// - Use `remake_free_chain` to rebuild the free-chain after bulk mutations.
/// - Prefer `free_if` / `fully_free_if` for conditional freeing, they maintain
///   allocator invariants and update `num_allocated` correctly.
/// - Be careful when converting between raw unit pointers and `GenIndex` —
///   always handle `EntityAccessErr` and `SlicePtrErr` reported by the helpers.
///
/// 实体分配器详述（中文）:
/// - 管理一组可增长的 `Chunk<E, P>`，每个 `Chunk` 包含固定大小的单元、位图和
///   generation 信息，用于检测 use-after-free。
/// - 在 chunk 内部的分配与释放为 O(1)；当现有 chunk 全满时通过新增 chunk 增长容量。
/// - 维护 `free_head` 指向含有空闲单元的 chunk，以加速分配。
/// - `num_allocated` 跟踪当前分配数量；`remake_free_chain` 可在批量修改后重建空闲链表。
/// - 结构体使用 `RefCell`/`Cell` 实现内部可变性，默认不安全以跨线程共享，若需要
///   并发访问请自行加锁或在上层保证互斥。
///
/// 示例:
/// ```ignore
/// // 参见 crate 文档中的示例，通常通过 crate 提供的安全 API 分配/访问/释放实体。
/// ```
pub struct EntityAlloc<E, P: IAllocPolicy = AllocPolicy256> {
    pub(crate) chunks: RefCell<Vec<Chunk<E, P>>>,
    pub(crate) free_head: Cell<u32>,
    pub(crate) num_allocated: Cell<usize>,
}
impl<E, P: IAllocPolicy> Default for EntityAlloc<E, P> {
    fn default() -> Self {
        Self {
            chunks: RefCell::new(Vec::new()),
            free_head: Cell::new(0),
            num_allocated: Cell::new(0),
        }
    }
}
impl<E, P: IAllocPolicy> EntityAlloc<E, P> {
    /// Create a new `EntityAlloc` with no preallocated capacity.
    pub fn new() -> Self {
        Self::default()
    }

    /// Create a new `EntityAlloc` with preallocated capacity for `cap` entities.
    pub fn with_capacity(cap: usize) -> Self {
        let chunks = {
            let nchunks = cap.div_ceil(P::CHUNK_SIZE);
            let mut chunks = Vec::with_capacity(nchunks);
            for chunk_id in 0..(nchunks as u32) {
                let chunk = Chunk::new(chunk_id, chunk_id + 1);
                chunks.push(chunk);
            }
            RefCell::new(chunks)
        };
        Self {
            chunks,
            free_head: Cell::new(0),
            num_allocated: Cell::new(0),
        }
    }

    pub(crate) fn try_allocate_unit(&self, val: E) -> Result<*const Unit<E>, (BorrowMutError, E)> {
        let mut chunks = match self.chunks.try_borrow_mut() {
            Ok(c) => c,
            Err(e) => return Err((e, val)),
        };
        let chunk = Self::find_chunk(self.free_head.get(), &mut chunks);
        let Ok(ptr) = chunk.allocate(val) else {
            unreachable!("Internal error: Cannot allocate inside a non-full chunk");
        };
        self.num_allocated.set(self.num_allocated.get() + 1);
        if chunk.is_full() {
            self.free_head.set(chunk.free_next);
        }
        Ok(ptr)
    }
    fn find_chunk(free: u32, chunks: &mut Vec<Chunk<E, P>>) -> &mut Chunk<E, P> {
        assert_ne!(free, u32::MAX, "OOM detected");
        let free = free as usize;
        while chunks.len() <= free {
            let id = chunks.len() as u32;
            chunks.push(Chunk::new(id, id + 1));
        }
        &mut chunks[free]
    }
    pub(crate) fn next_index(&self) -> GenIndex {
        let chunks = self.chunks.borrow();
        let free = self.free_head.get() as usize;
        let Some(chunk) = chunks.get(free) else {
            return GenIndex::compose(free << P::CHUNK_SIZE_LOG2, GenIndex::GEN_1);
        };
        let unit_id = chunk.next_available().unwrap();
        chunk.indexed_id_of_unit(unit_id)
    }

    pub(crate) fn free_gen_index(&mut self, index: GenIndex) -> Option<E> {
        let (chunk_id, unit_id, gene) = P::tear(index);
        let chunks = self.chunks.get_mut();
        let chunk = chunks.get_mut(chunk_id as usize)?;
        let chunk_full = chunk.is_full();
        let val = chunk.deallocate_checked(unit_id, gene)?;
        self.num_allocated.set(self.num_allocated.get() - 1);
        if chunk_full {
            chunk.free_next = self.free_head.get();
            self.free_head.set(chunk_id);
        }
        Some(val)
    }
    pub(crate) fn free_unit_ptr(&mut self, ptr: *const Unit<E>) -> Option<E> {
        let indexed = self.index_of_unit_ptr(ptr).ok()?;
        self.free_gen_index(indexed)
    }
    pub(crate) fn check_unit_validity(&self, ptr: *const Unit<E>) -> SlicePtrRes<GenIndex> {
        let chunks = self.chunks.borrow();
        for chunk in chunks.iter() {
            let unit_id = match chunk.unit_of_ptr(ptr) {
                Ok(id) => id,
                Err(SlicePtrErr::OutOfRange) => continue,
                Err(SlicePtrErr::NotAlignedWith) => return Err(SlicePtrErr::NotAlignedWith),
            };
            let unit = chunk.unit(unit_id);
            return if unit.indexed.is_invalid() {
                Err(SlicePtrErr::OutOfRange)
            } else {
                Ok(unit.indexed)
            };
        }
        Err(SlicePtrErr::OutOfRange)
    }
    pub(crate) fn index_of_unit_ptr(&self, ptr: *const Unit<E>) -> SlicePtrRes<GenIndex> {
        if cfg!(debug_assertions) {
            self.check_unit_validity(ptr)
        } else {
            let indexed = unsafe { ptr.as_ref().ok_or(SlicePtrErr::OutOfRange)?.indexed };
            if indexed.is_invalid() {
                Err(SlicePtrErr::OutOfRange)
            } else {
                Ok(indexed)
            }
        }
    }
    pub(crate) fn unit_ptr_of_indexed(
        &self,
        indexed: GenIndex,
    ) -> EntityAccessRes<NonNull<Unit<E>>> {
        if indexed.is_invalid() {
            return Err(EntityAccessErr::InvalidReference);
        }
        let (chunk_id, unit_id, _) = P::tear(indexed);
        let chunks = self.chunks.borrow();
        let Some(chunk) = chunks.get(chunk_id as usize) else {
            return Err(EntityAccessErr::UseAfterFree);
        };
        let unit = chunk.unit(unit_id);
        if unit.indexed != indexed {
            return Err(EntityAccessErr::UseAfterFree);
        }
        Ok(NonNull::new(unit as *const _ as *mut Unit<E>).unwrap())
    }
    pub(crate) fn unit_mut_of_indexed(
        &mut self,
        indexed: GenIndex,
    ) -> EntityAccessRes<&mut Unit<E>> {
        if indexed.is_invalid() {
            return Err(EntityAccessErr::InvalidReference);
        }
        let (chunk_id, unit_id, _) = P::tear(indexed);
        let chunks = self.chunks.get_mut();
        let Some(chunk) = chunks.get_mut(chunk_id as usize) else {
            return Err(EntityAccessErr::UseAfterFree);
        };
        let unit = chunk.unit_mut(unit_id);
        if unit.indexed != indexed {
            return Err(EntityAccessErr::UseAfterFree);
        }
        Ok(unit)
    }

    /// Rebuild the free-chain of chunks with available slots.
    pub fn remake_free_chain(&mut self) {
        let chunks = self.chunks.get_mut();
        let free_chunks = {
            let mut free_chunks = Vec::with_capacity(chunks.len());
            for chunk in chunks.iter_mut() {
                if !chunk.is_full() {
                    free_chunks.push(chunk.chunk_id);
                }
            }
            free_chunks
        };
        for i in 1..free_chunks.len() {
            let chunk_id = free_chunks[i - 1];
            let next_chunk_id = free_chunks[i];
            let chunk = &mut chunks[chunk_id as usize];
            chunk.free_next = next_chunk_id;
        }
        if let Some(&chunk_id) = free_chunks.first() {
            self.free_head.set(chunk_id);
        } else {
            self.free_head.set(chunks.len() as u32);
        }
    }
}

impl<E, P: IAllocPolicy> EntityAlloc<E, P> {
    /// Free entities that satisfy the given predicate.
    pub fn free_if(&mut self, pred: impl FnMut(&E, PtrID<E, P>, IndexedID<E, P>) -> bool) {
        self.fully_free_if(pred, drop);
    }

    /// Free entities that satisfy the given predicate, consuming them with the provided closure.
    ///
    /// This method automatically chooses between dense and sparse freeing strategies
    /// based on the chunk's allocation density for efficiency.
    pub fn fully_free_if(
        &mut self,
        mut should_free: impl FnMut(&E, PtrID<E, P>, IndexedID<E, P>) -> bool,
        mut consume: impl FnMut(E),
    ) {
        let chunks = self.chunks.get_mut();
        let mut num_last = self.num_allocated.get();
        for chunk in chunks.iter_mut() {
            if chunk.is_empty() {
                continue;
            }
            if num_last == 0 {
                break;
            }
            let chunk_full_previously = chunk.is_full();
            let chunk_num_allocated = chunk.num_allocated as usize;
            let num_freed = if chunk_num_allocated * 3 >= P::CHUNK_SIZE {
                Self::densely_free_chunk(&mut should_free, &mut consume, chunk)
            } else {
                Self::sparsely_free_chunk(&mut should_free, &mut consume, chunk)
            };

            num_last -= chunk_num_allocated;
            // 维护全局已分配计数
            if num_freed > 0 {
                self.num_allocated
                    .set(self.num_allocated.get().saturating_sub(num_freed));
            }
            // 如果之前是满的，现在有空位了，加入空闲链表
            if chunk_full_previously && !chunk.is_full() {
                chunk.free_next = self.free_head.get();
                self.free_head.set(chunk.chunk_id);
            }
        }
    }

    fn densely_free_chunk(
        should_free: &mut impl FnMut(&E, PtrID<E, P>, IndexedID<E, P>) -> bool,
        consume: &mut impl FnMut(E),
        chunk: &mut Chunk<E, P>,
    ) -> usize {
        let mut num_freed = 0;
        let mut count_down = chunk.num_allocated;
        for unit_id in 0..(P::CHUNK_SIZE as u16) {
            if count_down == 0 {
                break;
            }
            let index = chunk.indexed_id_of_unit(unit_id);
            if !chunk.allocated.is_allocated(unit_id) {
                continue;
            }
            count_down -= 1;

            let unit = chunk.unit_ptr_mut(unit_id);
            let ptr = PtrID::from(unit);
            let elem = unsafe { &*unit }.as_init_ref().unwrap();
            if !should_free(elem, ptr, IndexedID::from(index)) {
                continue;
            }
            let Some(e) = chunk.deallocate_unchecked(unit_id) else {
                unreachable!("Deallocation failed unexpectedly");
            };
            consume(e);
            num_freed += 1;
        }
        num_freed
    }
    fn sparsely_free_chunk(
        should_free: &mut impl FnMut(&E, PtrID<E, P>, IndexedID<E, P>) -> bool,
        consume: &mut impl FnMut(E),
        chunk: &mut Chunk<E, P>,
    ) -> usize {
        // 逐个枚举当前已分配位，避免全表扫描
        let mut num_freed = 0usize;
        let mut next = chunk.allocated.next_allocated(0);
        while let Some(unit_id) = next {
            let index = chunk.indexed_id_of_unit(unit_id);
            let unit = chunk.unit_ptr_mut(unit_id);
            let ptr = PtrID::from(unit);
            let elem = unsafe { &*unit }.as_init_ref().unwrap();
            if should_free(elem, ptr, IndexedID::from(index)) {
                if let Some(e) = chunk.deallocate_unchecked(unit_id) {
                    consume(e);
                    num_freed += 1;
                } else {
                    unreachable!("Deallocation failed unexpectedly");
                }
                // 删除了当前位后，从 unit_id+1 继续
                let next_start = unit_id.saturating_add(1);
                next = chunk.allocated.next_allocated(next_start);
            } else {
                // 未删除则从下一个位置开始
                let next_start = unit_id.saturating_add(1);
                next = chunk.allocated.next_allocated(next_start);
            }
        }
        num_freed
    }

    /// Get the number of currently allocated entities.
    ///
    /// Note that this is not the length of the internal chunk vector or the
    /// "allocated" sections in chunks, but the actual count of live entities.
    #[inline]
    pub fn len(&self) -> usize {
        self.num_allocated.get()
    }

    /// Get the current capacity in terms of number of entities.
    #[inline]
    pub fn capacity(&self) -> usize {
        self.chunks.borrow().len() * P::CHUNK_SIZE
    }

    /// Check if there are no allocated entities.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Clear all allocated entities, resetting the allocator to an empty state.
    pub fn clear(&mut self) {
        self.chunks.get_mut().clear();
        self.free_head.set(0);
        self.num_allocated.set(0);
    }

    /// Get an iterator over allocated entities for read-only access.
    ///
    /// The iterator yields tuples of `(IndexedID, PtrID, &E)`.
    #[inline]
    pub fn iter(&self) -> EntityAllocReadIter<'_, E, P> {
        EntityAllocReadIter::new(self)
    }
    /// Get an iterator over allocated entities for mutable access.
    ///
    /// The iterator yields tuples of `(IndexedID, PtrID, &mut E)`.
    #[inline]
    pub fn iter_mut(&mut self) -> EntityAllocEditIter<'_, E, P> {
        EntityAllocEditIter::new(self)
    }
}

impl<'alloc, E, P: IAllocPolicy> IntoIterator for &'alloc EntityAlloc<E, P> {
    type Item = (IndexedID<E, P>, PtrID<E, P>, &'alloc E);
    type IntoIter = EntityAllocReadIter<'alloc, E, P>;

    fn into_iter(self) -> Self::IntoIter {
        self.iter()
    }
}
impl<'alloc, E, P: IAllocPolicy> IntoIterator for &'alloc mut EntityAlloc<E, P> {
    type Item = (IndexedID<E, P>, PtrID<E, P>, &'alloc mut E);
    type IntoIter = EntityAllocEditIter<'alloc, E, P>;

    fn into_iter(self) -> Self::IntoIter {
        self.iter_mut()
    }
}
impl<E, P: IAllocPolicy> IntoIterator for EntityAlloc<E, P> {
    type Item = (GenIndex, E);
    type IntoIter = EntityAllocConsumeIter<E, P>;

    fn into_iter(self) -> Self::IntoIter {
        EntityAllocConsumeIter::new(self)
    }
}