mtb_entity_slab/

alloc.rs

use crate::{
    EntityAllocConsumeIter, EntityAllocEditIter, EntityAllocReadIter, PtrID,
    bitalloc::IBitAlloc,
    chunk::Chunk,
    id::IndexedID,
    policy::{AllocPolicy256, IAllocPolicy},
};
use std::cell::{BorrowMutError, Cell, RefCell};

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> AsRef<EntityAlloc<E, P>> for EntityAlloc<E, P> {
    fn as_ref(&self) -> &EntityAlloc<E, P> {
        self
    }
}
impl<E, P: IAllocPolicy> AsMut<EntityAlloc<E, P>> for EntityAlloc<E, P> {
    fn as_mut(&mut self) -> &mut EntityAlloc<E, P> {
        self
    }
}

impl<E, P: IAllocPolicy> EntityAlloc<E, P> {
    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 chunks = {
            let nchunks = cap.next_multiple_of(P::CHUNK_SIZE) / 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 fn allocate(&self, val: E) -> PtrID<E, P> {
        match self.try_allocate(val) {
            Ok(ptr) => ptr,
            Err(_) => panic!(
                "Cannot borrow EntityAlloc:{self:p} mutably for allocation. Are you traversing it now?"
            ),
        }
    }
    pub fn try_allocate(&self, val: E) -> Result<PtrID<E, P>, (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(unsafe { PtrID::from_raw_ptr(ptr as *mut _) })
    }
    /// 下一次分配时, 返回下一个可用的索引
    pub fn next_indexed(&self) -> IndexedID<E, P> {
        let chunks = self.chunks.borrow();
        let free_head = self.free_head.get() as usize;
        let Some(chunk) = chunks.get(free_head) else {
            return IndexedID::from(free_head << P::CHUNK_SIZE_LOG2);
        };
        let unit_id = chunk.next_available().unwrap();
        IndexedID::from(chunk.indexed_id_of_unit(unit_id))
    }
    fn find_chunk(free: u32, chunks: &mut Vec<Chunk<E, P>>) -> &mut Chunk<E, P> {
        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 = P::chunk_of_indexed_id(indexed);
        let unit_id = P::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, P>) -> Option<E> {
        let indexed = ptr.get_index(self)?;
        self.free_indexed(indexed)
    }

    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> {
    pub fn free_if(&mut self, pred: impl FnMut(&E, PtrID<E, P>, IndexedID<E, P>) -> bool) {
        self.fully_free_if(pred, drop);
    }
    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(unit_id);
            let ptr: PtrID<E, P> = PtrID::from_ref(unit);
            let elem = unit.as_init_ref().unwrap();
            if !should_free(elem, ptr, IndexedID::from(index)) {
                continue;
            }
            let Some(e) = chunk.deallocate_unit_id(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(unit_id);
            let ptr: PtrID<E, P> = PtrID::from_ref(unit);
            let elem = unit.as_init_ref().unwrap();
            if should_free(elem, ptr, IndexedID::from(index)) {
                if let Some(e) = chunk.deallocate_unit_id(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
    }

    #[inline]
    pub fn len(&self) -> usize {
        self.num_allocated.get()
    }
    #[inline]
    pub fn capacity(&self) -> usize {
        self.chunks.borrow().len() * P::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, P> {
        EntityAllocReadIter::new(self)
    }
    #[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 = (usize, 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 = (usize, 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 = (usize, E);
    type IntoIter = EntityAllocConsumeIter<E, P>;

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