dambi 0.1.2

use std::mem::MaybeUninit;

pub struct Slab<T> {
    slots: Vec<Slot<T>>,
    next_free: u32,
    len: u32,
}

#[derive(Clone, Copy)]
pub struct LiveKey(usize);

impl LiveKey {
    pub fn as_usize(self) -> usize {
        self.0
    }
}

struct Slot<T> {
    value: MaybeUninit<T>,
    next: u32,
    occupied: bool,
}

const NULL: u32 = u32::MAX;

impl<T> Slab<T> {
    fn to_slot_index(idx: usize) -> u32 {
        u32::try_from(idx).expect("Slab: index exceeds u32 range")
    }

    fn to_vec_index(idx: u32) -> usize {
        idx as usize
    }

    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    #[must_use]
    pub fn with_capacity(cap: usize) -> Self {
        Self {
            slots: Vec::with_capacity(cap),
            next_free: NULL,
            len: 0,
        }
    }
}

impl<T> Default for Slab<T> {
    fn default() -> Self {
        Self {
            slots: Vec::new(),
            next_free: NULL,
            len: 0,
        }
    }
}

impl<T> Slab<T> {
    pub fn len(&self) -> usize {
        self.len as usize
    }

    pub fn is_empty(&self) -> bool {
        self.len == 0
    }

    pub fn insert(&mut self, val: T) -> usize {
        self.insert_with(|_| val)
    }

    pub fn insert_with(&mut self, f: impl FnOnce(usize) -> T) -> usize {
        if self.next_free != NULL {
            let idx = Self::to_vec_index(self.next_free);
            let slot = &mut self.slots[idx];
            debug_assert!(!slot.occupied);
            let value = f(idx);
            let next = slot.next;
            slot.value.write(value);
            slot.occupied = true;
            self.next_free = next;
            self.len += 1;
            return idx;
        }
        let idx = self.slots.len();
        let _ = Self::to_slot_index(idx);
        let value = f(idx);
        self.slots.push(Slot {
            value: MaybeUninit::new(value),
            next: NULL,
            occupied: true,
        });
        self.len += 1;
        idx
    }

    pub fn get(&self, key: usize) -> Option<&T> {
        let slot = self.slots.get(key)?;
        if !slot.occupied {
            return None;
        }
        // SAFETY: `occupied == true` implies `value` is initialized.
        Some(unsafe { slot.value.assume_init_ref() })
    }

    pub fn get_mut(&mut self, key: usize) -> Option<&mut T> {
        let slot = self.slots.get_mut(key)?;
        if !slot.occupied {
            return None;
        }
        // SAFETY: see `get` — `occupied == true` implies `value` is initialized.
        Some(unsafe { slot.value.assume_init_mut() })
    }

    pub fn live_key(&self, key: usize) -> Option<LiveKey> {
        let slot = self.slots.get(key)?;
        if !slot.occupied {
            return None;
        }
        Some(LiveKey(key))
    }

    pub fn get_live(&mut self, key: LiveKey) -> &mut T {
        let slot = &mut self.slots[key.0];
        debug_assert!(slot.occupied);
        // SAFETY: LiveKey is only issued by `live_key` when `occupied == true`.
        unsafe { slot.value.assume_init_mut() }
    }

    pub fn remove(&mut self, key: usize) -> Option<T> {
        let slot = self.slots.get_mut(key)?;
        if !slot.occupied {
            return None;
        }
        slot.occupied = false;
        slot.next = self.next_free;
        self.next_free = Self::to_slot_index(key);
        self.len -= 1;
        // SAFETY: we just flipped `occupied` true → false, so this is the
        // sole read of the initialized value.
        Some(unsafe { slot.value.assume_init_read() })
    }

    pub fn remove_live(&mut self, key: LiveKey) -> T {
        let slot = &mut self.slots[key.0];
        debug_assert!(slot.occupied);
        slot.occupied = false;
        slot.next = self.next_free;
        self.next_free = Self::to_slot_index(key.0);
        self.len -= 1;
        // SAFETY: see `remove`; LiveKey is only issued for occupied slots.
        unsafe { slot.value.assume_init_read() }
    }
}

impl<T> Drop for Slab<T> {
    fn drop(&mut self) {
        for slot in &mut self.slots {
            if slot.occupied {
                // SAFETY: `occupied == true` ⇒ `value` is initialized.
                unsafe { slot.value.assume_init_drop() };
                slot.occupied = false;
            }
        }
    }
}