mtb-entity-slab 0.2.4

use std::{
    fmt::{Debug, Formatter, LowerHex, Pointer, UpperHex},
    hash::{Hash, Hasher},
    marker::PhantomData,
    num::NonZeroU16,
    ptr::NonNull,
};

use crate::{EntityAlloc, IAllocPolicy, chunk::Unit, gen_index::GenIndex};

/// Trait for ID wrappers that carry an associated object type and allocation policy.
///
/// This trait provides a uniform interface for ID types used by the crate's
/// APIs. An implementor represents an ID type that is bound to a concrete
/// `ObjectT` and `PolicyT`. The `BackID` associated type is a backend ID that
/// actually implements the conversion and dereference logic (for example
/// `PtrID` or `IndexedID`).
///
/// Implementations must be `Copy` and `Eq`. The trait exposes helpers to
/// convert to/from the backend representation and to dereference given an
/// `IDBoundAlloc` (an `EntityAlloc` specialized to the ID's object/policy).
///
/// 该 trait 用于将不同的 ID 表示（指针型或索引型）抽象为统一的接口，绑定对象类型与分配策略，
/// 便于上层 API 泛化处理 ID。
pub trait IPoliciedID: Copy + Eq + Debug {
    /// The object type this ID references.
    type ObjectT: Sized;
    /// The allocation policy this ID is associated with.
    type PolicyT: IAllocPolicy;
    /// The backend ID type that implements the actual logic.
    type BackID: IEntityAllocID<Self::ObjectT, Self::PolicyT>;

    /// Create this ID from its backend representation.
    fn from_backend(ptr: Self::BackID) -> Self;
    /// Convert this ID into its backend representation.
    fn into_backend(self) -> Self::BackID;

    /// Try to dereference the object this ID points to, returning `None` if invalid.
    fn try_deref_alloc(self, alloc: &IDBoundAlloc<Self>) -> Option<&Self::ObjectT> {
        self.into_backend().try_deref(alloc)
    }
    /// Try to mutably dereference the object this ID points to, returning `None` if invalid.
    fn try_deref_alloc_mut(self, alloc: &mut IDBoundAlloc<Self>) -> Option<&mut Self::ObjectT> {
        self.into_backend().try_deref_mut(alloc)
    }

    /// Dereference the object this ID points to, panicking if invalid.
    fn deref_alloc(self, alloc: &IDBoundAlloc<Self>) -> &Self::ObjectT {
        self.into_backend().deref(alloc)
    }
    /// Mutably dereference the object this ID points to, panicking if invalid.
    fn deref_alloc_mut(self, alloc: &mut IDBoundAlloc<Self>) -> &mut Self::ObjectT {
        self.into_backend().deref_mut(alloc)
    }
}

/// Alias for `EntityAlloc` specialized to the object and policy types of an ID.
///
/// 常用作函数签名中把 `ID` 与对应的 `EntityAlloc` 进行绑定，避免每次都写出复杂的关联类型。
pub type IDBoundAlloc<I> = EntityAlloc<<I as IPoliciedID>::ObjectT, <I as IPoliciedID>::PolicyT>;

/// Trait implemented by concrete ID representations that can reference
/// entities inside an `EntityAlloc`.
///
/// Responsibilities:
/// - Convert from/to `PtrID` and `IndexedID` where possible.
/// - Provide fallible dereference operations `try_deref` / `try_deref_mut`.
/// - Provide allocation (`allocate_from`) and deallocation (`free`) helpers
///   that tie the ID lifecycle to an `EntityAlloc` instance.
///
/// 中文说明：实现此 trait 的类型可作为对 `EntityAlloc` 中实体的引用，
/// 支持在指针/索引与自身之间转换并提供（可选）解引用与生命周期绑定方法。
pub trait IEntityAllocID<E, P: IAllocPolicy>: Sized + Copy {
    /// Create an ID from a pointer-backed ID, validating generation/index.
    fn from_ptr(alloc: &EntityAlloc<E, P>, ptr: PtrID<E, P>) -> Option<Self>;

