small_slot_map 0.3.1

use std::{
    fmt::Debug,
    hash::Hash,
    marker::PhantomData,
    ops::{Index, IndexMut},
};

#[cfg(test)]
mod tests {
    use crate::{DefaultU8Key, InvalidateEmptySlots, ReuseEmptySlots, SmallSlotMap};

    #[test]
    fn slot_reuse_works() {
        let mut reuses_slots = SmallSlotMap::<DefaultU8Key, i32, ReuseEmptySlots>::new();

        for _ in 0..u8::MAX {
            let first_key = reuses_slots.insert(0);
            reuses_slots.remove(first_key);

            for _ in 0..10 {
                let key = reuses_slots.insert(0);
                assert_eq!(key, first_key);
                reuses_slots.remove(key);
            }

            reuses_slots.insert(0);
        }

        assert!(reuses_slots.try_insert(0).is_none());
    }

    #[test]
    fn slot_invalidation_works() {
        let mut invalidates_slots = SmallSlotMap::<DefaultU8Key, i32, InvalidateEmptySlots>::new();

        let mut used_keys = Vec::new();

        for _ in 0..u8::MAX {
            let key = invalidates_slots.insert(0);
            invalidates_slots.remove(key);

            used_keys.push(key);
        }

        for used_key in used_keys {
            assert!(invalidates_slots.get(used_key).is_none());
        }

        assert!(invalidates_slots.try_insert(0).is_none());
    }
}

#[derive(Clone, Debug)]
pub struct SmallSlotMap<K, V, S>
where
    K: SmallKey,
    S: SlotBehavior,
{
    data: Vec<Option<V>>,
    first_free: usize,
    len: usize,
    _phantom: PhantomData<(K, S)>,
}

pub trait SmallKey: Sized + Copy + Eq + Ord + Hash + Debug {
    fn try_from_usize(value: usize) -> Option<Self>;

    fn into_usize(self) -> usize;
}

mod private {
    pub trait Sealed {}
}

pub trait SlotBehavior: private::Sealed {
    const DECREASE_FIRST_FREE: bool;
}

#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
pub struct ReuseEmptySlots;

impl private::Sealed for ReuseEmptySlots {}

impl SlotBehavior for ReuseEmptySlots {
    const DECREASE_FIRST_FREE: bool = true;
}

#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
pub struct InvalidateEmptySlots;

impl private::Sealed for InvalidateEmptySlots {}

impl SlotBehavior for InvalidateEmptySlots {
    const DECREASE_FIRST_FREE: bool = false;
}

impl<K, V, S> Default for SmallSlotMap<K, V, S>
where
    K: SmallKey,
    S: SlotBehavior,
{
    fn default() -> Self {
        Self::new()
    }
}

impl<K, V, S> SmallSlotMap<K, V, S>
where
    K: SmallKey,
    S: SlotBehavior,
{
    pub fn new() -> Self {
        Self {
            data: Vec::new(),
            first_free: 0,
            len: 0,
            _phantom: PhantomData,
        }
    }

    pub fn len(&self) -> usize {
        self.len
    }

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

    pub fn insert(&mut self, value: V) -> K {
        self.try_insert(value).expect("Too many elements")
    }

    pub fn try_insert(&mut self, value: V) -> Option<K> {
        assert!(self.first_free <= self.data.len());
        if self.first_free == self.data.len() {
            self.data.push(None);
        }

        let key = K::try_from_usize(self.first_free)?;
        self.data[self.first_free] = Some(value);

        while let Some(Some(_)) = self.data.get(self.first_free) {
            self.first_free += 1;
        }

        self.len += 1;

        Some(key)
    }

    pub fn get(&self, key: K) -> Option<&V> {
        self.data.get(key.into_usize())?.as_ref()
    }

    pub fn get_mut(&mut self, key: K) -> Option<&mut V> {
        self.data.get_mut(key.into_usize())?.as_mut()
    }

    /// Indexing with this key after this operation will result in unpredictable behavior if this
    /// `SmallMap` is set to `ReuseEmptySlots`.
    pub fn remove(&mut self, key: K) -> Option<V> {
        let index = key.into_usize();

        let value = self.data.get_mut(index)?.take();

        if value.is_some() {
            if S::DECREASE_FIRST_FREE {
                self.first_free = self.first_free.min(index);
            }
            self.len -= 1;
        }

        value
    }

    pub fn clear(&mut self) {
        if S::DECREASE_FIRST_FREE {
            self.data.clear();
            self.len = 0;
            self.first_free = 0;
        } else {
            for slot in &mut self.data {
                *slot = None;
            }
        }
    }
}

impl<K, V, S> Index<K> for SmallSlotMap<K, V, S>
where
    K: SmallKey,
    S: SlotBehavior,
{
    type Output = V;

    fn index(&self, index: K) -> &Self::Output {
        self.get(index).unwrap()
    }
}

