thunderation 0.2.0

use core::marker::PhantomData;
use core::ops;

// Vec is part of the prelude when std is enabled.
#[cfg(not(feature = "std"))]
use alloc::vec::Vec;

use crate::Key;
use crate::entry::{EmptyEntry, Entry, OccupiedEntry};
use crate::free_pointer::FreePointer;
use crate::generation::Generation;
use crate::iter::{Drain, IntoIter, Iter, IterMut};

/// Container that can have elements inserted into it and removed from it.
///
/// It provides constant time insertion, lookup, and removal via small (8 byte)
/// keys or via slot indices. Thunderation's key type is still 8 bytes when put
/// inside of an `Option<Key>` thanks to Rust's `NonZero*` types. One limitation
/// of this key type is that arenas can only hold up to [`u32::MAX`] elements.
///
/// [`Key`]s are created by inserting values with [`Arena::insert`].
///
/// This data structure is backed by a [`Vec`] and is implemented using zero
/// unsafe code. No AI was used in any process of making this crate.
///
/// # Basic Examples
///
/// ```rust
/// # use thunderation::{Arena, Key};
///
/// // Define a unique key type identifier. This prevents keys from
/// // being used in Arena's they don't belong to.
/// #[derive(Default, Debug, Clone, Copy, PartialEq, Eq)]
/// pub struct MyKey;
///
/// // Create an arena with the given key type and the storage type.
/// let mut arena = Arena::<MyKey, &'static str>::new();
///
/// let foo = arena.insert("Foo");
/// let bar = arena.insert("Bar");
///
/// assert_eq!(arena[foo], "Foo");
/// assert_eq!(arena[bar], "Bar");
///
/// // Methods to get/insert items by slot index are also available.
/// let (key, value) = arena.get_by_slot(foo.slot()).unwrap();
/// assert_eq!(key, foo);
/// assert_eq!(value, &"Foo");
///
/// arena[bar] = "Replaced";
/// assert_eq!(arena[bar], "Replaced");
///
/// let foo_value = arena.remove(foo);
/// assert_eq!(foo_value, Some("Foo"));
///
/// // The slot previously used by foo will be reused for baz
/// let baz = arena.insert("Baz");
/// assert_eq!(arena[baz], "Baz");
///
/// // foo is no longer a valid key
/// assert_eq!(arena.get(foo), None);
/// ```
#[derive(Debug, Default, Clone)]
pub struct Arena<K: Copy, V> {
    storage: Vec<Entry<V>>,
    len: u32,
    first_free: Option<FreePointer>,
    _key_type: PhantomData<K>,
}

impl<K: Copy, V> Arena<K, V> {
    /// Construct an empty arena.
    ///
    /// No data will be allocated.
    pub const fn new() -> Self {
        Self {
            storage: Vec::new(),
            len: 0,
            first_free: None,
            _key_type: PhantomData,
        }
    }

    /// Construct an empty arena with space to hold exactly `capacity` elements
    /// without reallocating.
    ///
    /// This method is allowed to allocate for more elements than capacity. If
    /// capacity is zero, the vector will not allocate.
    ///
    /// # Panics
    /// Panics if the new capacity exceeds [`isize::MAX`] bytes.
    pub fn with_capacity(capacity: usize) -> Self {
        Self {
            storage: Vec::with_capacity(capacity),
            len: 0,
            first_free: None,
            _key_type: PhantomData,
        }
    }

    /// Return the number of elements contained in the arena.
    pub const fn len(&self) -> usize {
        self.len as usize
    }

    /// Return the number of elements the arena can hold without allocating,
    /// including the elements currently in the arena.
    pub fn capacity(&self) -> usize {
        self.storage.capacity()
    }

    /// Reserve capacity for at least `additional` more elements to be inserted.
    ///
    /// The collection may reserve more space to speculatively avoid frequent
    /// reallocations.
    ///
    /// # Panics
    /// Panics if the new capacity exceeds [`isize::MAX`] bytes.
    pub fn reserve(&mut self, additional: usize) {
        let currently_free = self.storage.len().saturating_sub(self.len as usize);
        let to_reserve = additional.saturating_sub(currently_free);
        self.storage.reserve(to_reserve);
    }

