Skip to content

Commit

Permalink
Implement ACP 429: add Lazy{Cell,Lock}::get[_mut] and force_mut
Browse files Browse the repository at this point in the history 
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.
  • Loading branch information
@ChayimFriedman2 @workingjubilee
ChayimFriedman2 authored and workingjubilee committed Sep 17, 2024
1 parent 5601d14 commit d0a2ca4
Show file tree
Hide file tree
Showing 8 changed files with 333 additions and 13 deletions.
137 changes: 131 additions & 6 deletions library/core/src/cell/lazy.rs
View file
Original file line number Diff line number Diff line change
@@ -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<T, F> {
Expand Down Expand Up @@ -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(),
}
}

Expand Down Expand Up @@ -114,7 +115,75 @@ 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.
///
/// 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<T, F>) -> &mut T {
#[cold]
/// # Safety
/// May only be called when the state is `Uninit`.
unsafe fn really_init<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(state) },
State::Poisoned => panic_poisoned(),
}
}

Expand Down Expand Up @@ -152,13 +221,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]
#[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]
fn get(&self) -> Option<&T> {
#[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,
Expand All @@ -175,6 +286,14 @@ impl<T, F: FnOnce() -> T> Deref for LazyCell<T, F> {
}
}

#[stable(feature = "lazy_deref_mut", since = "CURRENT_RUSTC_VERSION")]
impl<T, F: FnOnce() -> T> DerefMut for LazyCell<T, F> {
#[inline]
fn deref_mut(&mut self) -> &mut T {
LazyCell::force_mut(self)
}
}

#[stable(feature = "lazy_cell", since = "1.80.0")]
impl<T: Default> Default for LazyCell<T> {
/// Creates a new lazy value using `Default` as the initializing function.
Expand All @@ -188,10 +307,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")
}
21 changes: 21 additions & 0 deletions library/core/tests/lazy.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down
130 changes: 123 additions & 7 deletions library/std/src/sync/lazy_lock.rs
View file
Original file line number Diff line number Diff line change
@@ -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};
Expand Down Expand Up @@ -121,7 +121,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();
Expand All @@ -134,6 +134,63 @@ impl<T, F: FnOnce() -> T> LazyLock<T, F> {
}
}

/// 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<T, F>) -> &mut T {
#[cold]
/// # Safety
/// May only be called when the state is `Incomplete`.
unsafe fn really_init<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(this) },
}
}

/// Forces the evaluation of this lazy value and returns a reference to
/// result. This is equivalent to the `Deref` impl, but is explicit.
///
Expand Down Expand Up @@ -174,13 +231,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
}
Expand Down Expand Up @@ -215,6 +317,14 @@ impl<T, F: FnOnce() -> T> Deref for LazyLock<T, F> {
}
}

#[stable(feature = "lazy_deref_mut", since = "CURRENT_RUSTC_VERSION")]
impl<T, F: FnOnce() -> T> DerefMut for LazyLock<T, F> {
#[inline]
fn deref_mut(&mut self) -> &mut T {
LazyLock::force_mut(self)
}
}

#[stable(feature = "lazy_cell", since = "1.80.0")]
impl<T: Default> Default for LazyLock<T> {
/// Creates a new lazy value using `Default` as the initializing function.
Expand All @@ -228,14 +338,20 @@ 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>")),
};
d.finish()
}
}

#[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")]
Expand Down
Loading

0 comments on commit d0a2ca4

Please sign in to comment.