diff --git a/library/core/src/cell/lazy.rs b/library/core/src/cell/lazy.rs index 64a6ce51b2eca..44adcfa1a9463 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,65 @@ 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"), - }) + // 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] + 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. + 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 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), + _ => None, + } } } @@ -105,6 +166,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() } }