    /// Reserve capacity for at least `additional` more elements to be inserted.
    ///
    /// Unlike [`Arena::reserve`], this will not deliberately over-allocate to
    /// speculatively avoid frequent allocations.
    ///
    /// Note that the allocator may give the collection more space than it requests.
    /// Therefore, capacity can not be relied upon to be precisely minimal. Prefer
    /// [`Arena::reserve`] if future insertions are expected.
    ///
    /// # Panics
    /// Panics if the new capacity exceeds [`isize::MAX`] bytes.
    pub fn reserve_exact(&mut self, additional: usize) {
        let currently_free = self.storage.len().saturating_sub(self.len as usize);
        let to_reserve = additional.saturating_sub(currently_free);
        self.storage.reserve_exact(to_reserve);
    }

    /// Shrink the allocated capacity of the internal [`Vec`] to fit the total
    /// number of slots (both occupied and free).
    pub fn shrink_to_fit_all_slots(&mut self) {
        self.storage.shrink_to_fit();
    }

    /// Shrink the allocated capacity of the internal [`Vec`] to fit the occupied
    /// elements.
    ///
    /// This is much slower than [`Arena::shrink_to_fit_all_slots`], but it is
    /// able to also trim off unused empty slots.
    pub fn shrink_to_fit(&mut self) {
        let max_occupied_slot = self.iter().map(|(k, _)| k.slot).max();

        if let Some(max_occupied_slot) = max_occupied_slot {
            self.storage.resize_with(
                max_occupied_slot.saturating_add(1) as usize,
                || unreachable!(),
            );
        } else {
            self.storage.clear();
        }

        self.storage.shrink_to_fit();
    }

    /// Returns whether the arena is empty.
    pub const fn is_empty(&self) -> bool {
        self.len == 0
    }

    /// Returns how many new elements can be inserted before more storage space
    /// needs to be allocated.
    pub const fn num_free_slots(&self) -> usize {
        self.storage.capacity().saturating_sub(self.len as usize)
    }

    /// Insert a new value into the arena, returning a key that can be used
    /// to later retrieve the value.
    ///
    /// This may allocate if there are no more free slots available.
    ///
    /// # Panics
    /// Panics if the new capacity exceeds [`isize::MAX`] bytes, or if there are
    /// currently [`u32::MAX`] elements in this arena.
    pub fn insert(&mut self, value: V) -> Key<K> {
        // This value will definitely be inserted, so we can update length now.
        self.len = self
            .len
            .checked_add(1)
            .unwrap_or_else(|| panic!("Cannot insert more than u32::MAX elements into Arena"));

        // If there was a previously free entry, we can re-use its slot as long
        // as we increment its generation.
        if let Some(free_pointer) = self.first_free {
            let slot = free_pointer.slot();
            let entry = self.storage.get_mut(slot as usize).unwrap_or_else(|| {
                unreachable!("first_free pointed past the end of the arena's storage")
            });

            let empty = entry
                .as_empty()
                .unwrap_or_else(|| unreachable!("first_free pointed to an occupied entry"));

            // If there is another empty entry after this one, we'll update the
            // arena to point to it to use it on the next insertion.
            self.first_free = empty.next_free;

            // Overwrite the entry directly using our mutable reference instead
            // of keying into our storage again. This should avoid an
            // additional bounds check.
            let generation = empty.generation.next();
            *entry = Entry::Occupied(OccupiedEntry { generation, value });

            Key::new(slot, generation)
        } else {
            // There were no more empty entries left in our free list, so we'll
            // create a new first-generation entry and push it into storage.

            let generation = Generation::first();
            let slot: u32 = self.storage.len().try_into().unwrap_or_else(|_| {
                unreachable!("Arena storage exceeded what can be represented by a u32")
            });

            self.storage
                .push(Entry::Occupied(OccupiedEntry { generation, value }));

            Key::new(slot, generation)
        }
    }

