diff --git a/src/imp_pl.rs b/src/imp_pl.rs index d838371..ed426b1 100644 --- a/src/imp_pl.rs +++ b/src/imp_pl.rs @@ -1,8 +1,7 @@ use std::{ cell::UnsafeCell, - mem::{self, MaybeUninit}, + hint, panic::{RefUnwindSafe, UnwindSafe}, - ptr, sync::atomic::{AtomicBool, Ordering}, }; @@ -11,7 +10,7 @@ use parking_lot::Mutex; pub(crate) struct OnceCell { mutex: Mutex<()>, is_initialized: AtomicBool, - value: UnsafeCell>, + value: UnsafeCell>, } // Why do we need `T: Send`? @@ -30,7 +29,7 @@ impl OnceCell { OnceCell { mutex: parking_lot::const_mutex(()), is_initialized: AtomicBool::new(false), - value: UnsafeCell::new(MaybeUninit::uninit()), + value: UnsafeCell::new(None), } } @@ -60,7 +59,9 @@ impl OnceCell { let value = f()?; // Safe b/c we have a unique access and no panic may happen // until the cell is marked as initialized. - unsafe { self.as_mut_ptr().write(value) }; + let slot: &mut Option = unsafe { &mut *self.value.get() }; + debug_assert!(slot.is_none()); + *slot = Some(value); self.is_initialized.store(true, Ordering::Release); } Ok(()) @@ -75,58 +76,29 @@ impl OnceCell { /// the contents are acquired by (synchronized to) this thread. pub(crate) unsafe fn get_unchecked(&self) -> &T { debug_assert!(self.is_initialized()); - &*self.as_ptr() + let slot: &Option = &*self.value.get(); + match slot { + Some(value) => value, + // This unsafe does improve performance, see `examples/bench`. + None => { + debug_assert!(false); + hint::unreachable_unchecked() + } + } } /// 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> { - if self.is_initialized() { - // Safe b/c we have a unique access and value is initialized. - Some(unsafe { &mut *self.as_mut_ptr() }) - } else { - None - } + // Safe b/c we have an exclusive access + let slot: &mut Option = unsafe { &mut *self.value.get() }; + slot.as_mut() } /// Consumes this `OnceCell`, returning the wrapped value. /// Returns `None` if the cell was empty. pub(crate) fn into_inner(self) -> Option { - if !self.is_initialized() { - return None; - } - - // Safe b/c we have a unique access and value is initialized. - let value: T = unsafe { ptr::read(self.as_ptr()) }; - - // It's OK to `mem::forget` without dropping, because both `self.mutex` - // and `self.is_initialized` are not heap-allocated. - mem::forget(self); - - Some(value) - } - - fn as_ptr(&self) -> *const T { - unsafe { - let slot: &MaybeUninit = &*self.value.get(); - slot.as_ptr() - } - } - - fn as_mut_ptr(&self) -> *mut T { - unsafe { - let slot: &mut MaybeUninit = &mut *self.value.get(); - slot.as_mut_ptr() - } - } -} - -impl Drop for OnceCell { - fn drop(&mut self) { - if self.is_initialized() { - // Safe b/c we have a unique access and value is initialized. - unsafe { ptr::drop_in_place(self.as_mut_ptr()) }; - } + self.value.into_inner() } } diff --git a/tests/test.rs b/tests/test.rs index c1351a7..18c86fe 100644 --- a/tests/test.rs +++ b/tests/test.rs @@ -186,6 +186,16 @@ mod unsync { }); eprintln!("use after free: {:?}", dangling_ref.get().unwrap()); } + + #[test] + // https://github.com/rust-lang/rust/issues/34761#issuecomment-256320669 + fn arrrrrrrrrrrrrrrrrrrrrr() { + let cell = OnceCell::new(); + { + let s = String::new(); + cell.set(&s).unwrap(); + } + } } #[cfg(feature = "std")] @@ -572,4 +582,14 @@ mod sync { .unwrap(); assert_eq!(cell.get(), Some(&"hello".to_string())); } + + #[test] + // https://github.com/rust-lang/rust/issues/34761#issuecomment-256320669 + fn arrrrrrrrrrrrrrrrrrrrrr() { + let cell = OnceCell::new(); + { + let s = String::new(); + cell.set(&s).unwrap(); + } + } }