xuko-core 0.6.1

//! Sparse set data structure

/// Sparse set data structure
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
pub struct SparseSet<T> {
    sparse: Vec<usize>,
    dense: Vec<T>,
}

impl<T> SparseSet<T> {
    const NULL_INDEX: usize = usize::MAX;

    /// Create a [`SparseSet`]
    pub fn new() -> Self {
        Self {
            sparse: Vec::new(),
            dense: Vec::new(),
        }
    }

    /// Inserts an item into the [`SparseSet`], returning the index
    pub fn insert(&mut self, value: T) -> usize {
        let index = self.sparse.len();
        if self.sparse.len() + index > self.sparse.capacity() {
            self.sparse.reserve(index.max(4));
        }
        self.sparse.insert(index, self.dense.len());
        self.dense.push(value);
        index
    }

    /// Get a reference to an item in the [`SparseSet`]
    pub fn get(&self, index: usize) -> Option<&T> {
        if index == Self::NULL_INDEX {
            None
        } else if let Some(i) = self.sparse.get(index).copied() {
            self.dense.get(i)
        } else {
            None
        }
    }

    /// Get a mutable reference to an item in the [`SparseSet`]
    pub fn get_mut(&mut self, index: usize) -> Option<&mut T> {
        if index == Self::NULL_INDEX {
            None
        } else if let Some(i) = self.sparse.get(index).copied() {
            self.dense.get_mut(i)
        } else {
            None
        }
    }

    /// Removes an element from the [`SparseSet`]
    pub fn remove(&mut self, index: usize) {
        if self.dense.is_empty() {
            return;
        }

        let Some(dense_index) = self.sparse.get(index).copied() else {
            return;
        };

        let back = self.dense.len() - 1;
        self.dense.swap(dense_index, back);
        self.dense.remove(back);
        self.sparse[index] = Self::NULL_INDEX;
        self.sparse[back] = index;
    }

    /// Return an iterator over the [`SparseSet`]
    pub fn iter(&self) -> iter::Iter<'_, T> {
        iter::Iter::new(self)
    }

    /// Return a mutable iterator over the [`SparseSet`]
    pub fn iter_mut(&mut self) -> iter::IterMut<'_, T> {
        iter::IterMut::new(self)
    }
}

impl<T> Default for SparseSet<T> {
    fn default() -> Self {
        Self {
            sparse: Vec::new(),
            dense: Vec::new(),
        }
    }
}

impl<T> IntoIterator for SparseSet<T> {
    type Item = T;
    type IntoIter = iter::IntoIter<T>;

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

impl<'a, T> IntoIterator for &'a SparseSet<T> {
    type Item = &'a T;
    type IntoIter = iter::Iter<'a, T>;

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

impl<'a, T> IntoIterator for &'a mut SparseSet<T> {
    type Item = &'a mut T;
    type IntoIter = iter::IterMut<'a, T>;

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

/// Iterators for [`SparseSet`]
pub mod iter {
    use std::marker::PhantomData;

    use super::SparseSet;

    /// Iterator over [`SparseSet`]
    pub struct Iter<'a, T> {
        _marker: PhantomData<&'a T>,
        ptr: *const T,
        end: *const T,
    }

    impl<'a, T> Iter<'a, T> {
        pub(crate) fn new(set: &'a SparseSet<T>) -> Self {
            let ptr = set.dense.as_ptr();
            unsafe {
                Self {
                    _marker: PhantomData,
                    ptr,
                    end: ptr.add(set.dense.len()),
                }
            }
        }
    }

    impl<'a, T> Iterator for Iter<'a, T> {
        type Item = &'a T;

        fn next(&mut self) -> Option<Self::Item> {
            if self.ptr == self.end {
                None
            } else {
                unsafe {
                    let ptr = self.ptr;
                    self.ptr = self.ptr.add(1);
                    Some(&*ptr)
                }
            }
        }
    }

    /// Mutable iterator over [`SparseSet`]
    pub struct IterMut<'a, T> {
        _marker: PhantomData<&'a T>,
        ptr: *mut T,
        end: *mut T,
    }

    impl<'a, T> IterMut<'a, T> {
        pub(crate) fn new(set: &'a mut SparseSet<T>) -> Self {
            let ptr = set.dense.as_ptr().cast_mut();
            unsafe {
                Self {
                    _marker: PhantomData,
                    ptr,
                    end: ptr.add(set.dense.len()),
                }
            }
        }
    }

    impl<'a, T> Iterator for IterMut<'a, T> {
        type Item = &'a mut T;

        fn next(&mut self) -> Option<Self::Item> {
            if self.ptr == self.end {
                None
            } else {
                unsafe {
                    let ptr = self.ptr;
                    self.ptr = self.ptr.add(1);
                    Some(&mut *ptr)
                }
            }
        }
    }

    /// Owned [`SparseSet`] iterator
    pub struct IntoIter<T> {
        _set: SparseSet<T>,
        ptr: *const T,
        end: *const T,
    }

    impl<T> IntoIter<T> {
        pub(crate) fn new(set: SparseSet<T>) -> Self {
            unsafe {
                let ptr = set.dense.as_ptr();
                let end = ptr.add(set.dense.len());

                Self {
                    _set: set,
                    ptr,
                    end,
                }
            }
        }
    }

    impl<T> Iterator for IntoIter<T> {
        type Item = T;
        fn next(&mut self) -> Option<Self::Item> {
            if self.ptr == self.end {
                None
            } else {
                unsafe {
                    let ptr = self.ptr;
                    self.ptr = self.ptr.add(1);
                    Some(ptr.read())
                }
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use std::time::SystemTime;

    use super::*;

    #[test]
    fn basic() {
        const N: usize = 1000000;
        let time = SystemTime::now();
        let mut set = SparseSet::new();

        let mut remove = vec![];
        for i in 0..=N {
            let i = set.insert(format!("Hello #{i}"));

            if i % 153 == 0 {
                remove.push(i);
            }
        }

        for i in &remove {
            set.remove(*i);
        }

        let dur = time.elapsed().unwrap();
        println!("Time took: {}ms", dur.as_millis());
        remove.clear();

        let mut hash = std::collections::HashMap::new();
        let time = SystemTime::now();
        for i in 0..=N {
            hash.insert(i, format!("Hello #{i}"));

            if i % 153 == 0 {
                remove.push(i);
            }
        }

        for i in &remove {
            hash.remove(i);
        }

        let dur = time.elapsed().unwrap();
        println!("Time took: {}ms", dur.as_millis());
    }
}