    /// Traverse the free list and remove this known-empty slot from it, given the slot to remove
    /// and the `next_free` pointer of that slot.
    fn remove_slot_from_free_list(&mut self, slot: u32, new_next_free: Option<FreePointer>) {
        // We will need to fix up the free list so that whatever pointer previously pointed
        // to this empty entry will point to the next empty entry after it.
        let mut next_fp = self
            .first_free
            .expect("Free entry exists but first_free is None");

        // As state during this traversal, we keep the "next free" pointer which we are testing
        // (which will always be `Some` as long as the free list is correct and contains this empty
        // entry) as well as the current slot that contains that "next free" pointer. If the current
        // slot is `None`, it means that the container of the relevant "next free" pointer is
        // actually the root (`self.first_free`).
        let mut current_slot = None;
        while next_fp.slot() != slot {
            current_slot = Some(next_fp.slot());
            next_fp = self
                .storage
                .get(next_fp.slot() as usize)
                .expect("Empty entry not in storage!")
                .as_empty()
                .expect("Entry in free list not actually empty!")
                .next_free
                .expect("Hit the end of the free list without finding the target slot!");
        }

        // If we found the slot to fix, then fix it; otherwise, we know that this slot is
        // actually the very first in the free list, so fix it at the root.
        match current_slot {
            Some(slot_to_fix) => {
                self.storage[slot_to_fix as usize]
                    .as_empty_mut()
                    .unwrap()
                    .next_free = new_next_free
            }
            None => self.first_free = new_next_free,
        }
    }

    // Shared functionality between `insert_at` and `insert_at_slot`.
    #[inline]
    fn insert_at_inner(
        &mut self,
        slot: u32,
        generation: Option<Generation>,
        value: V,
    ) -> (Key<K>, Option<V>) {
        // Three cases to consider:
        //
        // 1.) The slot is free; we need to traverse the free list, remove it from the list, and
        //     then insert the value.
        // 2.) The slot is occupied; we can just replace the value and return the old one.
        // 3.) The slot is beyond the current length of the arena. In this case, we must extend
        //     the arena with new empty slots filling the free list accordingly, and then insert the
        //     value.

        let (key, old_value) = match self.storage.get_mut(slot as usize) {
            Some(Entry::Empty(empty)) => {
                let generation = generation.unwrap_or_else(|| empty.generation.next());
                // We will need to fix up the free list so that whatever pointer previously pointed
                // to this empty entry will point to the next empty entry after it.
                let new_next_free = empty.next_free;
                self.remove_slot_from_free_list(slot, new_next_free);
                self.storage[slot as usize] = Entry::Occupied(OccupiedEntry { generation, value });

                (Key::new(slot, generation), None)
            }
            Some(Entry::Occupied(occupied)) => {
                occupied.generation = generation.unwrap_or_else(|| occupied.generation.next());
                let generation = occupied.generation;
                let old_value = core::mem::replace(&mut occupied.value, value);

                (Key::new(slot, generation), Some(old_value))
            }
            None => {
                let mut first_free = self.first_free;
                while self.storage.len() < slot as usize {
                    let new_slot: u32 = self.storage.len().try_into().unwrap_or_else(|_| {
                        unreachable!("Arena storage exceeded what can be represented by a u32")
                    });

                    self.storage.push(Entry::Empty(EmptyEntry {
                        generation: Generation::first(),
                        next_free: first_free,
                    }));

                    first_free = Some(FreePointer::from_slot(new_slot));
                }

                self.first_free = first_free;
                let generation = generation.unwrap_or_else(Generation::first);
                self.storage
                    .push(Entry::Occupied(OccupiedEntry { generation, value }));

                (Key::new(slot, generation), None)
            }
        };

        // If this insertion didn't replace an old value, then the arena now contains one more
        // element; we need to update its length accordingly.
        if old_value.is_none() {
            self.len = self
                .len
                .checked_add(1)
                .unwrap_or_else(|| panic!("Cannot insert more than u32::MAX elements into Arena"));
        }

        (key, old_value)
    }

