From 62ae6ca17bd0d214d92400adeeaf3d6d92701ebf Mon Sep 17 00:00:00 2001
From: Markus Reiter <me@reitermark.us>
Date: Wed, 7 Sep 2022 06:50:24 +0200
Subject: [PATCH] Fix `is_initialized` and use atomics if available.

---
 src/imp_cs.rs     | 36 ++++++++++++++++++++++++++++--------
 src/lib.rs        |  6 ++++++
 xtask/src/main.rs |  7 +++++++
 3 files changed, 41 insertions(+), 8 deletions(-)

diff --git a/src/imp_cs.rs b/src/imp_cs.rs
index 56aadb0..a1abede 100644
--- a/src/imp_cs.rs
+++ b/src/imp_cs.rs
@@ -1,10 +1,14 @@
 use core::panic::{RefUnwindSafe, UnwindSafe};
 
+#[cfg(feature = "atomic-polyfill")]
+use atomic_polyfill::{AtomicBool, Ordering};
 use critical_section::{CriticalSection, Mutex};
 
 use crate::unsync::OnceCell as UnsyncOnceCell;
 
 pub(crate) struct OnceCell<T> {
+    #[cfg(feature = "atomic-polyfill")]
+    initialized: AtomicBool,
     value: Mutex<UnsyncOnceCell<T>>,
 }
 
@@ -21,18 +25,30 @@ impl<T: UnwindSafe> UnwindSafe for OnceCell<T> {}
 
 impl<T> OnceCell<T> {
     pub(crate) const fn new() -> OnceCell<T> {
-        OnceCell { value: Mutex::new(UnsyncOnceCell::new()) }
+        OnceCell {
+            #[cfg(feature = "atomic-polyfill")]
+            initialized: AtomicBool::new(false),
+            value: Mutex::new(UnsyncOnceCell::new()),
+        }
     }
 
     pub(crate) const fn with_value(value: T) -> OnceCell<T> {
-        OnceCell { value: Mutex::new(UnsyncOnceCell::with_value(value)) }
+        OnceCell {
+            #[cfg(feature = "atomic-polyfill")]
+            initialized: AtomicBool::new(true),
+            value: Mutex::new(UnsyncOnceCell::with_value(value)),
+        }
     }
 
     #[inline]
     pub(crate) fn is_initialized(&self) -> bool {
-        // The only write access is synchronized in `initialize` with a
-        // critical section, so read-only access is valid everywhere else.
-        unsafe { self.value.borrow(CriticalSection::new()).get().is_some() }
+        #[cfg(feature = "atomic-polyfill")]
+        {
+            self.initialized.load(Ordering::Acquire)
+        }
+
+        #[cfg(not(feature = "atomic-polyfill"))]
+        critical_section::with(|cs| self.value.borrow(cs).get().is_some())
     }
 
     #[cold]
@@ -42,7 +58,12 @@ impl<T> OnceCell<T> {
     {
         critical_section::with(|cs| {
             let cell = self.value.borrow(cs);
-            cell.get_or_try_init(f).map(|_| ())
+            cell.get_or_try_init(f).map(|_| {
+                #[cfg(feature = "atomic-polyfill")]
+                self.initialized.store(true, Ordering::Release);
+
+                ()
+            })
         })
     }
 
@@ -55,8 +76,7 @@ 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());
-        // The only write access is synchronized in `initialize` with a
-        // critical section, so read-only access is valid everywhere else.
+        // SAFETY: The caller ensures that the value is initialized and access synchronized.
         self.value.borrow(CriticalSection::new()).get_unchecked()
     }
 
diff --git a/src/lib.rs b/src/lib.rs
index 64f09e5..8d1f56c 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -453,6 +453,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()
@@ -474,6 +475,7 @@ 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()
@@ -486,6 +488,7 @@ pub mod unsync {
         ///
         /// 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())
         }
@@ -985,6 +988,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()
         }
@@ -996,6 +1000,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()
         }
@@ -1181,6 +1186,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()
         }
diff --git a/xtask/src/main.rs b/xtask/src/main.rs
index 076c782..7ef0642 100644
--- a/xtask/src/main.rs
+++ b/xtask/src/main.rs
@@ -46,6 +46,13 @@ fn try_main() -> Result<()> {
             "cargo test --features 'unstable std critical-section' --no-default-features --release"
         )
         .run()?;
+
+        cmd!("cargo test --features 'unstable std atomic-polyfill critical-section' --no-default-features")
+                .run()?;
+        cmd!(
+                "cargo test --features 'unstable std atomic-polyfill critical-section' --no-default-features --release"
+            )
+            .run()?;
     }
 
     {