Skip to content

Commit

Permalink
Fix is_initialized and use atomics if available.
Browse files Browse the repository at this point in the history
  • Loading branch information
reitermarkus committed Sep 7, 2022
1 parent 7f05b37 commit 62ae6ca
Show file tree
Hide file tree
Showing 3 changed files with 41 additions and 8 deletions.
36 changes: 28 additions & 8 deletions 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>>,
}

Expand All @@ -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]
Expand All @@ -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);

()
})
})
}

Expand All @@ -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()
}

Expand Down
6 changes: 6 additions & 0 deletions src/lib.rs
Expand Up @@ -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()
Expand All @@ -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()
Expand All @@ -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())
}
Expand Down Expand Up @@ -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()
}
Expand All @@ -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()
}
Expand Down Expand Up @@ -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()
}
Expand Down
7 changes: 7 additions & 0 deletions xtask/src/main.rs
Expand Up @@ -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()?;
}

{
Expand Down

0 comments on commit 62ae6ca

Please sign in to comment.