wasmtime_internal_slab/
lib.rs

1//! A very simple, uniformly-typed slab arena that supports deallocation and
2//! reusing deallocated entries' space.
3//!
4//! > **⚠️ Warning ⚠️**: this crate is an internal-only crate for the Wasmtime
5//! > project and is not intended for general use. APIs are not strictly
6//! > reviewed for safety and usage outside of Wasmtime may have bugs. If
7//! > you're interested in using this feel free to file an issue on the
8//! > Wasmtime repository to start a discussion about doing so, but otherwise
9//! > be aware that your usage of this crate is not supported.
10//!
11//! The free list of vacant entries in the slab are stored inline in the slab's
12//! existing storage.
13//!
14//! # Example
15//!
16//! ```
17//! use wasmtime_internal_slab::{Id, Slab};
18//!
19//! let mut slab = Slab::new();
20//!
21//! // Insert some values into the slab.
22//! let rza = slab.alloc("Robert Fitzgerald Diggs");
23//! let gza = slab.alloc("Gary Grice");
24//! let bill = slab.alloc("Bill Gates");
25//!
26//! // Allocated elements can be accessed infallibly via indexing (and missing and
27//! // deallocated entries will panic).
28//! assert_eq!(slab[rza], "Robert Fitzgerald Diggs");
29//!
30//! // Alternatively, the `get` and `get_mut` methods provide fallible lookup.
31//! if let Some(genius) = slab.get(gza) {
32//!     println!("The gza gza genius: {}", genius);
33//! }
34//! if let Some(val) = slab.get_mut(bill) {
35//!     *val = "Bill Gates doesn't belong in this set...";
36//! }
37//!
38//! // We can remove values from the slab.
39//! slab.dealloc(bill);
40//!
41//! // Allocate a new entry.
42//! let bill = slab.alloc("Bill Murray");
43//! ```
44//!
45//! # Using `Id`s with the Wrong `Slab`
46//!
47//! `Slab` does NOT check that `Id`s used to access previously-allocated values
48//! came from the current `Slab` instance (as opposed to a different `Slab`
49//! instance). Using `Id`s from a different `Slab` is safe, but will yield an
50//! unrelated value, if any at all.
51//!
52//! If you desire checking that an `Id` came from the correct `Slab` instance,
53//! it should be easy to layer that functionality on top of this crate by
54//! wrapping `Slab` and `Id` in types that additionally maintain a slab instance
55//! identifier.
56//!
57//! # The ABA Problem
58//!
59//! This `Slab` type does NOT protect against ABA bugs, such as the following
60//! sequence:
61//!
62//! * Value `A` is allocated into the slab, yielding id `i`.
63//!
64//! * `A` is deallocated, and so `i`'s associated entry is added to the slab's
65//!   free list.
66//!
67//! * Value `B` is allocated into the slab, reusing `i`'s associated entry,
68//!   yielding id `i`.
69//!
70//! * The "original" id `i` is used to access the arena, expecting the
71//!   deallocated value `A`, but getting the new value `B`.
72//!
73//! That is, it does not detect and prevent against the memory-safe version of
74//! use-after-free bugs.
75//!
76//! If you need to protect against ABA bugs, it should be easy to layer that
77//! functionality on top of this crate by wrapping `Slab` with something like
78//! the following:
79//!
80//! ```rust
81//! pub struct GenerationalId {
82//!     id: wasmtime_internal_slab::Id,
83//!     generation: u32,
84//! }
85//!
86//! struct GenerationalEntry<T> {
87//!     value: T,
88//!     generation: u32,
89//! }
90//!
91//! pub struct GenerationalSlab<T> {
92//!     slab: wasmtime_internal_slab::Slab<GenerationalEntry<T>>,
93//!     generation: u32,
94//! }
95//!
96//! impl<T> GenerationalSlab<T> {
97//!     pub fn alloc(&mut self, value: T) -> GenerationalId {
98//!         let generation = self.generation;
99//!         let id = self.slab.alloc(GenerationalEntry { value, generation });
100//!         GenerationalId { id, generation }
101//!     }
102//!
103//!     pub fn get(&self, id: GenerationalId) -> Option<&T> {
104//!         let entry = self.slab.get(id.id)?;
105//!
106//!         // Check that the entry's generation matches the id's generation,
107//!         // else we have an ABA bug. (Alternatively, return `None` instead
108//!         // of panicking.)
109//!         assert_eq!(id.generation, entry.generation);
110//!
111//!         Some(&entry.value)
112//!     }
113//!
114//!     pub fn dealloc(&mut self, id: GenerationalId) {
115//!         // Check that the entry's generation matches the id's generation,
116//!         // else we have an ABA bug. (Alternatively, silently return on
117//!         // double-free instead of panicking.)
118//!         assert_eq!(id.generation, self.slab[id.id].generation);
119//!
120//!         self.slab.dealloc(id.id);
121//!
122//!         // Increment our generation whenever we deallocate so that any new
123//!         // value placed in this same entry will have a different generation
124//!         // and we can detect ABA bugs.
125//!         self.generation += 1;
126//!     }
127//! }
128//! ```
129
130#![no_std]
131#![forbid(unsafe_code)]
132#![deny(missing_docs, missing_debug_implementations)]
133
134extern crate alloc;
135
136use alloc::vec::Vec;
137use core::fmt;
138use core::num::NonZeroU32;
139
140/// An identifier for an allocated value inside a `slab`.
141#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
142#[repr(transparent)]
143pub struct Id(EntryIndex);
144
145impl fmt::Debug for Id {
146    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
147        f.debug_tuple("Id").field(&self.0.index()).finish()
148    }
149}
150
151impl Id {
152    /// Get the raw underlying representation of this `Id`.
153    #[inline]
154    pub fn into_raw(self) -> u32 {
155        u32::try_from(self.0.index()).unwrap()
156    }
157
158    /// Construct an `Id` from its raw underlying representation.
159    ///
160    /// `raw` should be a value that was previously created via
161    /// `Id::into_raw`. May panic if given arbitrary values.
162    #[inline]
163    pub fn from_raw(raw: u32) -> Self {
164        let raw = usize::try_from(raw).unwrap();
165        Self(EntryIndex::new(raw))
166    }
167}
168
169/// A simple, uni-typed slab arena.
170pub struct Slab<T> {
171    /// The slab's entries, each is either occupied and holding a `T` or vacant
172    /// and is a link the free list.
173    entries: Vec<Entry<T>>,
174
175    /// The index of the first free entry in the free list.
176    free: Option<EntryIndex>,
177
178    /// The number of occupied entries is this slab.
179    len: u32,
180}
181
182impl<T> fmt::Debug for Slab<T>
183where
184    T: fmt::Debug,
185{
186    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
187        f.debug_map().entries(self.iter()).finish()
188    }
189}
190
191enum Entry<T> {
192    /// An occupied entry holding a `T`.
193    Occupied(T),
194
195    /// A vacant entry.
196    Free {
197        /// A link in the slab's free list, pointing to the next free entry, if
198        /// any.
199        next_free: Option<EntryIndex>,
200    },
201}
202
203#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
204#[repr(transparent)]
205struct EntryIndex(NonZeroU32);
206
207impl EntryIndex {
208    #[inline]
209    fn new(index: usize) -> Self {
210        assert!(index <= Slab::<()>::MAX_CAPACITY);
211        let x = u32::try_from(index + 1).unwrap();
212        Self(NonZeroU32::new(x).unwrap())
213    }
214
215    #[inline]
216    fn index(&self) -> usize {
217        let index = self.0.get() - 1;
218        usize::try_from(index).unwrap()
219    }
220}
221
222impl<T> Default for Slab<T> {
223    #[inline]
224    fn default() -> Self {
225        Self {
226            entries: Vec::default(),
227            free: None,
228            len: 0,
229        }
230    }
231}
232
233impl<T> core::ops::Index<Id> for Slab<T> {
234    type Output = T;
235
236    #[inline]
237    fn index(&self, id: Id) -> &Self::Output {
238        self.get(id)
239            .expect("id from different slab or value was deallocated")
240    }
241}
242
243impl<T> core::ops::IndexMut<Id> for Slab<T> {
244    #[inline]
245    fn index_mut(&mut self, id: Id) -> &mut Self::Output {
246        self.get_mut(id)
247            .expect("id from different slab or value was deallocated")
248    }
249}
250
251impl<T> Slab<T> {
252    /// The maximum capacity any `Slab` can have: `u32::MAX - 1`.
253    pub const MAX_CAPACITY: usize = (u32::MAX - 1) as usize;
254
255    /// Construct a new, empty slab.
256    #[inline]
257    pub fn new() -> Self {
258        Slab::default()
259    }
260
261    /// Construct a new, empty slab, pre-reserving space for at least `capacity`
262    /// elements.
263    #[inline]
264    pub fn with_capacity(capacity: usize) -> Self {
265        let mut slab = Self::new();
266        slab.reserve(capacity);
267        slab
268    }
269
270    /// Ensure that there is space for at least `additional` elements in this
271    /// slab.
272    ///
273    /// # Panics
274    ///
275    /// Panics if the new capacity exceeds `Self::MAX_CAPACITY`.
276    pub fn reserve(&mut self, additional: usize) {
277        let cap = self.capacity();
278        let len = self.len();
279        assert!(cap >= len);
280        if cap - len >= additional {
281            // Already have `additional` capacity available.
282            return;
283        }
284
285        self.entries.reserve(additional);
286
287        // Maintain the invariant that `i <= MAX_CAPACITY` for all indices `i`
288        // in `self.entries`.
289        assert!(self.entries.capacity() <= Self::MAX_CAPACITY);
290    }
291
292    fn double_capacity(&mut self) {
293        // Double our capacity to amortize the cost of resizing. But make sure
294        // we add some amount of minimum additional capacity, since doubling
295        // zero capacity isn't useful.
296        const MIN_CAPACITY: usize = 16;
297        let additional = core::cmp::max(self.entries.capacity(), MIN_CAPACITY);
298        self.reserve(additional);
299    }
300
301    /// What is the capacity of this slab? That is, how many entries can it
302    /// contain within its current underlying storage?
303    #[inline]
304    pub fn capacity(&self) -> usize {
305        self.entries.capacity()
306    }
307
308    /// How many values are currently allocated within this slab?
309    #[inline]
310    pub fn len(&self) -> usize {
311        usize::try_from(self.len).unwrap()
312    }
313
314    /// Are there zero allocated values within this slab?
315    #[inline]
316    pub fn is_empty(&self) -> bool {
317        self.len() == 0
318    }
319
320    /// Try to allocate a `T` value within this slab.
321    ///
322    /// If there is no available capacity, ownership of the given value is
323    /// returned via `Err(value)`.
324    #[inline]
325    pub fn try_alloc(&mut self, value: T) -> Result<Id, T> {
326        if let Some(index) = self.try_alloc_index() {
327            let next_free = match self.entries[index.index()] {
328                Entry::Free { next_free } => next_free,
329                Entry::Occupied { .. } => unreachable!(),
330            };
331            self.free = next_free;
332            self.entries[index.index()] = Entry::Occupied(value);
333            self.len += 1;
334            Ok(Id(index))
335        } else {
336            Err(value)
337        }
338    }
339
340    #[inline]
341    fn try_alloc_index(&mut self) -> Option<EntryIndex> {
342        self.free.take().or_else(|| {
343            if self.entries.len() < self.entries.capacity() {
344                let index = EntryIndex::new(self.entries.len());
345                self.entries.push(Entry::Free { next_free: None });
346                Some(index)
347            } else {
348                None
349            }
350        })
351    }
352
353    /// Allocate a `T` value within this slab, allocating additional underlying
354    /// storage if there is no available capacity.
355    ///
356    /// # Panics
357    ///
358    /// Panics if allocating this value requires reallocating the underlying
359    /// storage, and the new capacity exceeds `Slab::MAX_CAPACITY`.
360    #[inline]
361    pub fn alloc(&mut self, value: T) -> Id {
362        self.try_alloc(value)
363            .unwrap_or_else(|value| self.alloc_slow(value))
364    }
365
366    /// Get the `Id` that will be returned for the next allocation in this slab.
367    #[inline]
368    pub fn next_id(&self) -> Id {
369        let index = self.free.unwrap_or_else(|| EntryIndex::new(self.len()));
370        Id(index)
371    }
372
373    #[inline(never)]
374    #[cold]
375    fn alloc_slow(&mut self, value: T) -> Id {
376        // Reserve additional capacity, since we didn't have space for the
377        // allocation.
378        self.double_capacity();
379        // After which the allocation will succeed.
380        self.try_alloc(value).ok().unwrap()
381    }
382
383    /// Get a shared borrow of the value associated with `id`.
384    ///
385    /// Returns `None` if the value has since been deallocated.
386    ///
387    /// If `id` comes from a different `Slab` instance, this method may panic,
388    /// return `None`, or return an arbitrary value.
389    #[inline]
390    pub fn get(&self, id: Id) -> Option<&T> {
391        match self
392            .entries
393            .get(id.0.index())
394            .expect("id from different slab")
395        {
396            Entry::Occupied(x) => Some(x),
397            Entry::Free { .. } => None,
398        }
399    }
400
401    /// Get an exclusive borrow of the value associated with `id`.
402    ///
403    /// Returns `None` if the value has since been deallocated.
404    ///
405    /// If `id` comes from a different `Slab` instance, this method may panic,
406    /// return `None`, or return an arbitrary value.
407    #[inline]
408    pub fn get_mut(&mut self, id: Id) -> Option<&mut T> {
409        match self
410            .entries
411            .get_mut(id.0.index())
412            .expect("id from different slab")
413        {
414            Entry::Occupied(x) => Some(x),
415            Entry::Free { .. } => None,
416        }
417    }
418
419    /// Does this slab contain an allocated value for `id`?
420    #[inline]
421    pub fn contains(&self, id: Id) -> bool {
422        match self.entries.get(id.0.index()) {
423            Some(Entry::Occupied(_)) => true,
424            None | Some(Entry::Free { .. }) => false,
425        }
426    }
427
428    /// Deallocate the value associated with the given `id`.
429    ///
430    /// If `id` comes from a different `Slab` instance, this method may panic or
431    /// deallocate an arbitrary value.
432    #[inline]
433    pub fn dealloc(&mut self, id: Id) -> T {
434        let entry = core::mem::replace(
435            self.entries
436                .get_mut(id.0.index())
437                .expect("id from a different slab"),
438            Entry::Free { next_free: None },
439        );
440        match entry {
441            Entry::Free { .. } => panic!("attempt to deallocate an entry that is already vacant"),
442            Entry::Occupied(value) => {
443                let next_free = core::mem::replace(&mut self.free, Some(id.0));
444                self.entries[id.0.index()] = Entry::Free { next_free };
445                self.len -= 1;
446                value
447            }
448        }
449    }
450
451    /// Iterate over all values currently allocated within this `Slab`.
452    ///
453    /// Yields pairs of an `Id` and the `Id`'s associated value.
454    ///
455    /// Iteration order is undefined.
456    #[inline]
457    pub fn iter(&self) -> impl Iterator<Item = (Id, &T)> + '_ {
458        assert!(self.entries.len() <= Self::MAX_CAPACITY);
459        self.entries
460            .iter()
461            .enumerate()
462            .filter_map(|(i, e)| match e {
463                Entry::Occupied(x) => Some((Id(EntryIndex::new(i)), x)),
464                Entry::Free { .. } => None,
465            })
466    }
467
468    /// Mutably iterate over all values currently allocated within this `Slab`.
469    ///
470    /// Yields pairs of an `Id` and the `Id`'s associated value.
471    ///
472    /// Iteration order is undefined.
473    #[inline]
474    pub fn iter_mut(&mut self) -> impl Iterator<Item = (Id, &mut T)> + '_ {
475        assert!(self.entries.len() <= Self::MAX_CAPACITY);
476        self.entries
477            .iter_mut()
478            .enumerate()
479            .filter_map(|(i, e)| match e {
480                Entry::Occupied(x) => Some((Id(EntryIndex::new(i)), x)),
481                Entry::Free { .. } => None,
482            })
483    }
484
485    /// Iterate over and remove all entries in this slab.
486    ///
487    /// The slab will be empty after calling this method.
488    ///
489    /// Yields pairs of an `Id` and the `Id`'s associated value.
490    ///
491    /// Iteration order is undefined.
492    #[inline]
493    pub fn drain(&mut self) -> impl Iterator<Item = (Id, T)> + '_ {
494        assert!(self.entries.len() <= Self::MAX_CAPACITY);
495        self.len = 0;
496        self.free = None;
497        self.entries
498            .drain(..)
499            .enumerate()
500            .filter_map(|(i, e)| match e {
501                Entry::Occupied(x) => Some((Id(EntryIndex::new(i)), x)),
502                Entry::Free { .. } => None,
503            })
504    }
505}