1 | // Based on unstable std::sync::OnceLock. |
2 | // |
3 | // Source: https://github.com/rust-lang/rust/blob/8e9c93df464b7ada3fc7a1c8ccddd9dcb24ee0a0/library/std/src/sync/once_lock.rs |
4 | |
5 | use core::cell::UnsafeCell; |
6 | use core::mem::MaybeUninit; |
7 | use std::sync::Once; |
8 | |
9 | pub(crate) struct OnceLock<T> { |
10 | once: Once, |
11 | value: UnsafeCell<MaybeUninit<T>>, |
12 | // Unlike std::sync::OnceLock, we don't need PhantomData here because |
13 | // we don't use #[may_dangle]. |
14 | } |
15 | |
16 | unsafe impl<T: Sync + Send> Sync for OnceLock<T> {} |
17 | unsafe impl<T: Send> Send for OnceLock<T> {} |
18 | |
19 | impl<T> OnceLock<T> { |
20 | /// Creates a new empty cell. |
21 | #[must_use ] |
22 | pub(crate) const fn new() -> Self { |
23 | Self { |
24 | once: Once::new(), |
25 | value: UnsafeCell::new(MaybeUninit::uninit()), |
26 | } |
27 | } |
28 | |
29 | /// Gets the contents of the cell, initializing it with `f` if the cell |
30 | /// was empty. |
31 | /// |
32 | /// Many threads may call `get_or_init` concurrently with different |
33 | /// initializing functions, but it is guaranteed that only one function |
34 | /// will be executed. |
35 | /// |
36 | /// # Panics |
37 | /// |
38 | /// If `f` panics, the panic is propagated to the caller, and the cell |
39 | /// remains uninitialized. |
40 | /// |
41 | /// It is an error to reentrantly initialize the cell from `f`. The |
42 | /// exact outcome is unspecified. Current implementation deadlocks, but |
43 | /// this may be changed to a panic in the future. |
44 | pub(crate) fn get_or_init<F>(&self, f: F) -> &T |
45 | where |
46 | F: FnOnce() -> T, |
47 | { |
48 | // Fast path check |
49 | if self.once.is_completed() { |
50 | // SAFETY: The inner value has been initialized |
51 | return unsafe { self.get_unchecked() }; |
52 | } |
53 | self.initialize(f); |
54 | |
55 | // SAFETY: The inner value has been initialized |
56 | unsafe { self.get_unchecked() } |
57 | } |
58 | |
59 | #[cold ] |
60 | fn initialize<F>(&self, f: F) |
61 | where |
62 | F: FnOnce() -> T, |
63 | { |
64 | let slot = self.value.get().cast::<T>(); |
65 | |
66 | self.once.call_once(|| { |
67 | let value = f(); |
68 | unsafe { |
69 | slot.write(value); |
70 | } |
71 | }); |
72 | } |
73 | |
74 | /// # Safety |
75 | /// |
76 | /// The value must be initialized |
77 | unsafe fn get_unchecked(&self) -> &T { |
78 | debug_assert!(self.once.is_completed()); |
79 | &*self.value.get().cast::<T>() |
80 | } |
81 | } |
82 | |
83 | impl<T> Drop for OnceLock<T> { |
84 | fn drop(&mut self) { |
85 | if self.once.is_completed() { |
86 | // SAFETY: The inner value has been initialized |
87 | unsafe { self.value.get().cast::<T>().drop_in_place() }; |
88 | } |
89 | } |
90 | } |
91 | |