forge/os/
once_lock.rs

1use core::{
2    cell::UnsafeCell,
3    mem::MaybeUninit,
4    sync::atomic::{AtomicU8, Ordering},
5};
6
7use sys::os::EventClearMode;
8
9use crate::os::sync::LightEvent;
10
11const UNINIT: u8 = 0;
12const INITING: u8 = 1;
13const INIT: u8 = 2;
14
15/// A cell that can be written exactly once and is safe to share across threads.
16///
17/// Equivalent to `std::sync::OnceLock`. The first call to [`get_or_init`](Self::get_or_init)
18/// runs the initializer; all subsequent calls block until initialization completes
19/// and then return a reference to the stored value.
20pub struct OnceLock<T> {
21    state: AtomicU8,
22    ready: LightEvent,
23    value: UnsafeCell<MaybeUninit<T>>,
24}
25
26unsafe impl<T: Send + Sync> Sync for OnceLock<T> {}
27unsafe impl<T: Send> Send for OnceLock<T> {}
28
29impl<T> OnceLock<T> {
30    /// Creates a new uninitialized `OnceLock`.
31    pub const fn new() -> Self {
32        Self {
33            state: AtomicU8::new(UNINIT),
34            ready: LightEvent::new(false, EventClearMode::Manual),
35            value: UnsafeCell::new(MaybeUninit::uninit()),
36        }
37    }
38
39    /// Returns a reference to the value if initialized, otherwise `None`.
40    pub fn get(&self) -> Option<&T> {
41        if self.state.load(Ordering::Acquire) == INIT {
42            Some(unsafe { (*self.value.get()).assume_init_ref() })
43        } else {
44            None
45        }
46    }
47
48    /// Returns a mutable reference to the value if initialized, otherwise `None`.
49    pub fn get_mut(&mut self) -> Option<&mut T> {
50        if self.state.load(Ordering::Acquire) == INIT {
51            Some(unsafe { (*self.value.get_mut()).assume_init_mut() })
52        } else {
53            None
54        }
55    }
56
57    /// Returns the value, initializing it with `f` if this is the first call.
58    ///
59    /// If another thread is already initializing the value, this call blocks until
60    /// initialization completes.
61    pub fn get_or_init(&self, f: impl FnOnce() -> T) -> &T {
62        if self.state.load(Ordering::Acquire) == INIT {
63            return unsafe { (*self.value.get()).assume_init_ref() };
64        }
65
66        match self
67            .state
68            .compare_exchange(UNINIT, INITING, Ordering::Acquire, Ordering::Acquire)
69        {
70            Ok(_) => {
71                unsafe { (*self.value.get()).write(f()) };
72                self.state.store(INIT, Ordering::Release);
73                self.ready.signal();
74            }
75            Err(_) => {
76                self.ready.wait();
77            }
78        }
79
80        unsafe { (*self.value.get()).assume_init_ref() }
81    }
82}
83
84impl<T> Drop for OnceLock<T> {
85    fn drop(&mut self) {
86        if *self.state.get_mut() == INIT {
87            unsafe { (*self.value.get()).assume_init_drop() };
88        }
89    }
90}
forge/os/once_lock.rs

forge/os/
once_lock.rs