    /// Create an ID from an index-backed ID, validating generation/index.
    fn from_index(alloc: &EntityAlloc<E, P>, indexed: IndexedID<E, P>) -> Option<Self>;

    /// Try to dereference the entity this ID points to, returning `None` if invalid.
    fn try_deref(self, alloc: &EntityAlloc<E, P>) -> Option<&E>;

    /// Try to mutably dereference the entity this ID points to, returning `None` if invalid.
    fn try_deref_mut(self, alloc: &mut EntityAlloc<E, P>) -> Option<&mut E>;

    /// Convert this ID to an index-backed ID, validating generation/index.
    fn to_index(self, alloc: &EntityAlloc<E, P>) -> Option<IndexedID<E, P>>;

    /// Convert this ID to a pointer-backed ID, validating generation/index.
    fn to_ptr(self, alloc: &EntityAlloc<E, P>) -> Option<PtrID<E, P>>;

    /// Directly dereference the entity this ID points to, panicking if invalid.
    #[inline]
    fn deref(self, alloc: &EntityAlloc<E, P>) -> &E {
        self.try_deref(alloc).expect("UAF detected!")
    }
    /// Directly mutably dereference the entity this ID points to, panicking if invalid.
    #[inline]
    fn deref_mut(self, alloc: &mut EntityAlloc<E, P>) -> &mut E {
        self.try_deref_mut(alloc).expect("UAF detected!")
    }

    /// Allocate a new entity in the given allocator and return its ID.
    /// NOTE that `MTB::Entity` library is ID-driven; allocation and deallocation
    /// must be done through IDs.
    fn allocate_from(alloc: &EntityAlloc<E, P>, val: E) -> Self;

    /// Free the referenced entity and return its owned value if successful.
    ///
    /// 释放当前 ID 指向的实体并返回其值（若释放成功）。
    fn free(self, alloc: &mut EntityAlloc<E, P>) -> Option<E>;
}

/// Pointer-backed ID referencing a `Unit<E>` inside an `EntityAlloc`.
///
/// `PtrID` is a thin, non-null wrapper around a raw `NonNull<Unit<E>>`. It is
/// intended for fast, pointer-based access to entities. The struct carries a
/// `PhantomData<P>` to associate the allocation policy at the type level.
///
/// Safety / 注意事项：
/// - `PtrID` assumes the pointer points to a properly initialized `Unit<E>`.
/// - Many operations are `unsafe` without validity checks; prefer the safe
///   conversion methods in `IEntityAllocID`/`EntityAlloc` that validate
///   generation/index when necessary.
pub struct PtrID<E, P> {
    pub(crate) ptr: NonNull<Unit<E>>,
    _marker: PhantomData<P>,
}
impl<E, P> From<NonNull<Unit<E>>> for PtrID<E, P> {
    fn from(ptr: NonNull<Unit<E>>) -> Self {
        Self {
            ptr,
            _marker: PhantomData,
        }
    }
}
impl<E, P> From<&mut Unit<E>> for PtrID<E, P> {
    fn from(unit: &mut Unit<E>) -> Self {
        Self {
            ptr: NonNull::from(unit),
            _marker: PhantomData,
        }
    }
}
impl<E, P> From<*mut Unit<E>> for PtrID<E, P> {
    fn from(raw: *mut Unit<E>) -> Self {
        Self {
            ptr: NonNull::new(raw).expect("PtrID cannot be null"),
            _marker: PhantomData,
        }
    }
}
impl<E, P> Copy for PtrID<E, P> {}
impl<E, P> Clone for PtrID<E, P> {
    fn clone(&self) -> Self {
        *self
    }
}
impl<E, P> Debug for PtrID<E, P> {
    #[inline]
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        Debug::fmt(&self.ptr, f)
    }
}
impl<E, P> Pointer for PtrID<E, P> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        Pointer::fmt(&self.ptr, f)
    }
}

