sparse.rs

#![no_std]
extern crate alloc;
use alloc::vec::Vec;
use primitive::Primitive;

/// Implementation of:
/// > An Efficient Representation for Sparse Sets
/// > by Preston Briggs and Linda Torczon
/// > <https://dl.acm.org/doi/pdf/10.1145/176454.176484>
/// > <https://research.swtch.com/sparse>
pub struct SparseSet<T> {
    sparse: Vec<T>,
    dense: Vec<T>,
}

impl<T> SparseSet<T>
where
    T: Primitive<usize> + PartialEq + Copy,
    usize: Primitive<T>,
{
    pub fn new(universe: T) -> Self {
        let universe = universe.primitive();
        let mut sparse = Vec::with_capacity(universe);
        unsafe {
            sparse.set_len(universe);
        }
        Self {
            sparse,
            dense: Vec::new(),
        }
    }

    #[inline(always)]
    pub fn len(&self) -> usize {
        self.dense.len()
    }

    #[inline(always)]
    pub fn is_empty(&self) -> bool {
        self.dense.is_empty()
    }

    #[inline(always)]
    pub fn add(&mut self, element: T) {
        if self.contains(element) {
            return;
        }
        self.sparse[element.primitive()] = self.dense.len().primitive();
        self.dense.push(element);
    }

    #[inline]
    pub fn remove(&mut self, element: T) {
        let index = self.sparse[element.primitive()].primitive();
        if index < self.dense.len() && self.dense[index] == element {
            let last_element = self.dense.pop().unwrap();
            if element != last_element {
                self.dense[index] = last_element;
                self.sparse[last_element.primitive()] = index.primitive();
            }
        }
    }

    #[inline(always)]
    pub fn clear(&mut self) {
        self.dense.clear();
    }

    #[inline(always)]
    pub fn contains(&self, element: T) -> bool {
        self.sparse[element.primitive()].primitive() < self.dense.len()
            && self.dense[self.sparse[element.primitive()].primitive()] == element
    }

    #[inline(always)]
    pub fn iter(&self) -> impl Iterator<Item = T> + '_ {
        self.dense.iter().copied()
    }
}

/// An associative map based on SparseSet
/// This allows O(1) operations
/// This version requires V to implement Copy so
/// We don't have to free it, otherwise we sould need a Vec<Option<V>> which
/// would require initializzation of the values vector.
pub struct SparseMap<K, V> {
    set: SparseSet<K>,
    values: Vec<V>,
}

impl<K, V> SparseMap<K, V>
where
    K: Primitive<usize> + PartialEq + Copy,
    usize: Primitive<K>,
    V: Copy
{

    pub fn new(universe: K) -> Self {
        let universe = universe.primitive();
        let mut values = Vec::with_capacity(universe);
        unsafe {
            values.set_len(universe);
        }
        Self {
            set: SparseSet::new(universe),
            values,
        }
    }

    pub fn get(&self, key: K) -> Option<&V> {
        if self.set.contains(key) {
            Some(&self.values[self.set.sparse[key.primitive()].primitive()])
        } else {
            None
        }
    }

    pub fn get_mut(&mut self, key: K) -> Option<&mut V> {
        if self.set.contains(key) {
            Some(&mut self.values[self.set.sparse[key.primitive()].primitive()])
        } else {
            None
        }
    }

    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
        if self.set.contains(key) {
            let index = self.set.sparse[key.primitive()].primitive();
            let old_value = self.values[index];
            self.values[index] = value;
            Some(old_value)
        } else {
            self.set.add(key);
            self.values.push(value);
            None
        }
    }


    pub fn remove(&mut self, key: K) -> Option<V> {
        if self.set.contains(key) {
            let index = self.set.sparse[key.primitive()].primitive();
            self.set.remove(key);
            Some(self.values.remove(index))
        } else {
            None
        }
    }

    pub fn clear(&mut self) {
        self.set.clear();
        self.values.clear();
    }

    pub fn contains_key(&self, key: K) -> bool {
        self.set.contains(key)
    }

    pub fn iter(&self) -> impl Iterator<Item = (K, V)> + '_ {
        self.set.iter().map(move |key| (key, self.values[self.set.sparse[key.primitive()].primitive()]))
    }

    pub fn iter_mut(&mut self) -> impl Iterator<Item = (&K, &mut V)> {
        self.set.iter().map(move |key| (&key, &mut self.values[self.set.sparse[key.primitive()].primitive()]))
    }

    pub fn keys(&self) -> impl Iterator<Item = K> + '_ {
        self.set.iter()
    }

    pub fn values(&self) -> impl Iterator<Item = &V> {
        self.set.iter().map(move |key| &self.values[self.set.sparse[key.primitive()].primitive()])
    }

    pub fn values_mut(&mut self) -> impl Iterator<Item = &mut V> {
        self.set.iter().map(move |key| &mut self.values[self.set.sparse[key.primitive()].primitive()])
    }
    
    pub fn len(&self) -> usize {
        self.set.len()
    }

    pub fn is_empty(&self) -> bool {
        self.set.is_empty()
    }
}