Split out synchronization from is_initialized

.travis.yml

@@ -32,6 +32,11 @@ matrix:
     script:
     - cargo build --no-default-features --test test
 
+  - rust: stable
+    script:
+    - cargo test --no-default-features --features "std maybe_uninit"
+    - cargo test --no-default-features --features "std maybe_uninit parking_lot"
+
   - rust: 1.31.1
     script:
     - mv Cargo.lock.min Cargo.lock

Cargo.toml

@@ -30,6 +30,8 @@ regex =  "1.2.0"
 default = ["std"]
 # Enables `once_cell::sync` module.
 std = []
+# Use `MaybeUninit` in the `sync` module.
+maybe_uninit = []
 
 [[example]]
 name = "bench"

src/imp_pl.rs

@@ -3,13 +3,14 @@ use std::{
     panic::{RefUnwindSafe, UnwindSafe},
     sync::atomic::{AtomicBool, Ordering},
 };
+use crate::maybe_uninit::MaybeUninit;
 
 use parking_lot::{lock_api::RawMutex as _RawMutex, RawMutex};
 
 pub(crate) struct OnceCell<T> {
     mutex: Mutex,
     is_initialized: AtomicBool,
-    pub(crate) value: UnsafeCell<Option<T>>,
+    pub(crate) value: UnsafeCell<MaybeUninit<T>>,
 }
 
 // Why do we need `T: Send`?
@@ -28,20 +29,33 @@ impl<T> OnceCell<T> {
         OnceCell {
             mutex: Mutex::new(),
             is_initialized: AtomicBool::new(false),
-            value: UnsafeCell::new(None),
+            value: UnsafeCell::new(MaybeUninit::uninit()),
         }
     }
 
-    /// Safety: synchronizes with store to value via Release/Acquire.
+    /// Safety: does no synchronization, only checks whether the `OnceCell` is initialized.
     #[inline]
     pub(crate) fn is_initialized(&self) -> bool {
-        self.is_initialized.load(Ordering::Acquire)
+        self.is_initialized.load(Ordering::Relaxed)
+    }
+
+    /// Safety: synchronizes with store to value via Release/Acquire.
+    #[inline]
+    pub(crate) fn sync_get(&self) -> Option<&T> {
+        if self.is_initialized.load(Ordering::Acquire) {
+            unsafe {
+                let slot: &MaybeUninit<T> = &*self.value.get();
+                Some(&*slot.as_ptr())
+            }
+        } else {
+            None
+        }
     }
 
     /// Safety: synchronizes with store to value via `is_initialized` or mutex
     /// lock/unlock, writes value only once because of the mutex.
     #[cold]
-    pub(crate) fn initialize<F, E>(&self, f: F) -> Result<(), E>
+    pub(crate) fn initialize<F, E>(&self, f: F) -> Result<&T, E>
     where
         F: FnOnce() -> Result<T, E>,
     {
@@ -56,12 +70,19 @@ impl<T> OnceCell<T> {
             // - finally, if it returns Ok, we store the value and store the flag with
             //   `Release`, which synchronizes with `Acquire`s.
             let value = f()?;
-            let slot: &mut Option<T> = unsafe { &mut *self.value.get() };
-            debug_assert!(slot.is_none());
-            *slot = Some(value);
+            unsafe {
+                let slot: &mut MaybeUninit<T> = &mut *self.value.get();
+                // FIXME: replace with `slot.as_mut_ptr().write(value)`
+                slot.write(value);
+            }
             self.is_initialized.store(true, Ordering::Release);
         }
-        Ok(())
+        unsafe {
+            let _acquire = self.is_initialized.load(Ordering::Acquire);
+            debug_assert!(_acquire);
+            let slot: &MaybeUninit<T> = &*self.value.get();
+            Ok(&*slot.as_ptr())
+        }
     }
 }
 
@@ -92,9 +113,8 @@ impl Drop for MutexGuard<'_> {
 }
 
 #[test]
-#[cfg(target_pointer_width = "64")]
 fn test_size() {
     use std::mem::size_of;
 
-    assert_eq!(size_of::<OnceCell<u32>>(), 3 * size_of::<u32>());
+    assert_eq!(size_of::<OnceCell<bool>>(), 3 * size_of::<u8>());
 }

src/imp_std.rs

@@ -10,18 +10,14 @@ use std::{
     sync::atomic::{AtomicBool, AtomicUsize, Ordering},
     thread::{self, Thread},
 };
+use crate::maybe_uninit::MaybeUninit;
 
