1#![allow(unsafe_code)]
19
20use core::convert::Infallible;
21use core::error::Error;
22use core::fmt;
23use core::future::Future;
24use core::mem::{self, MaybeUninit};
25use core::ptr;
26
27use super::contend::Closed;
28use super::internal::{AtomicPtr, Ordering, UnsafeCell, take_ptr};
29use super::raw_semaphore::{RawPermit, RawSemaphore, observed_closed};
30
31const fn inline_marker<T>() -> *mut T {
35 ptr::without_provenance_mut(usize::MAX)
36}
37
38#[derive(Debug, PartialEq, Eq)]
40pub enum SetError<T> {
41 AlreadyInitializedError(T),
43
44 InitializingError(T),
46}
47
48impl<T> fmt::Display for SetError<T> {
49 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
50 match self {
51 Self::AlreadyInitializedError(_) => f.write_str("AlreadyInitializedError"),
52 Self::InitializingError(_) => f.write_str("InitializingError"),
53 }
54 }
55}
56
57impl<T: fmt::Debug> Error for SetError<T> {}
58
59impl<T> SetError<T> {
60 pub fn is_already_init_err(&self) -> bool {
62 matches!(self, Self::AlreadyInitializedError(_))
63 }
64
65 pub fn is_initializing_err(&self) -> bool {
67 matches!(self, Self::InitializingError(_))
68 }
69}
70
71pub struct OnceCell<T> {
96 value: AtomicPtr<T>,
99 slot: UnsafeCell<MaybeUninit<T>>,
102 init: RawSemaphore,
105}
106
107unsafe impl<T: Send + Sync> Sync for OnceCell<T> {}
110
111impl<T> OnceCell<T> {
112 loom_const_fn! {
113 pub fn new() -> Self {
115 Self {
116 value: AtomicPtr::new(ptr::null_mut()),
117 slot: UnsafeCell::new(MaybeUninit::uninit()),
118 init: RawSemaphore::without_spin(1),
121 }
122 }
123 }
124
125 loom_const_fn! {
126 pub fn const_new() -> Self {
128 Self::new()
129 }
130 }
131
132 pub fn new_with(value: Option<T>) -> Self {
134 match value {
135 Some(value) => Self::from(value),
136 None => Self::new(),
137 }
138 }
139
140 loom_const_fn! {
141 pub fn const_new_with(value: T) -> Self {
143 Self {
144 value: AtomicPtr::new(inline_marker()),
145 slot: UnsafeCell::new(MaybeUninit::new(value)),
146 init: RawSemaphore::new_closed(),
147 }
148 }
149 }
150
151 pub fn initialized(&self) -> bool {
153 !self.value.load(Ordering::Acquire).is_null()
155 }
156
157 fn resolve(&self, stored: *mut T) -> &T {
159 if stored == inline_marker() {
160 unsafe { &*self.slot.with(|slot| slot.cast::<T>()) }
163 } else {
164 unsafe { &*stored }
167 }
168 }
169
170 pub fn get(&self) -> Option<&T> {
172 let stored = self.value.load(Ordering::Acquire);
173 (!stored.is_null()).then(|| self.resolve(stored))
174 }
175
176 pub fn get_mut(&mut self) -> Option<&mut T> {
178 let stored = self.value.load(Ordering::Acquire);
179 if stored.is_null() {
180 None
181 } else if stored == inline_marker() {
182 Some(unsafe { &mut *self.slot.with_mut(|slot| slot.cast::<T>()) })
185 } else {
186 Some(unsafe { &mut *stored })
188 }
189 }
190
191 fn take_value(&mut self) -> Option<T> {
194 let stored = take_ptr(&mut self.value);
195 if stored.is_null() {
196 None
197 } else if stored == inline_marker() {
198 Some(unsafe { self.slot.with_mut(|slot| (*slot).assume_init_read()) })
201 } else {
202 Some(*unsafe { Box::from_raw(stored) })
205 }
206 }
207
208 pub fn into_inner(mut self) -> Option<T> {
210 self.take_value()
211 }
212
213 pub fn take(&mut self) -> Option<T> {
216 mem::take(self).into_inner()
217 }
218
219 pub fn set(&self, value: T) -> Result<(), SetError<T>> {
229 if self.initialized() {
230 return Err(SetError::AlreadyInitializedError(value));
231 }
232 match self.init.try_acquire_permit(1) {
233 Ok(turn) => self
234 .publish_inline(value, turn)
235 .map(drop)
236 .map_err(|(value, _)| SetError::AlreadyInitializedError(value)),
237 Err(seen) if observed_closed(seen) => Err(SetError::AlreadyInitializedError(value)),
239 Err(_) => Err(SetError::InitializingError(value)),
241 }
242 }
243
244 pub fn get_or_init_racy(&self, init: impl FnOnce() -> T) -> &T {
250 match self.get() {
251 Some(value) => value,
252 None => self.publish_boxed(init()),
253 }
254 }
255
256 pub async fn get_or_init<F, Fut>(&self, f: F) -> &T
265 where
266 F: FnOnce() -> Fut,
267 Fut: Future<Output = T>,
268 {
269 match self
270 .get_or_try_init(|| async move { Ok::<T, Infallible>(f().await) })
271 .await
272 {
273 Ok(value) => value,
274 Err(never) => match never {},
275 }
276 }
277
278 pub async fn get_or_try_init<E, F, Fut>(&self, f: F) -> Result<&T, E>
287 where
288 F: FnOnce() -> Fut,
289 Fut: Future<Output = Result<T, E>>,
290 {
291 if let Some(value) = self.get() {
292 return Ok(value);
293 }
294 match self.init.acquire(1).await {
297 Ok(turn) => {
298 if let Some(value) = self.get() {
300 return Ok(value);
301 }
302 let value = f().await?;
303 Ok(match self.publish_inline(value, turn) {
306 Ok(value) | Err((_, value)) => value,
307 })
308 }
309 Err(Closed) => match self.get() {
314 Some(value) => Ok(value),
315 None => Ok(self.publish_boxed(f().await?)),
316 },
317 }
318 }
319
320 fn publish_inline(&self, value: T, turn: RawPermit<'_>) -> Result<&T, (T, &T)> {
326 self.slot.with_mut(|slot| unsafe {
329 (*slot).write(value);
330 });
331 match self.value.compare_exchange(
332 ptr::null_mut(),
333 inline_marker(),
334 Ordering::AcqRel,
335 Ordering::Acquire,
336 ) {
337 Ok(_) => {
338 turn.forget();
339 self.init.close();
340 Ok(self.resolve(inline_marker()))
341 }
342 Err(winner) => Err((
343 unsafe { self.slot.with_mut(|slot| (*slot).assume_init_read()) },
346 self.resolve(winner),
347 )),
348 }
349 }
350
351 fn publish_boxed(&self, value: T) -> &T {
354 let fresh = Box::into_raw(Box::new(value));
355 match self.value.compare_exchange(
356 ptr::null_mut(),
357 fresh,
358 Ordering::AcqRel,
359 Ordering::Acquire,
360 ) {
361 Ok(_) => {
362 self.init.close();
363 unsafe { &*fresh }
365 }
366 Err(winner) => {
367 drop(unsafe { Box::from_raw(fresh) });
369 self.resolve(winner)
370 }
371 }
372 }
373}
374
375impl<T> Drop for OnceCell<T> {
376 fn drop(&mut self) {
377 drop(self.take_value());
378 }
379}
380
381impl<T> Default for OnceCell<T> {
382 fn default() -> Self {
383 Self::new()
384 }
385}
386
387impl<T: fmt::Debug> fmt::Debug for OnceCell<T> {
388 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
389 f.debug_struct("OnceCell")
390 .field("value", &self.get())
391 .finish()
392 }
393}
394
395impl<T: Clone> Clone for OnceCell<T> {
396 fn clone(&self) -> Self {
397 Self::new_with(self.get().cloned())
398 }
399}
400
401impl<T: PartialEq> PartialEq for OnceCell<T> {
402 fn eq(&self, other: &Self) -> bool {
403 self.get() == other.get()
404 }
405}
406
407impl<T: Eq> Eq for OnceCell<T> {}
408
409impl<T> From<T> for OnceCell<T> {
410 fn from(value: T) -> Self {
411 Self::const_new_with(value)
412 }
413}
414
415#[cfg(all(test, not(loom)))]
416mod tests {
417 use super::super::queue::Policy;
418 use super::*;
419 use crate::sync::Event;
420 use crate::sync::test_support::{Polled, block_on};
421 use std::sync::Arc;
422 use std::sync::atomic::{AtomicUsize, Ordering as StdOrdering};
423
424 #[test]
429 fn a_waiter_for_the_initializer_turn_never_spins() {
430 let cell = OnceCell::<u8>::new();
431 assert_eq!(cell.init.queue().spin.rounds(), 0);
432 let opened = Event::new();
433 let mut first = Polled::new(cell.get_or_init(|| async {
434 opened.wait().await;
435 1
436 }));
437 first.pending();
438 let mut second = Polled::new(cell.get_or_init(|| async { 2 }));
439 second.pending();
440 opened.set();
441 assert_eq!(*first.ready(), 1);
442 assert_eq!(*second.ready(), 1);
443 assert_eq!(cell.init.queue().spin.rounds(), 0, "the turn never spins");
444 }
445
446 #[test]
447 fn set_publishes_once() {
448 let cell = OnceCell::new();
449 assert_eq!(cell.get(), None);
450 assert!(!cell.initialized());
451 assert_eq!(cell.set(1), Ok(()));
452 assert!(cell.initialized());
453 assert_eq!(cell.set(2), Err(SetError::AlreadyInitializedError(2)));
454 assert_eq!(cell.get(), Some(&1));
455 assert_eq!(cell.into_inner(), Some(1));
456 }
457
458 #[test]
459 fn set_fails_while_an_initializer_holds_the_turn() {
460 let cell = OnceCell::new();
461 let mut running = Polled::new(cell.get_or_init(core::future::pending::<u8>));
462 running.pending();
463 let err = cell.set(3).unwrap_err();
464 assert!(err.is_initializing_err() && !err.is_already_init_err());
465 assert_eq!(err, SetError::InitializingError(3));
466 assert_eq!(err.to_string(), "InitializingError");
467 assert_eq!(cell.get(), None);
468 drop(running);
469 assert_eq!(cell.set(3), Ok(()), "the turn was given up");
470 let err = cell.set(4).unwrap_err();
471 assert!(err.is_already_init_err());
472 assert_eq!(err.to_string(), "AlreadyInitializedError");
473 }
474
475 #[test]
476 fn a_racing_builder_publishes_beside_a_running_initializer() {
477 let cell = OnceCell::new();
478 let mut stuck = Polled::new(cell.get_or_init(core::future::pending));
479 stuck.pending();
480 let mut waiting = Polled::new(cell.get_or_init(|| async { 2 }));
481 waiting.pending();
482 assert_eq!(*cell.get_or_init_racy(|| 1), 1);
483 assert_eq!(waiting.wakes(), 1, "the publisher wakes the queue");
484 assert_eq!(*waiting.ready(), 1);
485 let mut late = Polled::new(cell.get_or_init(|| async { 3 }));
486 assert_eq!(*late.ready(), 1);
487 assert_eq!(
488 cell.set(9),
489 Err(SetError::AlreadyInitializedError(9)),
490 "the cell is full, not being initialized"
491 );
492 drop(stuck);
493 }
494
495 #[test]
496 fn an_initializer_that_lost_to_a_racing_builder_drops_its_value() {
497 let cell = OnceCell::new();
498 let opened = Event::new();
499 let mut slow = Polled::new(cell.get_or_init(|| async {
500 opened.wait().await;
501 Arc::new(5)
502 }));
503 slow.pending();
504 let racy = Arc::new(7);
505 assert_eq!(**cell.get_or_init_racy(|| Arc::clone(&racy)), 7);
506 opened.set();
507 assert_eq!(**slow.ready(), 7, "the racer's value won");
508 assert_eq!(Arc::strong_count(&racy), 2, "the cell holds the winner");
509 }
510
511 #[test]
512 fn racing_builders_publish_exactly_one_value() {
513 let cell = Arc::new(OnceCell::new());
514 let built = Arc::new(AtomicUsize::new(0));
515 let seen: Vec<usize> = (0..4)
516 .map(|i| {
517 let (cell, built) = (Arc::clone(&cell), Arc::clone(&built));
518 std::thread::spawn(move || {
519 *cell.get_or_init_racy(|| {
520 built.fetch_add(1, StdOrdering::SeqCst);
521 i
522 })
523 })
524 })
525 .map(|thread| thread.join().expect("racer"))
526 .collect();
527 assert!(seen.iter().all(|v| *v == seen[0]), "{seen:?}");
528 assert!(built.load(StdOrdering::SeqCst) >= 1);
529 }
530
531 #[test]
532 fn one_initializer_runs_while_the_others_wait() {
533 let cell = OnceCell::new();
534 let opened = Event::new();
535 let mut first = Polled::new(cell.get_or_init(|| async {
536 opened.wait().await;
537 1
538 }));
539 first.pending();
540 let mut second = Polled::new(cell.get_or_init(|| async { 2 }));
541 second.pending();
542 opened.set();
543 assert_eq!(*first.ready(), 1);
544 assert_eq!(second.wakes(), 1);
545 assert_eq!(*second.ready(), 1);
546 }
547
548 #[test]
549 fn a_closed_turn_queue_over_an_empty_cell_still_answers() {
550 let cell = OnceCell::<u8>::new();
553 cell.init.close();
554 assert_eq!(*block_on(cell.get_or_init(|| async { 5 })), 5);
555 assert_eq!(cell.get(), Some(&5));
556 let failed = OnceCell::<u8>::new();
557 failed.init.close();
558 assert_eq!(
559 block_on(failed.get_or_try_init(|| async { Err::<u8, &str>("no") })),
560 Err("no")
561 );
562 assert_eq!(failed.get(), None);
563 }
564
565 #[test]
566 fn the_initializer_is_not_called_when_the_cell_is_full() {
567 let cell = OnceCell::from(1);
568 let mut called = false;
569 let value = block_on(cell.get_or_init(|| {
570 called = true;
571 async { 2 }
572 }));
573 assert_eq!(*value, 1);
574 assert!(!called);
575 let tried = block_on(cell.get_or_try_init(|| async { Err::<u8, ()>(()) }));
576 assert_eq!(tried, Ok(&1));
577 }
578
579 #[test]
580 fn a_dropped_initializer_passes_the_turn_on() {
581 let cell = OnceCell::new();
582 let mut first = Polled::new(cell.get_or_init(core::future::pending));
583 first.pending();
584 let mut second = Polled::new(cell.get_or_init(|| async { 2 }));
585 second.pending();
586 drop(first);
587 assert_eq!(second.wakes(), 1);
588 assert_eq!(*second.ready(), 2);
589 }
590
591 #[test]
592 fn a_failed_initializer_lets_the_next_one_try() {
593 let cell = OnceCell::new();
594 let opened = Event::new();
595 let mut failing = Polled::new(cell.get_or_try_init(|| async {
596 opened.wait().await;
597 Err::<u32, &str>("no")
598 }));
599 failing.pending();
600 let mut next = Polled::new(cell.get_or_try_init(|| async { Ok::<u32, &str>(3) }));
601 next.pending();
602 opened.set();
603 assert_eq!(failing.ready(), Err("no"));
604 assert_eq!(next.ready(), Ok(&3));
605 }
606
607 #[test]
608 fn constructors_conversions_and_traits() {
609 static FULL: OnceCell<u32> = OnceCell::const_new_with(4);
610 static EMPTY: OnceCell<u32> = OnceCell::const_new();
611 assert_eq!(FULL.get(), Some(&4));
612 assert_eq!(EMPTY.get(), None);
613 assert_eq!(FULL.set(5), Err(SetError::AlreadyInitializedError(5)));
614 assert_eq!(*block_on(FULL.get_or_init(|| async { 6 })), 4);
615 assert_eq!(OnceCell::new_with(Some(1)).get(), Some(&1));
616 assert_eq!(OnceCell::<u8>::new_with(None).get(), None);
617 assert_eq!(OnceCell::<u8>::default().get(), None);
618 assert_eq!(OnceCell::from(2).into_inner(), Some(2));
619 assert_eq!(
620 format!("{:?}", OnceCell::from(2)),
621 "OnceCell { value: Some(2) }"
622 );
623 assert_eq!(
624 format!("{:?}", OnceCell::<u8>::new()),
625 "OnceCell { value: None }"
626 );
627
628 let cell = OnceCell::from(String::from("a"));
629 let copy = cell.clone();
630 assert_eq!(copy, cell);
631 assert_eq!(copy.get().map(String::as_str), Some("a"));
632 assert_ne!(cell, OnceCell::new());
633 assert_eq!(OnceCell::<u8>::new(), OnceCell::new());
634 assert_eq!(OnceCell::<String>::new().clone().get(), None);
635 }
636
637 #[test]
638 fn get_mut_and_take_work_for_either_home_of_the_value() {
639 let mut inline = OnceCell::new();
640 assert_eq!(inline.get_mut(), None);
641 assert_eq!(inline.set(vec![1]), Ok(()));
642 inline.get_mut().expect("full").push(2);
643 assert_eq!(inline.take(), Some(vec![1, 2]));
644 assert_eq!(inline.get(), None);
645 assert_eq!(inline.set(vec![3]), Ok(()), "a taken cell is set again");
646
647 let mut boxed = OnceCell::new();
648 assert_eq!(*boxed.get_or_init_racy(|| vec![1]), vec![1]);
649 boxed.get_mut().expect("full").push(2);
650 assert_eq!(boxed.take(), Some(vec![1, 2]));
651 assert_eq!(boxed.take(), None);
652 assert_eq!(*boxed.get_or_init_racy(|| vec![9]), vec![9]);
653
654 let mut const_full = OnceCell::const_new_with(String::from("c"));
655 const_full.get_mut().expect("full").push('d');
656 assert_eq!(const_full.into_inner().as_deref(), Some("cd"));
657 }
658
659 #[test]
660 fn every_value_is_dropped_exactly_once() {
661 let marker = Arc::new(());
662 let racy = OnceCell::new();
663 racy.get_or_init_racy(|| Arc::clone(&marker));
664 let inline = OnceCell::new();
665 assert!(inline.set(Arc::clone(&marker)).is_ok());
666 let lost = OnceCell::new();
667 assert!(lost.set(Arc::clone(&marker)).is_ok());
668 assert!(lost.set(Arc::clone(&marker)).is_err());
669 let from = OnceCell::from(Arc::clone(&marker));
670 assert_eq!(Arc::strong_count(&marker), 5);
671 drop((racy, inline, lost, from));
672 assert_eq!(Arc::strong_count(&marker), 1);
673 let zst = OnceCell::new();
674 assert!(zst.set(()).is_ok());
675 assert_eq!(zst.get(), Some(&()));
676 let zst_racy = OnceCell::<()>::new();
677 assert_eq!(zst_racy.get_or_init_racy(|| ()), &());
678 }
679
680 #[test]
681 fn threads_setting_and_initializing_agree_on_one_value() {
682 let rounds = if cfg!(miri) { 3 } else { 200 };
683 for _ in 0..rounds {
684 let cell = Arc::new(OnceCell::new());
685 let setter = {
686 let cell = Arc::clone(&cell);
687 std::thread::spawn(move || cell.set(1).is_ok())
688 };
689 let init = {
690 let cell = Arc::clone(&cell);
691 std::thread::spawn(move || *block_on(cell.get_or_init(|| async { 2 })))
692 };
693 let set_won = setter.join().expect("setter");
694 let got = init.join().expect("init");
695 assert_eq!(got, *cell.get().expect("full"));
696 if set_won {
697 assert_eq!(got, 1);
698 }
699 }
700 }
701}