Skip to main content

xtk_core/
sparse_set.rs

1//! Sparse Set
2
3use smallvec::SmallVec;
4
5/// An entry in the [SparseSet]
6#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
7pub struct Entry<T> {
8    sparse_index: usize,
9
10    /// Value contained in the entry
11    pub value: T,
12}
13
14impl<T> Entry<T> {
15    /// Get the sparse index associated with the entry
16    pub fn index(&self) -> usize {
17        self.sparse_index
18    }
19}
20
21/// See [module-level documentation](self)
22pub struct SparseSet<T> {
23    new_index: usize,
24    sparse: SmallVec<[usize; 32]>,
25    dense: Vec<Entry<T>>,
26}
27
28impl<T> SparseSet<T> {
29    /// Initialize a new empty sparse set
30    #[must_use]
31    pub fn new() -> Self {
32        Self {
33            new_index: 0,
34            sparse: SmallVec::new(),
35            dense: Vec::new(),
36        }
37    }
38
39    /// Create a sparse set with dense capacity `n`
40    #[must_use]
41    pub fn with_capacity(n: usize) -> Self {
42        Self {
43            new_index: 0,
44            sparse: SmallVec::new(),
45            dense: Vec::with_capacity(n),
46        }
47    }
48
49    /// Insert an item into the sparse set
50    pub fn insert(&mut self, value: T) -> usize {
51        let dense_index = self.dense.len();
52        let sparse_index = self.new_index;
53        self.new_index += 1;
54
55        self.dense.push(Entry {
56            sparse_index,
57            value,
58        });
59        self.sparse.push(dense_index);
60
61        return sparse_index;
62    }
63
64    /// Remove an element from the sparse set
65    pub fn remove(&mut self, index: usize) -> Option<T> {
66        if index >= self.sparse.len() {
67            return None;
68        }
69
70        let dense_index = self.sparse.get(index).copied()?;
71        if self.dense.get(dense_index).is_none() {
72            return None;
73        }
74
75        self.sparse[index] = self.dense.len();
76
77        Some(self.dense.swap_remove(dense_index).value)
78    }
79
80    /// Get a reference to element at `index`
81    #[must_use]
82    pub fn get(&self, index: usize) -> Option<&T> {
83        if index >= self.sparse.len() {
84            return None;
85        }
86
87        let dense_index = self.sparse.get(index).copied()?;
88        self.dense.get(dense_index).map(|entry| &entry.value)
89    }
90
91    /// Get a mutable reference to element at `index`
92    #[must_use]
93    pub fn get_mut(&mut self, index: usize) -> Option<&mut T> {
94        if index >= self.sparse.len() {
95            return None;
96        }
97
98        let dense_index = self.sparse.get(index).copied()?;
99        self.dense
100            .get_mut(dense_index)
101            .map(|entry| &mut entry.value)
102    }
103
104    /// Create an immutable iterator for the sparse set
105    #[must_use]
106    pub fn iter<'a>(&'a self) -> iter::Iter<'a, T> {
107        iter::Iter::new(self)
108    }
109
110    /// Create a mutable iterator for the sparse set
111    #[must_use]
112    pub fn iter_mut<'a>(&'a mut self) -> iter::IterMut<'a, T> {
113        iter::IterMut::new(self)
114    }
115}
116
117impl<T> Default for SparseSet<T> {
118    fn default() -> Self {
119        Self::new()
120    }
121}
122
123impl<'a, T> IntoIterator for &'a SparseSet<T> {
124    type Item = &'a Entry<T>;
125    type IntoIter = iter::Iter<'a, T>;
126
127    fn into_iter(self) -> Self::IntoIter {
128        self.iter()
129    }
130}
131
132impl<'a, T> IntoIterator for &'a mut SparseSet<T> {
133    type Item = &'a mut Entry<T>;
134    type IntoIter = iter::IterMut<'a, T>;
135
136    fn into_iter(self) -> Self::IntoIter {
137        self.iter_mut()
138    }
139}
140
141impl<T> IntoIterator for SparseSet<T> {
142    type Item = Entry<T>;
143    type IntoIter = iter::IntoIter<T>;
144
145    fn into_iter(self) -> Self::IntoIter {
146        iter::IntoIter::new(self)
147    }
148}
149
150#[cfg(feature = "serde")]
151impl<T: serde::Serialize> serde::Serialize for SparseSet<T> {
152    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
153    where
154        S: serde::Serializer,
155    {
156        let v = self.dense.iter().map(|e| &e.value).collect::<Vec<_>>();
157        v.serialize(serializer)
158    }
159}
160
161#[cfg(feature = "serde")]
162impl<'de, T: serde::Deserialize<'de>> serde::Deserialize<'de> for SparseSet<T> {
163    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
164    where
165        D: serde::Deserializer<'de>,
166    {
167        let de: Vec<T> = Vec::deserialize(deserializer)?;
168        let mut s = SparseSet::new();
169        for e in de {
170            s.insert(e);
171        }
172        Ok(s)
173    }
174}
175
176/// Iterators
177pub mod iter {
178    use std::mem;
179
180    use super::{Entry, SparseSet};
181
182    /// Immutable iterator to a [SparseSet]
183    pub struct Iter<'a, T> {
184        sparse: &'a SparseSet<T>,
185        index: usize,
186    }
187
188    impl<'a, T> Iter<'a, T> {
189        pub(crate) fn new(sparse: &'a SparseSet<T>) -> Self {
190            Self { sparse, index: 0 }
191        }
192    }
193
194    impl<'a, T> Iterator for Iter<'a, T> {
195        type Item = &'a Entry<T>;
196
197        fn next(&mut self) -> Option<Self::Item> {
198            let index = self.index;
199            self.index += 1;
200            self.sparse.dense.get(index)
201        }
202    }
203
204    /// Mutable iterator to a [SparseSet]
205    pub struct IterMut<'a, T> {
206        sparse: &'a mut SparseSet<T>,
207        index: usize,
208    }
209
210    impl<'a, T> IterMut<'a, T> {
211        pub(crate) fn new(sparse: &'a mut SparseSet<T>) -> Self {
212            Self { sparse, index: 0 }
213        }
214    }
215
216    impl<'a, T> Iterator for IterMut<'a, T> {
217        type Item = &'a mut Entry<T>;
218
219        fn next(&mut self) -> Option<Self::Item> {
220            let index = self.index;
221            self.index += 1;
222            self.sparse
223                .dense
224                .get_mut(index)
225                .map(|v| unsafe { &mut *(v as *mut Entry<T>) })
226        }
227    }
228
229    /// [SparseSet] iterator that owns the structure
230    pub struct IntoIter<T> {
231        sparse: SparseSet<T>,
232        index: usize,
233    }
234
235    impl<T> IntoIter<T> {
236        pub(crate) fn new(sparse: SparseSet<T>) -> Self {
237            Self { sparse, index: 0 }
238        }
239    }
240
241    impl<T> Iterator for IntoIter<T> {
242        type Item = Entry<T>;
243
244        fn next(&mut self) -> Option<Self::Item> {
245            let index = self.index;
246            if index == self.sparse.dense.len() {
247                None
248            } else {
249                self.index += 1;
250                let mut entry = unsafe { mem::zeroed() };
251                mem::swap(&mut self.sparse.dense[index], &mut entry);
252                Some(entry)
253            }
254        }
255    }
256}
257
258#[cfg(test)]
259mod tests {
260    use super::SparseSet;
261
262    #[test]
263    fn insert() {
264        let mut set = SparseSet::new();
265
266        set.insert(1);
267        let index = set.insert(2);
268        assert_eq!(set.remove(index), Some(2));
269        assert_eq!(set.insert(3), 2);
270    }
271
272    #[test]
273    fn iter() {
274        let mut set = SparseSet::new();
275
276        set.insert("static string");
277        set.insert("another static string");
278        let i = set.insert("abcdefghi");
279        set.remove(i);
280        assert_eq!(set.insert("hello"), 3);
281        set.insert("feijfoej");
282
283        for entry in set.iter() {
284            println!("entry #{}: {}", entry.index(), entry.value);
285        }
286        println!("{:?}, {:#?}", set.sparse, set.dense);
287    }
288}