From f015e6fe499f0d6dc2c3ea7422ac8efd1ddb3920 Mon Sep 17 00:00:00 2001
From: joboet <jonasboettiger@icloud.com>
Date: Wed, 22 Mar 2023 13:12:51 +0100
Subject: [PATCH 1/4] core: optimize `LazyCell` size
---
library/core/src/cell/lazy.rs | 70 +++++++++++++++++++++++++++++------
1 file changed, 59 insertions(+), 11 deletions(-)
diff --git a/library/core/src/cell/lazy.rs b/library/core/src/cell/lazy.rs
index 64a6ce51b2eca..13a4c297e8afe 100644
--- a/library/core/src/cell/lazy.rs
+++ b/library/core/src/cell/lazy.rs
@@ -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.
///
@@ -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> {
@@ -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
@@ -77,10 +83,47 @@ 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"),
- })
+ let state = unsafe { &*this.state.get() };
+ match state {
+ State::Init(data) => data,
+ 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]
+ unsafe fn really_init(this: &LazyCell<T, F>) -> &T {
+ let state = unsafe { &mut *this.state.get() };
+ // 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();
+
+ // If the closure accessed the cell, the mutable borrow will be
+ // invalidated, so create a new one here.
+ let state = unsafe { &mut *this.state.get() };
+ *state = State::Init(data);
+
+ // A reference obtained by downcasting from the mutable borrow
+ // would become stale if other references are created in `force`.
+ // Borrow the state directly 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> {
+ let state = unsafe { &*self.state.get() };
+ match state {
+ State::Init(data) => Some(data),
+ _ => None,
+ }
}
}
@@ -105,6 +148,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()
}
}
From c7f9739bada7c54b0d848cd029c4faa4665c9adc Mon Sep 17 00:00:00 2001
From: joboet <jonasboettiger@icloud.com>
Date: Tue, 28 Mar 2023 11:33:44 +0200
Subject: [PATCH 2/4] core: improve code documentation for `LazyCell`
---
library/core/src/cell/lazy.rs | 29 ++++++++++++++++++++++++-----
1 file changed, 24 insertions(+), 5 deletions(-)
diff --git a/library/core/src/cell/lazy.rs b/library/core/src/cell/lazy.rs
index 13a4c297e8afe..4039cc268120f 100644
--- a/library/core/src/cell/lazy.rs
+++ b/library/core/src/cell/lazy.rs
@@ -83,9 +83,15 @@ impl<T, F: FnOnce() -> T> LazyCell<T, F> {
#[inline]
#[unstable(feature = "lazy_cell", issue = "109736")]
pub fn force(this: &LazyCell<T, F>) -> &T {
+ // 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 invalidates in `really_init` (in the
+ // uninitialized case).
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"),
}
@@ -95,6 +101,10 @@ impl<T, F: FnOnce() -> T> LazyCell<T, F> {
/// May only be called when the state is `Uninit`.
#[cold]
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() };
// Temporarily mark the state as poisoned. This prevents reentrant
// accesses and correctly poisons the cell if the closure panicked.
@@ -102,14 +112,19 @@ impl<T, F: FnOnce() -> T> LazyCell<T, F> {
let data = f();
- // If the closure accessed the cell, the mutable borrow will be
- // invalidated, so create a new one here.
+ // 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 create a new
+ // one here.
let state = unsafe { &mut *this.state.get() };
*state = State::Init(data);
- // A reference obtained by downcasting from the mutable borrow
- // would become stale if other references are created in `force`.
- // Borrow the state directly instead.
+ // SAFETY:
+ // A reference obtained by downcasting from the mutable borrow would
+ // become stale the next time `force` is called (since there is a conflict
+ // between the mutable reference here and the shared reference there).
+ // Do a new shared borrow of the state instead.
let state = unsafe { &*this.state.get() };
let State::Init(data) = state else { unreachable!() };
data
@@ -119,6 +134,10 @@ impl<T, F: FnOnce() -> T> LazyCell<T, F> {
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),
From 97d49fcf60e0da214d8e08cb914821400d421caf Mon Sep 17 00:00:00 2001
From: joboet <jonasboettiger@icloud.com>
Date: Tue, 28 Mar 2023 11:39:09 +0200
Subject: [PATCH 3/4] core: use `pointer::write` to cleanup `LazyCell`
initialization
---
library/core/src/cell/lazy.rs | 15 +++++++--------
1 file changed, 7 insertions(+), 8 deletions(-)
diff --git a/library/core/src/cell/lazy.rs b/library/core/src/cell/lazy.rs
index 4039cc268120f..aff3388529a75 100644
--- a/library/core/src/cell/lazy.rs
+++ b/library/core/src/cell/lazy.rs
@@ -115,16 +115,15 @@ impl<T, F: FnOnce() -> T> LazyCell<T, 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 create a new
- // one here.
- let state = unsafe { &mut *this.state.get() };
- *state = State::Init(data);
+ // 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:
- // A reference obtained by downcasting from the mutable borrow would
- // become stale the next time `force` is called (since there is a conflict
- // between the mutable reference here and the shared reference there).
- // Do a new shared borrow of the state instead.
+ // 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
From af080bf04be2a82273c646d807d8411518089199 Mon Sep 17 00:00:00 2001
From: joboet <jonas.boettiger@icloud.com>
Date: Tue, 28 Mar 2023 13:55:24 +0200
Subject: [PATCH 4/4] fix typo and adjust comment
Co-authored-by: Ralf Jung <post@ralfj.de>
---
library/core/src/cell/lazy.rs | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/library/core/src/cell/lazy.rs b/library/core/src/cell/lazy.rs
index aff3388529a75..44adcfa1a9463 100644
--- a/library/core/src/cell/lazy.rs
+++ b/library/core/src/cell/lazy.rs
@@ -86,8 +86,8 @@ impl<T, F: FnOnce() -> T> LazyCell<T, F> {
// 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 invalidates in `really_init` (in the
- // uninitialized case).
+ // 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,