From 771a98b52503007fbec5cf9e4e1569ee2c40ef57 Mon Sep 17 00:00:00 2001 From: Markus Reiter Date: Tue, 6 Sep 2022 08:52:11 +0200 Subject: [PATCH] Add `no_std` implementation based on `critical-section`. --- Cargo.toml | 10 +++- src/imp_cs.rs | 131 ++++++++++++++++++++++++++++++++++++++++++++++ src/imp_pl.rs | 12 +---- src/imp_std.rs | 16 ++---- src/lib.rs | 78 ++++++++++++++++----------- src/race.rs | 1 + xtask/src/main.rs | 7 +++ 7 files changed, 200 insertions(+), 55 deletions(-) create mode 100644 src/imp_cs.rs diff --git a/Cargo.toml b/Cargo.toml index b7ebb35..bb40a01 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -31,15 +31,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 = { version = "1.1", optional = true } + [dev-dependencies] lazy_static = "1.0.0" crossbeam-utils = "0.8.7" regex = "1.2.0" +critical-section = { version = "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. @@ -48,8 +52,12 @@ race = [] # At the moment, this feature is unused. unstable = [] +sync = [] + parking_lot = ["parking_lot_core"] +critical-section = ["dep:critical-section", "sync"] + [[example]] name = "bench" required-features = ["std"] diff --git a/src/imp_cs.rs b/src/imp_cs.rs new file mode 100644 index 0000000..53848f8 --- /dev/null +++ b/src/imp_cs.rs @@ -0,0 +1,131 @@ +#[cfg(feature = "atomic-polyfill")] +use atomic_polyfill as atomic; +#[cfg(not(feature = "atomic-polyfill"))] +use core::sync::atomic; + +use atomic::{AtomicU8, Ordering}; + +use core::panic::{RefUnwindSafe, UnwindSafe}; + +use crate::unsync::OnceCell as UnsyncOnceCell; + +pub(crate) struct OnceCell { + state: AtomicU8, + value: UnsyncOnceCell, +} + +const INCOMPLETE: u8 = 0; +const RUNNING: u8 = 1; +const COMPLETE: u8 = 2; + +// 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 Sync for OnceCell {} +unsafe impl Send for OnceCell {} + +impl RefUnwindSafe for OnceCell {} +impl UnwindSafe for OnceCell {} + +impl OnceCell { + pub(crate) const fn new() -> OnceCell { + OnceCell { state: AtomicU8::new(INCOMPLETE), value: UnsyncOnceCell::new() } + } + + pub(crate) const fn with_value(value: T) -> OnceCell { + OnceCell { state: AtomicU8::new(COMPLETE), value: UnsyncOnceCell::with_value(value) } + } + + #[inline] + pub(crate) fn is_initialized(&self) -> bool { + self.state.load(Ordering::Acquire) == COMPLETE + } + + #[cold] + pub(crate) fn initialize(&self, f: F) -> Result<(), E> + where + F: FnOnce() -> Result, + { + let mut f = Some(f); + let mut res: Result<(), E> = Ok(()); + let slot: &UnsyncOnceCell = &self.value; + initialize_inner(&self.state, &mut || { + let f = unsafe { crate::take_unchecked(&mut f) }; + match f() { + Ok(value) => { + unsafe { crate::unwrap_unchecked(slot.set(value).ok()) }; + true + } + Err(err) => { + res = Err(err); + false + } + } + }); + res + } + + #[cold] + pub(crate) fn wait(&self) { + while !self.is_initialized() { + core::hint::spin_loop(); + } + } + + /// 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()); + self.value.get_unchecked() + } + + #[inline] + pub(crate) fn get_mut(&mut self) -> Option<&mut T> { + self.value.get_mut() + } + + #[inline] + pub(crate) fn into_inner(self) -> Option { + self.value.into_inner() + } +} + +struct Guard<'a> { + state: &'a AtomicU8, + new_state: u8, +} + +impl<'a> Drop for Guard<'a> { + fn drop(&mut self) { + self.state.store(self.new_state, Ordering::Release); + } +} + +#[inline(never)] +fn initialize_inner(state: &AtomicU8, init: &mut dyn FnMut() -> bool) { + loop { + match state.compare_exchange_weak(INCOMPLETE, RUNNING, Ordering::Acquire, Ordering::Acquire) + { + Ok(_) => { + let mut guard = Guard { state, new_state: INCOMPLETE }; + if init() { + guard.new_state = COMPLETE; + } + return; + } + Err(COMPLETE) => return, + Err(RUNNING) | Err(INCOMPLETE) => core::hint::spin_loop(), + Err(_) => { + debug_assert!(false); + unsafe { core::hint::unreachable_unchecked() } + } + } + } +} diff --git a/src/imp_pl.rs b/src/imp_pl.rs index d80ca5e..b723587 100644 --- a/src/imp_pl.rs +++ b/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 OnceCell { /// 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 = &*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. diff --git a/src/imp_std.rs b/src/imp_std.rs index 4d5b5fd..f828d45 100644 --- a/src/imp_std.rs +++ b/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 { // This `queue` field is the core of the implementation. It encodes two @@ -81,7 +78,7 @@ impl OnceCell { 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 OnceCell { /// 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 = &*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. diff --git a/src/lib.rs b/src/lib.rs index dcd64c9..4ab3b63 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -329,29 +329,38 @@ #![cfg_attr(not(feature = "std"), no_std)] +#[cfg(not(feature = "std"))] +use core as lib; +#[cfg(feature = "std")] +use std as lib; + +use lib::{ + cell::{Cell, UnsafeCell}, + fmt, mem, + ops::{Deref, DerefMut}, + panic::{RefUnwindSafe, UnwindSafe}, +}; + #[cfg(feature = "alloc")] extern crate alloc; -#[cfg(feature = "std")] +#[cfg(feature = "critical-section")] +#[path = "imp_cs.rs"] +mod imp; + +#[cfg(all(feature = "std", not(feature = "critical-section")))] #[cfg(feature = "parking_lot")] #[path = "imp_pl.rs"] mod imp; -#[cfg(feature = "std")] +#[cfg(all(feature = "std", not(feature = "critical-section")))] #[cfg(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, - ops::{Deref, DerefMut}, - }; - - #[cfg(feature = "std")] - use std::panic::{RefUnwindSafe, UnwindSafe}; + use super::*; /// A cell which can be written to only once. It is not thread safe. /// @@ -382,9 +391,7 @@ pub mod unsync { // `&unsync::OnceCell` to sneak a `T` through `catch_unwind`, // by initializing the cell in closure and extracting the value in the // `Drop`. - #[cfg(feature = "std")] impl RefUnwindSafe for OnceCell {} - #[cfg(feature = "std")] impl UnwindSafe for OnceCell {} impl Default for OnceCell { @@ -472,6 +479,17 @@ pub mod unsync { 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")] + 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 @@ -514,16 +532,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` @@ -680,7 +696,6 @@ pub mod unsync { init: Cell>, } - #[cfg(feature = "std")] impl RefUnwindSafe for Lazy where OnceCell: RefUnwindSafe {} impl fmt::Debug for Lazy { @@ -819,16 +834,11 @@ pub mod unsync { } /// Thread-safe, blocking version of `OnceCell`. -#[cfg(feature = "std")] +#[cfg(feature = "sync")] pub mod sync { - use std::{ - cell::Cell, - fmt, mem, - ops::{Deref, DerefMut}, - panic::RefUnwindSafe, - }; + use super::*; - use crate::{imp::OnceCell as Imp, take_unchecked}; + use imp::OnceCell as Imp; /// A thread-safe cell which can be written to only once. /// @@ -947,7 +957,7 @@ 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(); /// ``` pub fn wait(&self) -> &T { if !self.0.is_initialized() { @@ -1114,6 +1124,7 @@ pub mod sync { if let Some(value) = self.get() { return Ok(value); } + self.0.initialize(f)?; // Safe b/c value is initialized. @@ -1225,7 +1236,6 @@ pub mod sync { unsafe impl Sync for Lazy where OnceCell: Sync {} // auto-derived `Send` impl is OK. - #[cfg(feature = "std")] impl RefUnwindSafe for Lazy where OnceCell: RefUnwindSafe {} impl Lazy { @@ -1362,13 +1372,19 @@ pub mod sync { #[cfg(feature = "race")] pub mod race; -#[cfg(feature = "std")] +#[cfg(feature = "sync")] +#[inline] unsafe fn take_unchecked(val: &mut Option) -> T { - match val.take() { - Some(it) => it, + unwrap_unchecked(val.take()) +} + +#[inline] +unsafe fn unwrap_unchecked(val: Option) -> T { + match val { + Some(value) => value, None => { debug_assert!(false); - std::hint::unreachable_unchecked() + core::hint::unreachable_unchecked() } } } diff --git a/src/race.rs b/src/race.rs index e83e0b9..28acf62 100644 --- a/src/race.rs +++ b/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 }) } diff --git a/xtask/src/main.rs b/xtask/src/main.rs index 14da09a..076c782 100644 --- a/xtask/src/main.rs +++ b/xtask/src/main.rs @@ -39,6 +39,13 @@ fn try_main() -> Result<()> { cmd!("cargo test --features 'unstable alloc' --no-default-features --test it").run()?; cmd!("cargo test --features 'unstable std parking_lot alloc' --no-default-features") .run()?; + + cmd!("cargo test --features 'unstable std critical-section' --no-default-features") + .run()?; + cmd!( + "cargo test --features 'unstable std critical-section' --no-default-features --release" + ) + .run()?; } {