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 LazyCell size #109483

Merged
merged 4 commits into from
Apr 1, 2023
Merged
Changes from all commits
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
88 changes: 77 additions & 11 deletions library/core/src/cell/lazy.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,13 @@
use crate::cell::{Cell, OnceCell};
use crate::fmt;
use crate::ops::Deref;
use crate::{fmt, mem};

use super::UnsafeCell;

enum State<T, F> {
Uninit(F),
Init(T),
Poisoned,
}

/// A value which is initialized on the first access.
///
Expand Down Expand Up @@ -31,8 +38,7 @@ use crate::ops::Deref;
/// ```
#[unstable(feature = "lazy_cell", issue = "109736")]
pub struct LazyCell<T, F = fn() -> T> {
cell: OnceCell<T>,
init: Cell<Option<F>>,
state: UnsafeCell<State<T, F>>,
}

impl<T, F: FnOnce() -> T> LazyCell<T, F> {
Expand All @@ -53,8 +59,8 @@ impl<T, F: FnOnce() -> T> LazyCell<T, F> {
/// ```
#[inline]
#[unstable(feature = "lazy_cell", issue = "109736")]
pub const fn new(init: F) -> LazyCell<T, F> {
LazyCell { cell: OnceCell::new(), init: Cell::new(Some(init)) }
pub const fn new(f: F) -> LazyCell<T, F> {
LazyCell { state: UnsafeCell::new(State::Uninit(f)) }
}

/// Forces the evaluation of this lazy value and returns a reference to
Expand All @@ -77,10 +83,65 @@ impl<T, F: FnOnce() -> T> LazyCell<T, F> {
#[inline]
#[unstable(feature = "lazy_cell", issue = "109736")]
pub fn force(this: &LazyCell<T, F>) -> &T {
this.cell.get_or_init(|| match this.init.take() {
Some(f) => f(),
None => panic!("`Lazy` instance has previously been poisoned"),
})
// SAFETY:
// This invalidates any mutable references to the data. The resulting
// reference lives either until the end of the borrow of `this` (in the
// initialized case) or is invalidated in `really_init` (in the
// uninitialized case; `really_init` will create and return a fresh reference).
let state = unsafe { &*this.state.get() };
match state {
State::Init(data) => data,
// SAFETY: The state is uninitialized.
State::Uninit(_) => unsafe { LazyCell::really_init(this) },
State::Poisoned => panic!("LazyCell has previously been poisoned"),
}
}

/// # Safety
/// May only be called when the state is `Uninit`.
#[cold]
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am not entirely sure if a #[cold] marker is correct here. This function is virtually guaranteed to be called (just... once).

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since OnceCell marks the closure as #[cold], I didn't want to change it just yet. But you are probably right that it is unnecessary here.

unsafe fn really_init(this: &LazyCell<T, F>) -> &T {
// SAFETY:
// This function is only called when the state is uninitialized,
// so no references to `state` can exist except for the reference
// in `force`, which is invalidated here and not accessed again.
let state = unsafe { &mut *this.state.get() };
joboet marked this conversation as resolved.
Show resolved Hide resolved
// Temporarily mark the state as poisoned. This prevents reentrant
// accesses and correctly poisons the cell if the closure panicked.
let State::Uninit(f) = mem::replace(state, State::Poisoned) else { unreachable!() };

let data = f();

// SAFETY:
// If the closure accessed the cell through something like a reentrant
// mutex, but caught the panic resulting from the state being poisoned,
// the mutable borrow for `state` will be invalidated, so we need to
// go through the `UnsafeCell` pointer here. The state can only be
// poisoned at this point, so using `write` to skip the destructor
// of `State` should help the optimizer.
unsafe { this.state.get().write(State::Init(data)) };

// SAFETY:
// The previous references were invalidated by the `write` call above,
// so do a new shared borrow of the state instead.
let state = unsafe { &*this.state.get() };
let State::Init(data) = state else { unreachable!() };
data
}
}

impl<T, F> LazyCell<T, F> {
#[inline]
fn get(&self) -> Option<&T> {
// SAFETY:
// This is sound for the same reason as in `force`: once the state is
// initialized, it will not be mutably accessed again, so this reference
// will stay valid for the duration of the borrow to `self`.
let state = unsafe { &*self.state.get() };
match state {
State::Init(data) => Some(data),
_ => None,
}
}
}

Expand All @@ -105,6 +166,11 @@ impl<T: Default> Default for LazyCell<T> {
#[unstable(feature = "lazy_cell", issue = "109736")]
impl<T: fmt::Debug, F> fmt::Debug for LazyCell<T, F> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Lazy").field("cell", &self.cell).field("init", &"..").finish()
let mut d = f.debug_tuple("LazyCell");
match self.get() {
Some(data) => d.field(data),
None => d.field(&format_args!("<uninit>")),
};
d.finish()
}
}