impl<E, P> PartialEq for PtrID<E, P> {
    #[inline]
    fn eq(&self, other: &Self) -> bool {
        self.ptr == other.ptr
    }
}
impl<E, P> Eq for PtrID<E, P> {}
impl<E, P> PartialOrd for PtrID<E, P> {
    #[inline]
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}
impl<E, P> Ord for PtrID<E, P> {
    #[inline]
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        self.ptr.cmp(&other.ptr)
    }
}
impl<E, P> Hash for PtrID<E, P> {
    #[inline]
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.ptr.hash(state);
    }
}
unsafe impl<E: Send, P> Send for PtrID<E, P> {}
unsafe impl<E: Sync, P> Sync for PtrID<E, P> {}
impl<E, P: IAllocPolicy> IEntityAllocID<E, P> for PtrID<E, P> {
    fn from_ptr(_: &EntityAlloc<E, P>, ptr: PtrID<E, P>) -> Option<Self> {
        Some(ptr)
    }
    fn from_index(alloc: &EntityAlloc<E, P>, indexed: IndexedID<E, P>) -> Option<Self> {
        let unit_ptr = alloc.unit_ptr_of_indexed(indexed.indexed).ok()?;
        Some(Self::from(unit_ptr))
    }
    fn try_deref(self, alloc: &EntityAlloc<E, P>) -> Option<&E> {
        if cfg!(debug_assertions) {
            // 在 debug 模式下检查指针有效性
            alloc.check_unit_validity(self.ptr.as_ptr()).ok()?;
        }
        unsafe { self.ptr.as_ref().as_init_ref() }
    }
    fn try_deref_mut(mut self, alloc: &mut EntityAlloc<E, P>) -> Option<&mut E> {
        if cfg!(debug_assertions) {
            // 在 debug 模式下检查指针有效性
            alloc.check_unit_validity(self.ptr.as_ptr()).ok()?;
        }
        unsafe { self.ptr.as_mut().as_init_mut() }
    }
    fn to_index(self, alloc: &EntityAlloc<E, P>) -> Option<IndexedID<E, P>> {
        let gen_index = alloc.index_of_unit_ptr(self.ptr.as_ptr()).ok()?;
        Some(IndexedID::from(gen_index))
    }
    fn to_ptr(self, _alloc: &EntityAlloc<E, P>) -> Option<PtrID<E, P>> {
        Some(self)
    }

    fn allocate_from(alloc: &EntityAlloc<E, P>, val: E) -> Self {
        let ptr = match alloc.try_allocate_unit(val) {
            Ok(ptr) => NonNull::new(ptr as *mut _).unwrap(),
            Err(..) => panic!("Allocation failed in IEntityAllocID::allocate_from"),
        };
        PtrID {
            ptr,
            _marker: PhantomData,
        }
    }
    fn free(self, alloc: &mut EntityAlloc<E, P>) -> Option<E> {
        alloc.free_unit_ptr(self.ptr.as_ptr())
    }
}
impl<E, P: IAllocPolicy> IPoliciedID for PtrID<E, P> {
    type ObjectT = E;
    type PolicyT = P;
    type BackID = PtrID<E, P>;

    fn from_backend(ptr: Self::BackID) -> Self {
        ptr
    }
    fn into_backend(self) -> Self::BackID {
        self
    }
}
impl<E, P: IAllocPolicy> PtrID<E, P> {
    /// Check whether this pointer currently refers to a valid allocated unit.
    ///
    /// Returns `true` if the pointer can be resolved to a live `GenIndex` in
    /// the provided `alloc`. This performs the same generation/index check used
    /// by other conversion helpers.
    ///
    /// 检查此指针在给定的 `EntityAlloc` 中是否仍然有效（未被释放或越界）。
    pub fn check_validity(&self, alloc: &EntityAlloc<E, P>) -> bool {
        alloc.index_of_unit_ptr(self.ptr.as_ptr()).is_ok()
    }

    /// Directly dereference the underlying unit, without validity check.
    ///
    /// # Safety
    /// The caller must ensure the pointer is valid and the unit is initialized.
    /// 不做有效性检查，调用者必须保证指针有效且单元已初始化。
    pub unsafe fn direct_deref<'a>(self) -> Option<&'a E> {
        unsafe { self.ptr.as_ref().as_init_ref() }
    }

