diff --git a/src/imp.rs b/src/imp.rs index bfecc86..dff407e 100644 --- a/src/imp.rs +++ b/src/imp.rs @@ -1,5 +1,6 @@ use core::{ cell::UnsafeCell, + mem::MaybeUninit, panic::{RefUnwindSafe, UnwindSafe}, sync::atomic::{AtomicU8, Ordering}, }; @@ -8,7 +9,7 @@ use crate::error::ConcurrentInitialization; pub(crate) struct OnceCell { state: AtomicU8, - value: UnsafeCell>, + value: UnsafeCell>, } const INCOMPLETE: u8 = 0x0; @@ -28,11 +29,11 @@ impl UnwindSafe for OnceCell {} impl OnceCell { pub(crate) const fn new() -> OnceCell { - OnceCell { state: AtomicU8::new(INCOMPLETE), value: UnsafeCell::new(None) } + OnceCell { state: AtomicU8::new(INCOMPLETE), value: UnsafeCell::new(MaybeUninit::uninit()) } } pub(crate) const fn with_value(value: T) -> OnceCell { - OnceCell { state: AtomicU8::new(COMPLETE), value: UnsafeCell::new(Some(value)) } + OnceCell { state: AtomicU8::new(COMPLETE), value: UnsafeCell::new(MaybeUninit::new(value)) } } /// Safety: synchronizes with store to value via Release/Acquire. @@ -53,7 +54,7 @@ impl OnceCell { { let mut f = Some(f); let mut res: Result<(), E> = Ok(()); - let slot: *mut Option = self.value.get(); + let slot: *mut MaybeUninit = self.value.get(); try_initialize_inner(&self.state, &mut || { // We are calling user-supplied function and need to be careful. // - if it returns Err, we unlock mutex and return without touching anything @@ -68,8 +69,8 @@ impl OnceCell { Ok(value) => unsafe { // Safe b/c we have a unique access and no panic may happen // until the cell is marked as initialized. - debug_assert!((*slot).is_none()); - *slot = Some(value); + debug_assert!(self.state.load(Ordering::Relaxed) == RUNNING); + *slot = MaybeUninit::new(value); true }, Err(err) => { @@ -91,21 +92,37 @@ impl OnceCell { pub(crate) unsafe fn get_unchecked(&self) -> &T { debug_assert!(self.is_initialized()); let slot = &*self.value.get(); - slot.as_ref().unwrap_unchecked() + slot.assume_init_ref() } /// Gets the mutable reference to the underlying value. /// Returns `None` if the cell is empty. pub(crate) fn get_mut(&mut self) -> Option<&mut T> { - // Safe b/c we have an exclusive access - let slot: &mut Option = unsafe { &mut *self.value.get() }; - slot.as_mut() + match *self.state.get_mut() { + COMPLETE => unsafe { Some(self.value.get_mut().assume_init_mut()) }, + _ => None, + } } /// Consumes this `OnceCell`, returning the wrapped value. /// Returns `None` if the cell was empty. - pub(crate) fn into_inner(self) -> Option { - self.value.into_inner() + pub(crate) fn into_inner(mut self) -> Option { + match *self.state.get_mut() { + COMPLETE => { + *self.state.get_mut() = INCOMPLETE; + unsafe { Some(self.value.get_mut().assume_init_read()) } + } + _ => None, + } + } +} + +impl Drop for OnceCell { + fn drop(&mut self) { + match *self.state.get_mut() { + COMPLETE => unsafe { self.value.get_mut().assume_init_drop() }, + _ => {} + } } }