inttable/

lib.rs

//! This crate serves as a replacement for `intmap` in cases where the keys you are working with
//! are already hashed. This is essentially just a wrapper around a `hashbrown::HashTable<V>` (with
//! the added detail that the keys are stored alongside the values).
//!
//! This is significantly faster than `intmap` or the built-in [`std::collections::HashMap`] type,
//! as it performs _no_ hashing. Thus, if your keys are not randomly distributed, performance
//! can take a nosedive.
//!
//! And yes, I realise this crate is basically a "leftpad", but it can potentially save a tonne of
//! boilerplate. I had to write it anyway since, so I might as well clean it up and publish it.
//!
//! A lot of surrounding code (i.e., tests and benchmarks) was taken as-is from `intmap`.
#[cfg(feature = "serde")]
mod serde;

use std::ops::Index;

pub use hashbrown::HashTable;
use hashbrown::TryReserveError;

/// A Map associating unique integers and randomly distributed integers to values.
///
/// Tries to match [`std::collections::HashMap`]'s API, but leaves out functions that either don't
/// apply (anything to do with hashers), or that just don't make sense (like
/// [`std::collections::HashMap::get_key_value`], as the key can just be copied).
///
/// This is also the reason why this doesn't include any code examples. In fact, I could leave this
/// entire crate undocumented and just tell you to look at the built-in hashmap.
pub struct IntTable<V> {
    /// The hashtable actually working
    inner: HashTable<(u64, V)>,
}

impl<V> IntTable<V> {
    /// Creates a new, zero-capacity hashtable that will grow with insertions.
    pub const fn new() -> Self {
        Self {
            inner: HashTable::new(),
        }
    }

    /// Creates a new IntTable with the given capacity.
    pub fn with_capacity(capacity: usize) -> Self {
        Self {
            inner: HashTable::with_capacity(capacity),
        }
    }

    /// Removes the given key from the table, returning its associated value.
    pub fn remove(&mut self, key: u64) -> Option<V> {
        self.inner
            .find_entry(key, eq(key))
            .ok()
            .map(|v| v.remove().0 .1)
    }

    /// Shrinks the capacity of the table with a lower limit. It will drop down no lower than the
    /// supplied limit.
    pub fn shrink_to(&mut self, min_capacity: usize) {
        self.inner.shrink_to(min_capacity, hasher)
    }

    /// Shrinks the capacity of the underlying table to the number of elements it contains.
    pub fn shrink_to_fit(&mut self) {
        self.inner.shrink_to_fit(hasher)
    }

    /// Reserves at least `additional` spots in the backing allocation. This might allocate more.
    /// You know the drill.
    pub fn reserve(&mut self, additional: usize) {
        self.inner.reserve(additional, hasher)
    }

    /// Returns the map's capacity
    pub fn capacity(&self) -> usize {
        self.inner.capacity()
    }

    /// Removes all elements from self
    pub fn clear(&mut self) {
        self.inner.clear()
    }

    /// Insertns the given key-value pair into the table, returning the old value at `key` if there
    /// was one.
    pub fn insert(&mut self, key: u64, value: V) -> Option<V> {
        use hashbrown::hash_table::Entry;
        match self.inner.entry(key, eq(key), hasher) {
            Entry::Occupied(o) => {
                let ((_, old_val), empty) = o.remove();
                empty.insert((key, value));
                Some(old_val)
            }
            Entry::Vacant(v) => {
                v.insert((key, value));
                None
            }
        }
    }

    /// Attempts to insert `value` at `key`, but if `key` was occupied, this function doesn't
    /// replace the old value.
    pub fn try_insert(&mut self, key: u64, value: V) -> Result<&mut V, V> {
        use hashbrown::hash_table::Entry;
        match self.inner.entry(key, eq(key), hasher) {
            Entry::Occupied(_) => Err(value),
            Entry::Vacant(v) => {
                let entry = v.insert((key, value));
                Ok(&mut entry.into_mut().1)
            }
        }
    }

    /// Tries to reserve addional capacity, but doesn't abort the programme on failure.
    pub fn try_reserve(&mut self, additional: usize) -> Result<(), TryReserveError> {
        self.inner.try_reserve(additional, hasher)
    }

    /// Returns the value at `key`
    pub fn get(&self, key: u64) -> Option<&V> {
        self.inner.find(key, eq(key)).map(|(_, v)| v)
    }

