mtb-entity-slab 0.2.4

#[doc(hidden)]
pub trait IBitAlloc: Copy + Default {
    const BITSET_LEN: usize;

    fn allocate(&mut self) -> Option<u16>;
    fn is_allocated(&self, id: u16) -> bool;
    fn next_available(&self) -> Option<u16>;
    fn deallocate(&mut self, index: u16);

    /// 如果 unit 分配了则返回 unit, 否则返回 unit 的下一个已分配的 unit
    fn next_allocated(&self, unit: u16) -> Option<u16>;

    fn cells(&self) -> &[u64];

    fn count_allocated(&self) -> usize {
        let mut count = 0;
        for &cell in self.cells() {
            count += cell.count_ones() as usize;
        }
        count
    }
}

#[derive(Clone, Copy)]
pub struct BitAlloc<const LEN: usize> {
    cells: [u64; LEN],
}
impl<const LEN: usize> Default for BitAlloc<LEN> {
    fn default() -> Self {
        Self { cells: [0u64; LEN] }
    }
}
impl<const LEN: usize> IBitAlloc for BitAlloc<LEN> {
    const BITSET_LEN: usize = LEN;

    fn allocate(&mut self) -> Option<u16> {
        for i in 0..Self::BITSET_LEN {
            let cell = self.cells[i];
            if cell == u64::MAX {
                continue;
            }
            let pos = cell.trailing_ones();
            self.cells[i] = cell | 1 << pos;
            return Some(i as u16 * 64 + pos as u16);
        }
        None
    }
    fn is_allocated(&self, id: u16) -> bool {
        let (pos, bit) = (id >> 6, id & 0b11_1111);
        self.cells[pos as usize] & (1 << bit as u32) != 0
    }
    fn next_available(&self) -> Option<u16> {
        for i in 0..Self::BITSET_LEN {
            let cell = self.cells[i];
            if cell == u64::MAX {
                continue;
            }
            let pos = cell.trailing_ones();
            return Some(i as u16 * 64 + pos as u16);
        }
        None
    }
    fn deallocate(&mut self, index: u16) {
        let (pos, bit) = (index >> 6, index & 0b11_1111);
        let cell = self.cells[pos as usize];
        self.cells[pos as usize] = cell & !(1 << bit as u32);
    }
    fn cells(&self) -> &[u64] {
        &self.cells
    }
    fn next_allocated(&self, unit: u16) -> Option<u16> {
        array_next_allocated(&self.cells, unit)
    }
}

#[derive(Clone, Copy)]
pub struct SummaryAlloc<const LEN: usize> {
    cells: [u64; LEN],
    summary: u64,
}
impl<const LEN: usize> Default for SummaryAlloc<LEN> {
    fn default() -> Self {
        Self {
            cells: [0u64; LEN],
            summary: 0,
        }
    }
}
impl<const LEN: usize> IBitAlloc for SummaryAlloc<LEN> {
    const BITSET_LEN: usize = LEN;

    fn allocate(&mut self) -> Option<u16> {
        if self.summary == Self::BITMASK_FULL {
            return None;
        }
        let free_cell = self.summary.trailing_ones();
        let mut cell = self.cells[free_cell as usize];
        let pos = cell.trailing_ones();
        cell |= 1 << pos;
        self.cells[free_cell as usize] = cell;
        if cell == u64::MAX {
            self.summary |= 1 << free_cell;
        }
        Some((free_cell * 64 + pos) as u16)
    }
    fn is_allocated(&self, id: u16) -> bool {
        let (pos, bit) = (id >> 6, id & 0b11_1111);
        self.cells[pos as usize] & (1 << bit as u32) != 0
    }
    fn next_available(&self) -> Option<u16> {
        if self.summary == Self::BITMASK_FULL {
            return None;
        }
        let free_cell = self.summary.trailing_ones();
        let cell = self.cells[free_cell as usize];
        let pos = cell.trailing_ones();
        Some((free_cell * 64 + pos) as u16)
    }
    fn deallocate(&mut self, index: u16) {
        let (pos, bit) = (index >> 6, index & 0b11_1111);
        let cell = self.cells[pos as usize];
        self.cells[pos as usize] = cell & !(1 << bit as u32);
        self.summary &= !(1 << pos);
    }
    fn cells(&self) -> &[u64] {
        &self.cells
    }
    fn next_allocated(&self, unit: u16) -> Option<u16> {
        array_next_allocated(&self.cells, unit)
    }
}
impl<const LEN: usize> SummaryAlloc<LEN> {
    const BITMASK_FULL: u64 = if LEN == 64 {
        u64::MAX
    } else if LEN > 64 {
        panic!("SummaryAlloc supports up to 4096 units");
    } else {
        (1u64 << LEN) - 1
    };
}

/// 如果 start_unit 已分配则返回 start_unit, 否则返回下一个已分配的 unit
fn array_next_allocated<const LEN: usize>(cells: &[u64; LEN], start_unit: u16) -> Option<u16> {
    let (mut pos, bit) = (start_unit >> 6, start_unit & 0b11_1111);
    let mut cell = cells[pos as usize];
    // 先检查当前单元
    cell &= !((1 << bit as u32) - 1);
    if cell != 0 {
        let next_bit = cell.trailing_zeros();
        return Some(pos * 64 + next_bit as u16);
    }
    // 检查后续单元
    pos += 1;
    while (pos as usize) < LEN {
        let cell = cells[pos as usize];
        if cell != 0 {
            let next_bit = cell.trailing_zeros();
            return Some(pos * 64 + next_bit as u16);
        }
        pos += 1;
    }
    None
}