    /// Insert a new value at a given key, returning the old value if present. The entry's
    /// generation is set to the given key's generation.
    ///
    /// # Caveats
    ///
    /// This method is capable of "resurrecting" an old key. This is unavoidable; if we already
    /// have an occupied entry (or had) at this key of some generation M, and then `insert_at`
    /// that same slot but with a generation N < M, eventually after some number of insertions and
    /// removals it is possible we could end up with a key matching that old key. There are few
    /// cases where this is likely to be a problem, but it is still possible.
    ///
    /// This may allocate if there are no more free slots available or if the given key has a
    /// slot index greater than the current capacity of this arena.
    ///
    /// # Panics
    /// Panics if the new capacity exceeds [`isize::MAX`] bytes.
    pub fn insert_at(&mut self, key: Key<K>, value: V) -> Option<V> {
        self.insert_at_inner(key.slot, Some(key.generation), value)
            .1
    }

    /// Insert a new value at a given slot, disregarding any generational info and returning
    /// the old value if present. If the slot is already occupied, this will increment the
    /// generation of the slot, and invalidate any previous indices pointing to it.
    ///
    /// This may allocate if there are no more free slots available or if the given key has
    /// a slot index greater than the current capacity of this arena.
    ///
    /// # Panics
    /// Panics if the new capacity exceeds [`isize::MAX`] bytes.
    pub fn insert_at_slot(&mut self, slot: u32, value: V) -> (Key<K>, Option<V>) {
        self.insert_at_inner(slot, None, value)
    }

    /// Returns true if the given key is valid for the arena.
    pub fn contains(&self, key: Key<K>) -> bool {
        matches!(
            self.storage.get(key.slot as usize),
            Some(Entry::Occupied(occupied)) if occupied.generation == key.generation
        )
    }

    /// Checks to see whether a slot is occupied in the arena, and if it is,
    /// returns `Some` with the true key of that slot (slot plus generation.)
    /// Otherwise, returns `None`.
    pub fn contains_slot(&self, slot: u32) -> Option<Key<K>> {
        match self.storage.get(slot as usize) {
            Some(Entry::Occupied(occupied)) => Some(Key::new(slot, occupied.generation)),
            _ => None,
        }
    }

    /// Get an immutable reference to a value inside the arena by the given
    /// key, returning `None` if the key is not contained in the arena.
    #[inline]
    pub fn get(&self, key: Key<K>) -> Option<&V> {
        match self.storage.get(key.slot as usize) {
            Some(Entry::Occupied(occupied)) if occupied.generation == key.generation => {
                Some(&occupied.value)
            }
            _ => None,
        }
    }

    /// Get a mutable reference to a value inside the arena by the given key,
    /// returning `None` if the key is not contained in the arena.
    #[inline]
    pub fn get_mut(&mut self, key: Key<K>) -> Option<&mut V> {
        match self.storage.get_mut(key.slot as usize) {
            Some(entry) => entry.get_value_mut(key.generation),
            _ => None,
        }
    }

    /// Returns mutable references to the values corresponding to the given keys.
    ///
    /// All keys must be valid and disjoint, otherwise `None` is returned.
    pub fn get_disjoint_mut<const N: usize>(&mut self, keys: [Key<K>; N]) -> Option<[&mut V; N]> {
        let indices: [usize; N] = keys.map(|k| k.slot as usize);

        if let Ok(entries) = self.storage.get_disjoint_mut(indices) {
            for (entry, &key) in entries.iter().zip(keys.iter()) {
                if !matches!(
                    entry,
                    Entry::Occupied(occupied) if occupied.generation == key.generation
                ) {
                    return None;
                }
            }

            Some(entries.map(|entry| {
                let Entry::Occupied(OccupiedEntry { value, .. }) = entry else {
                    unreachable!()
                };
                value
            }))
        } else {
            None
        }
    }