    /// Returns a mutable reference to the value at `key`
    pub fn get_mut(&mut self, key: u64) -> Option<&mut V> {
        self.inner.find_mut(key, eq(key)).map(|(_, v)| v)
    }

    /// Attempts to get mutable references to `N` values in the map at once.
    ///
    /// Returns an array of length `N` with the results of each query. For soundness, at most one
    /// mutable reference will be returned to any value. `None` will be returned if any of the keys
    /// are duplicates or missing.
    pub fn get_many_mut<const N: usize>(&mut self, ks: [u64; N]) -> Option<[&mut V; N]> {
        self.inner
            .get_many_mut(ks, |idx, (k, _)| ks[idx] == *k)
            .map(|a| a.map(|(_, v)| v))
    }

    /// Attempts to get mutable references to `N` values in the map at once, without validating
    /// that the values are unique.
    ///
    /// Returns an array of length N with the resultsn of each query.
    /// `None` will be returned if one of the keys are missing.
    ///
    /// For a safe alternative see `[get_many_mut]`
    /// # Safety
    /// Calling this method with overlapping keys is undefined behaviour even if the
    /// resulting references are not used.
    pub unsafe fn get_many_unchecked_mut<const N: usize>(
        &mut self,
        ks: [u64; N],
    ) -> Option<[&mut V; N]> {
        self.inner
            .get_many_unchecked_mut(ks, |idx, (k, _)| ks[idx] == *k)
            .map(|a| a.map(|(_, v)| v))
    }

    /// Removes all values for which `p` evaluates to `false`.
    pub fn retain<F>(&mut self, mut p: F)
    where
        F: FnMut(&u64, &mut V) -> bool,
    {
        self.inner.retain(|(key, val)| p(&*key, val))
    }

    /// Returns an iterator over the keys
    pub fn keys(&self) -> impl Iterator<Item = &u64> {
        self.inner.iter().map(|(key, _)| key)
    }

    /// Returns an iterator over the values
    pub fn values(&self) -> impl Iterator<Item = &V> {
        self.inner.iter().map(|(_, val)| val)
    }

    /// Returns an iterator over mutable references to the stored values
    pub fn values_mut(&mut self) -> impl Iterator<Item = &mut V> {
        self.inner.iter_mut().map(|(_, val)| val)
    }

    /// Self-documenting.
    pub fn contains_key(&self, key: u64) -> bool {
        self.get(key).is_some()
    }

    /// Gets the given key's corresponding entry in the map for in-place manipulation.
    pub fn entry(&mut self, key: u64) -> Entry<V> {
        match self.inner.entry(key, eq(key), hasher) {
            hashbrown::hash_table::Entry::Occupied(o) => Entry::Occupied(OccupiedEntry(o)),
            hashbrown::hash_table::Entry::Vacant(v) => Entry::Vacant(VacantEntry(v, key)),
        }
    }

    /// Returns the number of elements in the map
    pub fn len(&self) -> usize {
        self.inner.len()
    }

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

    /// Returns an iterator over the mutable values, but provides a key. For API consistency, it
    /// returns a reference to the key as well (instead of passing it by value).
    pub fn iter_mut(&mut self) -> impl Iterator<Item = (&u64, &mut V)> {
        self.inner.iter_mut().map(|(key, val)| (&*key, val))
    }

    /// Removes all values from self and returns an iterator over all (key, value) pairs.
    pub fn drain(&mut self) -> impl Iterator<Item = (u64, V)> + '_ {
        self.inner.drain()
    }

    /// Returns an iterator over all (key, value) pairs.
    pub fn iter(&self) -> impl Iterator<Item = (&u64, &V)> {
        self.inner().iter().map(|(k, v)| (k, v))
    }

    /// Like [`Self::retain`], but allows you to reuse the values removed in this fashion. Also,
    /// this one removes the value if the predicate evaluates to `true`.
    ///
    /// Note that this lets the filter closure mutate every value, regardless of whether or not it
    /// shall be kept in the map.
    ///
    /// If the returned iterator is not exhausted, like if it was dropped without iterating or the
    /// iteration short-circuits, then the remaining elements will be retained.
    pub fn extract_if<F>(&mut self, mut f: F) -> ExtractIf<'_, V, impl FnMut(&mut (u64, V)) -> bool>
    where
        F: FnMut(&u64, &mut V) -> bool,
    {
        ExtractIf(self.inner.extract_if(move |(k, v)| f(&*k, v)))
    }

    /// Consumes self and returns an iterator over its keys.
    pub fn into_keys(self) -> impl Iterator<Item = u64> {
        self.into_iter().map(|(k, _)| k)
    }

    /// Consumes self and returns an iterator over its values.
    pub fn into_values(self) -> impl Iterator<Item = V> {
        self.into_iter().map(|(_, v)| v)
    }
}

