From a0f4a4b32819aa413dba3809981714213a4f6cf7 Mon Sep 17 00:00:00 2001 From: Chayim Refael Friedman Date: Wed, 21 Aug 2024 00:10:13 +0300 Subject: [PATCH] Implement ACP 429: add `Lazy{Cell,Lock}::get[_mut]` and `force_mut` In the implementation of `force_mut`, I chose performance over safety. For `LazyLock` this isn't really a choice; the code has to be unsafe. But for `LazyCell`, we can have a full-safe implementation, but it will be a bit less performant, so I went with the unsafe approach. --- core/src/cell/lazy.rs | 137 ++++++++++++++++++++++++++-- core/tests/lazy.rs | 21 +++++ std/src/sync/lazy_lock.rs | 130 ++++++++++++++++++++++++-- std/src/sync/lazy_lock/tests.rs | 21 +++++ std/src/sync/once.rs | 10 ++ std/src/sys/sync/once/futex.rs | 9 ++ std/src/sys/sync/once/no_threads.rs | 9 ++ std/src/sys/sync/once/queue.rs | 9 ++ 8 files changed, 333 insertions(+), 13 deletions(-) diff --git a/core/src/cell/lazy.rs b/core/src/cell/lazy.rs index 6ec1d2a33bef..6407aba7c811 100644 --- a/core/src/cell/lazy.rs +++ b/core/src/cell/lazy.rs @@ -1,5 +1,6 @@ use super::UnsafeCell; -use crate::ops::Deref; +use crate::hint::unreachable_unchecked; +use crate::ops::{Deref, DerefMut}; use crate::{fmt, mem}; enum State { @@ -82,7 +83,7 @@ impl T> LazyCell { match this.state.into_inner() { State::Init(data) => Ok(data), State::Uninit(f) => Err(f), - State::Poisoned => panic!("LazyCell instance has previously been poisoned"), + State::Poisoned => panic_poisoned(), } } @@ -114,7 +115,75 @@ impl T> LazyCell { State::Init(data) => data, // SAFETY: The state is uninitialized. State::Uninit(_) => unsafe { LazyCell::really_init(this) }, - State::Poisoned => panic!("LazyCell has previously been poisoned"), + State::Poisoned => panic_poisoned(), + } + } + + /// Forces the evaluation of this lazy value and returns a mutable reference to + /// the result. + /// + /// This is equivalent to the `DerefMut` impl, but is explicit. + /// + /// # Examples + /// + /// ``` + /// use std::cell::LazyCell; + /// + /// let mut lazy = LazyCell::new(|| 92); + /// + /// let p = LazyCell::force_mut(&mut lazy); + /// assert_eq!(*p, 92); + /// *p = 44; + /// assert_eq!(*lazy, 44); + /// *lazy = 55; // Using `DerefMut` + /// assert_eq!(*lazy, 55); + /// ``` + #[inline] + #[stable(feature = "lazy_deref_mut", since = "CURRENT_RUSTC_VERSION")] + pub fn force_mut(this: &mut LazyCell) -> &mut T { + #[cold] + /// # Safety + /// May only be called when the state is `Uninit`. + unsafe fn really_init T>(state: &mut State) -> &mut T { + // INVARIANT: Always valid, but the value may not be dropped. + struct PoisonOnPanic(*mut State); + impl Drop for PoisonOnPanic { + #[inline] + fn drop(&mut self) { + // SAFETY: Invariant states it is valid, and we don't drop the old value. + unsafe { + self.0.write(State::Poisoned); + } + } + } + + let State::Uninit(f) = state else { + // `unreachable!()` here won't optimize out because the function is cold. + // SAFETY: Precondition. + unsafe { unreachable_unchecked() }; + }; + // SAFETY: We never drop the state after we read `f`, and we write a valid value back + // in any case, panic or success. `f` can't access the `LazyCell` because it is mutably + // borrowed. + let f = unsafe { core::ptr::read(f) }; + // INVARIANT: Initiated from mutable reference, don't drop because we read it. + let guard = PoisonOnPanic(state); + let data = f(); + // SAFETY: `PoisonOnPanic` invariant, and we don't drop the old value. + unsafe { + core::ptr::write(guard.0, State::Init(data)); + } + core::mem::forget(guard); + let State::Init(data) = state else { unreachable!() }; + data + } + + let state = this.state.get_mut(); + match state { + State::Init(data) => data, + // SAFETY: `state` is `Uninit`. + State::Uninit(_) => unsafe { really_init(state) }, + State::Poisoned => panic_poisoned(), } } @@ -152,13 +221,55 @@ impl T> LazyCell { } impl LazyCell { + /// Returns a reference to the value if initialized, or `None` if not. + /// + /// # Examples + /// + /// ``` + /// #![feature(lazy_get)] + /// + /// use std::cell::LazyCell; + /// + /// let mut lazy = LazyCell::new(|| 92); + /// + /// assert_eq!(LazyCell::get_mut(&mut lazy), None); + /// let _ = LazyCell::force(&lazy); + /// *LazyCell::get_mut(&mut lazy).unwrap() = 44; + /// assert_eq!(*lazy, 44); + /// ``` + #[inline] + #[unstable(feature = "lazy_get", issue = "129333")] + pub fn get_mut(this: &mut LazyCell) -> Option<&mut T> { + let state = this.state.get_mut(); + match state { + State::Init(data) => Some(data), + _ => None, + } + } + + /// Returns a mutable reference to the value if initialized, or `None` if not. + /// + /// # Examples + /// + /// ``` + /// #![feature(lazy_get)] + /// + /// use std::cell::LazyCell; + /// + /// let lazy = LazyCell::new(|| 92); + /// + /// assert_eq!(LazyCell::get(&lazy), None); + /// let _ = LazyCell::force(&lazy); + /// assert_eq!(LazyCell::get(&lazy), Some(&92)); + /// ``` #[inline] - fn get(&self) -> Option<&T> { + #[unstable(feature = "lazy_get", issue = "129333")] + pub fn get(this: &LazyCell) -> 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() }; + let state = unsafe { &*this.state.get() }; match state { State::Init(data) => Some(data), _ => None, @@ -175,6 +286,14 @@ impl T> Deref for LazyCell { } } +#[stable(feature = "lazy_deref_mut", since = "CURRENT_RUSTC_VERSION")] +impl T> DerefMut for LazyCell { + #[inline] + fn deref_mut(&mut self) -> &mut T { + LazyCell::force_mut(self) + } +} + #[stable(feature = "lazy_cell", since = "1.80.0")] impl Default for LazyCell { /// Creates a new lazy value using `Default` as the initializing function. @@ -188,10 +307,16 @@ impl Default for LazyCell { impl fmt::Debug for LazyCell { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { let mut d = f.debug_tuple("LazyCell"); - match self.get() { + match LazyCell::get(self) { Some(data) => d.field(data), None => d.field(&format_args!("")), }; d.finish() } } + +#[cold] +#[inline(never)] +fn panic_poisoned() -> ! { + panic!("LazyCell instance has previously been poisoned") +} diff --git a/core/tests/lazy.rs b/core/tests/lazy.rs index a3b89f15b3f8..711511eaf4a5 100644 --- a/core/tests/lazy.rs +++ b/core/tests/lazy.rs @@ -113,6 +113,27 @@ fn lazy_type_inference() { let _ = LazyCell::new(|| ()); } +#[test] +#[should_panic = "LazyCell instance has previously been poisoned"] +fn lazy_force_mut_panic() { + let mut lazy = LazyCell::::new(|| panic!()); + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + let _ = LazyCell::force_mut(&mut lazy); + })) + .unwrap_err(); + let _ = &*lazy; +} + +#[test] +fn lazy_force_mut() { + let s = "abc".to_owned(); + let mut lazy = LazyCell::new(move || s); + LazyCell::force_mut(&mut lazy); + let p = LazyCell::force_mut(&mut lazy); + p.clear(); + LazyCell::force_mut(&mut lazy); +} + #[test] fn aliasing_in_get() { let x = OnceCell::new(); diff --git a/std/src/sync/lazy_lock.rs b/std/src/sync/lazy_lock.rs index 953aef40e7b7..afdfec43afd3 100644 --- a/std/src/sync/lazy_lock.rs +++ b/std/src/sync/lazy_lock.rs @@ -1,7 +1,7 @@ use super::once::ExclusiveState; use crate::cell::UnsafeCell; use crate::mem::ManuallyDrop; -use crate::ops::Deref; +use crate::ops::{Deref, DerefMut}; use crate::panic::{RefUnwindSafe, UnwindSafe}; use crate::sync::Once; use crate::{fmt, ptr}; @@ -121,7 +121,7 @@ impl T> LazyLock { pub fn into_inner(mut this: Self) -> Result { let state = this.once.state(); match state { - ExclusiveState::Poisoned => panic!("LazyLock instance has previously been poisoned"), + ExclusiveState::Poisoned => panic_poisoned(), state => { let this = ManuallyDrop::new(this); let data = unsafe { ptr::read(&this.data) }.into_inner(); @@ -134,6 +134,63 @@ impl T> LazyLock { } } + /// Forces the evaluation of this lazy value and returns a mutable reference to + /// the result. + /// + /// This is equivalent to the `DerefMut` impl, but is explicit. + /// + /// # Examples + /// + /// ``` + /// use std::sync::LazyLock; + /// + /// let mut lazy = LazyLock::new(|| 92); + /// + /// let p = LazyLock::force_mut(&mut lazy); + /// assert_eq!(*p, 92); + /// *p = 44; + /// assert_eq!(*lazy, 44); + /// *lazy = 55; // Using `DerefMut` + /// assert_eq!(*lazy, 55); + /// ``` + #[inline] + #[stable(feature = "lazy_deref_mut", since = "CURRENT_RUSTC_VERSION")] + pub fn force_mut(this: &mut LazyLock) -> &mut T { + #[cold] + /// # Safety + /// May only be called when the state is `Incomplete`. + unsafe fn really_init T>(this: &mut LazyLock) -> &mut T { + struct PoisonOnPanic<'a, T, F>(&'a mut LazyLock); + impl Drop for PoisonOnPanic<'_, T, F> { + #[inline] + fn drop(&mut self) { + self.0.once.set_state(ExclusiveState::Poisoned); + } + } + + // SAFETY: We always poison if the initializer panics (then we never check the data), + // or set the data on success. + let f = unsafe { ManuallyDrop::take(&mut this.data.get_mut().f) }; + // INVARIANT: Initiated from mutable reference, don't drop because we read it. + let guard = PoisonOnPanic(this); + let data = f(); + guard.0.data.get_mut().value = ManuallyDrop::new(data); + guard.0.once.set_state(ExclusiveState::Complete); + core::mem::forget(guard); + // SAFETY: We put the value there above. + unsafe { &mut this.data.get_mut().value } + } + + let state = this.once.state(); + match state { + ExclusiveState::Poisoned => panic_poisoned(), + // SAFETY: The `Once` states we completed the initialization. + ExclusiveState::Complete => unsafe { &mut this.data.get_mut().value }, + // SAFETY: The state is `Incomplete`. + ExclusiveState::Incomplete => unsafe { really_init(this) }, + } + } + /// Forces the evaluation of this lazy value and returns a reference to /// result. This is equivalent to the `Deref` impl, but is explicit. /// @@ -174,13 +231,58 @@ impl T> LazyLock { } impl LazyLock { - /// Gets the inner value if it has already been initialized. - fn get(&self) -> Option<&T> { - if self.once.is_completed() { + /// Returns a reference to the value if initialized, or `None` if not. + /// + /// # Examples + /// + /// ``` + /// #![feature(lazy_get)] + /// + /// use std::sync::LazyLock; + /// + /// let mut lazy = LazyLock::new(|| 92); + /// + /// assert_eq!(LazyLock::get_mut(&mut lazy), None); + /// let _ = LazyLock::force(&lazy); + /// *LazyLock::get_mut(&mut lazy).unwrap() = 44; + /// assert_eq!(*lazy, 44); + /// ``` + #[inline] + #[unstable(feature = "lazy_get", issue = "129333")] + pub fn get_mut(this: &mut LazyLock) -> Option<&mut T> { + // `state()` does not perform an atomic load, so prefer it over `is_complete()`. + let state = this.once.state(); + match state { + // SAFETY: + // The closure has been run successfully, so `value` has been initialized. + ExclusiveState::Complete => Some(unsafe { &mut this.data.get_mut().value }), + _ => None, + } + } + + /// Returns a mutable reference to the value if initialized, or `None` if not. + /// + /// # Examples + /// + /// ``` + /// #![feature(lazy_get)] + /// + /// use std::sync::LazyLock; + /// + /// let lazy = LazyLock::new(|| 92); + /// + /// assert_eq!(LazyLock::get(&lazy), None); + /// let _ = LazyLock::force(&lazy); + /// assert_eq!(LazyLock::get(&lazy), Some(&92)); + /// ``` + #[inline] + #[unstable(feature = "lazy_get", issue = "129333")] + pub fn get(this: &LazyLock) -> Option<&T> { + if this.once.is_completed() { // SAFETY: // The closure has been run successfully, so `value` has been initialized // and will not be modified again. - Some(unsafe { &*(*self.data.get()).value }) + Some(unsafe { &(*this.data.get()).value }) } else { None } @@ -215,6 +317,14 @@ impl T> Deref for LazyLock { } } +#[stable(feature = "lazy_deref_mut", since = "CURRENT_RUSTC_VERSION")] +impl T> DerefMut for LazyLock { + #[inline] + fn deref_mut(&mut self) -> &mut T { + LazyLock::force_mut(self) + } +} + #[stable(feature = "lazy_cell", since = "1.80.0")] impl Default for LazyLock { /// Creates a new lazy value using `Default` as the initializing function. @@ -228,7 +338,7 @@ impl Default for LazyLock { impl fmt::Debug for LazyLock { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { let mut d = f.debug_tuple("LazyLock"); - match self.get() { + match LazyLock::get(self) { Some(v) => d.field(v), None => d.field(&format_args!("")), }; @@ -236,6 +346,12 @@ impl fmt::Debug for LazyLock { } } +#[cold] +#[inline(never)] +fn panic_poisoned() -> ! { + panic!("LazyLock instance has previously been poisoned") +} + // We never create a `&F` from a `&LazyLock` so it is fine // to not impl `Sync` for `F`. #[stable(feature = "lazy_cell", since = "1.80.0")] diff --git a/std/src/sync/lazy_lock/tests.rs b/std/src/sync/lazy_lock/tests.rs index 8a6ab4ac4fd9..94044368305b 100644 --- a/std/src/sync/lazy_lock/tests.rs +++ b/std/src/sync/lazy_lock/tests.rs @@ -142,3 +142,24 @@ fn is_sync_send() { fn assert_traits() {} assert_traits::>(); } + +#[test] +#[should_panic = "has previously been poisoned"] +fn lazy_force_mut_panic() { + let mut lazy = LazyLock::::new(|| panic!()); + crate::panic::catch_unwind(crate::panic::AssertUnwindSafe(|| { + let _ = LazyLock::force_mut(&mut lazy); + })) + .unwrap_err(); + let _ = &*lazy; +} + +#[test] +fn lazy_force_mut() { + let s = "abc".to_owned(); + let mut lazy = LazyLock::new(move || s); + LazyLock::force_mut(&mut lazy); + let p = LazyLock::force_mut(&mut lazy); + p.clear(); + LazyLock::force_mut(&mut lazy); +} diff --git a/std/src/sync/once.rs b/std/src/sync/once.rs index bf595fdea2d2..5a1cd7d0b8b1 100644 --- a/std/src/sync/once.rs +++ b/std/src/sync/once.rs @@ -314,6 +314,16 @@ impl Once { pub(crate) fn state(&mut self) -> ExclusiveState { self.inner.state() } + + /// Sets 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 set_state(&mut self, new_state: ExclusiveState) { + self.inner.set_state(new_state); + } } #[stable(feature = "std_debug", since = "1.16.0")] diff --git a/std/src/sys/sync/once/futex.rs b/std/src/sys/sync/once/futex.rs index 2c8a054282b0..25588a4217b6 100644 --- a/std/src/sys/sync/once/futex.rs +++ b/std/src/sys/sync/once/futex.rs @@ -91,6 +91,15 @@ impl Once { } } + #[inline] + pub(crate) fn set_state(&mut self, new_state: ExclusiveState) { + *self.state_and_queued.get_mut() = match new_state { + ExclusiveState::Incomplete => INCOMPLETE, + ExclusiveState::Poisoned => POISONED, + ExclusiveState::Complete => COMPLETE, + }; + } + #[cold] #[track_caller] pub fn wait(&self, ignore_poisoning: bool) { diff --git a/std/src/sys/sync/once/no_threads.rs b/std/src/sys/sync/once/no_threads.rs index 12c1d9f5a6c9..cdcffe790f54 100644 --- a/std/src/sys/sync/once/no_threads.rs +++ b/std/src/sys/sync/once/no_threads.rs @@ -55,6 +55,15 @@ impl Once { } } + #[inline] + pub(crate) fn set_state(&mut self, new_state: ExclusiveState) { + self.state.set(match new_state { + ExclusiveState::Incomplete => State::Incomplete, + ExclusiveState::Poisoned => State::Poisoned, + ExclusiveState::Complete => State::Complete, + }); + } + #[cold] #[track_caller] pub fn wait(&self, _ignore_poisoning: bool) { diff --git a/std/src/sys/sync/once/queue.rs b/std/src/sys/sync/once/queue.rs index 86f72c82008b..17abaf0bf267 100644 --- a/std/src/sys/sync/once/queue.rs +++ b/std/src/sys/sync/once/queue.rs @@ -140,6 +140,15 @@ impl Once { } } + #[inline] + pub(crate) fn set_state(&mut self, new_state: ExclusiveState) { + *self.state_and_queue.get_mut() = match new_state { + ExclusiveState::Incomplete => ptr::without_provenance_mut(INCOMPLETE), + ExclusiveState::Poisoned => ptr::without_provenance_mut(POISONED), + ExclusiveState::Complete => ptr::without_provenance_mut(COMPLETE), + }; + } + #[cold] #[track_caller] pub fn wait(&self, ignore_poisoning: bool) {