io2/
slab.rs

1use std::alloc::Allocator;
2
3pub struct Slab<T, A: Allocator = std::alloc::Global> {
4    elems: Vec<Entry<T>, A>,
5    first_free_entry: u32,
6    current_generation: u32,
7}
8
9impl<T, A: Allocator> Slab<T, A> {
10    pub fn with_capacity_in(capacity: usize, alloc: A) -> Self {
11        let capacity = if capacity > 0 {
12            capacity.next_power_of_two()
13        } else {
14            0
15        };
16
17        let mut elems = Vec::with_capacity_in(capacity, alloc);
18
19        for i in 0..u32::try_from(capacity).unwrap() {
20            elems.push(Entry::Free { next_free: i + 1 });
21        }
22
23        Self {
24            elems,
25            first_free_entry: 0,
26            current_generation: 0,
27        }
28    }
29
30    pub fn insert(&mut self, val: T) -> Key {
31        let key_idx = usize::try_from(self.first_free_entry).unwrap();
32        let entry = match self.elems.get_mut(key_idx) {
33            Some(entry) => entry,
34            None => {
35                assert_eq!(key_idx, self.elems.len());
36                let initial_len = u32::try_from(self.elems.len()).unwrap();
37                let extend_len = initial_len.max(16);
38                self.elems.reserve(usize::try_from(extend_len).unwrap());
39                for i in initial_len..initial_len + extend_len {
40                    self.elems.push(Entry::Free { next_free: i + 1 });
41                }
42                self.elems.get_mut(key_idx).unwrap()
43            }
44        };
45
46        match entry {
47            Entry::Free { next_free } => {
48                self.first_free_entry = *next_free;
49                *entry = Entry::Occupied {
50                    generation: self.current_generation,
51                    val,
52                };
53            }
54            _ => unreachable!(),
55        }
56
57        Key {
58            generation: self.current_generation,
59            index: u32::try_from(key_idx).unwrap(),
60        }
61    }
62
63    pub fn get(&self, key: Key) -> Option<&T> {
64        match self.elems.get(usize::try_from(key.index).unwrap()) {
65            Some(entry) => match entry {
66                Entry::Occupied { generation, val } => {
67                    if *generation > key.generation {
68                        None
69                    } else {
70                        Some(val)
71                    }
72                }
73                Entry::Free { .. } => None,
74            },
75            None => None,
76        }
77    }
78
79    pub fn get_mut(&mut self, key: Key) -> Option<&mut T> {
80        match self.elems.get_mut(usize::try_from(key.index).unwrap()) {
81            Some(entry) => match entry {
82                Entry::Occupied { generation, val } => {
83                    if *generation > key.generation {
84                        None
85                    } else {
86                        Some(val)
87                    }
88                }
89                Entry::Free { .. } => None,
90            },
91            None => None,
92        }
93    }
94
95    pub fn remove(&mut self, key: Key) -> Option<T> {
96        match self.elems.get_mut(usize::try_from(key.index).unwrap()) {
97            Some(entry) => match entry {
98                Entry::Occupied { generation, .. } => {
99                    if *generation > key.generation {
100                        None
101                    } else {
102                        let entry = std::mem::replace(
103                            entry,
104                            Entry::Free {
105                                next_free: self.first_free_entry,
106                            },
107                        );
108                        self.first_free_entry = key.index;
109                        self.current_generation = self.current_generation.wrapping_add(1);
110                        match entry {
111                            Entry::Occupied { val, .. } => Some(val),
112                            _ => unreachable!(),
113                        }
114                    }
115                }
116                Entry::Free { .. } => None,
117            },
118            None => None,
119        }
120    }
121}
122
123enum Entry<T> {
124    Occupied { generation: u32, val: T },
125    Free { next_free: u32 },
126}
127
128#[derive(Clone, Copy, PartialEq, Eq, Hash)]
129pub struct Key {
130    index: u32,
131    generation: u32,
132}
133
134impl From<u64> for Key {
135    fn from(val: u64) -> Self {
136        let index = val as u32;
137        let generation = (val >> 32) as u32;
138
139        Self { index, generation }
140    }
141}
142
143impl From<Key> for u64 {
144    fn from(key: Key) -> Self {
145        key.index as u64 | ((key.generation as u64) << 32)
146    }
147}
io2/slab.rs

io2/
slab.rs