impl<V> FromIterator<(u64, V)> for IntTable<V> {
    /// Constructs an IntTable from an iterator. Note that this deduplicates the values and keeps
    /// the *last* value.
    fn from_iter<T: IntoIterator<Item = (u64, V)>>(iter: T) -> Self {
        let mut res = Self::new();
        res.extend(iter);
        res
    }
}

impl<V> Default for IntTable<V> {
    fn default() -> Self {
        Self::new()
    }
}

impl<V> IntoIterator for IntTable<V> {
    type Item = (u64, V);

    type IntoIter = IntTableIter<Self::Item>;

    fn into_iter(self) -> Self::IntoIter {
        let it = self.inner.into_iter();
        IntTableIter(it)
    }
}

impl<V> Index<u64> for IntTable<V> {
    type Output = V;

    fn index(&self, index: u64) -> &Self::Output {
        self.get(index).expect("Key not in map")
    }
}

impl<V> IntTable<V> {
    /// Returns the inner hashtable for your viewing pleasure.
    pub fn inner(&self) -> &HashTable<(u64, V)> {
        &self.inner
    }

    /// Unwraps the inner hashtable.
    pub fn into_inner(self) -> HashTable<(u64, V)> {
        self.inner
    }
}

/// Iterator over extracted values.
pub struct ExtractIf<'a, T, F: FnMut(&mut (u64, T)) -> bool>(
    ::hashbrown::hash_table::ExtractIf<'a, (u64, T), F>,
);

impl<V, F> Iterator for ExtractIf<'_, V, F>
where
    F: FnMut(&mut (u64, V)) -> bool,
{
    type Item = (u64, V);

    fn next(&mut self) -> Option<Self::Item> {
        self.0.next()
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        self.0.size_hint()
    }
}

/// An IntTable entry
#[derive(Debug)]
pub enum Entry<'a, V> {
    /// Entry was occupied
    Occupied(OccupiedEntry<'a, V>),
    /// Entry was empty
    Vacant(VacantEntry<'a, V>),
}

impl<'a, V> Entry<'a, V> {
    /// Ensures a value is in the entry by inserting the default if empty, and returns a mutable
    /// reference to the value in the entry.
    pub fn or_insert(self, default: V) -> &'a mut V {
        match self {
            Entry::Occupied(o) => o.into_mut(),
            Entry::Vacant(v) => v.insert(default),
        }
    }

    /// Ensures a value is in the entry by inserting the default if empty, and returns a mutable
    /// reference to the value in the entry. Only calls the closure if the entry is empty.
    pub fn or_insert_with<F>(self, default: F) -> &'a mut V
    where
        F: FnOnce() -> V,
    {
        match self {
            Entry::Occupied(o) => o.into_mut(),
            Entry::Vacant(v) => v.insert(default()),
        }
    }

    /// Ensures a value is in the entry by inserting the default if empty, and returns a mutable
    /// reference to the value in the entry. Only calls the closure if the entry is empty.
    pub fn or_insert_with_key<F>(self, default: F) -> &'a mut V
    where
        F: FnOnce(&u64) -> V,
    {
        match self {
            Entry::Occupied(o) => o.into_mut(),
            Entry::Vacant(v) => {
                let key = v.1;
                v.insert(default(&key))
            }
        }
    }

    /// Returns the key for this entry
    pub fn key(&self) -> &u64 {
        match self {
            Entry::Occupied(o) => &o.0.get().0,
            Entry::Vacant(v) => &v.1,
        }
    }

    /// Provides in-place mutable access to an occupied entry before any potential inserts into the
    /// map
    pub fn and_modify<F>(mut self, f: F) -> Self
    where
        F: FnOnce(&mut V),
    {
        if let Entry::Occupied(o) = &mut self {
            f(o.get_mut())
        };
        self
    }

    /// Sets the value of the entry and returns an [`OccupiedEntry`].
    pub fn insert_entry(self, value: V) -> OccupiedEntry<'a, V> {
        match self {
            Entry::Occupied(mut o) => {
                o.insert(value);
                o
            }
            Entry::Vacant(v) => v.insert_entry(value),
        }
    }
}

