Rollup merge of #130476 - workingjubilee:more-lazy-methods-take-2, r=Amanieu

Implement ACP 429: add `Lazy{Cell,Lock}::get[_mut]` and `force_mut`

Tracking issue for `lazy_get`: https://github.com/rust-lang/rust/issues/129333
This commit is contained in:
Jubilee 2024-09-18 14:32:26 -07:00 committed by GitHub
commit 12b59e52bc
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
11 changed files with 312 additions and 11 deletions

View File

@ -1,4 +1,5 @@
use super::UnsafeCell;
use crate::hint::unreachable_unchecked;
use crate::ops::Deref;
use crate::{fmt, mem};
@ -82,7 +83,7 @@ impl<T, F: FnOnce() -> T> LazyCell<T, F> {
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,72 @@ impl<T, F: FnOnce() -> T> LazyCell<T, F> {
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.
///
/// # Examples
///
/// ```
/// #![feature(lazy_get)]
/// 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);
/// ```
#[inline]
#[unstable(feature = "lazy_get", issue = "129333")]
pub fn force_mut(this: &mut LazyCell<T, F>) -> &mut T {
#[cold]
/// # Safety
/// May only be called when the state is `Uninit`.
unsafe fn really_init_mut<T, F: FnOnce() -> T>(state: &mut State<T, F>) -> &mut T {
// INVARIANT: Always valid, but the value may not be dropped.
struct PoisonOnPanic<T, F>(*mut State<T, F>);
impl<T, F> Drop for PoisonOnPanic<T, F> {
#[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_mut(state) },
State::Poisoned => panic_poisoned(),
}
}
@ -152,13 +218,55 @@ impl<T, F: FnOnce() -> T> LazyCell<T, F> {
}
impl<T, F> LazyCell<T, F> {
/// 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]
fn get(&self) -> Option<&T> {
#[unstable(feature = "lazy_get", issue = "129333")]
pub fn get_mut(this: &mut LazyCell<T, F>) -> 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]
#[unstable(feature = "lazy_get", issue = "129333")]
pub fn get(this: &LazyCell<T, F>) -> 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,
@ -188,10 +296,16 @@ impl<T: Default> Default for LazyCell<T> {
impl<T: fmt::Debug, F> fmt::Debug for LazyCell<T, F> {
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!("<uninit>")),
};
d.finish()
}
}
#[cold]
#[inline(never)]
fn panic_poisoned() -> ! {
panic!("LazyCell instance has previously been poisoned")
}

View File

@ -170,6 +170,7 @@
#![feature(is_ascii_octdigit)]
#![feature(is_val_statically_known)]
#![feature(isqrt)]
#![feature(lazy_get)]
#![feature(link_cfg)]
#![feature(offset_of_enum)]
#![feature(panic_internals)]

View File

@ -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::<String>::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();

View File

@ -76,6 +76,7 @@
#![feature(iterator_try_collect)]
#![feature(iterator_try_reduce)]
#![feature(layout_for_ptr)]
#![feature(lazy_get)]
#![feature(maybe_uninit_fill)]
#![feature(maybe_uninit_uninit_array_transpose)]
#![feature(maybe_uninit_write_slice)]

View File

@ -337,6 +337,7 @@
#![feature(hasher_prefixfree_extras)]
#![feature(hashmap_internals)]
#![feature(ip)]
#![feature(lazy_get)]
#![feature(maybe_uninit_slice)]
#![feature(maybe_uninit_write_slice)]
#![feature(panic_can_unwind)]

View File

@ -119,7 +119,7 @@ impl<T, F: FnOnce() -> T> LazyLock<T, F> {
pub fn into_inner(mut this: Self) -> Result<T, F> {
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();
@ -132,6 +132,60 @@ impl<T, F: FnOnce() -> T> LazyLock<T, F> {
}
}
/// Forces the evaluation of this lazy value and returns a mutable reference to
/// the result.
///
/// # Examples
///
/// ```
/// #![feature(lazy_get)]
/// 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);
/// ```
#[inline]
#[unstable(feature = "lazy_get", issue = "129333")]
pub fn force_mut(this: &mut LazyLock<T, F>) -> &mut T {
#[cold]
/// # Safety
/// May only be called when the state is `Incomplete`.
unsafe fn really_init_mut<T, F: FnOnce() -> T>(this: &mut LazyLock<T, F>) -> &mut T {
struct PoisonOnPanic<'a, T, F>(&'a mut LazyLock<T, F>);
impl<T, F> 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_mut(this) },
}
}
/// Forces the evaluation of this lazy value and returns a reference to
/// result. This is equivalent to the `Deref` impl, but is explicit.
///
@ -172,13 +226,58 @@ impl<T, F: FnOnce() -> T> LazyLock<T, F> {
}
impl<T, F> LazyLock<T, F> {
/// 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<T, F>) -> 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<T, F>) -> 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
}
@ -226,7 +325,7 @@ impl<T: Default> Default for LazyLock<T> {
impl<T: fmt::Debug, F> fmt::Debug for LazyLock<T, F> {
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!("<uninit>")),
};
@ -234,6 +333,12 @@ impl<T: fmt::Debug, F> fmt::Debug for LazyLock<T, F> {
}
}
#[cold]
#[inline(never)]
fn panic_poisoned() -> ! {
panic!("LazyLock instance has previously been poisoned")
}
// We never create a `&F` from a `&LazyLock<T, F>` so it is fine
// to not impl `Sync` for `F`.
#[stable(feature = "lazy_cell", since = "1.80.0")]

View File

@ -142,3 +142,24 @@ fn is_sync_send() {
fn assert_traits<T: Send + Sync>() {}
assert_traits::<LazyLock<String>>();
}
#[test]
#[should_panic = "has previously been poisoned"]
fn lazy_force_mut_panic() {
let mut lazy = LazyLock::<String>::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);
}

View File

@ -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")]

View File

@ -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) {

View File

@ -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) {

View File

@ -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) {