Add no_std implementation based on critical-section.

Cargo.toml

@@ -32,15 +32,19 @@ parking_lot_core = { version = "0.9.3", optional = true, default_features = fals
 # and make sure you understand all the implications
 atomic-polyfill = { version = "1", optional = true }
 
+# Uses `critical-section` to implement `once_cell::sync::OnceCell`.
+critical_section = { package = "critical-section", version = "1", optional = true }
+
 [dev-dependencies]
 lazy_static = "1.0.0"
 crossbeam-utils = "0.8.7"
 regex =  "1.2.0"
+critical_section = { package = "critical-section", version = "1.1.1", features = ["std"] }
 
 [features]
 default = ["std"]
 # Enables `once_cell::sync` module.
-std = ["alloc"]
+std = ["alloc", "sync"]
 # Enables `once_cell::race::OnceBox` type.
 alloc = ["race"]
 # Enables `once_cell::race` module.
@@ -49,8 +53,12 @@ race = []
 # At the moment, this feature is unused.
 unstable = []
 
+sync = []
+
 parking_lot = ["parking_lot_core"]
 
+critical-section = ["critical_section", "atomic-polyfill", "sync"]
+
 [[example]]
 name = "bench"
 required-features = ["std"]

src/imp_cs.rs

@@ -0,0 +1,78 @@
+use core::panic::{RefUnwindSafe, UnwindSafe};
+
+use atomic_polyfill::{AtomicBool, Ordering};
+use critical_section::{CriticalSection, Mutex};
+
+use crate::unsync;
+
+pub(crate) struct OnceCell<T> {
+    initialized: AtomicBool,
+    // Use `unsync::OnceCell` internally since `Mutex` does not provide
+    // interior mutability and to be able to re-use `get_or_try_init`.
+    value: Mutex<unsync::OnceCell<T>>,
+}
+
+// Why do we need `T: Send`?
+// Thread A creates a `OnceCell` and shares it with
+// scoped thread B, which fills the cell, which is
+// then destroyed by A. That is, destructor observes
+// a sent value.
+unsafe impl<T: Sync + Send> Sync for OnceCell<T> {}
+unsafe impl<T: Send> Send for OnceCell<T> {}
+
+impl<T: RefUnwindSafe + UnwindSafe> RefUnwindSafe for OnceCell<T> {}
+impl<T: UnwindSafe> UnwindSafe for OnceCell<T> {}
+
+impl<T> OnceCell<T> {
+    pub(crate) const fn new() -> OnceCell<T> {
+        OnceCell { initialized: AtomicBool::new(false), value: Mutex::new(unsync::OnceCell::new()) }
+    }
+
+    pub(crate) const fn with_value(value: T) -> OnceCell<T> {
+        OnceCell {
+            initialized: AtomicBool::new(true),
+            value: Mutex::new(unsync::OnceCell::with_value(value)),
+        }
+    }
+
+    #[inline]
+    pub(crate) fn is_initialized(&self) -> bool {
+        self.initialized.load(Ordering::Acquire)
+    }
+
+    #[cold]
+    pub(crate) fn initialize<F, E>(&self, f: F) -> Result<(), E>
+    where
+        F: FnOnce() -> Result<T, E>,
+    {
+        critical_section::with(|cs| {
+            let cell = self.value.borrow(cs);
+            cell.get_or_try_init(f).map(|_| {
+                self.initialized.store(true, Ordering::Release);
+            })
+        })
+    }
+
+    /// Get the reference to the underlying value, without checking if the cell
+    /// is initialized.
+    ///
+    /// # Safety
+    ///
+    /// Caller must ensure that the cell is in initialized state, and that
+    /// the contents are acquired by (synchronized to) this thread.
+    pub(crate) unsafe fn get_unchecked(&self) -> &T {
+        debug_assert!(self.is_initialized());
+        // SAFETY: The caller ensures that the value is initialized and access synchronized.
+        self.value.borrow(CriticalSection::new()).get_unchecked()
+    }
+
+    #[inline]
+    pub(crate) fn get_mut(&mut self) -> Option<&mut T> {
+        self.value.get_mut().get_mut()
+    }
+
+    #[inline]
+    pub(crate) fn into_inner(self) -> Option<T> {
+        self.value.into_inner().into_inner()
+    }
+}

src/imp_pl.rs

@@ -1,6 +1,5 @@
 use std::{
     cell::UnsafeCell,
-    hint,
     panic::{RefUnwindSafe, UnwindSafe},
     sync::atomic::{AtomicU8, Ordering},
 };
@@ -101,15 +100,8 @@ impl<T> OnceCell<T> {
     /// the contents are acquired by (synchronized to) this thread.
     pub(crate) unsafe fn get_unchecked(&self) -> &T {
         debug_assert!(self.is_initialized());
-        let slot: &Option<T> = &*self.value.get();
-        match slot {
-            Some(value) => value,
-            // This unsafe does improve performance, see `examples/bench`.
-            None => {
-                debug_assert!(false);
-                hint::unreachable_unchecked()
-            }
-        }
+        let slot = &*self.value.get();
+        crate::unwrap_unchecked(slot.as_ref())
     }
 
     /// Gets the mutable reference to the underlying value.

src/imp_std.rs

@@ -5,15 +5,12 @@
 
 use std::{
     cell::{Cell, UnsafeCell},
-    hint::unreachable_unchecked,
     marker::PhantomData,
     panic::{RefUnwindSafe, UnwindSafe},
     sync::atomic::{AtomicBool, AtomicPtr, Ordering},
     thread::{self, Thread},
 };
 
-use crate::take_unchecked;
-
 #[derive(Debug)]
 pub(crate) struct OnceCell<T> {
     // This `queue` field is the core of the implementation. It encodes two
@@ -81,7 +78,7 @@ impl<T> OnceCell<T> {
         initialize_or_wait(
             &self.queue,
             Some(&mut || {
-                let f = unsafe { take_unchecked(&mut f) };
+                let f = unsafe { crate::take_unchecked(&mut f) };
                 match f() {
                     Ok(value) => {
                         unsafe { *slot = Some(value) };
@@ -111,15 +108,8 @@ impl<T> OnceCell<T> {
     /// the contents are acquired by (synchronized to) this thread.
     pub(crate) unsafe fn get_unchecked(&self) -> &T {
         debug_assert!(self.is_initialized());
-        let slot: &Option<T> = &*self.value.get();
-        match slot {
-            Some(value) => value,
-            // This unsafe does improve performance, see `examples/bench`.
-            None => {
-                debug_assert!(false);
-                unreachable_unchecked()
-            }
-        }
+        let slot = &*self.value.get();
+        crate::unwrap_unchecked(slot.as_ref())
     }
 
     /// Gets the mutable reference to the underlying value.

src/lib.rs

@@ -332,25 +332,29 @@
 #[cfg(feature = "alloc")]
 extern crate alloc;
 
-#[cfg(feature = "std")]
-#[cfg(feature = "parking_lot")]
+#[cfg(all(feature = "critical-section", not(feature = "std")))]
+#[path = "imp_cs.rs"]
+mod imp;
+
+#[cfg(all(feature = "std", feature = "parking_lot"))]
 #[path = "imp_pl.rs"]
 mod imp;
 
-#[cfg(feature = "std")]
-#[cfg(not(feature = "parking_lot"))]
+#[cfg(all(feature = "std", not(feature = "parking_lot")))]
 #[path = "imp_std.rs"]
 mod imp;
 
 /// Single-threaded version of `OnceCell`.
 pub mod unsync {
     use core::{
         cell::{Cell, UnsafeCell},
-        fmt, hint, mem,
+        fmt, mem,
         ops::{Deref, DerefMut},
         panic::{RefUnwindSafe, UnwindSafe},
     };
 
+    use super::unwrap_unchecked;
+
     /// A cell which can be written to only once. It is not thread safe.
     ///
     /// Unlike [`std::cell::RefCell`], a `OnceCell` provides simple `&`
@@ -442,6 +446,7 @@ pub mod unsync {
         /// Gets a reference to the underlying value.
         ///
         /// Returns `None` if the cell is empty.
+        #[inline]
         pub fn get(&self) -> Option<&T> {
             // Safe due to `inner`'s invariant
             unsafe { &*self.inner.get() }.as_ref()
@@ -463,11 +468,24 @@ pub mod unsync {
         /// *cell.get_mut().unwrap() = 93;
         /// assert_eq!(cell.get(), Some(&93));
         /// ```
+        #[inline]
         pub fn get_mut(&mut self) -> Option<&mut T> {
             // Safe because we have unique access
             unsafe { &mut *self.inner.get() }.as_mut()
         }
 
+        /// Get the reference to the underlying value, without checking if the
+        /// cell is initialized.
+        ///
+        /// # Safety
+        ///
+        /// Caller must ensure that the cell is in initialized state.
+        #[cfg(feature = "critical-section")]
+        #[inline]
+        pub(crate) unsafe fn get_unchecked(&self) -> &T {
+            crate::unwrap_unchecked(self.get())
+        }
+
         /// Sets the contents of this cell to `value`.
         ///
         /// Returns `Ok(())` if the cell was empty and `Err(value)` if it was
@@ -510,16 +528,14 @@ pub mod unsync {
             if let Some(old) = self.get() {
                 return Err((old, value));
             }
+
             let slot = unsafe { &mut *self.inner.get() };
             // This is the only place where we set the slot, no races
             // due to reentrancy/concurrency are possible, and we've
             // checked that slot is currently `None`, so this write
             // maintains the `inner`'s invariant.
             *slot = Some(value);
-            Ok(match &*slot {
-                Some(value) => value,
-                None => unsafe { hint::unreachable_unchecked() },
-            })
+            Ok(unsafe { unwrap_unchecked(slot.as_ref()) })
         }
 
         /// Gets the contents of the cell, initializing it with `f`
@@ -592,7 +608,7 @@ pub mod unsync {
             // `assert`, while keeping `set/get` would be sound, but it seems
             // better to panic, rather than to silently use an old value.
             assert!(self.set(val).is_ok(), "reentrant init");
-            Ok(self.get().unwrap())
+            Ok(unsafe { unwrap_unchecked(self.get()) })
         }
 
         /// Takes the value out of this `OnceCell`, moving it back to an uninitialized state.
@@ -814,16 +830,16 @@ pub mod unsync {
 }
 
 /// Thread-safe, blocking version of `OnceCell`.
-#[cfg(feature = "std")]
+#[cfg(feature = "sync")]
 pub mod sync {
-    use std::{
+    use core::{
         cell::Cell,
         fmt, mem,
         ops::{Deref, DerefMut},
         panic::RefUnwindSafe,
     };
 
-    use crate::{imp::OnceCell as Imp, take_unchecked};
+    use super::{imp::OnceCell as Imp, take_unchecked};
 
     /// A thread-safe cell which can be written to only once.
     ///
@@ -942,8 +958,9 @@ pub mod sync {
         /// // Will return 92, but might block until the other thread does `.set`.
         /// let value: &u32 = cell.wait();
         /// assert_eq!(*value, 92);
-        /// t.join().unwrap();;
+        /// t.join().unwrap();
         /// ```
+        #[cfg(not(feature = "critical-section"))]
         pub fn wait(&self) -> &T {
             if !self.0.is_initialized() {
                 self.0.wait()
@@ -969,6 +986,7 @@ pub mod sync {
         /// cell.set(92).unwrap();
         /// cell = OnceCell::new();
         /// ```
+        #[inline]
         pub fn get_mut(&mut self) -> Option<&mut T> {
             self.0.get_mut()
         }
@@ -980,6 +998,7 @@ pub mod sync {
         ///
         /// Caller must ensure that the cell is in initialized state, and that
         /// the contents are acquired by (synchronized to) this thread.
+        #[inline]
         pub unsafe fn get_unchecked(&self) -> &T {
             self.0.get_unchecked()
         }
@@ -1109,6 +1128,7 @@ pub mod sync {
             if let Some(value) = self.get() {
                 return Ok(value);
             }
+
             self.0.initialize(f)?;
 
             // Safe b/c value is initialized.
@@ -1164,6 +1184,7 @@ pub mod sync {
         /// cell.set("hello".to_string()).unwrap();
         /// assert_eq!(cell.into_inner(), Some("hello".to_string()));
         /// ```
+        #[inline]
         pub fn into_inner(self) -> Option<T> {
             self.0.into_inner()
         }
@@ -1356,13 +1377,19 @@ pub mod sync {
 #[cfg(feature = "race")]
 pub mod race;
 
-#[cfg(feature = "std")]
+#[cfg(feature = "sync")]
+#[inline]
 unsafe fn take_unchecked<T>(val: &mut Option<T>) -> T {
-    match val.take() {
-        Some(it) => it,
+    unwrap_unchecked(val.take())
+}
+
+#[inline]
+unsafe fn unwrap_unchecked<T>(val: Option<T>) -> T {
+    match val {
+        Some(value) => value,
         None => {
             debug_assert!(false);
-            std::hint::unreachable_unchecked()
+            core::hint::unreachable_unchecked()
         }
     }
 }

src/race.rs

@@ -165,6 +165,7 @@ impl OnceBool {
     fn from_usize(value: NonZeroUsize) -> bool {
         value.get() == 1
     }
+
     #[inline]
     fn to_usize(value: bool) -> NonZeroUsize {
         unsafe { NonZeroUsize::new_unchecked(if value { 1 } else { 2 }) }