    /// Directly dereference the underlying unit mutably, without validity check.
    ///
    /// # Safety
    /// The caller must ensure exclusive access and that the pointer is valid.
    /// 不做有效性检查，调用者必须保证对该内存拥有唯一可变访问权。
    pub unsafe fn direct_deref_mut<'a>(mut self) -> Option<&'a mut E> {
        unsafe { self.ptr.as_mut().as_init_mut() }
    }
}

/// Index-backed ID storing a `GenIndex` (real index + generation).
///
/// `IndexedID` is a compact representation suitable for serialization,
/// hashing, and comparisons. It stores the `GenIndex` produced by the allocator
/// and keeps a `PhantomData` to tie the ID to the element and policy types.
///
/// 中文说明：索引型 ID，包含 `GenIndex` 用于表达槽位与世代信息，适合哈希、比较与持久化场景。
#[repr(C)]
pub struct IndexedID<E, P> {
    /// The underlying `GenIndex` representing the indexed entity.
    pub indexed: GenIndex,
    _marker: PhantomData<(E, P)>,
}
impl<E, P> From<GenIndex> for IndexedID<E, P> {
    fn from(indexed: GenIndex) -> Self {
        Self {
            indexed,
            _marker: PhantomData,
        }
    }
}
impl<E, P> Copy for IndexedID<E, P> {}
impl<E, P> Clone for IndexedID<E, P> {
    fn clone(&self) -> Self {
        *self
    }
}
impl<E, P> Debug for IndexedID<E, P> {
    #[inline]
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        let (real, gene) = self.indexed.tear();
        write!(f, "IndexedID({real:x} gen {gene})")
    }
}
impl<E, P> Pointer for IndexedID<E, P> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "{:#x}", u64::from(self.indexed))
    }
}
impl<E, P> LowerHex for IndexedID<E, P> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        LowerHex::fmt(&u64::from(self.indexed), f)
    }
}
impl<E, P> UpperHex for IndexedID<E, P> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        UpperHex::fmt(&u64::from(self.indexed), f)
    }
}
impl<E, P> PartialEq for IndexedID<E, P> {
    #[inline]
    fn eq(&self, other: &Self) -> bool {
        self.indexed == other.indexed
    }
}
impl<E, P> Eq for IndexedID<E, P> {}
impl<E, P> PartialOrd for IndexedID<E, P> {
    #[inline]
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}
impl<E, P> Ord for IndexedID<E, P> {
    #[inline]
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        self.indexed.cmp(&other.indexed)
    }
}
impl<E, P> Hash for IndexedID<E, P> {
    #[inline]
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.indexed.hash(state);
    }
}
#[cfg(feature = "serde")]
/// Indexed ID Serialize Syntax: string literal `{ID:x}:{Gen:x}`
impl<E, P> serde_core::Serialize for IndexedID<E, P> {
    fn serialize<S: serde_core::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        use std::fmt::Write;
        #[derive(Default)]
        struct Buffer {
            buf: [u8; 32],
            len: usize,
        }
        impl Buffer {
            fn as_str(&self) -> &str {
                // We only ever write valid UTF-8 via `write_str`, so this is safe.
                unsafe { std::str::from_utf8_unchecked(&self.buf[..self.len]) }
            }
        }
        impl Write for Buffer {
            fn write_str(&mut self, s: &str) -> std::fmt::Result {
                let bytes = s.as_bytes();
                let avail = self.buf.len().saturating_sub(self.len);
                if bytes.len() > avail {
                    return Err(std::fmt::Error);
                }
                let dst = &mut self.buf[self.len..self.len + bytes.len()];
                dst.copy_from_slice(bytes);
                self.len += bytes.len();
                Ok(())
            }
        }

        let (real, gene) = self.indexed.tear();
        let mut buf = Buffer::default();
        write!(&mut buf, "{real:x}:{gene:x}").unwrap();
        serializer.serialize_str(buf.as_str())
    }
}
#[cfg(feature = "serde")]
impl<'de, E, P> serde_core::Deserialize<'de> for IndexedID<E, P> {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde_core::Deserializer<'de>,
    {
        use serde_core::de::Error as _;

