mtb-entity-slab 0.2.4

use crate::{
    bitalloc::{BitAlloc, IBitAlloc, SummaryAlloc},
    chunk::Unit,
    gen_index::GenIndex,
};
use std::{num::NonZeroU16, ops::IndexMut};

/// Errors that can occur when working with slice pointers.
///
/// This is used in pointer ID validation in `EntityAlloc`.
#[derive(Debug)]
pub enum SlicePtrErr {
    /// The pointer is out of the valid address range of the slice.
    OutOfRange,

    /// The pointer is not aligned with the start of any unit in the slice.
    NotAlignedWith,
}
impl std::fmt::Display for SlicePtrErr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let msg = match self {
            SlicePtrErr::OutOfRange => "Pointer is out of the valid address range of the slice.",
            SlicePtrErr::NotAlignedWith => {
                "Pointer is not aligned with the start of any unit in the slice."
            }
        };
        f.write_str(msg)
    }
}
impl std::error::Error for SlicePtrErr {}

/// Result type for operations involving slice pointers.
pub type SlicePtrRes<T = ()> = Result<T, SlicePtrErr>;

pub trait IUnitArray<E>: Sized + IndexMut<usize, Output = Unit<E>> {
    const LEN: usize;
    unsafe fn boxed_uninit() -> Box<Self>;
    fn indexof_unit_ptr(&self, ptr: *const Unit<E>) -> Result<usize, SlicePtrErr>;
}
impl<E, const N: usize> IUnitArray<E> for [Unit<E>; N] {
    const LEN: usize = N;

    unsafe fn boxed_uninit() -> Box<Self> {
        let mut ret: Box<Self> = unsafe { Box::new_uninit().assume_init() };
        let init_id = initial_id::initial_id();
        for i in 0..Self::LEN {
            ret[i].indexed = init_id;
        }
        ret
    }

    fn indexof_unit_ptr(&self, ptr: *const Unit<E>) -> Result<usize, SlicePtrErr> {
        use std::ops::Range;
        let unit_size = std::mem::size_of::<Unit<E>>();
        let Range { start, end } = self.as_ptr_range();
        let start_addr = start as usize;
        let end_addr = end as usize;
        let ptr_addr = ptr as usize;
        if ptr_addr < start_addr || ptr_addr >= end_addr {
            return Err(SlicePtrErr::OutOfRange);
        }
        let offset = ptr_addr - start_addr;
        if !offset.is_multiple_of(unit_size) {
            return Err(SlicePtrErr::NotAlignedWith);
        }
        Ok(offset / unit_size)
    }
}

mod initial_id {
    use crate::gen_index::GenIndex;

    #[cfg(feature = "random-generation")]
    pub(super) fn initial_id() -> GenIndex {
        use std::num::NonZeroU16;

        let generation = NonZeroU16::new(fastrand::u16(1..)).unwrap();
        GenIndex::compose(GenIndex::INVALID_INDEX, generation)
    }
    #[cfg(not(feature = "random-generation"))]
    pub(super) const fn initial_id() -> GenIndex {
        GenIndex::compose(GenIndex::INVALID_INDEX, GenIndex::GEN_1)
    }
}

