xuko_core/
sparse_set.rs

1//! Sparse set data structure
2
3/// Sparse set data structure
4#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
5pub struct SparseSet<T> {
6    sparse: Vec<usize>,
7    dense: Vec<T>,
8}
9
10impl<T> SparseSet<T> {
11    const NULL_INDEX: usize = usize::MAX;
12
13    /// Create a [`SparseSet`]
14    pub fn new() -> Self {
15        Self {
16            sparse: Vec::new(),
17            dense: Vec::new(),
18        }
19    }
20
21    /// Inserts an item into the [`SparseSet`], returning the index
22    pub fn insert(&mut self, value: T) -> usize {
23        let index = self.sparse.len();
24        if self.sparse.len() + index > self.sparse.capacity() {
25            self.sparse.reserve(index.max(4));
26        }
27        self.sparse.insert(index, self.dense.len());
28        self.dense.push(value);
29        index
30    }
31
32    /// Get a reference to an item in the [`SparseSet`]
33    pub fn get(&self, index: usize) -> Option<&T> {
34        if index == Self::NULL_INDEX {
35            None
36        } else if let Some(i) = self.sparse.get(index).copied() {
37            self.dense.get(i)
38        } else {
39            None
40        }
41    }
42
43    /// Get a mutable reference to an item in the [`SparseSet`]
44    pub fn get_mut(&mut self, index: usize) -> Option<&mut T> {
45        if index == Self::NULL_INDEX {
46            None
47        } else if let Some(i) = self.sparse.get(index).copied() {
48            self.dense.get_mut(i)
49        } else {
50            None
51        }
52    }
53
54    /// Removes an element from the [`SparseSet`]
55    pub fn remove(&mut self, index: usize) {
56        if self.dense.is_empty() {
57            return;
58        }
59
60        let Some(dense_index) = self.sparse.get(index).copied() else {
61            return;
62        };
63
64        let back = self.dense.len() - 1;
65        self.dense.swap(dense_index, back);
66        self.dense.remove(back);
67        self.sparse[index] = Self::NULL_INDEX;
68        self.sparse[back] = index;
69    }
70
71    /// Return an iterator over the [`SparseSet`]
72    pub fn iter(&self) -> iter::Iter<'_, T> {
73        iter::Iter::new(self)
74    }
75
76    /// Return a mutable iterator over the [`SparseSet`]
77    pub fn iter_mut(&mut self) -> iter::IterMut<'_, T> {
78        iter::IterMut::new(self)
79    }
80}
81
82impl<T> Default for SparseSet<T> {
83    fn default() -> Self {
84        Self {
85            sparse: Vec::new(),
86            dense: Vec::new(),
87        }
88    }
89}
90
91impl<T> IntoIterator for SparseSet<T> {
92    type Item = T;
93    type IntoIter = iter::IntoIter<T>;
94
95    fn into_iter(self) -> Self::IntoIter {
96        iter::IntoIter::new(self)
97    }
98}
99
100impl<'a, T> IntoIterator for &'a SparseSet<T> {
101    type Item = &'a T;
102    type IntoIter = iter::Iter<'a, T>;
103
104    fn into_iter(self) -> Self::IntoIter {
105        self.iter()
106    }
107}
108
109impl<'a, T> IntoIterator for &'a mut SparseSet<T> {
110    type Item = &'a mut T;
111    type IntoIter = iter::IterMut<'a, T>;
112
113    fn into_iter(self) -> Self::IntoIter {
114        self.iter_mut()
115    }
116}
117
118/// Iterators for [`SparseSet`]
119pub mod iter {
120    use std::marker::PhantomData;
121
122    use super::SparseSet;
123
124    /// Iterator over [`SparseSet`]
125    pub struct Iter<'a, T> {
126        _marker: PhantomData<&'a T>,
127        ptr: *const T,
128        end: *const T,
129    }
130
131    impl<'a, T> Iter<'a, T> {
132        pub(crate) fn new(set: &'a SparseSet<T>) -> Self {
133            let ptr = set.dense.as_ptr();
134            unsafe {
135                Self {
136                    _marker: PhantomData,
137                    ptr,
138                    end: ptr.add(set.dense.len()),
139                }
140            }
141        }
142    }
143
144    impl<'a, T> Iterator for Iter<'a, T> {
145        type Item = &'a T;
146
147        fn next(&mut self) -> Option<Self::Item> {
148            if self.ptr == self.end {
149                None
150            } else {
151                unsafe {
152                    let ptr = self.ptr;
153                    self.ptr = self.ptr.add(1);
154                    Some(&*ptr)
155                }
156            }
157        }
158    }
159
160    /// Mutable iterator over [`SparseSet`]
161    pub struct IterMut<'a, T> {
162        _marker: PhantomData<&'a T>,
163        ptr: *mut T,
164        end: *mut T,
165    }
166
167    impl<'a, T> IterMut<'a, T> {
168        pub(crate) fn new(set: &'a mut SparseSet<T>) -> Self {
169            let ptr = set.dense.as_ptr().cast_mut();
170            unsafe {
171                Self {
172                    _marker: PhantomData,
173                    ptr,
174                    end: ptr.add(set.dense.len()),
175                }
176            }
177        }
178    }
179
180    impl<'a, T> Iterator for IterMut<'a, T> {
181        type Item = &'a mut T;
182
183        fn next(&mut self) -> Option<Self::Item> {
184            if self.ptr == self.end {
185                None
186            } else {
187                unsafe {
188                    let ptr = self.ptr;
189                    self.ptr = self.ptr.add(1);
190                    Some(&mut *ptr)
191                }
192            }
193        }
194    }
195
196    /// Owned [`SparseSet`] iterator
197    pub struct IntoIter<T> {
198        _set: SparseSet<T>,
199        ptr: *const T,
200        end: *const T,
201    }
202
203    impl<T> IntoIter<T> {
204        pub(crate) fn new(set: SparseSet<T>) -> Self {
205            unsafe {
206                let ptr = set.dense.as_ptr();
207                let end = ptr.add(set.dense.len());
208
209                Self {
210                    _set: set,
211                    ptr,
212                    end,
213                }
214            }
215        }
216    }
217
218    impl<T> Iterator for IntoIter<T> {
219        type Item = T;
220        fn next(&mut self) -> Option<Self::Item> {
221            if self.ptr == self.end {
222                None
223            } else {
224                unsafe {
225                    let ptr = self.ptr;
226                    self.ptr = self.ptr.add(1);
227                    Some(ptr.read())
228                }
229            }
230        }
231    }
232}
233
234#[cfg(test)]
235mod tests {
236    use std::time::SystemTime;
237
238    use super::*;
239
240    #[test]
241    fn basic() {
242        const N: usize = 1000000;
243        let time = SystemTime::now();
244        let mut set = SparseSet::new();
245
246        let mut remove = vec![];
247        for i in 0..=N {
248            let i = set.insert(format!("Hello #{i}"));
249
250            if i % 153 == 0 {
251                remove.push(i);
252            }
253        }
254
255        for i in &remove {
256            set.remove(*i);
257        }
258
259        let dur = time.elapsed().unwrap();
260        println!("Time took: {}ms", dur.as_millis());
261        remove.clear();
262
263        let mut hash = std::collections::HashMap::new();
264        let time = SystemTime::now();
265        for i in 0..=N {
266            hash.insert(i, format!("Hello #{i}"));
267
268            if i % 153 == 0 {
269                remove.push(i);
270            }
271        }
272
273        for i in &remove {
274            hash.remove(i);
275        }
276
277        let dur = time.elapsed().unwrap();
278        println!("Time took: {}ms", dur.as_millis());
279    }
280}
xuko_core/sparse_set.rs

xuko_core/
sparse_set.rs