-#[derive(Debug)]
 pub(crate) struct OnceCell<T> {
     // This `state` word is actually an encoded version of just a pointer to a
     // `Waiter`, so we add the `PhantomData` appropriately.
     state_and_queue: AtomicUsize,
     _marker: PhantomData<*mut Waiter>,
-    // FIXME: switch to `std::mem::MaybeUninit` once we are ready to bump MSRV
-    // that far. It was stabilized in 1.36.0, so, if you are reading this and
-    // it's higher than 1.46.0 outside, please send a PR! ;) (and do the same
-    // for `Lazy`, while we are at it).
-    pub(crate) value: UnsafeCell<Option<T>>,
+    pub(crate) value: UnsafeCell<MaybeUninit<T>>,
 }
 
 // Why do we need `T: Send`?
@@ -66,25 +62,38 @@ impl<T> OnceCell<T> {
         OnceCell {
             state_and_queue: AtomicUsize::new(INCOMPLETE),
             _marker: PhantomData,
-            value: UnsafeCell::new(None),
+            value: UnsafeCell::new(MaybeUninit::uninit()),
         }
     }
 
-    /// Safety: synchronizes with store to value via Release/(Acquire|SeqCst).
+    /// Safety: does no synchronization, only checks whether the `OnceCell` is initialized.
     #[inline]
     pub(crate) fn is_initialized(&self) -> bool {
+        self.state_and_queue.load(Ordering::Relaxed) == COMPLETE
+    }
+
+    /// Safety: synchronizes with store to value via Release/Acquire.
+    #[inline]
+    pub(crate) fn sync_get(&self) -> Option<&T> {
         // An `Acquire` load is enough because that makes all the initialization
         // operations visible to us, and, this being a fast path, weaker
         // ordering helps with performance. This `Acquire` synchronizes with
         // `SeqCst` operations on the slow path.
-        self.state_and_queue.load(Ordering::Acquire) == COMPLETE
+        if self.state_and_queue.load(Ordering::Acquire) == COMPLETE {
+            unsafe {
+                let slot: &MaybeUninit<T> = &*self.value.get();
+                Some(&*slot.as_ptr())
+            }
+        } else {
+            None
+        }
     }
 
     /// Safety: synchronizes with store to value via SeqCst read from state,
     /// writes value only once because we never get to INCOMPLETE state after a
     /// successful write.
     #[cold]
-    pub(crate) fn initialize<F, E>(&self, f: F) -> Result<(), E>
+    pub(crate) fn initialize<F, E>(&self, f: F) -> Result<&T, E>
     where
         F: FnOnce() -> Result<T, E>,
     {
@@ -95,7 +104,11 @@ impl<T> OnceCell<T> {
             let f = f.take().unwrap();
             match f() {
                 Ok(value) => {
-                    unsafe { *slot.get() = Some(value) };
+                    unsafe {
+                        let slot: &mut MaybeUninit<T> = &mut *slot.get();
+                        // FIXME: replace with `slot.as_mut_ptr().write(value)`
+                        slot.write(value);
+                    }
                     true
                 }
                 Err(e) => {
@@ -104,7 +117,13 @@ impl<T> OnceCell<T> {
                 }
             }
         });
-        res
+        res?;
+        unsafe {
+            let _acquire = self.state_and_queue.load(Ordering::Acquire);
+            debug_assert!(_acquire == COMPLETE);
+            let slot: &MaybeUninit<T> = &*self.value.get();
+            Ok(&*slot.as_ptr())
+        }
     }
 }
 
@@ -315,10 +334,9 @@ mod tests {
     }
 
     #[test]
-    #[cfg(target_pointer_width = "64")]
     fn test_size() {
         use std::mem::size_of;
 
-        assert_eq!(size_of::<OnceCell<u32>>(), 4 * size_of::<u32>());
+        assert_eq!(size_of::<OnceCell<&usize>>(), 2 * size_of::<usize>());
     }
 }

src/lib.rs

@@ -233,6 +233,8 @@ mod imp;
 #[path = "imp_std.rs"]
 mod imp;
 
