kernel/sync/
once.rs

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

kernel/sync/
once.rs