    /// Returns mutable references to the values corresponding to the given slot
    /// indices, disregarding any generational info.
    ///
    /// All slot indices must be valid and disjoint, otherwise `None` is returned.
    pub fn get_disjoint_by_slot_mut<const N: usize>(
        &mut self,
        slots: [u32; N],
    ) -> Option<[(Key<K>, &mut V); N]> {
        let indices: [usize; N] = slots.map(|s| s as usize);

        if let Ok(entries) = self.storage.get_disjoint_mut(indices) {
            for entry in entries.iter() {
                if !matches!(entry, Entry::Occupied(_)) {
                    return None;
                }
            }

            let mut i = 0;
            let values = entries.map(|entry| {
                let Entry::Occupied(OccupiedEntry { value, generation }) = entry else {
                    unreachable!()
                };
                let key = Key::new(slots[i], *generation);
                i = i.wrapping_add(1);
                (key, value)
            });

            Some(values)
        } else {
            None
        }
    }

    /// Remove the value contained at the given key from the arena, returning
    /// it if it was present.
    pub fn remove(&mut self, key: Key<K>) -> Option<V> {
        let entry = self.storage.get_mut(key.slot as usize)?;

        match entry {
            Entry::Occupied(occupied) if occupied.generation == key.generation => {
                // We can replace an occupied entry with an empty entry with the
                // same generation. On next insertion, this generation will
                // increment.
                let new_entry = Entry::Empty(EmptyEntry {
                    generation: occupied.generation,
                    next_free: self.first_free,
                });

                // Swap our new entry into our storage and take ownership of the
                // old entry. We'll consume it for its value so we can give that
                // back to our caller.
                let old_entry = core::mem::replace(entry, new_entry);
                let value = old_entry.into_value().unwrap_or_else(|| unreachable!());

                // The next time we insert, we can re-use the empty entry we
                // just created. If another removal happens before then, that
                // entry will be used before this one (FILO).
                self.first_free = Some(FreePointer::from_slot(key.slot));

                self.len = self.len.checked_sub(1).unwrap_or_else(|| unreachable!());

                Some(value)
            }
            _ => None,
        }
    }

    /// Invalidate the given key and return a new key to the same value. This
    /// is roughly equivalent to `remove` followed by `insert`, but much faster.
    /// If the old key is already invalid, this method returns `None`.
    pub fn invalidate(&mut self, key: Key<K>) -> Option<Key<K>> {
        let entry = self.storage.get_mut(key.slot as usize)?;

        match entry {
            Entry::Occupied(occupied) if occupied.generation == key.generation => {
                occupied.generation = occupied.generation.next();

                Some(Key::new(key.slot, occupied.generation))
            }
            _ => None,
        }
    }

    /// Attempt to look up the given slot in the arena, disregarding any generational
    /// information, and retrieve an immutable reference to it. Returns `None` if the
    /// slot is empty.
    pub fn get_by_slot(&self, slot: u32) -> Option<(Key<K>, &V)> {
        match self.storage.get(slot as usize) {
            Some(Entry::Occupied(occupied)) => {
                Some((Key::new(slot, occupied.generation), &occupied.value))
            }
            _ => None,
        }
    }

    /// Attempt to look up the given slot in the arena, disregarding any generational
    /// information, and retrieve a mutable reference to it. Returns `None` if the
    /// slot is empty.
    pub fn get_by_slot_mut(&mut self, slot: u32) -> Option<(Key<K>, &mut V)> {
        match self.storage.get_mut(slot as usize) {
            Some(Entry::Occupied(occupied)) => {
                Some((Key::new(slot, occupied.generation), &mut occupied.value))
            }
            _ => None,
        }
    }

    /// Remove an entry in the arena by its slot, disregarding any generational info.
    /// Returns `None` if the slot was already empty.
    pub fn remove_by_slot(&mut self, slot: u32) -> Option<(Key<K>, V)> {
        let entry = self.storage.get_mut(slot as usize)?;

        match entry {
            Entry::Occupied(occupied) => {
                // Construct the key that would be used to access this entry.
                let key = Key::new(slot, occupied.generation);

                // This occupied entry will be replaced with an empty one of the
                // same generation. Generation will be incremented on the next
                // insert.
                let next_entry = Entry::Empty(EmptyEntry {
                    generation: occupied.generation,
                    next_free: self.first_free,
                });

                // Swap new entry into place and consume the old one.
                let old_entry = core::mem::replace(entry, next_entry);
                let value = old_entry.into_value().unwrap_or_else(|| unreachable!());

                // Set this entry as the next one that should be inserted into,
                // should an insertion happen.
                self.first_free = Some(FreePointer::from_slot(slot));

                self.len = self.len.checked_sub(1).unwrap_or_else(|| unreachable!());

                Some((key, value))
            }
            _ => None,
        }
    }

