1 // Based on unstable std::sync::OnceLock. 2 // 3 // Source: https://github.com/rust-lang/rust/blob/8e9c93df464b7ada3fc7a1c8ccddd9dcb24ee0a0/library/std/src/sync/once_lock.rs 4 5 use core::cell::UnsafeCell; 6 use core::mem::MaybeUninit; 7 use core::sync::atomic::{AtomicBool, Ordering}; 8 use std::sync::Once; 9 10 pub(crate) struct OnceLock<T> { 11 once: Once, 12 // Once::is_completed requires Rust 1.43, so use this to track of whether they have been initialized. 13 is_initialized: AtomicBool, 14 value: UnsafeCell<MaybeUninit<T>>, 15 // Unlike std::sync::OnceLock, we don't need PhantomData here because 16 // we don't use #[may_dangle]. 17 } 18 19 unsafe impl<T: Sync + Send> Sync for OnceLock<T> {} 20 unsafe impl<T: Send> Send for OnceLock<T> {} 21 22 impl<T> OnceLock<T> { 23 /// Creates a new empty cell. 24 #[must_use] new() -> Self25 pub(crate) const fn new() -> Self { 26 Self { 27 once: Once::new(), 28 is_initialized: AtomicBool::new(false), 29 value: UnsafeCell::new(MaybeUninit::uninit()), 30 } 31 } 32 33 /// Gets the contents of the cell, initializing it with `f` if the cell 34 /// was empty. 35 /// 36 /// Many threads may call `get_or_init` concurrently with different 37 /// initializing functions, but it is guaranteed that only one function 38 /// will be executed. 39 /// 40 /// # Panics 41 /// 42 /// If `f` panics, the panic is propagated to the caller, and the cell 43 /// remains uninitialized. 44 /// 45 /// It is an error to reentrantly initialize the cell from `f`. The 46 /// exact outcome is unspecified. Current implementation deadlocks, but 47 /// this may be changed to a panic in the future. get_or_init<F>(&self, f: F) -> &T where F: FnOnce() -> T,48 pub(crate) fn get_or_init<F>(&self, f: F) -> &T 49 where 50 F: FnOnce() -> T, 51 { 52 // Fast path check 53 if self.is_initialized() { 54 // SAFETY: The inner value has been initialized 55 return unsafe { self.get_unchecked() }; 56 } 57 self.initialize(f); 58 59 debug_assert!(self.is_initialized()); 60 61 // SAFETY: The inner value has been initialized 62 unsafe { self.get_unchecked() } 63 } 64 65 #[inline] is_initialized(&self) -> bool66 fn is_initialized(&self) -> bool { 67 self.is_initialized.load(Ordering::Acquire) 68 } 69 70 #[cold] initialize<F>(&self, f: F) where F: FnOnce() -> T,71 fn initialize<F>(&self, f: F) 72 where 73 F: FnOnce() -> T, 74 { 75 let slot = self.value.get().cast::<T>(); 76 let is_initialized = &self.is_initialized; 77 78 self.once.call_once(|| { 79 let value = f(); 80 unsafe { 81 slot.write(value); 82 } 83 is_initialized.store(true, Ordering::Release); 84 }); 85 } 86 87 /// # Safety 88 /// 89 /// The value must be initialized get_unchecked(&self) -> &T90 unsafe fn get_unchecked(&self) -> &T { 91 debug_assert!(self.is_initialized()); 92 &*self.value.get().cast::<T>() 93 } 94 } 95 96 impl<T> Drop for OnceLock<T> { drop(&mut self)97 fn drop(&mut self) { 98 if self.is_initialized() { 99 // SAFETY: The inner value has been initialized 100 unsafe { self.value.get().cast::<T>().drop_in_place() }; 101 } 102 } 103 } 104