Skip to content

Commit

Permalink
Auto merge of #107329 - joboet:optimize_lazylock, r=m-ou-se
Browse files Browse the repository at this point in the history
Optimize `LazyLock` size

The initialization function was unnecessarily stored separately from the data to be initialized. Since both cannot exist at the same time, a `union` can be used, with the `Once` acting as discriminant. This unfortunately requires some extra methods on `Once` so that `Drop` can be implemented correctly and efficiently.

`@rustbot` label +T-libs +A-atomic
  • Loading branch information
bors committed Feb 18, 2023
2 parents 6d819a4 + 6520488 commit 3701bdc
Show file tree
Hide file tree
Showing 6 changed files with 120 additions and 17 deletions.
86 changes: 70 additions & 16 deletions library/std/src/sync/lazy_lock.rs
Original file line number Diff line number Diff line change
@@ -1,8 +1,21 @@
use crate::cell::Cell;
use crate::cell::UnsafeCell;
use crate::fmt;
use crate::mem::ManuallyDrop;
use crate::ops::Deref;
use crate::panic::{RefUnwindSafe, UnwindSafe};
use crate::sync::OnceLock;
use crate::sync::Once;

use super::once::ExclusiveState;

// We use the state of a Once as discriminant value. Upon creation, the state is
// "incomplete" and `f` contains the initialization closure. In the first call to
// `call_once`, `f` is taken and run. If it succeeds, `value` is set and the state
// is changed to "complete". If it panics, the Once is poisoned, so none of the
// two fields is initialized.
union Data<T, F> {
value: ManuallyDrop<T>,
f: ManuallyDrop<F>,
}

/// A value which is initialized on the first access.
///
Expand Down Expand Up @@ -43,16 +56,17 @@ use crate::sync::OnceLock;
/// ```
#[unstable(feature = "once_cell", issue = "74465")]
pub struct LazyLock<T, F = fn() -> T> {
cell: OnceLock<T>,
init: Cell<Option<F>>,
once: Once,
data: UnsafeCell<Data<T, F>>,
}

impl<T, F: FnOnce() -> T> LazyLock<T, F> {
/// Creates a new lazy value with the given initializing
/// function.
#[inline]
#[unstable(feature = "once_cell", issue = "74465")]
pub const fn new(f: F) -> LazyLock<T, F> {
LazyLock { cell: OnceLock::new(), init: Cell::new(Some(f)) }
LazyLock { once: Once::new(), data: UnsafeCell::new(Data { f: ManuallyDrop::new(f) }) }
}

/// Forces the evaluation of this lazy value and
Expand All @@ -74,10 +88,50 @@ impl<T, F: FnOnce() -> T> LazyLock<T, F> {
#[inline]
#[unstable(feature = "once_cell", issue = "74465")]
pub fn force(this: &LazyLock<T, F>) -> &T {
this.cell.get_or_init(|| match this.init.take() {
Some(f) => f(),
None => panic!("Lazy instance has previously been poisoned"),
})
this.once.call_once(|| {
// SAFETY: `call_once` only runs this closure once, ever.
let data = unsafe { &mut *this.data.get() };
let f = unsafe { ManuallyDrop::take(&mut data.f) };
let value = f();
data.value = ManuallyDrop::new(value);
});

// SAFETY:
// There are four possible scenarios:
// * the closure was called and initialized `value`.
// * the closure was called and panicked, so this point is never reached.
// * the closure was not called, but a previous call initialized `value`.
// * the closure was not called because the Once is poisoned, so this point
// is never reached.
// So `value` has definitely been initialized and will not be modified again.
unsafe { &*(*this.data.get()).value }
}
}

impl<T, F> LazyLock<T, F> {
/// Get the inner value if it has already been initialized.
fn get(&self) -> Option<&T> {
if self.once.is_completed() {
// SAFETY:
// The closure has been run successfully, so `value` has been initialized
// and will not be modified again.
Some(unsafe { &*(*self.data.get()).value })
} else {
None
}
}
}

#[unstable(feature = "once_cell", issue = "74465")]
impl<T, F> Drop for LazyLock<T, F> {
fn drop(&mut self) {
match self.once.state() {
ExclusiveState::Incomplete => unsafe { ManuallyDrop::drop(&mut self.data.get_mut().f) },
ExclusiveState::Complete => unsafe {
ManuallyDrop::drop(&mut self.data.get_mut().value)
},
ExclusiveState::Poisoned => {}
}
}
}

Expand All @@ -103,23 +157,23 @@ impl<T: Default> Default for LazyLock<T> {
#[unstable(feature = "once_cell", issue = "74465")]
impl<T: fmt::Debug, F> fmt::Debug for LazyLock<T, F> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Lazy").field("cell", &self.cell).finish_non_exhaustive()
match self.get() {
Some(v) => f.debug_tuple("LazyLock").field(v).finish(),
None => f.write_str("LazyLock(Uninit)"),
}
}
}

