avian3d/data_structures/
stable_vec.rs

1//! A vector data structure that maintains stable indices for its elements.
2//!
3//! # Why Not `slab`?
4//!
5//! The `slab` crate provides a similar `Slab` data structure.
6//! However, its insertion order is not necessarily preserved when elements are removed.
7//! For example, after clearing the slab, the next element might not be inserted at index 0,
8//! but rather it reuses some other slot. This can lead to seemingly non-deterministic behavior
9//! when clearing and repopulating scenes, for example.
10//!
11//! This `StableVec` implementation instead always pushes elements to the first available slot
12//! with the lowest index, ensuring that insertion order is preserved across removals and clear operations.
13//!
14//! # Why Not `stable_vec`?
15//!
16//! The `stable_vec` crate provides a similar `StableVec` data structure.
17//! However, it doesn't actually reuse slots of removed elements, leading to unbounded memory growth
18//! if elements are frequently added and removed. This makes it unsuitable for general-purpose use
19//! for gameplay scenarios.
20//!
21//! This `StableVec` implementation reuses slots of removed elements, using the [`IdPool`] type
22//! to track free indices.
23
24use crate::data_structures::id_pool::IdPool;
25
26/// A [`Vec<T>`]-like collection that maintains stable indices for its elements.
27///
28/// Unlike with a standard [`Vec<T>`], removing elements from a [`StableVec<T>`] is O(1),
29/// and it does not move other elements or invalidate their indices.
30/// This is achieved by internally storing each element as an [`Option<T>`],
31/// and reusing freed slots for new elements.
32///
33/// # Overiew of Important Methods
34///
35/// - [`push`](Self::push) adds a new element and returns its stable index (O(1)).
36/// - [`remove`](Self::remove) removes an element at a given index and returns it (O(1)).
37/// - [`try_remove`](Self::try_remove) attempts to remove an element at a given index, returning `None` if it doesn't exist (O(1)).
38/// - [`get`](Self::get) and [`get_mut`](Self::get_mut) provide access to elements by index (O(1)).
39/// - [`clear`](Self::clear) removes all elements without deallocating memory (O(1)).
40#[derive(Clone, Debug)]
41pub struct StableVec<T> {
42    data: Vec<Option<T>>,
43    indices: IdPool,
44}
45
46impl<T> StableVec<T> {
47    /// Creates a new empty [`StableVec`].
48    #[inline(always)]
49    pub const fn new() -> Self {
50        Self {
51            data: Vec::new(),
52            indices: IdPool::new(),
53        }
54    }
55
56    /// Creates a new [`StableVec`] with the given initial capacity.
57    ///
58    /// This is useful for preallocating space to avoid reallocations.
59    #[inline(always)]
60    pub fn with_capacity(capacity: usize) -> Self {
61        Self {
62            data: Vec::with_capacity(capacity),
63            indices: IdPool::with_capacity(capacity),
64        }
65    }
66
67    /// Pushes a new element to the first available slot, returning its index.
68    ///
69    /// This may reuse a previously freed slot.
70    ///
71    /// # Time Complexity
72    ///
73    /// O(1)
74    #[inline(always)]
75    pub fn push(&mut self, element: T) -> usize {
76        let index = self.indices.alloc() as usize;
77
78        if index >= self.data.len() {
79            self.data.push(Some(element));
80        } else {
81            self.data[index] = Some(element);
82        }
83
84        index
85    }
86
87    /// Returns the next index that will be used for a push.
88    ///
89    /// This index may reuse a previously freed slot.
90    ///
91    /// # Time Complexity
92    ///
93    /// O(1)
94    #[inline(always)]
95    pub fn next_push_index(&self) -> usize {
96        self.indices.next_id() as usize
97    }
98
99    /// Removes the element at the given index, returning it.
100    ///
101    /// # Time Complexity
102    ///
103    /// O(1)
104    ///
105    /// # Panics
106    ///
107    /// Panics if there is no element at the given index.
108    #[inline(always)]
109    pub fn remove(&mut self, index: usize) -> T {
110        self.try_remove(index).unwrap_or_else(|| {
111            panic!("no element at index {} in StableVec::remove", index);
112        })
113    }
114
115    /// Tries to remove the element at the given index, returning it if it existed.
116    ///
117    /// # Time Complexity
118    ///
119    /// O(1)
120    #[inline(always)]
121    pub fn try_remove(&mut self, index: usize) -> Option<T> {
122        if index >= self.data.len() {
123            return None;
124        }
125        let element = self.data[index].take();
126        if element.is_some() {
127            self.indices.free(index as u32);
128        }
129        element
130    }
131
132    /// Removes all elements from the [`StableVec`].
133    ///
134    /// No memory is deallocated, so the capacity remains the same.
135    ///
136    /// # Time Complexity
137    ///
138    /// O(1)
139    #[inline(always)]
140    pub fn clear(&mut self) {
141        self.data.clear();
142        self.indices.clear();
143    }
144
145    /// Returns a reference to the element at the given index, if it exists.
146    ///
147    /// # Time Complexity
148    ///
149    /// O(1)
150    #[inline(always)]
151    pub fn get(&self, index: usize) -> Option<&T> {
152        self.data.get(index)?.as_ref()
153    }
154
155    /// Returns a mutable reference to the element at the given index, if it exists.
156    ///
157    /// # Time Complexity
158    ///
159    /// O(1)
160    #[inline(always)]
161    pub fn get_mut(&mut self, index: usize) -> Option<&mut T> {
162        self.data.get_mut(index)?.as_mut()
163    }
164
165    /// Returns a reference to the element at the given index without bounds checking.
166    ///
167    /// # Time Complexity
168    ///
169    /// O(1)
170    ///
171    /// # Safety
172    ///
173    /// The caller must ensure that the index is in bounds
174    /// and that there is an element at that index.
175    #[inline(always)]
176    pub unsafe fn get_unchecked(&self, index: usize) -> &T {
177        unsafe { self.data.get_unchecked(index).as_ref().unwrap_unchecked() }
178    }
179
180    /// Returns a mutable reference to the element at the given index without bounds checking.
181    ///
182    /// # Time Complexity
183    ///
184    /// O(1)
185    ///
186    /// # Safety
187    ///
188    /// The caller must ensure that the index is in bounds
189    /// and that there is an element at that index.
190    #[inline(always)]
191    pub unsafe fn get_unchecked_mut(&mut self, index: usize) -> &mut T {
192        unsafe {
193            self.data
194                .get_unchecked_mut(index)
195                .as_mut()
196                .unwrap_unchecked()
197        }
198    }
199
200    /// Returns mutable references to two disjoint elements at the given indices, if they exist.
201    ///
202    /// If the indices are the same, or if either index is out of bounds or has no element,
203    /// `None` is returned.
204    ///
205    /// # Time Complexity
206    ///
207    /// O(1)
208    #[inline(always)]
209    pub fn get_disjoint_mut2(&mut self, index1: usize, index2: usize) -> Option<[&mut T; 2]> {
210        // TODO: Return a `Result`.
211        let [first, second] = self.data.get_disjoint_mut([index1, index2]).ok()?;
212        Some([first.as_mut()?, second.as_mut()?])
213    }
214
215    /// Returns mutable references to two disjoint elements at the given indices without bounds checking.
216    ///
217    /// If the indices are the same, or if either index has no element,
218    /// the behavior is undefined.
219    ///
220    /// # Time Complexity
221    ///
222    /// O(1)
223    ///
224    /// # Safety
225    ///
226    /// The caller must ensure that the indices are disjoint
227    /// and that there are elements at both indices.
228    #[inline(always)]
229    pub unsafe fn get_disjoint_mut_unchecked(
230        &mut self,
231        index1: usize,
232        index2: usize,
233    ) -> [&mut T; 2] {
234        // TODO: Return a `Result`.
235        unsafe {
236            let [first, second] = self.data.get_disjoint_unchecked_mut([index1, index2]);
237            [
238                first.as_mut().unwrap_unchecked(),
239                second.as_mut().unwrap_unchecked(),
240            ]
241        }
242    }
243
244    /// Returns the number of elements in the [`StableVec`].
245    ///
246    /// # Time Complexity
247    ///
248    /// O(1)
249    #[inline(always)]
250    pub fn len(&self) -> usize {
251        self.indices.len()
252    }
253
254    /// Returns `true` if the [`StableVec`] is empty.
255    ///
256    /// # Time Complexity
257    ///
258    /// O(1)
259    #[inline(always)]
260    pub fn is_empty(&self) -> bool {
261        self.len() == 0
262    }
263
264    /// Returns the capacity of the [`StableVec`].
265    ///
266    /// # Time Complexity
267    ///
268    /// O(1)
269    #[inline(always)]
270    pub fn capacity(&self) -> usize {
271        self.data.capacity()
272    }
273
274    /// Returns an iterator over the elements and indices of the [`StableVec`].
275    #[inline(always)]
276    pub fn iter(&self) -> impl Iterator<Item = (usize, &T)> {
277        self.data
278            .iter()
279            .enumerate()
280            .filter_map(|(i, opt)| opt.as_ref().map(|v| (i, v)))
281    }
282
283    /// Returns a mutable iterator over the elements and indices of the [`StableVec`].
284    #[inline(always)]
285    pub fn iter_mut(&mut self) -> impl Iterator<Item = (usize, &mut T)> {
286        self.data
287            .iter_mut()
288            .enumerate()
289            .filter_map(|(i, opt)| opt.as_mut().map(|v| (i, v)))
290    }
291}
292
293impl<T> Default for StableVec<T> {
294    /// Creates a new empty [`StableVec`].
295    #[inline(always)]
296    fn default() -> Self {
297        Self::new()
298    }
299}
300
301impl<T> core::ops::Index<usize> for StableVec<T> {
302    type Output = T;
303
304    /// Returns a reference to the element at the given index.
305    ///
306    /// # Time Complexity
307    ///
308    /// O(1)
309    ///
310    /// # Panics
311    ///
312    /// Panics if there is no element at the given index.
313    #[inline(always)]
314    fn index(&self, index: usize) -> &Self::Output {
315        self.get(index).unwrap_or_else(|| {
316            panic!("no element at index {} in StableVec::index", index);
317        })
318    }
319}
320
321impl<T> core::ops::IndexMut<usize> for StableVec<T> {
322    /// Returns a mutable reference to the element at the given index.
323    ///
324    /// # Time Complexity
325    ///
326    /// O(1)
327    ///
328    /// # Panics
329    ///
330    /// Panics if there is no element at the given index.
331    #[inline(always)]
332    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
333        self.get_mut(index).unwrap_or_else(|| {
334            panic!("no element at index {} in StableVec::index_mut", index);
335        })
336    }
337}
338
339impl<T> IntoIterator for StableVec<T> {
340    type Item = T;
341    type IntoIter = StableVecIntoIterator<T>;
342
343    #[inline(always)]
344    fn into_iter(self) -> Self::IntoIter {
345        StableVecIntoIterator {
346            inner: self.data.into_iter(),
347        }
348    }
349}
350
351/// An iterator over the elements of a [`StableVec`].
352pub struct StableVecIntoIterator<T> {
353    inner: alloc::vec::IntoIter<Option<T>>,
354}
355
356impl<T> Iterator for StableVecIntoIterator<T> {
357    type Item = T;
358
359    #[inline(always)]
360    fn next(&mut self) -> Option<Self::Item> {
361        self.inner.by_ref().flatten().next()
362    }
363}
364
365impl<T> DoubleEndedIterator for StableVecIntoIterator<T> {
366    #[inline(always)]
367    fn next_back(&mut self) -> Option<Self::Item> {
368        self.inner.by_ref().flatten().next_back()
369    }
370}
371
372impl<T> core::iter::FusedIterator for StableVecIntoIterator<T> {}
373
374#[cfg(test)]
375mod tests {
376    use super::StableVec;
377
378    #[test]
379    fn test_stable_vec_push_remove() {
380        let mut sv = StableVec::new();
381        let idx1 = sv.push(10);
382        let idx2 = sv.push(20);
383
384        assert_eq!(sv.len(), 2);
385        assert_eq!(sv[idx1], 10);
386        assert_eq!(sv[idx2], 20);
387
388        let val = sv.remove(idx1);
389        assert_eq!(val, 10);
390        assert_eq!(sv.len(), 1);
391        assert!(sv.get(idx1).is_none());
392
393        let idx3 = sv.push(30);
394        assert_eq!(idx3, idx1); // Reused index
395        assert_eq!(sv[idx3], 30);
396    }
397
398    #[test]
399    fn test_stable_vec_clear() {
400        let mut sv = StableVec::new();
401        sv.push(1);
402        sv.push(2);
403        sv.clear();
404        assert_eq!(sv.len(), 0);
405    }
406
407    #[test]
408    fn test_stable_vec_iter() {
409        let mut sv = StableVec::new();
410        sv.push(1);
411        sv.push(2);
412        sv.push(3);
413
414        let collected: Vec<_> = sv.into_iter().collect();
415        assert_eq!(collected, vec![1, 2, 3]);
416    }
417}
avian3d/data_structures/stable_vec.rs

avian3d/data_structures/
stable_vec.rs
