moa_idalloc 0.1.4

//! 世代槽位表(id ↔ 对象)。

use alloc::vec::Vec;

/// 空闲链链尾哨兵,同时是 index 的非法值(故有效 index 上限为 `u32::MAX - 1`)。
const NIL: u32 = u32::MAX;

/// 世代化 id:`generation`(高 32 位)+ `index`(低 32 位)打包成 `u64`。
///
/// 同一槽位被复用时 `generation` 递增,因此旧 id 在 [`IdTable::get`] 等查找中会被拒绝。
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub struct Id(u64);

impl Id {
    #[inline(always)]
    const fn new(index: u32, generation: u32) -> Self {
        Self(((generation as u64) << 32) | index as u64)
    }

    // 取低 32 位 index(高 32 位是 generation),截断为预期行为。
    #[inline(always)]
    #[allow(clippy::cast_possible_truncation)]
    fn index(self) -> u32 {
        self.0 as u32
    }

    #[inline(always)]
    fn generation(self) -> u32 {
        (self.0 >> 32) as u32
    }

    /// 打包后的原始 `u64` 值(可作句柄/键存储或传出)。
    #[inline(always)]
    pub fn bits(self) -> u64 {
        self.0
    }

    /// 从 [`bits`](Id::bits) 得到的原始值重建。
    #[inline(always)]
    pub fn from_bits(bits: u64) -> Self {
        Self(bits)
    }
}

enum State<T> {
    Occupied(T),
    /// 空闲槽,串入空闲链;`NIL` 为链尾。
    Vacant {
        next_free: u32,
    },
}

struct Slot<T> {
    generation: u32,
    state: State<T>,
}

/// 世代槽位表:按 [`Id`] 存取对象,`O(1)` 分配/释放/查找。
///
/// 释放后槽位的 `generation` 递增并挂回空闲链,复用时返回带新 `generation` 的 [`Id`]。
/// 这样任何被缓存的旧 id 在查找时都会因 `generation` 不符被拒,从而免疫 ABA/复用问题。
/// 适合替代"裸数字句柄表"。
pub struct IdTable<T> {
    slots: Vec<Slot<T>>,
    /// 空闲链表头;`NIL` 表示无空闲槽。
    free_head: u32,
    len: usize,
}

impl<T> IdTable<T> {
    /// 创建空表(不预分配)。
    pub const fn new() -> Self {
        Self { slots: Vec::new(), free_head: NIL, len: 0 }
    }

    /// 当前存活对象数。
    #[inline(always)]
    pub fn len(&self) -> usize {
        self.len
    }

    /// 是否为空。
    #[inline(always)]
    pub fn is_empty(&self) -> bool {
        self.len == 0
    }

    /// 插入对象,返回其世代化 [`Id`]。优先复用空闲槽,否则追加。
    #[inline]
    pub fn insert(&mut self, value: T) -> Id {
        self.len += 1;
        if self.free_head == NIL {
            // 无空闲槽:追加。index 必须落在 [0, NIL) 内(NIL 是非法哨兵)。
            let index = u32::try_from(self.slots.len())
                .ok()
                .filter(|&i| i != NIL)
                .expect("IdTable index space exhausted");
            self.slots.push(Slot { generation: 0, state: State::Occupied(value) });
            Id::new(index, 0)
        } else {
            // 复用空闲链表头的槽。
            let index = self.free_head;
            let slot = &self.slots[index as usize];
            let State::Vacant { next_free } = slot.state else {
                unreachable!("free list points to an occupied slot");
            };
            let generation = slot.generation;
            self.free_head = next_free;
            self.slots[index as usize].state = State::Occupied(value);
            Id::new(index, generation)
        }
    }

    /// 按 [`Id`] 取对象引用;index 越界、`generation` 不符或槽已空均返回 `None`。
    #[inline]
    pub fn get(&self, id: Id) -> Option<&T> {
        match self.slots.get(id.index() as usize)? {
            Slot { generation, state: State::Occupied(v) } if *generation == id.generation() => {
                Some(v)
            },
            _ => None,
        }
    }

    /// 按 [`Id`] 取对象可变引用。
    #[inline]
    pub fn get_mut(&mut self, id: Id) -> Option<&mut T> {
        match self.slots.get_mut(id.index() as usize)? {
            Slot { generation, state: State::Occupied(v) } if *generation == id.generation() => {
                Some(v)
            },
            _ => None,
        }
    }

    /// 按 [`Id`] 移除并返回对象;旧/无效 id 返回 `None`。
    #[inline]
    pub fn remove(&mut self, id: Id) -> Option<T> {
        let index = id.index() as usize;
        match self.slots.get(index) {
            Some(Slot { generation, state: State::Occupied(_) })
                if *generation == id.generation() => {},
            _ => return None,
        }
        let free_head = self.free_head;
        let slot = &mut self.slots[index];
        // 提升 generation,使已发出的旧 id 失效;再把槽挂回空闲链。
        slot.generation = slot.generation.wrapping_add(1);
        let old = core::mem::replace(&mut slot.state, State::Vacant { next_free: free_head });
        self.free_head = id.index();
        self.len -= 1;
        match old {
            State::Occupied(v) => Some(v),
            State::Vacant { .. } => unreachable!(),
        }
    }
}

impl<T> Default for IdTable<T> {
    fn default() -> Self {
        Self::new()
    }
}