// We never create a `&F` from a `&LazyLock<T, F>` so it is fine
// to not impl `Sync` for `F`
// we do create a `&mut Option<F>` in `force`, but this is
// properly synchronized, so it only happens once
// so it also does not contribute to this impl.
#[unstable(feature = "once_cell", issue = "74465")]
unsafe impl<T, F: Send> Sync for LazyLock<T, F> where OnceLock<T>: Sync {}
unsafe impl<T: Sync + Send, F: Send> Sync for LazyLock<T, F> {}
// auto-derived `Send` impl is OK.

#[unstable(feature = "once_cell", issue = "74465")]
impl<T, F: UnwindSafe> RefUnwindSafe for LazyLock<T, F> where OnceLock<T>: RefUnwindSafe {}
impl<T: RefUnwindSafe + UnwindSafe, F: UnwindSafe> RefUnwindSafe for LazyLock<T, F> {}
#[unstable(feature = "once_cell", issue = "74465")]
impl<T, F: UnwindSafe> UnwindSafe for LazyLock<T, F> where OnceLock<T>: UnwindSafe {}
impl<T: UnwindSafe, F: UnwindSafe> UnwindSafe for LazyLock<T, F> {}

#[cfg(test)]
mod tests;
2 changes: 1 addition & 1 deletion library/std/src/sync/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -186,7 +186,7 @@ mod condvar;
mod lazy_lock;
mod mpmc;
mod mutex;
mod once;
pub(crate) mod once;
mod once_lock;
mod poison;
mod remutex;
Expand Down
16 changes: 16 additions & 0 deletions library/std/src/sync/once.rs
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,12 @@ pub struct OnceState {
pub(crate) inner: sys::OnceState,
}

pub(crate) enum ExclusiveState {
Incomplete,
Poisoned,
Complete,
}

/// Initialization value for static [`Once`] values.
///
/// # Examples
Expand Down Expand Up @@ -248,6 +254,16 @@ impl Once {
pub fn is_completed(&self) -> bool {
self.inner.is_completed()
}

/// Returns the current state of the `Once` instance.
///
/// Since this takes a mutable reference, no initialization can currently
/// be running, so the state must be either "incomplete", "poisoned" or
/// "complete".
#[inline]
pub(crate) fn state(&mut self) -> ExclusiveState {
self.inner.state()
}
}

#[stable(feature = "std_debug", since = "1.16.0")]
Expand Down
11 changes: 11 additions & 0 deletions library/std/src/sys/unsupported/once.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
use crate::cell::Cell;
use crate::sync as public;
use crate::sync::once::ExclusiveState;

pub struct Once {
state: Cell<State>,
Expand Down Expand Up @@ -44,6 +45,16 @@ impl Once {
self.state.get() == State::Complete
}

#[inline]
pub(crate) fn state(&mut self) -> ExclusiveState {
match self.state.get() {
State::Incomplete => ExclusiveState::Incomplete,
State::Poisoned => ExclusiveState::Poisoned,
State::Complete => ExclusiveState::Complete,
_ => unreachable!("invalid Once state"),
}
}

#[cold]
#[track_caller]
pub fn call(&self, ignore_poisoning: bool, f: &mut impl FnMut(&public::OnceState)) {
Expand Down
11 changes: 11 additions & 0 deletions library/std/src/sys_common/once/futex.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ use crate::sync::atomic::{
AtomicU32,
Ordering::{Acquire, Relaxed, Release},
};
use crate::sync::once::ExclusiveState;
use crate::sys::futex::{futex_wait, futex_wake_all};

// On some platforms, the OS is very nice and handles the waiter queue for us.
Expand Down Expand Up @@ -78,6 +79,16 @@ impl Once {
self.state.load(Acquire) == COMPLETE
}

#[inline]
pub(crate) fn state(&mut self) -> ExclusiveState {
match *self.state.get_mut() {
INCOMPLETE => ExclusiveState::Incomplete,
POISONED => ExclusiveState::Poisoned,
COMPLETE => ExclusiveState::Complete,
_ => unreachable!("invalid Once state"),
}
}

// This uses FnMut to match the API of the generic implementation. As this
// implementation is quite light-weight, it is generic over the closure and
// so avoids the cost of dynamic dispatch.
Expand Down
11 changes: 11 additions & 0 deletions library/std/src/sys_common/once/queue.rs
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,7 @@ use crate::fmt;
use crate::ptr;
use crate::sync as public;
use crate::sync::atomic::{AtomicBool, AtomicPtr, Ordering};
use crate::sync::once::ExclusiveState;
use crate::thread::{self, Thread};

type Masked = ();
Expand Down Expand Up @@ -121,6 +122,16 @@ impl Once {
self.state_and_queue.load(Ordering::Acquire).addr() == COMPLETE
}

#[inline]
pub(crate) fn state(&mut self) -> ExclusiveState {
match self.state_and_queue.get_mut().addr() {
INCOMPLETE => ExclusiveState::Incomplete,
POISONED => ExclusiveState::Poisoned,
COMPLETE => ExclusiveState::Complete,
_ => unreachable!("invalid Once state"),
}
}

// This is a non-generic function to reduce the monomorphization cost of
// using `call_once` (this isn't exactly a trivial or small implementation).
//
Expand Down

0 comments on commit 3701bdc

Please sign in to comment.