nordic-dev.net/rust
nordic-dev.net/rust
2
0
Fork 0
mirror of https://github.com/rust-lang/rust.git synced 2024-10-30 05:51:58 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
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.
This commit is contained in:
Chayim Refael Friedman 2024-08-21 00:10:13 +03:00 committed by Jubilee Young
parent 5601d14249
commit d0a2ca4867
8 changed files with 333 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

137
library/core/src/cell/lazy.rs
View File

@ -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> {
@ -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,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(),
}
}
@ -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]
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,
@ -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.
@ -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
library/core/tests/lazy.rs
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();

130
library/std/src/sync/lazy_lock.rs
View File

@ -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, 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();
@ -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.
///
@ -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
}
@ -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.
@ -228,7 +338,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>")),
};
@ -236,6 +346,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")]

21
library/std/src/sync/lazy_lock/tests.rs
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);
}

10
library/std/src/sync/once.rs
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")]

9
library/std/src/sys/sync/once/futex.rs
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) {

9
library/std/src/sys/sync/once/no_threads.rs
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) {

9
library/std/src/sys/sync/once/queue.rs
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) {