1use core::fmt;
2use core::hash::{Hash, Hasher};
3use core::marker::PhantomData;
4
5use o3::collections::{CellSlab, Slab, SlabGeneration, SlabKey, SlabKeyParts, SlabVacantEntry};
6
7use super::{EPOCH_MASK, SLOT_MASK, SlotIndex, Token, TokenTag};
8
9const GENERATION_LIMIT: u32 = EPOCH_MASK as u32;
10
11#[repr(transparent)]
12pub struct Key<Tag>(SlabKey<Tag, GENERATION_LIMIT>);
13
14#[repr(transparent)]
15pub struct KeyParts<Tag> {
16 parts: SlabKeyParts<GENERATION_LIMIT>,
17 tag: PhantomData<*mut Tag>,
18}
19
20#[repr(transparent)]
21pub struct TokenCellSlab<T, Tag>(CellSlab<T, Tag, GENERATION_LIMIT>);
22
23#[derive(Clone, Copy, PartialEq, Eq, Debug, Hash)]
24#[repr(transparent)]
25pub struct TokenCapacity(usize);
26
27#[repr(transparent)]
28pub struct TokenSlab<T, Tag>(Slab<T, Tag, GENERATION_LIMIT>);
29
30#[repr(transparent)]
31pub struct TokenSlabVacantEntry<'a, T, Tag>(SlabVacantEntry<'a, T, Tag, GENERATION_LIMIT>);
32
33impl<Tag> Clone for Key<Tag> {
34 fn clone(&self) -> Self {
35 *self
36 }
37}
38
39impl<Tag> Copy for Key<Tag> {}
40
41impl<Tag> PartialEq for Key<Tag> {
42 fn eq(&self, other: &Self) -> bool {
43 self.0 == other.0
44 }
45}
46
47impl<Tag> Eq for Key<Tag> {}
48
49impl<Tag> Hash for Key<Tag> {
50 fn hash<H: Hasher>(&self, state: &mut H) {
51 self.0.hash(state);
52 }
53}
54
55impl<Tag> fmt::Debug for Key<Tag> {
56 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
57 self.0.fmt(f)
58 }
59}
60
61impl<Tag> Key<Tag> {
62 const fn from_raw(raw: SlabKey<Tag, GENERATION_LIMIT>) -> Self {
63 Self(raw)
64 }
65
66 const fn raw(self) -> SlabKey<Tag, GENERATION_LIMIT> {
67 self.0
68 }
69
70 pub const fn index(self) -> u32 {
71 self.0.index()
72 }
73
74 pub const fn slot(self) -> SlotIndex {
75 SlotIndex::from_bounded(self.0.index())
76 }
77
78 pub const fn generation(self) -> SlabGeneration<GENERATION_LIMIT> {
79 self.0.generation()
80 }
81}
82
83impl<Tag> Clone for KeyParts<Tag> {
84 fn clone(&self) -> Self {
85 *self
86 }
87}
88
89impl<Tag> Copy for KeyParts<Tag> {}
90
91impl<Tag> KeyParts<Tag> {
92 pub(super) const fn new(index: u32, generation: u32) -> Option<Self> {
93 match SlabKeyParts::new(index, generation) {
94 Some(parts) => Some(Self::from_raw(parts)),
95 None => None,
96 }
97 }
98
99 const fn from_raw(parts: SlabKeyParts<GENERATION_LIMIT>) -> Self {
100 Self {
101 parts,
102 tag: PhantomData,
103 }
104 }
105
106 const fn raw(self) -> SlabKeyParts<GENERATION_LIMIT> {
107 self.parts
108 }
109
110 pub const fn index(self) -> u32 {
111 self.parts.index()
112 }
113
114 pub const fn slot(self) -> SlotIndex {
115 SlotIndex::from_bounded(self.parts.index())
116 }
117
118 pub(super) const fn generation(self) -> SlabGeneration<GENERATION_LIMIT> {
119 self.parts.generation()
120 }
121}
122
123impl TokenCapacity {
124 pub const EMPTY: Self = Self(0);
125
126 pub const fn new(capacity: usize) -> Option<Self> {
127 if capacity <= SLOT_MASK as usize + 1 {
128 Some(Self(capacity))
129 } else {
130 None
131 }
132 }
133
134 pub const fn get(self) -> usize {
135 self.0
136 }
137
138 pub const fn slot(self, index: usize) -> Option<SlotIndex> {
139 if index < self.0 {
140 Some(SlotIndex::from_bounded(index as u32))
141 } else {
142 None
143 }
144 }
145
146 pub const fn sentinel(self) -> Option<SlotIndex> {
147 SlotIndex::try_new(self.0 as u32)
148 }
149
150 pub fn slots(self) -> impl ExactSizeIterator<Item = SlotIndex> + DoubleEndedIterator {
151 (0..self.0 as u32).map(SlotIndex::from_bounded)
152 }
153}
154
155impl<T, Tag> TokenSlab<T, Tag> {
156 #[must_use]
157 pub fn with_capacity(capacity: TokenCapacity) -> Self {
158 Self(Slab::with_capacity(capacity.get()))
159 }
160
161 pub fn capacity(&self) -> TokenCapacity {
162 TokenCapacity(self.0.capacity())
163 }
164
165 pub fn len(&self) -> usize {
166 self.0.len()
167 }
168
169 pub fn is_empty(&self) -> bool {
170 self.0.is_empty()
171 }
172
173 pub fn insert(&mut self, value: T) -> Result<Key<Tag>, T> {
174 self.0.insert(value).map(Key::from_raw)
175 }
176
177 pub fn insert_entry(&mut self, value: T) -> Result<(Key<Tag>, &mut T), T> {
178 self.0
179 .insert_entry(value)
180 .map(|(key, value)| (Key::from_raw(key), value))
181 }
182
183 pub fn vacant_entry(&mut self) -> Option<TokenSlabVacantEntry<'_, T, Tag>> {
184 self.0.vacant_entry().map(TokenSlabVacantEntry)
185 }
186
187 pub fn vacant_entry_at(&mut self, index: u32) -> Option<TokenSlabVacantEntry<'_, T, Tag>> {
188 self.0.vacant_entry_at(index).map(TokenSlabVacantEntry)
189 }
190
191 pub fn get(&self, key: Key<Tag>) -> Option<&T> {
192 self.0.get(key.raw())
193 }
194
195 pub fn get_parts(&self, parts: KeyParts<Tag>) -> Option<&T> {
196 self.0.get_parts(parts.raw())
197 }
198
199 pub fn get_parts_mut(&mut self, parts: KeyParts<Tag>) -> Option<&mut T> {
200 self.0.get_parts_mut(parts.raw())
201 }
202
203 pub fn remove(&mut self, key: Key<Tag>) -> Option<T> {
204 self.0.remove(key.raw())
205 }
206
207 pub fn remove_parts(&mut self, parts: KeyParts<Tag>) -> Option<T> {
208 self.0.remove_parts(parts.raw())
209 }
210
211 pub fn remove_index_with<R>(
212 &mut self,
213 index: u32,
214 f: impl FnOnce(&mut T, Key<Tag>) -> Option<R>,
215 ) -> Option<(T, R)> {
216 self.0
217 .remove_index_with(index, |value, key| f(value, Key::from_raw(key)))
218 }
219
220 pub fn get_index(&self, index: u32) -> Option<(&T, Key<Tag>)> {
221 self.0
222 .get_index(index)
223 .map(|(value, key)| (value, Key::from_raw(key)))
224 }
225
226 pub fn get_index_mut(&mut self, index: u32) -> Option<(&mut T, Key<Tag>)> {
227 self.0
228 .get_index_mut(index)
229 .map(|(value, key)| (value, Key::from_raw(key)))
230 }
231
232 pub fn key(&self, index: u32) -> Option<Key<Tag>> {
233 self.0.key(index).map(Key::from_raw)
234 }
235
236 pub fn values(&self) -> impl Iterator<Item = &T> {
237 self.0.values()
238 }
239
240 pub fn values_mut(&mut self) -> impl Iterator<Item = &mut T> {
241 self.0.values_mut()
242 }
243}
244
245impl<T, Tag> TokenSlabVacantEntry<'_, T, Tag> {
246 pub fn insert(self, value: T) {
247 let _ = self.0.insert(value);
248 }
249}
250
251impl<T, Tag: TokenTag> TokenSlabVacantEntry<'_, T, Tag> {
252 pub fn token(&self) -> Token {
253 Token::from_key(Key::from_raw(self.0.key()))
254 }
255}
256
257impl<T, Tag> TokenCellSlab<T, Tag> {
258 #[must_use]
259 pub fn with_capacity(capacity: TokenCapacity) -> Self {
260 Self(CellSlab::with_capacity(capacity.get()))
261 }
262
263 pub fn is_empty(&self) -> bool {
264 self.0.is_empty()
265 }
266
267 pub fn keys(&self) -> impl Iterator<Item = Key<Tag>> + '_ {
268 self.0.keys().map(Key::from_raw)
269 }
270
271 pub fn insert(&self, value: T) -> Result<Key<Tag>, T> {
272 self.0.insert(value).map(Key::from_raw)
273 }
274
275 pub fn update<R>(&self, key: Key<Tag>, f: impl FnOnce(&mut T) -> R) -> Option<R> {
276 self.0.update(key.raw(), f)
277 }
278
279 pub fn update_parts<R>(&self, parts: KeyParts<Tag>, f: impl FnOnce(&mut T) -> R) -> Option<R> {
280 self.0.update_parts(parts.raw(), f)
281 }
282
283 pub fn remove(&self, key: Key<Tag>) -> Option<T> {
284 self.0.remove(key.raw())
285 }
286
287 pub fn remove_parts(&self, parts: KeyParts<Tag>) -> Option<T> {
288 self.0.remove_parts(parts.raw())
289 }
290
291 pub fn remove_parts_with<R>(
292 &self,
293 parts: KeyParts<Tag>,
294 f: impl FnOnce(&mut T) -> Option<R>,
295 ) -> Option<(T, R)> {
296 self.0.remove_parts_with(parts.raw(), f)
297 }
298}