Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Optimize LazyLock size #107329

Merged
merged 3 commits into from
Feb 18, 2023
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
82 changes: 66 additions & 16 deletions library/std/src/sync/lazy_lock.rs
Original file line number Diff line number Diff line change
@@ -1,8 +1,22 @@
use crate::cell::Cell;
use crate::mem::ManuallyDrop;

joboet marked this conversation as resolved.
Show resolved Hide resolved
use crate::cell::UnsafeCell;
use crate::fmt;
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 +57,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 +89,43 @@ 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() { Some(unsafe { &*(*self.data.get()).value }) } else { None }
joboet marked this conversation as resolved.
Show resolved Hide resolved
}
}

#[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 +151,25 @@ 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()
f.debug_tuple("LazyLock")
.field(match self.get() {
Some(v) => v,
None => &"Uninit",
})
.finish()
joboet marked this conversation as resolved.
Show resolved Hide resolved
}
}

// 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> {}
m-ou-se marked this conversation as resolved.
Show resolved Hide resolved

#[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,
m-ou-se marked this conversation as resolved.
Show resolved Hide resolved
_ => 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