impl<'a, V> Entry<'a, V>
where
    V: Default,
{
    /// Sets the value to the default
    pub fn or_default(self) -> &'a mut V {
        #[allow(clippy::unwrap_or_default)] // Thank you clippy, very cool
        self.or_insert(V::default())
    }
}

/// Allows for in-place manipulation of existing table entries
#[derive(Debug)]
pub struct OccupiedEntry<'a, V>(::hashbrown::hash_table::OccupiedEntry<'a, (u64, V)>);
/// Allows for in-place manipulation of vacant table entries
#[derive(Debug)]
pub struct VacantEntry<'a, V>(::hashbrown::hash_table::VacantEntry<'a, (u64, V)>, u64);

impl<'a, V> OccupiedEntry<'a, V> {
    /// Returns the value pointed to by this entry
    pub fn get(&self) -> &V {
        &self.0.get().1
    }

    /// Returns a mutable reference to the value pointed to by this entry
    pub fn get_mut(&mut self) -> &mut V {
        &mut self.0.get_mut().1
    }

    /// Inserts a new value into the entry, returning the old one.
    pub fn insert(&mut self, new_val: V) -> V {
        std::mem::replace(self.get_mut(), new_val)
    }

    /// Consumes this entry and returns the raw mutable reference of the value.
    pub fn into_mut(self) -> &'a mut V {
        &mut self.0.into_mut().1
    }

    /// Removes the value, leaving it empty.
    pub fn remove(self) -> V {
        self.0.remove().0 .1
    }

    /// Take the ownership of the key and value from the map.
    pub fn remove_entry(self) -> (u64, V) {
        self.0.remove().0
    }

    /// Returns the entry's key
    pub fn key(&self) -> &u64 {
        &self.0.get().0
    }
}

impl<'a, V> VacantEntry<'a, V> {
    /// Sets the value of the entry, and returns a mutable reference to it.
    pub fn insert(self, value: V) -> &'a mut V {
        let key = self.1;
        &mut self.0.insert((key, value)).into_mut().1
    }

    /// Sets the value of the entry and returns the [`OccupiedEntry`]
    fn insert_entry(self, value: V) -> OccupiedEntry<'a, V> {
        OccupiedEntry(self.0.insert((self.1, value)))
    }

    /// Returns the key for this entry
    pub fn key(&self) -> &u64 {
        &self.1
    }
}

/// Convenience function to not have to retype this every time lol
fn eq<V>(key: u64) -> impl Fn(&(u64, V)) -> bool {
    move |(other_key, _)| *other_key == key
}

/// Convenience function
fn hasher<V>((key, _value): &(u64, V)) -> u64 {
    *key
}

/// Wrapper around an owned iterator from an [`IntTable`]
pub struct IntTableIter<V>(::hashbrown::hash_table::IntoIter<V>);
impl<V> ExactSizeIterator for IntTableIter<V> {}

impl<V> Iterator for IntTableIter<V> {
    fn size_hint(&self) -> (usize, Option<usize>) {
        let l = self.0.len();
        (l, Some(l))
    }

    fn count(self) -> usize
    where
        Self: Sized,
    {
        self.len()
    }

    type Item = V;

    fn next(&mut self) -> Option<Self::Item> {
        self.0.next()
    }
}

impl<V> Extend<(u64, V)> for IntTable<V> {
    fn extend<T: IntoIterator<Item = (u64, V)>>(&mut self, iter: T) {
        let iter = iter.into_iter();
        let (min, max) = iter.size_hint();
        self.reserve(max.unwrap_or(min));
        for (k, v) in iter {
            self.insert(k, v);
        }
    }
}

impl<V> PartialEq for IntTable<V>
where
    V: PartialEq,
{
    fn eq(&self, other: &IntTable<V>) -> bool {
        self.iter().eq(other.iter())
    }
}

impl<V> std::fmt::Debug for IntTable<V>
where
    V: std::fmt::Debug,
{
    fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
        let mut map_debug = fmt.debug_map();
        for (k, v) in self.iter() {
            map_debug.key(k);
            map_debug.value(v);
        }

        map_debug.finish()
    }
}