Skip to main content

odem_rs_util/pool/
backing.rs

1use super::Storage;
2use core::{
3	cell::{Cell, UnsafeCell},
4	mem::MaybeUninit,
5	pin::Pin,
6};
7
8/// A fixed-size array that can be used as a [`Storage`] backend for a [`Pool`].
9///
10/// [`Pool`]: super::Pool
11pub struct Array<A, const N: usize> {
12	/// The number of currently initialized elements.
13	len: Cell<usize>,
14	/// The raw memory buffer.
15	mem: [UnsafeCell<MaybeUninit<A>>; N],
16}
17
18impl<A, const N: usize> Array<A, N> {
19	/// Creates a new, empty `Array` storage.
20	pub const fn new() -> Self {
21		Array {
22			len: Cell::new(0),
23			// SAFETY: An array of `MaybeUninit` does not require initialization.
24			mem: [const { UnsafeCell::new(MaybeUninit::uninit()) }; N],
25		}
26	}
27}
28
29impl<A, const N: usize> Default for Array<A, N> {
30	fn default() -> Self {
31		Array::new()
32	}
33}
34
35impl<A, const N: usize> Storage for Array<A, N> {
36	type Slot = A;
37
38	fn try_alloc(self: Pin<&Self>, value: A) -> Option<Pin<&mut A>> {
39		// Check if there is space for a new allocation.
40		let res = self.mem.get(self.len.get())?;
41		self.len.set(self.len.get() + 1);
42
43		// Write the value into the uninitialized slot and return a pinned reference.
44		// SAFETY: Pinning is safe because the `Array` itself is pinned, and the
45		// memory location of the element will not change until it is dropped.
46		Some(unsafe { Pin::new_unchecked((*res.get()).write(value)) })
47	}
48
49	fn len(&self) -> usize {
50		self.len.get()
51	}
52}
53
54impl<A, const N: usize> Drop for Array<A, N> {
55	fn drop(&mut self) {
56		use core::{mem::transmute, ptr::drop_in_place};
57		let len = self.len.get();
58
59		// Manually drop the first `len` initialized elements in place.
60		// SAFETY: We are only dropping the elements that we have explicitly
61		// initialized. The slice is transmuted from `MaybeUninit<A>` to `A`
62		// because we know these `len` elements are valid.
63		unsafe {
64			drop_in_place(transmute::<
65				&mut [UnsafeCell<MaybeUninit<A>>],
66				&mut [UnsafeCell<A>],
67			>(&mut self.mem[..len]));
68		}
69	}
70}
71
72// Implements `Storage` for `typed_arena::Arena` when the `alloc` feature is enabled.
73#[cfg(feature = "alloc")]
74#[cfg_attr(docsrs, doc(cfg(feature = "alloc")))]
75impl<A> Storage for typed_arena::Arena<A> {
76	type Slot = A;
77
78	fn alloc(self: Pin<&Self>, value: A) -> Pin<&mut A> {
79		let value = self.get_ref().alloc(value);
80
81		// SAFETY: The arena guarantees stable addresses for its allocations,
82		// so creating a pinned reference is safe.
83		unsafe { Pin::new_unchecked(value) }
84	}
85
86	fn try_alloc(self: Pin<&Self>, value: A) -> Option<Pin<&mut A>> {
87		Some(Storage::alloc(self, value))
88	}
89
90	fn len(&self) -> usize {
91		self.len()
92	}
93}