mtb-entity-slab 0.2.4

use crate::{
    SlicePtrErr,
    bitalloc::IBitAlloc,
    gen_index::GenIndex,
    policy::{IAllocPolicy, IAllocPolicyPrivate, IUnitArray},
};
use std::{mem::MaybeUninit, num::NonZeroU16};

#[repr(C)]
pub struct Unit<E> {
    pub data: MaybeUninit<E>,
    pub indexed: GenIndex,
}

impl<E> Drop for Unit<E> {
    fn drop(&mut self) {
        if self.indexed.is_invalid() {
            return;
        }
        unsafe { self.data.assume_init_drop() };
        // 析构以后可能还有 ID 指向这块内存区域,
        // 所以把 indexed_id 设为 INVALID
        self.indexed = self.indexed.to_invalid();
    }
}

impl<E: Clone> Clone for Unit<E> {
    fn clone(&self) -> Self {
        if self.indexed.is_invalid() {
            Self {
                data: MaybeUninit::uninit(),
                indexed: GenIndex::to_invalid(self.indexed),
            }
        } else {
            Self {
                data: MaybeUninit::new(unsafe { self.data.assume_init_ref().clone() }),
                indexed: self.indexed,
            }
        }
    }
}

impl<E> Unit<E> {
    pub fn as_init_mut(&mut self) -> Option<&mut E> {
        if self.indexed.is_invalid() {
            None
        } else {
            unsafe { Some(self.data.assume_init_mut()) }
        }
    }
    pub fn as_init_ref(&self) -> Option<&E> {
        if self.indexed.is_invalid() {
            None
        } else {
            unsafe { Some(self.data.assume_init_ref()) }
        }
    }

    unsafe fn write(&mut self, data: E, real_index: usize) -> &mut E {
        if cfg!(debug_assertions) && !self.indexed.is_invalid() {
            let indexed = self.indexed.0;
            panic!("Unit{self:p} to be written should be invalid, but got id {indexed:x}");
        }
        let gene = GenIndex::do_generation_inc(self.indexed.generation());
        unsafe {
            self.data.write(data);
            self.indexed = GenIndex::compose(real_index, gene);
            self.data.assume_init_mut()
        }
    }
    pub(super) unsafe fn free(&mut self) -> Option<E> {
        if self.indexed.is_invalid() {
            return None;
        }
        let val = unsafe { self.data.assume_init_read() };
        self.indexed = self.indexed.to_invalid();
        Some(val)
    }
}

pub(crate) struct Chunk<E, P: IAllocPolicy> {
    pub units: *mut P::Units<E>,
    pub allocated: P::BitAlloc,
    /// 分配了多少元素
    pub num_allocated: u32,
    /// 该内存块的索引
    pub chunk_id: u32,
    /// 如果该内存块有空位, 那下一个有空位的内存块索引在哪儿
    pub free_next: u32,
}
unsafe impl<E: Send, P: IAllocPolicy> Send for Chunk<E, P> {}

impl<E, P: IAllocPolicy> Drop for Chunk<E, P> {
    fn drop(&mut self) {
        // # Safety
        //
        // 1. self.units 是通过 P::unit_array() 创建的, 其返回值是一个 Box<Self::Units<E>>,
        //    也就是一个 Box<[Unit<E>; CHUNK_SIZE]>.
        // 2. Box::from_raw(self.units) 创建的 Box<Self::Units<E>> 会在 drop 时正确地调用每个
        //    Unit<E> 的 drop 方法, 从而正确地析构每个 Unit<E> 中的 E.
        let _ = unsafe { Box::from_raw(self.units) };
    }
}

impl<E, P: IAllocPolicy> Chunk<E, P> {
    pub fn new(id: u32, free_next: u32) -> Self {
        Self {
            units: unsafe { Box::into_raw(P::unit_array()) },
            allocated: Default::default(),
            num_allocated: 0,
            chunk_id: id,
            free_next,
        }
    }

    pub fn is_full(&self) -> bool {
        self.num_allocated as usize == P::CHUNK_SIZE
    }
    pub fn is_empty(&self) -> bool {
        self.num_allocated == 0
    }

    #[inline]
    pub fn unit(&self, unit_id: u16) -> &Unit<E> {
        unsafe { &*(self.units as *const Unit<E>).add(unit_id as usize) }
    }
    #[inline]
    pub fn unit_mut(&mut self, unit_id: u16) -> &mut Unit<E> {
        unsafe { &mut *(self.units as *mut Unit<E>).add(unit_id as usize) }
    }
    #[inline]
    pub fn unit_ptr_mut(&self, unit_id: u16) -> *mut Unit<E> {
        unsafe { (self.units as *mut Unit<E>).add(unit_id as usize) }
    }

    pub fn allocate(&mut self, val: E) -> Result<*mut Unit<E>, E> {
        if self.is_full() {
            return Err(val);
        }
        let Some(unit_id) = self.allocated.allocate() else {
            unreachable!("Length maintainance error: full ID allocator when num_free == 0");
        };
        self.num_allocated += 1;
        let real_index = self.real_index_of_unit(unit_id);
        unsafe { self.unit_mut(unit_id).write(val, real_index) };
        Ok(self.unit_ptr_mut(unit_id))
    }
    pub fn deallocate_checked(&mut self, unit_id: u16, gene: NonZeroU16) -> Option<E> {
        debug_assert!(
            (unit_id as usize) < P::CHUNK_SIZE,
            "Found overflowing unit ID {unit_id}"
        );
        debug_assert!(
            self.allocated.is_allocated(unit_id),
            "Found unallocated unit ID {unit_id}"
        );
        if self.unit_mut(unit_id).indexed.generation() != gene {
            return None;
        }
        self.deallocate_unchecked(unit_id)
    }
    pub fn deallocate_unchecked(&mut self, unit_id: u16) -> Option<E> {
        let val = unsafe { self.unit_mut(unit_id).free() }?;
        self.allocated.deallocate(unit_id);
        // 维护已分配计数，避免 chunk 永远保持满状态
        self.num_allocated -= 1;
        Some(val)
    }

    #[inline]
    pub fn indexed_id_of_unit(&self, unit_id: u16) -> GenIndex {
        let gene = self.unit(unit_id).indexed.generation();
        P::compose(self.chunk_id, unit_id, gene)
    }
    #[inline]
    pub fn real_index_of_unit(&self, unit_id: u16) -> usize {
        let chunk_id = self.chunk_id as usize;
        let unit_id = unit_id as usize;
        (chunk_id << P::CHUNK_SIZE_LOG2) | unit_id
    }
    pub fn unit_of_ptr(&self, ptr: *const Unit<E>) -> Result<u16, SlicePtrErr> {
        unsafe { &*self.units }
            .indexof_unit_ptr(ptr)
            .map(|i| i as u16)
    }

    /// 获取下一个已分配的单元ID. 如果 unit_id 自己就是已分配的, 则返回自己.
    #[inline]
    pub fn next_allocated(&self, unit_id: u16) -> Option<u16> {
        self.allocated.next_allocated(unit_id)
    }
    #[inline]
    pub fn next_available(&self) -> Option<u16> {
        self.allocated.next_available()
    }
}