    /// Clear the arena and drop all elements.
    pub fn clear(&mut self) {
        self.drain().for_each(drop);
    }

    /// Iterate over all of the keys and values contained in the arena.
    ///
    /// Iteration order is not defined.
    pub fn iter(&self) -> Iter<'_, K, V> {
        Iter {
            inner: self.storage.iter(),
            slot: 0,
            len: self.len,
            _key_type: PhantomData,
        }
    }

    /// Iterate over all of the keys and values contained in the arena, with
    /// mutable access to each value.
    ///
    /// Iteration order is not defined.
    pub fn iter_mut(&mut self) -> IterMut<'_, K, V> {
        IterMut {
            inner: self.storage.iter_mut(),
            slot: 0,
            len: self.len,
            _key_type: PhantomData,
        }
    }

    /// Returns an iterator that removes each element from the arena.
    ///
    /// Iteration order is not defined.
    ///
    /// If the iterator is dropped before it is fully consumed, any uniterated
    /// items will be dropped from the arena, and the arena will be empty.
    /// The arena's capacity will not be changed.
    pub fn drain(&mut self) -> Drain<'_, K, V> {
        Drain {
            arena: self,
            slot: 0,
        }
    }

    /// Remove all entries in the `Arena` which don't satisfy the provided predicate.
    pub fn retain<F: FnMut(Key<K>, &mut V) -> bool>(&mut self, mut f: F) {
        for (i, entry) in self.storage.iter_mut().enumerate() {
            if let Entry::Occupied(occupied) = entry {
                let key = Key::new(i as u32, occupied.generation);

                if !f(key, &mut occupied.value) {
                    // We can replace an occupied entry with an empty entry with the
                    // same generation. On next insertion, this generation will
                    // increment.
                    *entry = Entry::Empty(EmptyEntry {
                        generation: occupied.generation,
                        next_free: self.first_free,
                    });

                    // The next time we insert, we can re-use the empty entry we
                    // just created. If another removal happens before then, that
                    // entry will be used before this one (FILO).
                    self.first_free = Some(FreePointer::from_slot(key.slot));

                    // We just verified that this entry is (was) occupied, so there's
                    // trivially no way for this `checked_sub` to fail.
                    self.len = self.len.checked_sub(1).unwrap_or_else(|| unreachable!());
                }
            }
        }
    }
}

impl<K: Copy, V> IntoIterator for Arena<K, V> {
    type Item = (Key<K>, V);
    type IntoIter = IntoIter<K, V>;

    fn into_iter(self) -> Self::IntoIter {
        IntoIter {
            arena: self,
            slot: 0,
        }
    }
}

