mtb_entity_slab/

alloc.rs

use crate::{
    EntityAllocConsumeIter, EntityAllocEditIter, EntityAllocReadIter, IEntityAllocPolicy,
    IndexedID, PtrID,
    chunk::{Chunk, NULL_INDEXED_ID},
    iter::EntitySlotIter,
};
use std::{
    cell::{BorrowMutError, Cell, RefCell},
    marker::PhantomData,
    ptr::NonNull,
};

pub trait IEntityAllocatable: Sized {
    /// Every allocatable item should define an allocation policy type.
    type AllocatePolicyT: IEntityAllocPolicy<Self>;

    /// Pointer ID type used by this allocatable item.
    /// Normally you can put `PtrID<Self>` here. While you have special needs,
    /// you can define your own pointer ID structure and hide the `PtrID` inside ir.
    type PtrID: Copy + Eq + From<PtrID<Self>> + Into<PtrID<Self>>;

    const CHUNK_SIZE: usize = Self::AllocatePolicyT::CHUNK_SIZE;
    const BITSET_LEN: usize = Self::AllocatePolicyT::BITSET_LEN;

    fn compose_indexed_id(chunk_id: u32, unit_id: u16) -> usize {
        Self::AllocatePolicyT::compose_indexed_id(chunk_id, unit_id)
    }
    fn split_indexed_id(indexed_id: usize) -> (u32, u16) {
        Self::AllocatePolicyT::split_indexed_id(indexed_id)
    }
    fn unit_of_indexed_id(indexed_id: usize) -> u16 {
        Self::AllocatePolicyT::unit_of_indexed_id(indexed_id)
    }
    fn chunk_of_indexed_id(indexed_id: usize) -> u32 {
        Self::AllocatePolicyT::chunk_of_indexed_id(indexed_id)
    }

    fn ptr_of_id(id: Self::PtrID) -> PtrID<Self> {
        id.into()
    }
    fn id_of_ptr(ptr: PtrID<Self>) -> Self::PtrID {
        Self::PtrID::from(ptr)
    }
}

pub struct EntityAlloc<E: IEntityAllocatable> {
    pub(crate) chunks: RefCell<Vec<Chunk<E, E::AllocatePolicyT>>>,
    pub(crate) free_head: Cell<u32>,
    pub(crate) num_allocated: Cell<usize>,
}

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

    #[inline]
    fn into_iter(self) -> Self::IntoIter {
        EntityAllocReadIter::new(self)
    }
}
impl<'alloc, E: IEntityAllocatable> IntoIterator for &'alloc mut EntityAlloc<E> {
    type Item = (usize, E::PtrID, &'alloc mut E);
    type IntoIter = EntityAllocEditIter<'alloc, E>;

    #[inline]
    fn into_iter(self) -> Self::IntoIter {
        EntityAllocEditIter::new(self)
    }
}
impl<E: IEntityAllocatable> IntoIterator for EntityAlloc<E> {
    type Item = (usize, E);
    type IntoIter = EntityAllocConsumeIter<E>;

    #[inline]
    fn into_iter(self) -> Self::IntoIter {
        EntityAllocConsumeIter::new(self)
    }
}

impl<E: IEntityAllocatable> EntityAlloc<E> {
    pub fn new() -> Self {
        Self {
            chunks: RefCell::new(Vec::new()),
            free_head: Cell::new(0),
            num_allocated: Cell::new(0),
        }
    }

