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