        let s = <&str>::deserialize(deserializer)?;
        let mut parts = s.splitn(2, ':');
        let real_str = parts
            .next()
            .ok_or_else(|| D::Error::custom("missing real index"))?;
        let gen_str = parts
            .next()
            .ok_or_else(|| D::Error::custom("missing generation"))?;

        let real_u64 = u64::from_str_radix(real_str, 16)
            .map_err(|e| D::Error::custom(format!("invalid real index: {e}")))?;
        let gen_u64 = u64::from_str_radix(gen_str, 16)
            .map_err(|e| D::Error::custom(format!("invalid generation: {e}")))?;

        if gen_u64 > u16::MAX as u64 || gen_u64 == 0 {
            return Err(D::Error::custom("generation out of range"));
        }

        let real_usize = real_u64 as usize;
        // Case 'gen_u64 == 0' is already handled above, so this unwrap is safe.
        let generation = NonZeroU16::new(gen_u64 as u16).unwrap();

        Ok(IndexedID::from(GenIndex::compose(real_usize, generation)))
    }
}

impl<E, P: IAllocPolicy> IEntityAllocID<E, P> for IndexedID<E, P> {
    fn from_ptr(alloc: &EntityAlloc<E, P>, ptr: PtrID<E, P>) -> Option<Self> {
        let gen_index = alloc.index_of_unit_ptr(ptr.ptr.as_ptr()).ok()?;
        Some(IndexedID::from(gen_index))
    }
    fn from_index(_: &EntityAlloc<E, P>, indexed: IndexedID<E, P>) -> Option<Self> {
        Some(indexed)
    }

    fn try_deref(self, alloc: &EntityAlloc<E, P>) -> Option<&E> {
        let unit_ptr = alloc.unit_ptr_of_indexed(self.indexed).ok()?;
        unsafe { unit_ptr.as_ref().as_init_ref() }
    }
    fn try_deref_mut(self, alloc: &mut EntityAlloc<E, P>) -> Option<&mut E> {
        alloc
            .unit_mut_of_indexed(self.indexed)
            .ok()
            .and_then(|p| p.as_init_mut())
    }
    fn to_index(self, _: &EntityAlloc<E, P>) -> Option<IndexedID<E, P>> {
        Some(self)
    }
    fn to_ptr(self, alloc: &EntityAlloc<E, P>) -> Option<PtrID<E, P>> {
        let unit_ptr = alloc.unit_ptr_of_indexed(self.indexed).ok()?;
        Some(PtrID::from(unit_ptr))
    }

    fn allocate_from(alloc: &EntityAlloc<E, P>, val: E) -> Self {
        let Ok(unit) = alloc.try_allocate_unit(val) else {
            panic!("Allocation failed in IEntityAllocID::allocate_from");
        };
        unsafe { Self::from(unit.as_ref().unwrap().indexed) }
    }
    fn free(self, alloc: &mut EntityAlloc<E, P>) -> Option<E> {
        alloc.free_gen_index(self.indexed)
    }
}
impl<E, P: IAllocPolicy> IPoliciedID for IndexedID<E, P> {
    type ObjectT = E;
    type PolicyT = P;
    type BackID = IndexedID<E, P>;

    fn from_backend(ptr: Self::BackID) -> Self {
        ptr
    }
    fn into_backend(self) -> Self::BackID {
        self
    }
}
impl<E, P: IAllocPolicy> IndexedID<E, P> {
    /// Create an `IndexedID` that represents the next index the allocator would allocate.
    ///
    /// 注意：这只是对分配器内部 `next_index` 的封装，实际分配仍需调用分配接口。
    pub fn next_to_alloc(alloc: &EntityAlloc<E, P>) -> Self {
        Self::from(alloc.next_index())
    }

    /// Return the generation stored in this `IndexedID`.
    ///
    /// 返回索引中的世代号（用于检测过期的引用）。
    pub fn get_generation(self) -> NonZeroU16 {
        self.indexed.generation()
    }

    /// Return the real-order/index portion of this `IndexedID`.
    ///
    /// 返回索引的真实位置，用于定位在分配器内部的槽位。
    pub fn get_order(self) -> usize {
        self.indexed.real_index()
    }
}