    pub fn with_capacity(cap: usize) -> Self {
        let cap = cap.next_multiple_of(E::AllocatePolicyT::CHUNK_SIZE);
        let chunks = {
            let nchunks = cap / E::AllocatePolicyT::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 fn allocate(&self, val: E) -> PtrID<E> {
        let Ok(ptr) = self.try_allocate(val) else {
            panic!(
                "Cannot borrow EntityAlloc {self:p} mutably for allocation. Are you traversing it now?"
            );
        };
        ptr
    }
    pub fn try_allocate(&self, val: E) -> Result<PtrID<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 {
            panic!("Allocation error!");
        };
        self.num_allocated.set(self.num_allocated.get() + 1);
        if chunk.is_full() {
            self.free_head.set(chunk.free_next);
        }
        Ok(PtrID(NonNull::new(ptr as *mut _).unwrap()))
    }
    /// 下一次分配时, 返回下一个可用的索引
    pub fn next_indexed(&self) -> IndexedID<E> {
        let chunks = self.chunks.borrow();
        let free_head = self.free_head.get() as usize;
        let Some(chunk) = chunks.get(free_head) else {
            return IndexedID(
                free_head << E::AllocatePolicyT::CHUNK_SIZE_LOG2,
                PhantomData,
            );
        };
        let unit_id = chunk.next_available().unwrap();
        IndexedID(chunk.indexed_id_of_unit(unit_id), PhantomData)
    }
    fn find_chunk(
        free: u32,
        chunks: &mut Vec<Chunk<E, E::AllocatePolicyT>>,
    ) -> &mut Chunk<E, E::AllocatePolicyT> {
        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 fn free_indexed(&mut self, indexed: usize) -> Option<E> {
        let chunks = self.chunks.get_mut();
        let chunk_id = E::chunk_of_indexed_id(indexed);
        let unit_id = E::unit_of_indexed_id(indexed);
        let chunk = chunks.get_mut(chunk_id as usize)?;
        let chunk_full = chunk.is_full();
        let e = chunk.deallocate_unit_id(unit_id)?;
        // 如果之前是满的，现在有空位了，加入空闲链表
        if chunk_full {
            chunk.free_next = self.free_head.get();
            self.free_head.set(chunk_id);
        }
        self.num_allocated.set(self.num_allocated.get() - 1);
        Some(e)
    }
    pub fn free_ptr(&mut self, ptr: PtrID<E>) -> Option<E> {
        if cfg!(debug_assertions) && !ptr.check_validity(self.chunks.get_mut()) {
            return None;
        }
        let indexed = unsafe { ptr.direct_get_indexed() };
        if indexed == NULL_INDEXED_ID {
            return None;
        }
        self.free_indexed(indexed)
    }
    pub fn free_if(&mut self, pred: impl Fn(&E, PtrID<E>, IndexedID<E>) -> bool) {
        self.fully_free_if(pred, drop);
    }
    pub fn fully_free_if(
        &mut self,
        pred: impl Fn(&E, PtrID<E>, IndexedID<E>) -> bool,
        mut consume: impl FnMut(E),
    ) {
        let mut num_allocated_dec = self.num_allocated.get();
        let mut num_allocated_inc = 0;
        for slot in EntitySlotIter::new(self) {
            let (ptr, indexed_id) = unsafe { slot.copy_release_all() };
            if pred(slot.get().unwrap(), ptr, indexed_id) {
                let elem = unsafe { slot.raw_free() };
                let Some(elem) = elem else {
                    continue;
                };
                consume(elem);
                num_allocated_dec -= 1;
            } else {
                num_allocated_inc += 1;
            }
        }
        assert_eq!(
            num_allocated_dec, num_allocated_inc,
            "Length maintenance error in fully_free_if"
        );
        self.num_allocated.set(num_allocated_dec);
    }

    #[inline]
    pub fn len(&self) -> usize {
        self.num_allocated.get()
    }
    #[inline]
    pub fn capacity(&self) -> usize {
        self.chunks.borrow().len() * E::AllocatePolicyT::CHUNK_SIZE
    }
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    pub fn clear(&mut self) {
        self.chunks.get_mut().clear();
        self.free_head.set(0);
        self.num_allocated.set(0);
    }

    #[inline]
    pub fn iter(&self) -> EntityAllocReadIter<'_, E> {
        EntityAllocReadIter::new(self)
    }
    #[inline]
    pub fn iter_mut(&mut self) -> EntityAllocEditIter<'_, E> {
        EntityAllocEditIter::new(self)
    }
}