From f015e6fe499f0d6dc2c3ea7422ac8efd1ddb3920 Mon Sep 17 00:00:00 2001 From: joboet 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 { + 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> { - cell: OnceCell, - init: Cell>, + state: UnsafeCell>, } impl T> LazyCell { @@ -53,8 +59,8 @@ impl T> LazyCell { /// ``` #[inline] #[unstable(feature = "lazy_cell", issue = "109736")] - pub const fn new(init: F) -> LazyCell { - LazyCell { cell: OnceCell::new(), init: Cell::new(Some(init)) } + pub const fn new(f: F) -> LazyCell { + 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> LazyCell { #[inline] #[unstable(feature = "lazy_cell", issue = "109736")] pub fn force(this: &LazyCell) -> &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 { + 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 LazyCell { + #[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 Default for LazyCell { #[unstable(feature = "lazy_cell", issue = "109736")] impl fmt::Debug for LazyCell { 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!("")), + }; + d.finish() } } From c7f9739bada7c54b0d848cd029c4faa4665c9adc Mon Sep 17 00:00:00 2001 From: joboet 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> LazyCell { #[inline] #[unstable(feature = "lazy_cell", issue = "109736")] pub fn force(this: &LazyCell) -> &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> LazyCell { /// May only be called when the state is `Uninit`. #[cold] unsafe fn really_init(this: &LazyCell) -> &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> LazyCell { 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> LazyCell { impl LazyCell { #[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 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> LazyCell { // 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 Date: Tue, 28 Mar 2023 13:55:24 +0200 Subject: [PATCH 4/4] fix typo and adjust comment Co-authored-by: Ralf Jung --- 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> LazyCell { // 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,