gpu_alloc/
slab.rs

1use {crate::unreachable_unchecked, alloc::vec::Vec, core::mem::replace};
2
3#[derive(Debug)]
4enum Entry<T> {
5    Vacant(usize),
6    Occupied(T),
7}
8#[derive(Debug)]
9pub(crate) struct Slab<T> {
10    next_vacant: usize,
11    entries: Vec<Entry<T>>,
12}
13
14impl<T> Slab<T> {
15    pub fn new() -> Self {
16        Slab {
17            next_vacant: !0,
18            entries: Vec::new(),
19        }
20    }
21
22    /// Inserts value into this linked vec and returns index
23    /// at which value can be accessed in constant time.
24    pub fn insert(&mut self, value: T) -> usize {
25        if self.next_vacant >= self.entries.len() {
26            self.entries.push(Entry::Occupied(value));
27            self.entries.len() - 1
28        } else {
29            match *unsafe { self.entries.get_unchecked(self.next_vacant) } {
30                Entry::Vacant(next_vacant) => {
31                    unsafe {
32                        *self.entries.get_unchecked_mut(self.next_vacant) = Entry::Occupied(value);
33                    }
34                    replace(&mut self.next_vacant, next_vacant)
35                }
36                _ => unsafe { unreachable_unchecked() },
37            }
38        }
39    }
40
41    pub fn len(&self) -> usize {
42        self.entries.len()
43    }
44
45    pub unsafe fn get_unchecked(&self, index: usize) -> &T {
46        debug_assert!(index < self.len());
47
48        match self.entries.get_unchecked(index) {
49            Entry::Occupied(value) => value,
50            _ => unreachable_unchecked(),
51        }
52    }
53
54    pub unsafe fn get_unchecked_mut(&mut self, index: usize) -> &mut T {
55        debug_assert!(index < self.len());
56
57        match self.entries.get_unchecked_mut(index) {
58            Entry::Occupied(value) => value,
59            _ => unreachable_unchecked(),
60        }
61    }
62
63    pub fn get(&self, index: usize) -> &T {
64        match self.entries.get(index) {
65            Some(Entry::Occupied(value)) => value,
66            _ => panic!("Invalid index"),
67        }
68    }
69
70    pub fn get_mut(&mut self, index: usize) -> &mut T {
71        match self.entries.get_mut(index) {
72            Some(Entry::Occupied(value)) => value,
73            _ => panic!("Invalid index"),
74        }
75    }
76
77    pub unsafe fn remove_unchecked(&mut self, index: usize) -> T {
78        let entry = replace(
79            self.entries.get_unchecked_mut(index),
80            Entry::Vacant(self.next_vacant),
81        );
82
83        self.next_vacant = index;
84
85        match entry {
86            Entry::Occupied(value) => value,
87            _ => unreachable_unchecked(),
88        }
89    }
90
91    pub fn remove(&mut self, index: usize) -> T {
92        match self.entries.get_mut(index) {
93            Some(Entry::Occupied(_)) => unsafe { self.remove_unchecked(index) },
94            _ => panic!("Invalid index"),
95        }
96    }
97}