1//! Synchronization primitive for initializing a value once, allowing others to await a reference to the value.
2
3use core::cell::Cell;
4use core::future::{poll_fn, Future};
5use core::mem::MaybeUninit;
6use core::sync::atomic::{AtomicBool, Ordering};
7use core::task::Poll;
8
9/// The `OnceLock` is a synchronization primitive that allows for
10/// initializing a value once, and allowing others to `.await` a
11/// reference to the value. This is useful for lazy initialization of
12/// a static value.
13///
14/// **Note**: this implementation uses a busy loop to poll the value,
15/// which is not as efficient as registering a dedicated `Waker`.
16/// However, if the usecase for it is to initialize a static variable
17/// relatively early in the program life cycle, it should be fine.
18///
19/// # Example
20/// ```
21/// use futures_executor::block_on;
22/// use embassy_sync::once_lock::OnceLock;
23///
24/// // Define a static value that will be lazily initialized
25/// static VALUE: OnceLock<u32> = OnceLock::new();
26///
27/// let f = async {
28///
29/// // Initialize the value
30/// let reference = VALUE.get_or_init(|| 20);
31/// assert_eq!(reference, &20);
32///
33/// // Wait for the value to be initialized
34/// // and get a static reference it
35/// assert_eq!(VALUE.get().await, &20);
36///
37/// };
38/// block_on(f)
39/// ```
40pub struct OnceLock<T> {
41 init: AtomicBool,
42 data: Cell<MaybeUninit<T>>,
43}
44
45unsafe impl<T> Sync for OnceLock<T> {}
46
47impl<T> OnceLock<T> {
48 /// Create a new uninitialized `OnceLock`.
49 pub const fn new() -> Self {
50 Self {
51 init: AtomicBool::new(false),
52 data: Cell::new(MaybeUninit::zeroed()),
53 }
54 }
55
56 /// Get a reference to the underlying value, waiting for it to be set.
57 /// If the value is already set, this will return immediately.
58 pub fn get(&self) -> impl Future<Output = &T> {
59 poll_fn(|cx| match self.try_get() {
60 Some(data) => Poll::Ready(data),
61 None => {
62 cx.waker().wake_by_ref();
63 Poll::Pending
64 }
65 })
66 }
67
68 /// Try to get a reference to the underlying value if it exists.
69 pub fn try_get(&self) -> Option<&T> {
70 if self.init.load(Ordering::Relaxed) {
71 Some(unsafe { self.get_ref_unchecked() })
72 } else {
73 None
74 }
75 }
76
77 /// Set the underlying value. If the value is already set, this will return an error with the given value.
78 pub fn init(&self, value: T) -> Result<(), T> {
79 // Critical section is required to ensure that the value is
80 // not simultaneously initialized elsewhere at the same time.
81 critical_section::with(|_| {
82 // If the value is not set, set it and return Ok.
83 if !self.init.load(Ordering::Relaxed) {
84 self.data.set(MaybeUninit::new(value));
85 self.init.store(true, Ordering::Relaxed);
86 Ok(())
87
88 // Otherwise return an error with the given value.
89 } else {
90 Err(value)
91 }
92 })
93 }
94
95 /// Get a reference to the underlying value, initializing it if it does not exist.
96 pub fn get_or_init<F>(&self, f: F) -> &T
97 where
98 F: FnOnce() -> T,
99 {
100 // Critical section is required to ensure that the value is
101 // not simultaneously initialized elsewhere at the same time.
102 critical_section::with(|_| {
103 // If the value is not set, set it.
104 if !self.init.load(Ordering::Relaxed) {
105 self.data.set(MaybeUninit::new(f()));
106 self.init.store(true, Ordering::Relaxed);
107 }
108 });
109
110 // Return a reference to the value.
111 unsafe { self.get_ref_unchecked() }
112 }
113
114 /// Consume the `OnceLock`, returning the underlying value if it was initialized.
115 pub fn into_inner(self) -> Option<T> {
116 if self.init.load(Ordering::Relaxed) {
117 Some(unsafe { self.data.into_inner().assume_init() })
118 } else {
119 None
120 }
121 }
122
123 /// Take the underlying value if it was initialized, uninitializing the `OnceLock` in the process.
124 pub fn take(&mut self) -> Option<T> {
125 // If the value is set, uninitialize the lock and return the value.
126 critical_section::with(|_| {
127 if self.init.load(Ordering::Relaxed) {
128 let val = unsafe { self.data.replace(MaybeUninit::zeroed()).assume_init() };
129 self.init.store(false, Ordering::Relaxed);
130 Some(val)
131
132 // Otherwise return None.
133 } else {
134 None
135 }
136 })
137 }
138
139 /// Check if the value has been set.
140 pub fn is_set(&self) -> bool {
141 self.init.load(Ordering::Relaxed)
142 }
143
144 /// Get a reference to the underlying value.
145 /// # Safety
146 /// Must only be used if a value has been set.
147 unsafe fn get_ref_unchecked(&self) -> &T {
148 (*self.data.as_ptr()).assume_init_ref()
149 }
150}
151
152#[cfg(test)]
153mod tests {
154 use super::*;
155
156 #[test]
157 fn once_lock() {
158 let lock = OnceLock::new();
159 assert_eq!(lock.try_get(), None);
160 assert_eq!(lock.is_set(), false);
161
162 let v = 42;
163 assert_eq!(lock.init(v), Ok(()));
164 assert_eq!(lock.is_set(), true);
165 assert_eq!(lock.try_get(), Some(&v));
166 assert_eq!(lock.try_get(), Some(&v));
167
168 let v = 43;
169 assert_eq!(lock.init(v), Err(v));
170 assert_eq!(lock.is_set(), true);
171 assert_eq!(lock.try_get(), Some(&42));
172 }
173
174 #[test]
175 fn once_lock_get_or_init() {
176 let lock = OnceLock::new();
177 assert_eq!(lock.try_get(), None);
178 assert_eq!(lock.is_set(), false);
179
180 let v = lock.get_or_init(|| 42);
181 assert_eq!(v, &42);
182 assert_eq!(lock.is_set(), true);
183 assert_eq!(lock.try_get(), Some(&42));
184
185 let v = lock.get_or_init(|| 43);
186 assert_eq!(v, &42);
187 assert_eq!(lock.is_set(), true);
188 assert_eq!(lock.try_get(), Some(&42));
189 }
190
191 #[test]
192 fn once_lock_static() {
193 static LOCK: OnceLock<i32> = OnceLock::new();
194
195 let v: &'static i32 = LOCK.get_or_init(|| 42);
196 assert_eq!(v, &42);
197
198 let v: &'static i32 = LOCK.get_or_init(|| 43);
199 assert_eq!(v, &42);
200 }
201
202 #[futures_test::test]
203 async fn once_lock_async() {
204 static LOCK: OnceLock<i32> = OnceLock::new();
205
206 assert!(LOCK.init(42).is_ok());
207
208 let v: &'static i32 = LOCK.get().await;
209 assert_eq!(v, &42);
210 }
211
212 #[test]
213 fn once_lock_into_inner() {
214 let lock: OnceLock<i32> = OnceLock::new();
215
216 let v = lock.get_or_init(|| 42);
217 assert_eq!(v, &42);
218
219 assert_eq!(lock.into_inner(), Some(42));
220 }
221
222 #[test]
223 fn once_lock_take_init() {
224 let mut lock: OnceLock<i32> = OnceLock::new();
225
226 assert_eq!(lock.get_or_init(|| 42), &42);
227 assert_eq!(lock.is_set(), true);
228
229 assert_eq!(lock.take(), Some(42));
230 assert_eq!(lock.is_set(), false);
231
232 assert_eq!(lock.get_or_init(|| 43), &43);
233 assert_eq!(lock.is_set(), true);
234 }
235}
236