impl<K, V, S> IndexMut<K> for SmallSlotMap<K, V, S>
where
    K: SmallKey,
    S: SlotBehavior,
{
    fn index_mut(&mut self, index: K) -> &mut Self::Output {
        self.get_mut(index).unwrap()
    }
}

pub struct IntoIter<K: SmallKey, V> {
    inner: Vec<Option<V>>,
    _phantom: PhantomData<K>,
}

impl<K, V, S> IntoIterator for SmallSlotMap<K, V, S>
where
    K: SmallKey,
    S: SlotBehavior,
{
    type Item = V;

    type IntoIter = IntoIter<K, V>;

    fn into_iter(self) -> Self::IntoIter {
        IntoIter {
            inner: self.data,
            _phantom: PhantomData,
        }
    }
}

impl<K: SmallKey, V> Iterator for IntoIter<K, V> {
    type Item = V;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            if let Some(element) = self.inner.pop()? {
                return Some(element);
            }
        }
    }
}

pub struct Iter<'a, K: SmallKey, V> {
    inner: &'a [Option<V>],
    _phantom: PhantomData<K>,
}

impl<K, V, S> SmallSlotMap<K, V, S>
where
    K: SmallKey,
    S: SlotBehavior,
{
    pub fn iter(&self) -> Iter<'_, K, V> {
        Iter {
            inner: &self.data,
            _phantom: PhantomData,
        }
    }
}

impl<'a, K, V, S> IntoIterator for &'a SmallSlotMap<K, V, S>
where
    K: SmallKey,
    S: SlotBehavior,
{
    type Item = (K, &'a V);

    type IntoIter = Iter<'a, K, V>;

    fn into_iter(self) -> Self::IntoIter {
        self.iter()
    }
}

impl<'a, K: SmallKey, V> Iterator for Iter<'a, K, V> {
    type Item = (K, &'a V);

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            if let Some(element) = self.inner.split_off_last()? {
                let key = K::try_from_usize(self.inner.len()).unwrap();
                return Some((key, element));
            }
        }
    }
}

pub struct IterMut<'a, K: SmallKey, V> {
    inner: &'a mut [Option<V>],
    _phantom: PhantomData<K>,
}

impl<K, V, S> SmallSlotMap<K, V, S>
where
    K: SmallKey,
    S: SlotBehavior,
{
    pub fn iter_mut(&mut self) -> IterMut<'_, K, V> {
        IterMut {
            inner: &mut self.data,
            _phantom: PhantomData,
        }
    }
}

impl<'a, K, V, S> IntoIterator for &'a mut SmallSlotMap<K, V, S>
where
    K: SmallKey,
    S: SlotBehavior,
{
    type Item = (K, &'a mut V);

    type IntoIter = IterMut<'a, K, V>;

    fn into_iter(self) -> Self::IntoIter {
        self.iter_mut()
    }
}

impl<'a, K: SmallKey, V> Iterator for IterMut<'a, K, V> {
    type Item = (K, &'a mut V);

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            if let Some(element) = self.inner.split_off_last_mut()? {
                let key = K::try_from_usize(self.inner.len()).unwrap();
                return Some((key, element));
            }
        }
    }
}

pub struct Keys<'a, K: SmallKey, V> {
    inner: &'a [Option<V>],
    _phantom: PhantomData<K>,
}

impl<K, V, S> SmallSlotMap<K, V, S>
where
    K: SmallKey,
    S: SlotBehavior,
{
    pub fn keys(&self) -> Keys<'_, K, V> {
        Keys {
            inner: &self.data,
            _phantom: PhantomData,
        }
    }
}

impl<'a, K: SmallKey, V> Iterator for Keys<'a, K, V> {
    type Item = K;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            if let Some(_) = self.inner.split_off_last()? {
                let key = K::try_from_usize(self.inner.len()).unwrap();
                return Some(key);
            }
        }
    }
}

/// CREDIT: Adapted from slotmap's `new_key_type`.
#[macro_export]
macro_rules! new_small_key_type {
    ( $(#[$outer:meta])* $vis:vis struct $name:ident($inner:ty); $($rest:tt)* ) => {
        $(#[$outer])*
        #[derive(Copy, Clone,
                 Eq, PartialEq, Ord, PartialOrd,
                 Hash, Debug)]
        #[repr(transparent)]
        $vis struct $name(std::num::NonZero<$inner>);

        // Make it a bit harder to accidentally misuse the macro
        const _: () = assert!(<$inner>::BITS <= u32::BITS);

        impl $crate::SmallKey for $name {
            fn try_from_usize(value: usize) -> Option<Self> {
                // If adding 1 overflows, it will result in 0, which will return an error
                Some(Self(std::num::NonZero::new(<$inner>::try_from(value).ok()?.wrapping_add(1))?))
            }

            fn into_usize(self) -> usize {
                (self.0.get() - 1) as usize
            }
        }

        $crate::new_small_key_type!($($rest)*);
    };

    () => {}
}

new_small_key_type! {
    pub struct DefaultU8Key(u8);
    pub struct DefaultU16Key(u16);
    pub struct DefaultU32Key(u32);
}