osiris/sync/
once.rs

1#![allow(dead_code)]
2
3use core::cell::UnsafeCell;
4use core::mem::MaybeUninit;
5use core::sync::atomic::AtomicU8;
6use core::sync::atomic::Ordering;
7
8/// A synchronization primitive that can be used to block a thread until a value is ready.
9/// The procedure is as follows:
10/// 1. The Caller calls step(NOT_READY) to indicate that it is about to start the initialization process.
11/// 2. The Caller initializes the value.
12/// 3. The Caller calls step(IN_TRANSIT) to indicate that the value is ready.
13///    If step 1 fails, the value is already being initialized and the Caller must wait until is() returns true.
14pub struct Ready {
15    ready: AtomicU8,
16}
17
18impl Ready {
19    const READY: u8 = 2;
20    const IN_TRANSIT: u8 = 1;
21    const NOT_READY: u8 = 0;
22
23    /// Initializes a new Ready.
24    pub const fn new() -> Self {
25        Self {
26            ready: AtomicU8::new(0),
27        }
28    }
29
30    /// Move the Ready to the next state, if it is in state `from`.
31    pub fn step(&self, from: u8) -> bool {
32        self.forward(from, from + 1)
33    }
34
35    /// Move the Ready to state `to` if it is in state `from`.
36    fn forward(&self, _from: u8, _to: u8) -> bool {
37        self.ready
38            .compare_exchange(_from, _to, Ordering::AcqRel, Ordering::Acquire)
39            .is_ok()
40    }
41
42    /// Returns true if the value is ready.
43    pub fn is(&self) -> bool {
44        self.ready.load(Ordering::Acquire) == Self::READY
45    }
46}
47
48/// A synchronization primitive that represents a value that is initialized at most once.
49pub struct OnceCell<T> {
50    value: UnsafeCell<MaybeUninit<T>>,
51    init: Ready,
52}
53
54/// Safety:
55/// 1. A `value` is only written to atomically and once.
56/// 2. A `value` is only readable from after the initialization process is finished.
57/// 3. A `init` is only written and read from atomically.
58unsafe impl<T> Sync for OnceCell<T> {}
59
60impl<T> OnceCell<T> {
61    /// Initializes a new OnceCell.
62    pub const fn new() -> Self {
63        Self {
64            value: UnsafeCell::new(MaybeUninit::uninit()),
65            init: Ready::new(),
66        }
67    }
68
69    /// Returns a reference to the value if it is initialized.
70    pub fn get(&self) -> Option<&T> {
71        if self.init.is() {
72            // Safety:
73            // 1. By contract, is the value initialized if init.is() returns true.
74            // 2. No writes are allowed to the value after the initialization process is finished.
75            Some(unsafe { self.get_unchecked() })
76        } else {
77            None
78        }
79    }
80
81    /// Sets the value if it is not already initialized, and returns a reference to the value.
82    pub fn set_or_get(&self, value: T) -> &T {
83        if let Some(value) = self.set(value) {
84            value
85        } else {
86            // If we reach this point, initialization is already in progress.
87            while !self.init.is() {
88                hal::asm::nop!();
89            }
90            // Safety:
91            // 1. By contract, is the value initialized if init.is() returns true.
92            // 2. No writes are allowed to the value after the initialization process is finished.
93            unsafe { self.get_unchecked() }
94        }
95    }
96
97    /// Sets the value if it is not already initialized, and returns a reference to the value.
98    pub fn do_or_get<F>(&self, f: F) -> &T
99    where
100        F: FnOnce() -> T,
101    {
102        self.set_or_get(f())
103    }
104
105    /// Sets the value if it is not already initialized, returns a reference to the value if it was not set previously.
106    pub fn set(&self, value: T) -> Option<&T> {
107        if self.init.is() {
108            return None;
109        }
110
111        if self.init.step(Ready::NOT_READY) {
112            // Safety: We are now in the IN_TRANSIT state, so we are the only ones that can write to the value.
113            // We are also the only ones that can read from the value.
114            unsafe {
115                self.value.get().write(MaybeUninit::new(value));
116            }
117
118            if self.init.step(Ready::IN_TRANSIT) {
119                // Safety: We are now in the READY state, so no writes can happen to the value.
120                // 1. It is safe to create a immutable reference to the value.
121                // 2. We initialized the value, so it is safe to return a reference to it.
122                return Some(unsafe { self.get_unchecked() });
123            }
124
125            // By contract, only the thread that started the initialization process can finish it.
126            unreachable!();
127        }
128
129        None
130    }
131
132    /// Returns a reference to the value, unchecked.
133    ///
134    /// # Safety
135    /// Preconditions: The value must be initialized.
136    /// Postconditions: The value is returned.
137    unsafe fn get_unchecked(&self) -> &T {
138        unsafe { (*self.value.get()).assume_init_ref() }
139    }
140}
osiris/sync/once.rs

osiris/sync/
once.rs