/// Entity allocation policy trait.
///
/// You are NOT ALLOWED to implement this trait for your own types.
/// You can only select the following pre-defined policies:
///
/// - `AllocPolicy128` - Chunk size of 128.
/// - `AllocPolicy256` - Chunk size of 256.
/// - `AllocPolicy512` - Chunk size of 512.
/// - `AllocPolicy1024` - Chunk size of 1024.
/// - `AllocPolicy2048` - Chunk size of 2048.
/// - `AllocPolicy4096` - Chunk size of 4096.
///
/// Entity 分配策略接口.
///
/// 你不允许为自己的类型实现此接口. 你只能选择以下预定义的策略:
///
/// - `AllocPolicy128` - 块大小为 128.
/// - `AllocPolicy256` - 块大小为 256.
/// - `AllocPolicy512` - 块大小为 512.
/// - `AllocPolicy1024` - 块大小为 1024.
/// - `AllocPolicy2048` - 块大小为 2048.
/// - `AllocPolicy4096` - 块大小为 4096.
pub trait IAllocPolicy: 'static {
    /// Chunk size log2 (e.g., 7 for 128, 8 for 256, etc.)
    const CHUNK_SIZE_LOG2: usize;

    /// Chunk size (e.g., 128, 256, etc.)
    const CHUNK_SIZE: usize;

    /// Bitset length (number of u64 units in the bitset)
    const BITSET_LEN: usize;

    /// Unit ID mask, usually CHUNK_SIZE - 1
    const UNIT_ID_MASK: usize;

    /// Bit allocator type. `MTB::Entity` chooses different bit allocators
    /// for different chunk sizes.
    type BitAlloc: IBitAlloc;

    /// Unit array type for the chunk.
    type Units<E>: IUnitArray<E>;

    /// Compose a `GenIndex` from chunk ID, unit ID, and generation.
    fn compose(chunk: u32, unit: u16, generation: NonZeroU16) -> GenIndex {
        let chunk_id_shift = Self::CHUNK_SIZE_LOG2;
        let unit_id_mask = Self::UNIT_ID_MASK;
        let chunk_id = chunk as u64;
        let unit_id = (unit as u64) & (unit_id_mask as u64);
        let real_index = (chunk_id << chunk_id_shift) | unit_id;
        GenIndex::compose(real_index as usize, generation)
    }
    /// Extract chunk ID from a real index.
    fn chunkof_real(real_index: usize) -> u32 {
        (real_index >> Self::CHUNK_SIZE_LOG2) as u32
    }
    /// Extract chunk ID from a `GenIndex`.
    fn chunk(index: GenIndex) -> u32 {
        Self::chunkof_real(index.real_index())
    }

    /// Extract unit ID from a real index.
    fn unitof_real(real_index: usize) -> u16 {
        (real_index & Self::UNIT_ID_MASK) as u16
    }
    /// Extract unit ID from a `GenIndex`.
    fn unit(index: GenIndex) -> u16 {
        (index.real_index() & Self::UNIT_ID_MASK) as u16
    }

    /// Tear down a `GenIndex` into chunk ID, unit ID, and generation.
    fn tear(index: GenIndex) -> (u32, u16, NonZeroU16) {
        let chunk = Self::chunk(index);
        let unit = Self::unit(index);
        let generation = index.generation();
        (chunk, unit, generation)
    }
}
pub(crate) trait IAllocPolicyPrivate: IAllocPolicy {
    /// # Safety
    ///
    /// This function returns a boxed array of uninitialized units.
    unsafe fn unit_array<E>() -> Box<Self::Units<E>> {
        unsafe { Self::Units::boxed_uninit() }
    }
}
impl<T: IAllocPolicy> IAllocPolicyPrivate for T {}

macro_rules! define_entity_alloc_policy {
    ($name:ident, $chunk_size_log2:expr, $bit_alloc:ty, $doc:expr) => {
        #[doc = $doc]
        pub struct $name;

        impl IAllocPolicy for $name {
            const CHUNK_SIZE_LOG2: usize = $chunk_size_log2;
            const CHUNK_SIZE: usize = 1 << $chunk_size_log2;
            const BITSET_LEN: usize = Self::CHUNK_SIZE / 64;
            const UNIT_ID_MASK: usize = (1 << Self::CHUNK_SIZE_LOG2) - 1;

            type BitAlloc = $bit_alloc;
            type Units<E> = [Unit<E>; Self::CHUNK_SIZE];
        }
    };
}

define_entity_alloc_policy! {AllocPolicy128,  7,  BitAlloc<2>, "Entity allocation policy for chunks of size 128."}
define_entity_alloc_policy! {AllocPolicy256,  8,  BitAlloc<4>, "Entity allocation policy for chunks of size 256."}
define_entity_alloc_policy! {AllocPolicy512,  9,  BitAlloc<8>, "Entity allocation policy for chunks of size 512."}
define_entity_alloc_policy! {AllocPolicy1024, 10, SummaryAlloc<16>, "Entity allocation policy for chunks of size 1024."}
define_entity_alloc_policy! {AllocPolicy2048, 11, SummaryAlloc<32>, "Entity allocation policy for chunks of size 2048."}
define_entity_alloc_policy! {AllocPolicy4096, 12, SummaryAlloc<64>, "Entity allocation policy for chunks of size 4096."}