1 | use core::cell::UnsafeCell; |
2 | use core::convert::Infallible; |
3 | use core::fmt; |
4 | use core::future::Future; |
5 | use core::mem::{forget, MaybeUninit}; |
6 | use core::ptr; |
7 | use core::sync::atomic::{AtomicUsize, Ordering}; |
8 | |
9 | #[cfg (all(feature = "std" , not(target_family = "wasm" )))] |
10 | use core::task::{Context, Poll, RawWaker, RawWakerVTable, Waker}; |
11 | |
12 | use event_listener::{Event, EventListener}; |
13 | use event_listener_strategy::{NonBlocking, Strategy}; |
14 | |
15 | /// The current state of the `OnceCell`. |
16 | #[derive (Copy, Clone, PartialEq, Eq)] |
17 | #[repr (usize)] |
18 | enum State { |
19 | /// The `OnceCell` is uninitialized. |
20 | Uninitialized = 0, |
21 | /// The `OnceCell` is being initialized. |
22 | Initializing = 1, |
23 | /// The `OnceCell` is initialized. |
24 | Initialized = 2, |
25 | } |
26 | |
27 | impl From<usize> for State { |
28 | fn from(val: usize) -> Self { |
29 | match val { |
30 | 0 => State::Uninitialized, |
31 | 1 => State::Initializing, |
32 | 2 => State::Initialized, |
33 | _ => unreachable!("Invalid state" ), |
34 | } |
35 | } |
36 | } |
37 | |
38 | impl From<State> for usize { |
39 | fn from(val: State) -> Self { |
40 | val as usize |
41 | } |
42 | } |
43 | |
44 | /// A memory location that can be written to at most once. |
45 | /// |
46 | /// A `OnceCell` can be used to store a single value, and only once. However, |
47 | /// once the value is stored, it can be accessed directly through a reference |
48 | /// instead of needing an RAII guard like `Mutex` or `RwLock`. |
49 | /// |
50 | /// # Examples |
51 | /// |
52 | /// This structure is useful for a variety of patterns, most notably for one-time |
53 | /// initialization. |
54 | /// |
55 | /// ```rust |
56 | /// use async_lock::OnceCell; |
57 | /// |
58 | /// # struct Foobar; |
59 | /// |
60 | /// async fn very_expensive_initialization() -> Foobar { |
61 | /// // Imagine this is very expensive to initialize, |
62 | /// // for instance, it requires a network request or |
63 | /// // a database call. |
64 | /// # Foobar |
65 | /// } |
66 | /// |
67 | /// struct LazyFoobar { |
68 | /// inner: OnceCell<Foobar>, |
69 | /// } |
70 | /// |
71 | /// impl LazyFoobar { |
72 | /// fn new() -> Self { |
73 | /// Self { |
74 | /// inner: OnceCell::new(), |
75 | /// } |
76 | /// } |
77 | /// |
78 | /// async fn load(&self) -> &Foobar { |
79 | /// self.inner.get_or_init(|| async { |
80 | /// very_expensive_initialization().await |
81 | /// }).await |
82 | /// } |
83 | /// } |
84 | /// ``` |
85 | pub struct OnceCell<T> { |
86 | /// Listeners waiting for a chance to initialize the cell. |
87 | /// |
88 | /// These are the users of get_or_init() and similar functions. |
89 | active_initializers: Event, |
90 | |
91 | /// Listeners waiting for the cell to be initialized. |
92 | /// |
93 | /// These are the users of wait(). |
94 | passive_waiters: Event, |
95 | |
96 | /// State associated with the cell. |
97 | state: AtomicUsize, |
98 | |
99 | /// The value of the cell. |
100 | value: UnsafeCell<MaybeUninit<T>>, |
101 | } |
102 | |
103 | unsafe impl<T: Send> Send for OnceCell<T> {} |
104 | unsafe impl<T: Send + Sync> Sync for OnceCell<T> {} |
105 | |
106 | impl<T> OnceCell<T> { |
107 | /// Create a new, uninitialized `OnceCell`. |
108 | /// |
109 | /// # Example |
110 | /// |
111 | /// ```rust |
112 | /// use async_lock::OnceCell; |
113 | /// |
114 | /// let cell = OnceCell::new(); |
115 | /// # cell.set_blocking(1); |
116 | /// ``` |
117 | pub const fn new() -> Self { |
118 | Self { |
119 | active_initializers: Event::new(), |
120 | passive_waiters: Event::new(), |
121 | state: AtomicUsize::new(State::Uninitialized as _), |
122 | value: UnsafeCell::new(MaybeUninit::uninit()), |
123 | } |
124 | } |
125 | |
126 | /// Tell whether or not the cell is initialized. |
127 | /// |
128 | /// This may not always be accurate. For instance, it is possible for |
129 | /// another thread to initialize the cell between the time when this |
130 | /// function is called and the time when the result is actually used. |
131 | /// |
132 | /// # Example |
133 | /// |
134 | /// ```rust |
135 | /// use async_lock::OnceCell; |
136 | /// |
137 | /// # futures_lite::future::block_on(async { |
138 | /// let cell = OnceCell::new(); |
139 | /// assert!(!cell.is_initialized()); |
140 | /// cell.set(1).await; |
141 | /// assert!(cell.is_initialized()); |
142 | /// # }); |
143 | /// ``` |
144 | pub fn is_initialized(&self) -> bool { |
145 | State::from(self.state.load(Ordering::Acquire)) == State::Initialized |
146 | } |
147 | |
148 | /// Get a reference to the inner value, or `None` if the value |
149 | /// is not yet initialized. |
150 | /// |
151 | /// # Example |
152 | /// |
153 | /// ```rust |
154 | /// use async_lock::OnceCell; |
155 | /// |
156 | /// # futures_lite::future::block_on(async { |
157 | /// let cell = OnceCell::new(); |
158 | /// assert!(cell.get().is_none()); |
159 | /// cell.set(1).await; |
160 | /// assert_eq!(cell.get(), Some(&1)); |
161 | /// # }); |
162 | /// ``` |
163 | pub fn get(&self) -> Option<&T> { |
164 | if self.is_initialized() { |
165 | // SAFETY: We know that the value is initialized, so it is safe to |
166 | // read it. |
167 | Some(unsafe { self.get_unchecked() }) |
168 | } else { |
169 | None |
170 | } |
171 | } |
172 | |
173 | /// Get a mutable reference to the inner value, or `None` if the value |
174 | /// is not yet initialized. |
175 | /// |
176 | /// This function is useful for initializing the value inside the cell |
177 | /// when we still have a mutable reference to the cell. |
178 | /// |
179 | /// # Example |
180 | /// |
181 | /// ```rust |
182 | /// use async_lock::OnceCell; |
183 | /// |
184 | /// # futures_lite::future::block_on(async { |
185 | /// let mut cell = OnceCell::new(); |
186 | /// assert!(cell.get_mut().is_none()); |
187 | /// cell.set(1).await; |
188 | /// assert_eq!(cell.get_mut(), Some(&mut 1)); |
189 | /// *cell.get_mut().unwrap() = 2; |
190 | /// assert_eq!(cell.get(), Some(&2)); |
191 | /// # }); |
192 | /// ``` |
193 | pub fn get_mut(&mut self) -> Option<&mut T> { |
194 | if State::from(*self.state.get_mut()) == State::Initialized { |
195 | // SAFETY: We know that the value is initialized, so it is safe to |
196 | // read it. |
197 | Some(unsafe { &mut *self.value.get().cast() }) |
198 | } else { |
199 | None |
200 | } |
201 | } |
202 | |
203 | /// Take the value out of this `OnceCell`, moving it back to the uninitialized |
204 | /// state. |
205 | /// |
206 | /// # Example |
207 | /// |
208 | /// ```rust |
209 | /// use async_lock::OnceCell; |
210 | /// |
211 | /// # futures_lite::future::block_on(async { |
212 | /// let mut cell = OnceCell::new(); |
213 | /// cell.set(1).await; |
214 | /// assert_eq!(cell.take(), Some(1)); |
215 | /// assert!(!cell.is_initialized()); |
216 | /// # }); |
217 | /// ``` |
218 | pub fn take(&mut self) -> Option<T> { |
219 | if State::from(*self.state.get_mut()) == State::Initialized { |
220 | // SAFETY: We know that the value is initialized, so it is safe to |
221 | // read it. |
222 | let value = unsafe { ptr::read(self.value.get().cast()) }; |
223 | *self.state.get_mut() = State::Uninitialized.into(); |
224 | Some(value) |
225 | } else { |
226 | None |
227 | } |
228 | } |
229 | |
230 | /// Convert this `OnceCell` into the inner value, if it is initialized. |
231 | /// |
232 | /// # Example |
233 | /// |
234 | /// ```rust |
235 | /// use async_lock::OnceCell; |
236 | /// |
237 | /// # futures_lite::future::block_on(async { |
238 | /// let cell = OnceCell::new(); |
239 | /// cell.set(1).await; |
240 | /// assert_eq!(cell.into_inner(), Some(1)); |
241 | /// # }); |
242 | /// ``` |
243 | pub fn into_inner(mut self) -> Option<T> { |
244 | self.take() |
245 | } |
246 | |
247 | /// Wait for the cell to be initialized, and then return a reference to the |
248 | /// inner value. |
249 | /// |
250 | /// # Example |
251 | /// |
252 | /// ```rust |
253 | /// use async_lock::OnceCell; |
254 | /// use std::sync::Arc; |
255 | /// use std::time::Duration; |
256 | /// use std::thread::{sleep, spawn}; |
257 | /// |
258 | /// let cell = Arc::new(OnceCell::new()); |
259 | /// let cell2 = cell.clone(); |
260 | /// |
261 | /// spawn(move || { |
262 | /// sleep(Duration::from_millis(5)); |
263 | /// cell2.set_blocking(1); |
264 | /// }); |
265 | /// |
266 | /// # futures_lite::future::block_on(async { |
267 | /// assert_eq!(cell.wait().await, &1); |
268 | /// # }); |
269 | /// ``` |
270 | pub async fn wait(&self) -> &T { |
271 | // Fast path: see if the value is already initialized. |
272 | if let Some(value) = self.get() { |
273 | return value; |
274 | } |
275 | |
276 | // Slow path: wait for the value to be initialized. |
277 | let listener = EventListener::new(); |
278 | pin!(listener); |
279 | listener.as_mut().listen(&self.passive_waiters); |
280 | |
281 | // Try again. |
282 | if let Some(value) = self.get() { |
283 | return value; |
284 | } |
285 | |
286 | listener.await; |
287 | debug_assert!(self.is_initialized()); |
288 | |
289 | // SAFETY: We know that the value is initialized, so it is safe to |
290 | // read it. |
291 | unsafe { self.get_unchecked() } |
292 | } |
293 | |
294 | /// Wait for the cell to be initialized, and then return a reference to the |
295 | /// inner value. |
296 | /// |
297 | /// # Blocking |
298 | /// |
299 | /// In contrast to the `wait` method, this method blocks the current thread of |
300 | /// execution instead of awaiting. |
301 | /// |
302 | /// This method should not be used in an asynchronous context. It is intended |
303 | /// to be used such that a `OnceCell` can be used in both asynchronous and synchronous contexts. |
304 | /// Calling this method in an asynchronous context may result in deadlocks. |
305 | /// |
306 | /// # Example |
307 | /// |
308 | /// ```rust |
309 | /// use async_lock::OnceCell; |
310 | /// use std::sync::Arc; |
311 | /// use std::time::Duration; |
312 | /// use std::thread::{sleep, spawn}; |
313 | /// |
314 | /// let cell = Arc::new(OnceCell::new()); |
315 | /// let cell2 = cell.clone(); |
316 | /// |
317 | /// spawn(move || { |
318 | /// sleep(Duration::from_millis(5)); |
319 | /// cell2.set_blocking(1); |
320 | /// }); |
321 | /// |
322 | /// assert_eq!(cell.wait_blocking(), &1); |
323 | /// ``` |
324 | #[cfg (all(feature = "std" , not(target_family = "wasm" )))] |
325 | pub fn wait_blocking(&self) -> &T { |
326 | // Fast path: see if the value is already initialized. |
327 | if let Some(value) = self.get() { |
328 | return value; |
329 | } |
330 | |
331 | // Slow path: wait for the value to be initialized. |
332 | let listener = EventListener::new(); |
333 | pin!(listener); |
334 | listener.as_mut().listen(&self.passive_waiters); |
335 | |
336 | // Try again. |
337 | if let Some(value) = self.get() { |
338 | return value; |
339 | } |
340 | |
341 | listener.wait(); |
342 | debug_assert!(self.is_initialized()); |
343 | |
344 | // SAFETY: We know that the value is initialized, so it is safe to |
345 | // read it. |
346 | unsafe { self.get_unchecked() } |
347 | } |
348 | |
349 | /// Either get the value or initialize it with the given closure. |
350 | /// |
351 | /// The cell will not be initialized if the closure returns an error. |
352 | /// |
353 | /// # Example |
354 | /// |
355 | /// ```rust |
356 | /// use async_lock::OnceCell; |
357 | /// # |
358 | /// # // Prevent explicit value errors. |
359 | /// # fn _explicit(_: &Result<&i32, ()>) {} |
360 | /// |
361 | /// # futures_lite::future::block_on(async { |
362 | /// let cell = OnceCell::new(); |
363 | /// |
364 | /// let result = cell.get_or_try_init(|| async { Err(()) }).await; |
365 | /// assert!(result.is_err()); |
366 | /// |
367 | /// let result = cell.get_or_try_init(|| async { Ok(1) }).await; |
368 | /// # _explicit(&result); |
369 | /// assert_eq!(result.unwrap(), &1); |
370 | /// |
371 | /// let result = cell.get_or_try_init(|| async { Err(()) }).await; |
372 | /// |
373 | /// assert_eq!(result.unwrap(), &1); |
374 | /// # }); |
375 | /// ``` |
376 | pub async fn get_or_try_init<E, Fut: Future<Output = Result<T, E>>>( |
377 | &self, |
378 | closure: impl FnOnce() -> Fut, |
379 | ) -> Result<&T, E> { |
380 | // Fast path: see if the value is already initialized. |
381 | if let Some(value) = self.get() { |
382 | return Ok(value); |
383 | } |
384 | |
385 | // Slow path: initialize the value. |
386 | self.initialize_or_wait(closure, &mut NonBlocking::default()) |
387 | .await?; |
388 | debug_assert!(self.is_initialized()); |
389 | |
390 | // SAFETY: We know that the value is initialized, so it is safe to |
391 | // read it. |
392 | Ok(unsafe { self.get_unchecked() }) |
393 | } |
394 | |
395 | /// Either get the value or initialize it with the given closure. |
396 | /// |
397 | /// The cell will not be initialized if the closure returns an error. |
398 | /// |
399 | /// # Blocking |
400 | /// |
401 | /// In contrast to the `get_or_try_init` method, this method blocks the current thread of |
402 | /// execution instead of awaiting. |
403 | /// |
404 | /// This method should not be used in an asynchronous context. It is intended |
405 | /// to be used such that a `OnceCell` can be used in both asynchronous and synchronous contexts. |
406 | /// Calling this method in an asynchronous context may result in deadlocks. |
407 | /// |
408 | /// # Example |
409 | /// |
410 | /// ```rust |
411 | /// use async_lock::OnceCell; |
412 | /// # |
413 | /// # // Prevent explicit type errors. |
414 | /// # fn _explicit(_: &Result<&i32, ()>) {} |
415 | /// |
416 | /// let cell = OnceCell::new(); |
417 | /// |
418 | /// let result = cell.get_or_try_init_blocking(|| Err(())); |
419 | /// assert!(result.is_err()); |
420 | /// |
421 | /// let result = cell.get_or_try_init_blocking(|| Ok(1)); |
422 | /// # _explicit(&result); |
423 | /// assert_eq!(result.unwrap(), &1); |
424 | /// |
425 | /// let result = cell.get_or_try_init_blocking(|| Err(())); |
426 | /// |
427 | /// assert_eq!(result.unwrap(), &1); |
428 | /// ``` |
429 | #[cfg (all(feature = "std" , not(target_family = "wasm" )))] |
430 | pub fn get_or_try_init_blocking<E>( |
431 | &self, |
432 | closure: impl FnOnce() -> Result<T, E>, |
433 | ) -> Result<&T, E> { |
434 | // Fast path: see if the value is already initialized. |
435 | if let Some(value) = self.get() { |
436 | return Ok(value); |
437 | } |
438 | |
439 | // Slow path: initialize the value. |
440 | // The futures provided should never block, so we can use `now_or_never`. |
441 | now_or_never(self.initialize_or_wait( |
442 | move || core::future::ready(closure()), |
443 | &mut event_listener_strategy::Blocking::default(), |
444 | ))?; |
445 | debug_assert!(self.is_initialized()); |
446 | |
447 | // SAFETY: We know that the value is initialized, so it is safe to |
448 | // read it. |
449 | Ok(unsafe { self.get_unchecked() }) |
450 | } |
451 | |
452 | /// Either get the value or initialize it with the given closure. |
453 | /// |
454 | /// Many tasks may call this function, but the value will only be set once |
455 | /// and only one closure will be invoked. |
456 | /// |
457 | /// # Example |
458 | /// |
459 | /// ```rust |
460 | /// use async_lock::OnceCell; |
461 | /// |
462 | /// # futures_lite::future::block_on(async { |
463 | /// let cell = OnceCell::new(); |
464 | /// assert_eq!(cell.get_or_init(|| async { 1 }).await, &1); |
465 | /// assert_eq!(cell.get_or_init(|| async { 2 }).await, &1); |
466 | /// # }); |
467 | /// ``` |
468 | pub async fn get_or_init<Fut: Future<Output = T>>(&self, closure: impl FnOnce() -> Fut) -> &T { |
469 | match self |
470 | .get_or_try_init(move || async move { |
471 | let result: Result<T, Infallible> = Ok(closure().await); |
472 | result |
473 | }) |
474 | .await |
475 | { |
476 | Ok(value) => value, |
477 | Err(infallible) => match infallible {}, |
478 | } |
479 | } |
480 | |
481 | /// Either get the value or initialize it with the given closure. |
482 | /// |
483 | /// Many tasks may call this function, but the value will only be set once |
484 | /// and only one closure will be invoked. |
485 | /// |
486 | /// # Blocking |
487 | /// |
488 | /// In contrast to the `get_or_init` method, this method blocks the current thread of |
489 | /// execution instead of awaiting. |
490 | /// |
491 | /// This method should not be used in an asynchronous context. It is intended |
492 | /// to be used such that a `OnceCell` can be used in both asynchronous and synchronous contexts. |
493 | /// Calling this method in an asynchronous context may result in deadlocks. |
494 | /// |
495 | /// # Example |
496 | /// |
497 | /// ```rust |
498 | /// use async_lock::OnceCell; |
499 | /// |
500 | /// let cell = OnceCell::new(); |
501 | /// assert_eq!(cell.get_or_init_blocking(|| 1), &1); |
502 | /// assert_eq!(cell.get_or_init_blocking(|| 2), &1); |
503 | /// ``` |
504 | #[cfg (all(feature = "std" , not(target_family = "wasm" )))] |
505 | pub fn get_or_init_blocking(&self, closure: impl FnOnce() -> T + Unpin) -> &T { |
506 | match self.get_or_try_init_blocking(move || { |
507 | let result: Result<T, Infallible> = Ok(closure()); |
508 | result |
509 | }) { |
510 | Ok(value) => value, |
511 | Err(infallible) => match infallible {}, |
512 | } |
513 | } |
514 | |
515 | /// Try to set the value of the cell. |
516 | /// |
517 | /// If the cell is already initialized, this method returns the original |
518 | /// value back. |
519 | /// |
520 | /// # Example |
521 | /// |
522 | /// ```rust |
523 | /// use async_lock::OnceCell; |
524 | /// |
525 | /// # futures_lite::future::block_on(async { |
526 | /// let cell = OnceCell::new(); |
527 | /// |
528 | /// assert_eq!(cell.set(1).await, Ok(&1)); |
529 | /// assert_eq!(cell.get(), Some(&1)); |
530 | /// assert_eq!(cell.set(2).await, Err(2)); |
531 | /// # }); |
532 | /// ``` |
533 | pub async fn set(&self, value: T) -> Result<&T, T> { |
534 | let mut value = Some(value); |
535 | self.get_or_init(|| async { value.take().unwrap() }).await; |
536 | |
537 | match value { |
538 | Some(value) => Err(value), |
539 | None => { |
540 | // SAFETY: value was taken, so we are initialized |
541 | Ok(unsafe { self.get_unchecked() }) |
542 | } |
543 | } |
544 | } |
545 | |
546 | /// Try to set the value of the cell. |
547 | /// |
548 | /// If the cell is already initialized, this method returns the original |
549 | /// value back. |
550 | /// |
551 | /// # Blocking |
552 | /// |
553 | /// In contrast to the `set` method, this method blocks the current thread of |
554 | /// execution instead of awaiting. |
555 | /// |
556 | /// This method should not be used in an asynchronous context. It is intended |
557 | /// to be used such that a `OnceCell` can be used in both asynchronous and synchronous contexts. |
558 | /// Calling this method in an asynchronous context may result in deadlocks. |
559 | /// |
560 | /// # Example |
561 | /// |
562 | /// ```rust |
563 | /// use async_lock::OnceCell; |
564 | /// |
565 | /// let cell = OnceCell::new(); |
566 | /// |
567 | /// assert_eq!(cell.set_blocking(1), Ok(&1)); |
568 | /// assert_eq!(cell.get(), Some(&1)); |
569 | /// assert_eq!(cell.set_blocking(2), Err(2)); |
570 | /// ``` |
571 | #[cfg (all(feature = "std" , not(target_family = "wasm" )))] |
572 | pub fn set_blocking(&self, value: T) -> Result<&T, T> { |
573 | let mut value = Some(value); |
574 | self.get_or_init_blocking(|| value.take().unwrap()); |
575 | |
576 | match value { |
577 | Some(value) => Err(value), |
578 | None => { |
579 | // SAFETY: value was taken, so we are initialized |
580 | Ok(unsafe { self.get_unchecked() }) |
581 | } |
582 | } |
583 | } |
584 | |
585 | /// Wait for the cell to be initialized, optionally using a closure |
586 | /// to initialize the cell if it is not initialized yet. |
587 | #[cold ] |
588 | async fn initialize_or_wait<E, Fut: Future<Output = Result<T, E>>, F: FnOnce() -> Fut>( |
589 | &self, |
590 | closure: F, |
591 | strategy: &mut impl for<'a> Strategy<'a>, |
592 | ) -> Result<(), E> { |
593 | // The event listener we're currently waiting on. |
594 | let event_listener = EventListener::new(); |
595 | pin!(event_listener); |
596 | |
597 | let mut closure = Some(closure); |
598 | |
599 | loop { |
600 | // Check the current state of the cell. |
601 | let state = self.state.load(Ordering::Acquire); |
602 | |
603 | // Determine what we should do based on our state. |
604 | match state.into() { |
605 | State::Initialized => { |
606 | // The cell is initialized now, so we can return. |
607 | return Ok(()); |
608 | } |
609 | State::Initializing => { |
610 | // The cell is currently initializing, or the cell is uninitialized |
611 | // but we do not have the ability to initialize it. |
612 | // |
613 | // We need to wait the initialization to complete. |
614 | if event_listener.is_listening() { |
615 | strategy.wait(event_listener.as_mut()).await; |
616 | } else { |
617 | event_listener.as_mut().listen(&self.active_initializers); |
618 | } |
619 | } |
620 | State::Uninitialized => { |
621 | // Try to move the cell into the initializing state. |
622 | if self |
623 | .state |
624 | .compare_exchange( |
625 | State::Uninitialized.into(), |
626 | State::Initializing.into(), |
627 | Ordering::AcqRel, |
628 | Ordering::Acquire, |
629 | ) |
630 | .is_err() |
631 | { |
632 | // The cell was initialized while we were trying to |
633 | // initialize it. |
634 | continue; |
635 | } |
636 | |
637 | // Now that we have an exclusive lock on the cell's value, |
638 | // we can try to initialize it. |
639 | let _guard = Guard(self); |
640 | let initializer = closure.take().unwrap(); |
641 | match (initializer)().await { |
642 | Ok(value) => { |
643 | // Write the value into the cell and update the state. |
644 | unsafe { |
645 | ptr::write(self.value.get().cast(), value); |
646 | } |
647 | forget(_guard); |
648 | self.state |
649 | .store(State::Initialized.into(), Ordering::Release); |
650 | |
651 | // Notify the listeners that the value is initialized. |
652 | self.active_initializers.notify_additional(core::usize::MAX); |
653 | self.passive_waiters.notify_additional(core::usize::MAX); |
654 | |
655 | return Ok(()); |
656 | } |
657 | Err(err) => { |
658 | // Update the state to indicate that the value is |
659 | // uninitialized. |
660 | drop(_guard); |
661 | |
662 | return Err(err); |
663 | } |
664 | } |
665 | } |
666 | } |
667 | } |
668 | |
669 | /// Set the cell's state back to `UNINITIALIZED on drop. |
670 | /// |
671 | /// If the closure panics, this ensures that the cell's state is set back to |
672 | /// `UNINITIALIZED` and that the next listener is notified. |
673 | struct Guard<'a, T>(&'a OnceCell<T>); |
674 | |
675 | impl<'a, T> Drop for Guard<'a, T> { |
676 | fn drop(&mut self) { |
677 | self.0 |
678 | .state |
679 | .store(State::Uninitialized.into(), Ordering::Release); |
680 | |
681 | // Notify the next initializer that it's their turn. |
682 | self.0.active_initializers.notify(1); |
683 | } |
684 | } |
685 | } |
686 | |
687 | /// Get a reference to the inner value. |
688 | /// |
689 | /// # Safety |
690 | /// |
691 | /// The caller must ensure that the cell is initialized. |
692 | /// |
693 | /// # Example |
694 | /// |
695 | /// ```rust |
696 | /// use async_lock::OnceCell; |
697 | /// |
698 | /// # futures_lite::future::block_on(async { |
699 | /// let cell = OnceCell::new(); |
700 | /// cell.set(1).await; |
701 | /// |
702 | /// // SAFETY: We know that the value is initialized, so it is safe to |
703 | /// // read it. |
704 | /// assert_eq!(unsafe { cell.get_unchecked() }, &1); |
705 | /// # }); |
706 | /// ``` |
707 | pub unsafe fn get_unchecked(&self) -> &T { |
708 | // SAFETY: The caller asserts that the value is initialized |
709 | &*self.value.get().cast() |
710 | } |
711 | } |
712 | |
713 | impl<T> From<T> for OnceCell<T> { |
714 | /// Create a new, initialized `OnceCell` from an existing value. |
715 | /// |
716 | /// # Example |
717 | /// |
718 | /// ```rust |
719 | /// use async_lock::OnceCell; |
720 | /// |
721 | /// let cell = OnceCell::from(42); |
722 | /// assert_eq!(cell.get(), Some(&42)); |
723 | /// ``` |
724 | fn from(value: T) -> Self { |
725 | Self { |
726 | active_initializers: Event::new(), |
727 | passive_waiters: Event::new(), |
728 | state: AtomicUsize::new(State::Initialized.into()), |
729 | value: UnsafeCell::new(MaybeUninit::new(val:value)), |
730 | } |
731 | } |
732 | } |
733 | |
734 | impl<T: fmt::Debug> fmt::Debug for OnceCell<T> { |
735 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { |
736 | struct Inner<'a, T>(&'a OnceCell<T>); |
737 | |
738 | impl<T: fmt::Debug> fmt::Debug for Inner<'_, T> { |
739 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { |
740 | match self.0.state.load(order:Ordering::Acquire).into() { |
741 | State::Uninitialized => f.write_str(data:"<uninitialized>" ), |
742 | State::Initializing => f.write_str(data:"<initializing>" ), |
743 | State::Initialized => { |
744 | // SAFETY: "value" is initialized. |
745 | let value: &T = unsafe { self.0.get_unchecked() }; |
746 | fmt::Debug::fmt(self:value, f) |
747 | } |
748 | } |
749 | } |
750 | } |
751 | |
752 | f.debug_tuple(name:"OnceCell" ).field(&Inner(self)).finish() |
753 | } |
754 | } |
755 | |
756 | impl<T> Drop for OnceCell<T> { |
757 | fn drop(&mut self) { |
758 | if State::from(*self.state.get_mut()) == State::Initialized { |
759 | // SAFETY: We know that the value is initialized, so it is safe to |
760 | // drop it. |
761 | unsafe { self.value.get().cast::<T>().drop_in_place() } |
762 | } |
763 | } |
764 | } |
765 | |
766 | impl<T> Default for OnceCell<T> { |
767 | // Calls `OnceCell::new`. |
768 | #[inline ] |
769 | fn default() -> Self { |
770 | Self::new() |
771 | } |
772 | } |
773 | |
774 | /// Either return the result of a future now, or panic. |
775 | #[cfg (all(feature = "std" , not(target_family = "wasm" )))] |
776 | fn now_or_never<T>(f: impl Future<Output = T>) -> T { |
777 | const NOOP_WAKER: RawWakerVTable = RawWakerVTable::new(clone, wake, wake_by_ref, drop); |
778 | |
779 | unsafe fn wake(_: *const ()) {} |
780 | unsafe fn wake_by_ref(_: *const ()) {} |
781 | unsafe fn clone(_: *const ()) -> RawWaker { |
782 | RawWaker::new(data:ptr::null(), &NOOP_WAKER) |
783 | } |
784 | unsafe fn drop(_: *const ()) {} |
785 | |
786 | pin!(f); |
787 | |
788 | let waker: Waker = unsafe { Waker::from_raw(waker:RawWaker::new(data:ptr::null(), &NOOP_WAKER)) }; |
789 | |
790 | // Poll the future exactly once. |
791 | let mut cx: Context<'_> = Context::from_waker(&waker); |
792 | |
793 | match f.poll(&mut cx) { |
794 | Poll::Ready(value: T) => value, |
795 | Poll::Pending => unreachable!("future not ready" ), |
796 | } |
797 | } |
798 | |