maitake_sync/spin/
once.rs

1//! Cells storing a value which must be initialized prior to use.
2//!
3//! This module provides:
4//!
5//! - [`InitOnce`]: a cell storing a [`MaybeUninit`](core::mem::MaybeUninit)
6//!   value which must be manually initialized prior to use.
7//! - [`Lazy`]: an [`InitOnce`] cell coupled with an initializer function. The
8//!   [`Lazy`] cell ensures the initializer is called to initialize the
9//!   value the first time it is accessed.
10use crate::{
11    loom::sync::atomic::{AtomicU8, Ordering},
12    util::{Backoff, CheckedMaybeUninit},
13};
14use core::{
15    any,
16    cell::UnsafeCell,
17    fmt,
18    ops::{Deref, DerefMut},
19};
20
21/// A cell which may be initialized a single time after it is created.
22///
23/// This can be used as a safer alternative to `static mut`.
24///
25/// For performance-critical use-cases, this type also has a [`get_unchecked`]
26/// method, which dereferences the cell **without** checking if it has been
27/// initialized. This method is unsafe and should be used with caution ---
28/// incorrect usage can result in reading uninitialized memory.
29///
30/// [`get_unchecked`]: Self::get_unchecked
31pub struct InitOnce<T> {
32    value: UnsafeCell<CheckedMaybeUninit<T>>,
33    state: AtomicU8,
34}
35
36/// A cell which will be lazily initialized by the provided function the first
37/// time it is accessed.
38///
39/// This can be used as a safer alternative to `static mut`.
40pub struct Lazy<T, F = fn() -> T> {
41    value: UnsafeCell<CheckedMaybeUninit<T>>,
42    state: AtomicU8,
43    initializer: F,
44}
45
46/// Errors returned by [`InitOnce::try_init`].
47///
48/// This contains the value that the caller was attempting to use to initialize
49/// the cell.
50pub struct TryInitError<T> {
51    value: T,
52    actual: u8,
53}
54
55const UNINITIALIZED: u8 = 0;
56const INITIALIZING: u8 = 1;
57const INITIALIZED: u8 = 2;
58
59// === impl InitOnce ===
60
61impl<T> InitOnce<T> {
62    loom_const_fn! {
63        /// Returns a new `InitOnce` in the uninitialized state.
64        #[must_use]
65        pub fn uninitialized() -> Self {
66            Self {
67                value: UnsafeCell::new(CheckedMaybeUninit::uninit()),
68                state: AtomicU8::new(UNINITIALIZED),
69            }
70        }
71    }
72
73    /// Initialize the cell to `value`, returning an error if it has already
74    /// been initialized.
75    ///
76    /// If the cell has already been initialized, the returned error contains
77    /// the value.
78    pub fn try_init(&self, value: T) -> Result<(), TryInitError<T>> {
79        if let Err(actual) = self.state.compare_exchange(
80            UNINITIALIZED,
81            INITIALIZING,
82            Ordering::AcqRel,
83            Ordering::Acquire,
84        ) {
85            return Err(TryInitError { value, actual });
86        };
87        unsafe {
88            *(self.value.get()) = CheckedMaybeUninit::new(value);
89        }
90        let _prev = self.state.swap(INITIALIZED, Ordering::AcqRel);
91        debug_assert_eq!(
92            _prev,
93            INITIALIZING,
94            "InitOnce<{}>: state changed while locked. This is a bug!",
95            any::type_name::<T>(),
96        );
97        Ok(())
98    }
99
100    /// Initialize the cell to `value`, panicking if it has already been
101    /// initialized.
102    ///
103    /// # Panics
104    ///
105    /// If the cell has already been initialized.
106    #[track_caller]
107    pub fn init(&self, value: T) -> &T {
108        self.try_init(value).unwrap();
109        self.get()
110    }
111
112    /// Borrow the contents of this `InitOnce` cell, if it has been
113    /// initialized. Otherwise, if the cell has not yet been initialized, this
114    /// returns [`None`].
115    #[inline]
116    #[must_use]
117    pub fn try_get(&self) -> Option<&T> {
118        if self.state.load(Ordering::Acquire) != INITIALIZED {
119            return None;
120        }
121        unsafe {
122            // Safety: we just checked if the value was initialized.
123            Some(&*((*self.value.get()).as_ptr()))
124        }
125    }
126
127    /// Borrow the contents of this `InitOnce` cell, or panic if it has not
128    /// been initialized.
129    ///
130    /// # Panics
131    ///
132    /// If the cell has not yet been initialized.
133    #[track_caller]
134    #[inline]
135    #[must_use]
136    pub fn get(&self) -> &T {
137        if self.state.load(Ordering::Acquire) != INITIALIZED {
138            panic!("InitOnce<{}> not yet initialized!", any::type_name::<T>());
139        }
140        unsafe {
141            // Safety: we just checked if the value was initialized.
142            &*((*self.value.get()).as_ptr())
143        }
144    }
145
146    /// Borrow the contents of this `InitOnce` cell, or initialize it with the
147    /// provided closure.
148    ///
149    /// If the cell has been initialized, this returns the current value.
150    /// Otherwise, it calls the closure, puts the returned value from the
151    /// closure in the cell, and borrows the current value.
152    #[must_use]
153    pub fn get_or_else(&self, f: impl FnOnce() -> T) -> &T {
154        if let Some(val) = self.try_get() {
155            return val;
156        }
157
158        let _ = self.try_init(f());
159        self.get()
160    }
161
162    /// Borrow the contents of this `InitOnce` cell, **without** checking
163    /// whether it has been initialized.
164    ///
165    /// # Safety
166    ///
167    /// The caller is responsible for ensuring that the value has already been
168    /// initialized.
169    ///
170    /// In debug mode, this still checks the state of the cell, so if it has yet
171    /// to be initialized, this will panic. However, in release mode builds,
172    /// this is completely unchecked. If the value has not yet been initialized,
173    /// this may return a pointer to uninitialized memory! It may also return a
174    /// pointer to memory that is currently being written to.
175    ///
176    /// If you see this method panic in debug mode, please, **please** re-check
177    /// your code.
178    #[cfg_attr(not(debug_assertions), inline(always))]
179    #[cfg_attr(debug_assertions, track_caller)]
180    #[must_use]
181    pub unsafe fn get_unchecked(&self) -> &T {
182        debug_assert_eq!(
183            INITIALIZED,
184            self.state.load(Ordering::Acquire),
185            "InitOnce<{}>: accessed before initialized!\n\
186            /!\\ EXTREMELY SERIOUS WARNING: /!\\ This is REAL BAD! If you were \
187            running in release mode, you would have just read uninitialized \
188            memory! That's bad news indeed, buddy. Double- or triple-check \
189            your assumptions, or consider Just Using A Goddamn Mutex --- it's \
190            much safer that way. Maybe this whole `InitOnce` thing was a \
191            mistake...
192            ",
193            any::type_name::<T>(),
194        );
195        unsafe {
196            // Safety: hahaha wheeee no rules! You can't stop meeeeee!
197            &*((*self.value.get()).as_ptr())
198        }
199    }
200}
201
202impl<T: fmt::Debug> fmt::Debug for InitOnce<T> {
203    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
204        match self.state.load(Ordering::Acquire) {
205            INITIALIZED => self.get().fmt(f),
206            INITIALIZING => f.pad("<initializing>"),
207            UNINITIALIZED => f.pad("<uninitialized>"),
208            _state => unsafe {
209                unreachable_unchecked!("unexpected state value {}, this is a bug!", _state)
210            },
211        }
212    }
213}
214
215unsafe impl<T: Send> Send for InitOnce<T> {}
216unsafe impl<T: Sync> Sync for InitOnce<T> {}
217
218// === impl Lazy ===
219
220impl<T, F> Lazy<T, F> {
221    loom_const_fn! {
222        /// Returns a new `Lazy` cell, initialized with the provided `initializer`
223        /// function.
224        #[must_use]
225        pub fn new(initializer: F) -> Self {
226            Self {
227                value: UnsafeCell::new(CheckedMaybeUninit::uninit()),
228                state: AtomicU8::new(UNINITIALIZED),
229                initializer,
230            }
231        }
232    }
233
234    /// Returns the value of the lazy cell, if it has already been initialized.
235    /// Otherwise, returns `None`.
236    #[inline]
237    #[must_use]
238    pub fn get_if_present(&self) -> Option<&T> {
239        if self.state.load(Ordering::Acquire) == INITIALIZED {
240            let value = unsafe {
241                // Safety: we just ensured the cell was initialized.
242                &*((*self.value.get()).as_ptr())
243            };
244            Some(value)
245        } else {
246            None
247        }
248    }
249}
250
251impl<T, F> Lazy<T, F>
252where
253    F: Fn() -> T,
254{
255    /// Borrow the value, or initialize it if it has not yet been initialized.
256    #[inline]
257    #[must_use]
258    pub fn get(&self) -> &T {
259        self.init();
260        unsafe {
261            // Safety: we just ensured the cell was initialized.
262            &*((*self.value.get()).as_ptr())
263        }
264    }
265
266    /// Borrow the value mutably, or initialize it if it has not yet been initialized.
267    #[inline]
268    #[must_use]
269    pub fn get_mut(&mut self) -> &mut T {
270        self.init();
271        unsafe {
272            // Safety: we just ensured the cell was initialized.
273            &mut *((*self.value.get()).as_mut_ptr())
274        }
275    }
276
277    /// Ensure that the cell has been initialized.
278    ///
279    /// If the cell has yet to be initialized, this initializes it. If it is
280    /// currently initializing, this spins until it has been fully initialized.
281    /// Otherwise, this returns immediately.
282    pub fn init(&self) {
283        let state = self.state.compare_exchange(
284            UNINITIALIZED,
285            INITIALIZING,
286            Ordering::AcqRel,
287            Ordering::Acquire,
288        );
289
290        match state {
291            Err(INITIALIZED) => {
292                // Already initialized! Just return the value.
293            }
294            Err(INITIALIZING) => {
295                // Wait for the cell to be initialized
296                let mut backoff = Backoff::new();
297                while self.state.load(Ordering::Acquire) != INITIALIZED {
298                    backoff.spin();
299                }
300            }
301            Ok(_) => {
302                // Now we have to actually initialize the cell.
303                unsafe {
304                    *(self.value.get()) = CheckedMaybeUninit::new((self.initializer)());
305                }
306                if let Err(actual) = self.state.compare_exchange(
307                    INITIALIZING,
308                    INITIALIZED,
309                    Ordering::AcqRel,
310                    Ordering::Acquire,
311                ) {
312                    unreachable!(
313                        "Lazy<{}>: state changed while locked. This is a bug! (state={})",
314                        any::type_name::<T>(),
315                        actual
316                    );
317                }
318            }
319            Err(_state) => unsafe {
320                unreachable_unchecked!(
321                    "Lazy<{}>: unexpected state {}!. This is a bug!",
322                    any::type_name::<T>(),
323                    _state
324                )
325            },
326        };
327    }
328}
329
330impl<T, F> Deref for Lazy<T, F>
331where
332    F: Fn() -> T,
333{
334    type Target = T;
335
336    fn deref(&self) -> &Self::Target {
337        self.get()
338    }
339}
340
341impl<T, F> DerefMut for Lazy<T, F>
342where
343    F: Fn() -> T,
344{
345    fn deref_mut(&mut self) -> &mut Self::Target {
346        self.get_mut()
347    }
348}
349
350impl<T, F> fmt::Debug for Lazy<T, F>
351where
352    T: fmt::Debug,
353{
354    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
355        match self.state.load(Ordering::Acquire) {
356            INITIALIZED => self
357                .get_if_present()
358                .expect("if state is `INITIALIZED`, value should be present")
359                .fmt(f),
360            INITIALIZING => f.pad("<initializing>"),
361            UNINITIALIZED => f.pad("<uninitialized>"),
362            _state => unsafe {
363                unreachable_unchecked!("unexpected state value {}, this is a bug!", _state)
364            },
365        }
366    }
367}
368
369unsafe impl<T: Send, F: Send> Send for Lazy<T, F> {}
370unsafe impl<T: Sync, F: Sync> Sync for Lazy<T, F> {}
371
372// === impl TryInitError ===
373
374impl<T> TryInitError<T> {
375    /// Returns the value that the caller attempted to initialize the cell with
376    #[must_use]
377    pub fn into_inner(self) -> T {
378        self.value
379    }
380}
381
382impl<T> fmt::Debug for TryInitError<T> {
383    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
384        f.debug_struct("TryInitError")
385            .field("type", &any::type_name::<T>())
386            .field("value", &format_args!("..."))
387            .field(
388                "state",
389                &format_args!("State::{}", match self.actual {
390                    UNINITIALIZED => "UNINITIALIZED",
391                    INITIALIZING => "INITIALIZING",
392                    INITIALIZED => unsafe { unreachable_unchecked!("an error should not be returned when InitOnce is in the initialized state, this is a bug!") },
393                    _state => unsafe { unreachable_unchecked!("unexpected state value {}, this is a bug!", _state) },
394                }),
395            )
396            .finish()
397    }
398}
399
400impl<T> fmt::Display for TryInitError<T> {
401    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
402        write!(f, "InitOnce<{}> already initialized", any::type_name::<T>())
403    }
404}
405
406#[cfg(feature = "core-error")]
407impl<T> core::error::Error for TryInitError<T> {}
maitake_sync/spin/once.rs

maitake_sync/spin/
once.rs