+mod maybe_uninit;
+
 pub mod unsync {
     use core::{
         cell::{Cell, UnsafeCell},
@@ -568,14 +570,15 @@ pub mod unsync {
 #[cfg(feature = "std")]
 pub mod sync {
     use std::{
-        cell::Cell,
+        cell::{Cell, UnsafeCell},
         fmt,
-        hint::unreachable_unchecked,
+        mem::{self, ManuallyDrop},
         ops::{Deref, DerefMut},
         panic::RefUnwindSafe,
     };
 
     use crate::imp::OnceCell as Imp;
+    use crate::maybe_uninit::MaybeUninit;
 
     /// A thread-safe cell which can be written to only once.
     ///
@@ -622,6 +625,12 @@ pub mod sync {
         }
     }
 
+    impl<T> Drop for OnceCell<T> {
+        fn drop(&mut self) {
+            self.take_inner();
+        }
+    }
+
     impl<T: Clone> Clone for OnceCell<T> {
         fn clone(&self) -> OnceCell<T> {
             let res = OnceCell::new();
@@ -662,20 +671,22 @@ pub mod sync {
         /// Returns `None` if the cell is empty, or being initialized. This
         /// method never blocks.
         pub fn get(&self) -> Option<&T> {
-            if self.0.is_initialized() {
-                // Safe b/c checked is_initialize
-                Some(unsafe { self.get_unchecked() })
-            } else {
-                None
-            }
+            self.0.sync_get()
         }
 
         /// Gets the mutable reference to the underlying value.
         ///
         /// Returns `None` if the cell is empty.
         pub fn get_mut(&mut self) -> Option<&mut T> {
-            // Safe b/c we have a unique access.
-            unsafe { &mut *self.0.value.get() }.as_mut()
+            // Safe because we have a unique access.
+            // We don't need to do any synchronization. Either we did the initialization ourselves,
+            // or whatever mechanism send us the unique reference from another thread took care of
+            // synchronization.
+            if self.0.is_initialized() {
+                Some(unsafe { &mut *(*self.0.value.get()).as_mut_ptr() })
+            } else {
+                None
+            }
         }
 
         /// Get the reference to the underlying value, without checking if the
@@ -687,15 +698,8 @@ pub mod sync {
         /// the contents are acquired by (synchronized to) this thread.
         pub unsafe fn get_unchecked(&self) -> &T {
             debug_assert!(self.0.is_initialized());
-            let slot: &Option<T> = &*self.0.value.get();
-            match slot {
-                Some(value) => value,
-                // This unsafe does improve performance, see `examples/bench`.
-                None => {
-                    debug_assert!(false);
-                    unreachable_unchecked()
-                }
-            }
+            let slot: &MaybeUninit<T> = &*self.0.value.get();
+            &*slot.as_ptr()
         }
 
         /// Sets the contents of this cell to `value`.
@@ -800,11 +804,7 @@ pub mod sync {
             if let Some(value) = self.get() {
                 return Ok(value);
             }
-            self.0.initialize(f)?;
-
-            // Safe b/c called initialize
-            debug_assert!(self.0.is_initialized());
-            Ok(unsafe { self.get_unchecked() })
+            self.0.initialize(f)
         }
 
         /// Consumes the `OnceCell`, returning the wrapped value. Returns
@@ -822,10 +822,28 @@ pub mod sync {
         /// cell.set("hello".to_string()).unwrap();
         /// assert_eq!(cell.into_inner(), Some("hello".to_string()));
         /// ```
-        pub fn into_inner(self) -> Option<T> {
-            // Because `into_inner` takes `self` by value, the compiler statically verifies
-            // that it is not currently borrowed. So it is safe to move out `Option<T>`.
-            self.0.value.into_inner()
+        pub fn into_inner(mut self) -> Option<T> {
+            let res = self.take_inner();
+            // Don't drop this `OnceCell`. We just moved out one of the fields, but didn't set
+            // the state to uninitialized.
+            ManuallyDrop::new(self);
+            res
+        }
+
+        /// Takes the wrapped value out of a `OnceCell`.
+        /// Afterwards the cell is no longer initialized, and must be free'd WITHOUT dropping.
+        /// Only used by `into_inner` and `drop`.
+        fn take_inner(&mut self) -> Option<T> {
+            // The mutable reference guarantees there are no other threads that can observe us
+            // taking out the wrapped value.
+            // Right after this function `self` is supposed to be freed, so it makes little sense
+            // to atomically set the state to uninitialized.
+            if self.0.is_initialized() {
+                let value = mem::replace(&mut self.0.value, UnsafeCell::new(MaybeUninit::uninit()));
+                Some(unsafe { value.into_inner().assume_init() })
+            } else {
+                None
+            }
         }
     }
 
@@ -863,6 +881,7 @@ pub mod sync {
     /// ```
     pub struct Lazy<T, F = fn() -> T> {
         cell: OnceCell<T>,
+        // FIXME: replace with MaybeUninit
         init: Cell<Option<F>>,
     }