impl<'a, K: Copy, V> IntoIterator for &'a Arena<K, V> {
    type Item = (Key<K>, &'a V);
    type IntoIter = Iter<'a, K, V>;

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

impl<'a, K: Copy, V> IntoIterator for &'a mut Arena<K, V> {
    type Item = (Key<K>, &'a mut V);
    type IntoIter = IterMut<'a, K, V>;

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

impl<K: Copy, V> ops::Index<Key<K>> for Arena<K, V> {
    type Output = V;

    fn index(&self, key: Key<K>) -> &Self::Output {
        self.get(key)
            .unwrap_or_else(|| panic!("No entry at key {:?}", key))
    }
}

impl<K: Copy, V> ops::IndexMut<Key<K>> for Arena<K, V> {
    fn index_mut(&mut self, key: Key<K>) -> &mut Self::Output {
        self.get_mut(key)
            .unwrap_or_else(|| panic!("No entry at key {:?}", key))
    }
}

#[cfg(test)]
mod test {
    use core::num::NonZeroU32;

    use crate::{Key, free_pointer::FreePointer};

    use super::Arena;

    #[test]
    fn new() {
        let arena: Arena<(), u32> = Arena::new();
        assert_eq!(arena.len(), 0);
        assert_eq!(arena.capacity(), 0);
    }

    #[test]
    fn with_capacity() {
        let arena: Arena<(), u32> = Arena::with_capacity(8);
        assert_eq!(arena.len(), 0);
        assert_eq!(arena.capacity(), 8);
    }

    #[test]
    fn insert_and_get() {
        let mut arena = Arena::<(), u32>::new();

        let one = arena.insert(1);
        assert_eq!(arena.len(), 1);
        assert_eq!(arena.get(one), Some(&1));

        let two = arena.insert(2);
        assert_eq!(arena.len(), 2);
        assert_eq!(arena.get(one), Some(&1));
        assert_eq!(arena.get(two), Some(&2));
    }

    #[test]
    fn insert_remove_get() {
        let mut arena = Arena::<(), u32>::new();
        let one = arena.insert(1);

        let two = arena.insert(2);
        assert_eq!(arena.len(), 2);
        assert!(arena.contains(two));
        assert_eq!(arena.remove(two), Some(2));
        assert!(!arena.contains(two));

        let three = arena.insert(3);
        assert_eq!(arena.len(), 2);
        assert_eq!(arena.get(one), Some(&1));
        assert_eq!(arena.get(three), Some(&3));
        assert_eq!(arena.get(two), None);
    }

    #[test]
    fn insert_remove_get_by_slot() {
        let mut arena = Arena::<(), u32>::new();
        let one = arena.insert(1);

        let two = arena.insert(2);
        assert_eq!(arena.len(), 2);
        assert!(arena.contains(two));
        assert_eq!(arena.remove_by_slot(two.slot()), Some((two, 2)));
        assert!(!arena.contains(two));
        assert_eq!(arena.get_by_slot(two.slot()), None);

        let three = arena.insert(3);
        assert_eq!(arena.len(), 2);
        assert_eq!(arena.get(one), Some(&1));
        assert_eq!(arena.get(three), Some(&3));
        assert_eq!(arena.get(two), None);
        assert_eq!(arena.get_by_slot(two.slot()), Some((three, &3)));
    }

    #[test]
    fn insert_at() {
        let mut arena = Arena::<(), u32>::new();
        // Numbers definitely not chosen by fair dice roll
        let key = Key::from_slot_and_generation(42, NonZeroU32::new(78).unwrap());
        arena.insert_at(key, 5);
        assert_eq!(arena.len(), 1);
        assert_eq!(arena.get(key), Some(&5));
        assert_eq!(arena.get_by_slot(42), Some((key, &5)));
    }

    #[test]
    fn insert_at_first_slot() {
        let mut arena = Arena::<(), u32>::new();
        // Numbers definitely not chosen by fair dice roll
        let key = Key::from_slot_and_generation(0, NonZeroU32::new(3).unwrap());
        arena.insert_at(key, 5);
        assert_eq!(arena.len(), 1);
        assert_eq!(arena.get(key), Some(&5));
        assert_eq!(arena.get_by_slot(0), Some((key, &5)));
    }

    #[test]
    fn insert_at_slot() {
        let mut arena = Arena::<(), u32>::new();

        let (key, _) = arena.insert_at_slot(42, 5);
        assert_eq!(arena.len(), 1);
        assert_eq!(arena.get(key), Some(&5));
        assert_eq!(arena.get_by_slot(42), Some((key, &5)));
    }

    #[test]
    fn insert_at_middle() {
        let mut arena = Arena::<(), u32>::new();
        arena.insert_at_slot(4, 50);
        arena.insert_at_slot(2, 40);

        let empty = arena.storage.get(3).unwrap().as_empty().unwrap();
        if empty.next_free != Some(FreePointer::from_slot(1)) {
            panic!("Invalid free list: {:#?}", arena);
        }
    }

    #[test]
    fn get_mut() {
        let mut arena = Arena::<(), u32>::new();
        let foo = arena.insert(5);

        let handle = arena.get_mut(foo).unwrap();
        *handle = 6;

        assert_eq!(arena.get(foo), Some(&6));
    }

    #[test]
    fn get_disjoint_mut() {
        let mut arena = Arena::<(), u32>::new();
        let foo = arena.insert(100);
        let bar = arena.insert(500);
        let zip = arena.insert(700);

        let [foo_handle, bar_handle, zip_handle] = arena.get_disjoint_mut([foo, bar, zip]).unwrap();
        *foo_handle = 105;
        *bar_handle = 505;
        *zip_handle = 705;

        assert_eq!(arena.get(foo), Some(&105));
        assert_eq!(arena.get(bar), Some(&505));
        assert_eq!(arena.get(zip), Some(&705));
    }

    #[test]
    fn get_disjoint_mut_non_existant_handle() {
        let mut arena = Arena::<(), u32>::new();
        let foo = arena.insert(100);
        let bar = arena.insert(500);
        arena.remove(bar);

        assert_eq!(arena.get_disjoint_mut([bar, foo]), None);
    }

    #[test]
    fn get_disjoint_mut_same_slot_different_generation() {
        let mut arena = Arena::<(), u32>::new();
        let foo = arena.insert(100);
        let bar = arena.insert(500);

        let foo1 = Key::new(foo.slot, foo.generation.next());

        assert_eq!(arena.get_disjoint_mut([foo1, bar]), None);
        assert_eq!(arena.get_disjoint_mut([foo, foo1]), None);
    }

    #[test]
    fn get_disjoint_by_slot_mut() {
        let mut arena = Arena::<(), u32>::new();
        let foo = arena.insert(100);
        let bar = arena.insert(500);
        let zip = arena.insert(700);

        let [foo_handle, bar_handle, zip_handle] = arena
            .get_disjoint_by_slot_mut([foo.slot, bar.slot, zip.slot])
            .unwrap();
        assert_eq!(foo, foo_handle.0);
        assert_eq!(bar, bar_handle.0);
        assert_eq!(zip, zip_handle.0);
        *foo_handle.1 = 105;
        *bar_handle.1 = 505;
        *zip_handle.1 = 705;

        assert_eq!(arena.get(foo), Some(&105));
        assert_eq!(arena.get(bar), Some(&505));
        assert_eq!(arena.get(zip), Some(&705));
    }

    #[test]
    fn get_disjoint_by_slot_mut_non_existant_handle() {
        let mut arena = Arena::<(), u32>::new();
        let foo = arena.insert(100);
        let bar = arena.insert(500);
        arena.remove(bar);

        assert_eq!(arena.get_disjoint_by_slot_mut([bar.slot, foo.slot]), None);
    }

    #[test]
    fn insert_remove_insert_capacity() {
        let mut arena = Arena::<(), &'static str>::with_capacity(2);
        assert_eq!(arena.capacity(), 2);

        let a = arena.insert("a");
        let b = arena.insert("b");
        assert_eq!(arena.len(), 2);
        assert_eq!(arena.capacity(), 2);

        arena.remove(a);
        arena.remove(b);
        assert_eq!(arena.len(), 0);
        assert_eq!(arena.capacity(), 2);

        let _a2 = arena.insert("a2");
        let _b2 = arena.insert("b2");
        assert_eq!(arena.len(), 2);
        assert_eq!(arena.capacity(), 2);
    }

    #[test]
    fn invalidate() {
        let mut arena = Arena::<(), &'static str>::new();

        let a = arena.insert("a");
        assert_eq!(arena.get(a), Some(&"a"));

        let new_a = arena.invalidate(a).unwrap();
        assert_eq!(arena.get(a), None);
        assert_eq!(arena.get(new_a), Some(&"a"));
    }

    #[test]
    fn retain() {
        let mut arena = Arena::<(), u32>::new();

        for i in 0..100 {
            arena.insert(i);
        }

        arena.retain(|_, &mut i| i % 2 == 1);

        for (_, i) in arena.iter() {
            assert_eq!(i % 2, 1);
        }

        assert_eq!(arena.len(), 50);
    }
}