Skip to content
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.

Commit a0f4a4b

Browse files
ChayimFriedman2workingjubilee
ChayimFriedman2
authored and
workingjubilee
committedSep 17, 2024
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.
1 parent ebe99f3 commit a0f4a4b

File tree

8 files changed

+333
-13
lines changed

8 files changed

+333
-13
lines changed
 

‎core/src/cell/lazy.rs

+131-6
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
use super::UnsafeCell;
2-
use crate::ops::Deref;
2+
use crate::hint::unreachable_unchecked;
3+
use crate::ops::{Deref, DerefMut};
34
use crate::{fmt, mem};
45

56
enum State<T, F> {
@@ -82,7 +83,7 @@ impl<T, F: FnOnce() -> T> LazyCell<T, F> {
8283
match this.state.into_inner() {
8384
State::Init(data) => Ok(data),
8485
State::Uninit(f) => Err(f),
85-
State::Poisoned => panic!("LazyCell instance has previously been poisoned"),
86+
State::Poisoned => panic_poisoned(),
8687
}
8788
}
8889

@@ -114,7 +115,75 @@ impl<T, F: FnOnce() -> T> LazyCell<T, F> {
114115
State::Init(data) => data,
115116
// SAFETY: The state is uninitialized.
116117
State::Uninit(_) => unsafe { LazyCell::really_init(this) },
117-
State::Poisoned => panic!("LazyCell has previously been poisoned"),
118+
State::Poisoned => panic_poisoned(),
119+
}
120+
}
121+
122+
/// Forces the evaluation of this lazy value and returns a mutable reference to
123+
/// the result.
124+
///
125+
/// This is equivalent to the `DerefMut` impl, but is explicit.
126+
///
127+
/// # Examples
128+
///
129+
/// ```
130+
/// use std::cell::LazyCell;
131+
///
132+
/// let mut lazy = LazyCell::new(|| 92);
133+
///
134+
/// let p = LazyCell::force_mut(&mut lazy);
135+
/// assert_eq!(*p, 92);
136+
/// *p = 44;
137+
/// assert_eq!(*lazy, 44);
138+
/// *lazy = 55; // Using `DerefMut`
139+
/// assert_eq!(*lazy, 55);
140+
/// ```
141+
#[inline]
142+
#[stable(feature = "lazy_deref_mut", since = "CURRENT_RUSTC_VERSION")]
143+
pub fn force_mut(this: &mut LazyCell<T, F>) -> &mut T {
144+
#[cold]
145+
/// # Safety
146+
/// May only be called when the state is `Uninit`.
147+
unsafe fn really_init<T, F: FnOnce() -> T>(state: &mut State<T, F>) -> &mut T {
148+
// INVARIANT: Always valid, but the value may not be dropped.
149+
struct PoisonOnPanic<T, F>(*mut State<T, F>);
150+
impl<T, F> Drop for PoisonOnPanic<T, F> {
151+
#[inline]
152+
fn drop(&mut self) {
153+
// SAFETY: Invariant states it is valid, and we don't drop the old value.
154+
unsafe {
155+
self.0.write(State::Poisoned);
156+
}
157+
}
158+
}
159+
160+
let State::Uninit(f) = state else {
161+
// `unreachable!()` here won't optimize out because the function is cold.
162+
// SAFETY: Precondition.
163+
unsafe { unreachable_unchecked() };
164+
};
165+
// SAFETY: We never drop the state after we read `f`, and we write a valid value back
166+
// in any case, panic or success. `f` can't access the `LazyCell` because it is mutably
167+
// borrowed.
168+
let f = unsafe { core::ptr::read(f) };
169+
// INVARIANT: Initiated from mutable reference, don't drop because we read it.
170+
let guard = PoisonOnPanic(state);
171+
let data = f();
172+
// SAFETY: `PoisonOnPanic` invariant, and we don't drop the old value.
173+
unsafe {
174+
core::ptr::write(guard.0, State::Init(data));
175+
}
176+
core::mem::forget(guard);
177+
let State::Init(data) = state else { unreachable!() };
178+
data
179+
}
180+
181+
let state = this.state.get_mut();
182+
match state {
183+
State::Init(data) => data,
184+
// SAFETY: `state` is `Uninit`.
185+
State::Uninit(_) => unsafe { really_init(state) },
186+
State::Poisoned => panic_poisoned(),
118187
}
119188
}
120189

@@ -152,13 +221,55 @@ impl<T, F: FnOnce() -> T> LazyCell<T, F> {
152221
}
153222

154223
impl<T, F> LazyCell<T, F> {
224+
/// Returns a reference to the value if initialized, or `None` if not.
225+
///
226+
/// # Examples
227+
///
228+
/// ```
229+
/// #![feature(lazy_get)]
230+
///
231+
/// use std::cell::LazyCell;
232+
///
233+
/// let mut lazy = LazyCell::new(|| 92);
234+
///
235+
/// assert_eq!(LazyCell::get_mut(&mut lazy), None);
236+
/// let _ = LazyCell::force(&lazy);
237+
/// *LazyCell::get_mut(&mut lazy).unwrap() = 44;
238+
/// assert_eq!(*lazy, 44);
239+
/// ```
240+
#[inline]
241+
#[unstable(feature = "lazy_get", issue = "129333")]
242+
pub fn get_mut(this: &mut LazyCell<T, F>) -> Option<&mut T> {
243+
let state = this.state.get_mut();
244+
match state {
245+
State::Init(data) => Some(data),
246+
_ => None,
247+
}
248+
}
249+
250+
/// Returns a mutable reference to the value if initialized, or `None` if not.
251+
///
252+
/// # Examples
253+
///
254+
/// ```
255+
/// #![feature(lazy_get)]
256+
///
257+
/// use std::cell::LazyCell;
258+
///
259+
/// let lazy = LazyCell::new(|| 92);
260+
///
261+
/// assert_eq!(LazyCell::get(&lazy), None);
262+
/// let _ = LazyCell::force(&lazy);
263+
/// assert_eq!(LazyCell::get(&lazy), Some(&92));
264+
/// ```
155265
#[inline]
156-
fn get(&self) -> Option<&T> {
266+
#[unstable(feature = "lazy_get", issue = "129333")]
267+
pub fn get(this: &LazyCell<T, F>) -> Option<&T> {
157268
// SAFETY:
158269
// This is sound for the same reason as in `force`: once the state is
159270
// initialized, it will not be mutably accessed again, so this reference
160271
// will stay valid for the duration of the borrow to `self`.
161-
let state = unsafe { &*self.state.get() };
272+
let state = unsafe { &*this.state.get() };
162273
match state {
163274
State::Init(data) => Some(data),
164275
_ => None,
@@ -175,6 +286,14 @@ impl<T, F: FnOnce() -> T> Deref for LazyCell<T, F> {
175286
}
176287
}
177288

289+
#[stable(feature = "lazy_deref_mut", since = "CURRENT_RUSTC_VERSION")]
290+
impl<T, F: FnOnce() -> T> DerefMut for LazyCell<T, F> {
291+
#[inline]
292+
fn deref_mut(&mut self) -> &mut T {
293+
LazyCell::force_mut(self)
294+
}
295+
}
296+
178297
#[stable(feature = "lazy_cell", since = "1.80.0")]
179298
impl<T: Default> Default for LazyCell<T> {
180299
/// Creates a new lazy value using `Default` as the initializing function.
@@ -188,10 +307,16 @@ impl<T: Default> Default for LazyCell<T> {
188307
impl<T: fmt::Debug, F> fmt::Debug for LazyCell<T, F> {
189308
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
190309
let mut d = f.debug_tuple("LazyCell");
191-
match self.get() {
310+
match LazyCell::get(self) {
192311
Some(data) => d.field(data),
193312
None => d.field(&format_args!("<uninit>")),
194313
};
195314
d.finish()
196315
}
197316
}
317+
318+
#[cold]
319+
#[inline(never)]
320+
fn panic_poisoned() -> ! {
321+
panic!("LazyCell instance has previously been poisoned")
322+
}

‎core/tests/lazy.rs

+21
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,27 @@ fn lazy_type_inference() {
113113
let _ = LazyCell::new(|| ());
114114
}
115115

116+
#[test]
117+
#[should_panic = "LazyCell instance has previously been poisoned"]
118+
fn lazy_force_mut_panic() {
119+
let mut lazy = LazyCell::<String>::new(|| panic!());
120+
std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
121+
let _ = LazyCell::force_mut(&mut lazy);
122+
}))
123+
.unwrap_err();
124+
let _ = &*lazy;
125+
}
126+
127+
#[test]
128+
fn lazy_force_mut() {
129+
let s = "abc".to_owned();
130+
let mut lazy = LazyCell::new(move || s);
131+
LazyCell::force_mut(&mut lazy);
132+
let p = LazyCell::force_mut(&mut lazy);
133+
p.clear();
134+
LazyCell::force_mut(&mut lazy);
135+
}
136+
116137
#[test]
117138
fn aliasing_in_get() {
118139
let x = OnceCell::new();

‎std/src/sync/lazy_lock.rs

+123-7
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
use super::once::ExclusiveState;
22
use crate::cell::UnsafeCell;
33
use crate::mem::ManuallyDrop;
4-
use crate::ops::Deref;
4+
use crate::ops::{Deref, DerefMut};
55
use crate::panic::{RefUnwindSafe, UnwindSafe};
66
use crate::sync::Once;
77
use crate::{fmt, ptr};
@@ -121,7 +121,7 @@ impl<T, F: FnOnce() -> T> LazyLock<T, F> {
121121
pub fn into_inner(mut this: Self) -> Result<T, F> {
122122
let state = this.once.state();
123123
match state {
124-
ExclusiveState::Poisoned => panic!("LazyLock instance has previously been poisoned"),
124+
ExclusiveState::Poisoned => panic_poisoned(),
125125
state => {
126126
let this = ManuallyDrop::new(this);
127127
let data = unsafe { ptr::read(&this.data) }.into_inner();
@@ -134,6 +134,63 @@ impl<T, F: FnOnce() -> T> LazyLock<T, F> {
134134
}
135135
}
136136

137+
/// Forces the evaluation of this lazy value and returns a mutable reference to
138+
/// the result.
139+
///
140+
/// This is equivalent to the `DerefMut` impl, but is explicit.
141+
///
142+
/// # Examples
143+
///
144+
/// ```
145+
/// use std::sync::LazyLock;
146+
///
147+
/// let mut lazy = LazyLock::new(|| 92);
148+
///
149+
/// let p = LazyLock::force_mut(&mut lazy);
150+
/// assert_eq!(*p, 92);
151+
/// *p = 44;
152+
/// assert_eq!(*lazy, 44);
153+
/// *lazy = 55; // Using `DerefMut`
154+
/// assert_eq!(*lazy, 55);
155+
/// ```
156+
#[inline]
157+
#[stable(feature = "lazy_deref_mut", since = "CURRENT_RUSTC_VERSION")]
158+
pub fn force_mut(this: &mut LazyLock<T, F>) -> &mut T {
159+
#[cold]
160+
/// # Safety
161+
/// May only be called when the state is `Incomplete`.
162+
unsafe fn really_init<T, F: FnOnce() -> T>(this: &mut LazyLock<T, F>) -> &mut T {
163+
struct PoisonOnPanic<'a, T, F>(&'a mut LazyLock<T, F>);
164+
impl<T, F> Drop for PoisonOnPanic<'_, T, F> {
165+
#[inline]
166+
fn drop(&mut self) {
167+
self.0.once.set_state(ExclusiveState::Poisoned);
168+
}
169+
}
170+
171+
// SAFETY: We always poison if the initializer panics (then we never check the data),
172+
// or set the data on success.
173+
let f = unsafe { ManuallyDrop::take(&mut this.data.get_mut().f) };
174+
// INVARIANT: Initiated from mutable reference, don't drop because we read it.
175+
let guard = PoisonOnPanic(this);
176+
let data = f();
177+
guard.0.data.get_mut().value = ManuallyDrop::new(data);
178+
guard.0.once.set_state(ExclusiveState::Complete);
179+
core::mem::forget(guard);
180+
// SAFETY: We put the value there above.
181+
unsafe { &mut this.data.get_mut().value }
182+
}
183+
184+
let state = this.once.state();
185+
match state {
186+
ExclusiveState::Poisoned => panic_poisoned(),
187+
// SAFETY: The `Once` states we completed the initialization.
188+
ExclusiveState::Complete => unsafe { &mut this.data.get_mut().value },
189+
// SAFETY: The state is `Incomplete`.
190+
ExclusiveState::Incomplete => unsafe { really_init(this) },
191+
}
192+
}
193+
137194
/// Forces the evaluation of this lazy value and returns a reference to
138195
/// result. This is equivalent to the `Deref` impl, but is explicit.
139196
///
@@ -174,13 +231,58 @@ impl<T, F: FnOnce() -> T> LazyLock<T, F> {
174231
}
175232

176233
impl<T, F> LazyLock<T, F> {
177-
/// Gets the inner value if it has already been initialized.
178-
fn get(&self) -> Option<&T> {
179-
if self.once.is_completed() {
234+
/// Returns a reference to the value if initialized, or `None` if not.
235+
///
236+
/// # Examples
237+
///
238+
/// ```
239+
/// #![feature(lazy_get)]
240+
///
241+
/// use std::sync::LazyLock;
242+
///
243+
/// let mut lazy = LazyLock::new(|| 92);
244+
///
245+
/// assert_eq!(LazyLock::get_mut(&mut lazy), None);
246+
/// let _ = LazyLock::force(&lazy);
247+
/// *LazyLock::get_mut(&mut lazy).unwrap() = 44;
248+
/// assert_eq!(*lazy, 44);
249+
/// ```
250+
#[inline]
251+
#[unstable(feature = "lazy_get", issue = "129333")]
252+
pub fn get_mut(this: &mut LazyLock<T, F>) -> Option<&mut T> {
253+
// `state()` does not perform an atomic load, so prefer it over `is_complete()`.
254+
let state = this.once.state();
255+
match state {
256+
// SAFETY:
257+
// The closure has been run successfully, so `value` has been initialized.
258+
ExclusiveState::Complete => Some(unsafe { &mut this.data.get_mut().value }),
259+
_ => None,
260+
}
261+
}
262+
263+
/// Returns a mutable reference to the value if initialized, or `None` if not.
264+
///
265+
/// # Examples
266+
///
267+
/// ```
268+
/// #![feature(lazy_get)]
269+
///
270+
/// use std::sync::LazyLock;
271+
///
272+
/// let lazy = LazyLock::new(|| 92);
273+
///
274+
/// assert_eq!(LazyLock::get(&lazy), None);
275+
/// let _ = LazyLock::force(&lazy);
276+
/// assert_eq!(LazyLock::get(&lazy), Some(&92));
277+
/// ```
278+
#[inline]
279+
#[unstable(feature = "lazy_get", issue = "129333")]
280+
pub fn get(this: &LazyLock<T, F>) -> Option<&T> {
281+
if this.once.is_completed() {
180282
// SAFETY:
181283
// The closure has been run successfully, so `value` has been initialized
182284
// and will not be modified again.
183-
Some(unsafe { &*(*self.data.get()).value })
285+
Some(unsafe { &(*this.data.get()).value })
184286
} else {
185287
None
186288
}
@@ -215,6 +317,14 @@ impl<T, F: FnOnce() -> T> Deref for LazyLock<T, F> {
215317
}
216318
}
217319

320+
#[stable(feature = "lazy_deref_mut", since = "CURRENT_RUSTC_VERSION")]
321+
impl<T, F: FnOnce() -> T> DerefMut for LazyLock<T, F> {
322+
#[inline]
323+
fn deref_mut(&mut self) -> &mut T {
324+
LazyLock::force_mut(self)
325+
}
326+
}
327+
218328
#[stable(feature = "lazy_cell", since = "1.80.0")]
219329
impl<T: Default> Default for LazyLock<T> {
220330
/// Creates a new lazy value using `Default` as the initializing function.
@@ -228,14 +338,20 @@ impl<T: Default> Default for LazyLock<T> {
228338
impl<T: fmt::Debug, F> fmt::Debug for LazyLock<T, F> {
229339
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
230340
let mut d = f.debug_tuple("LazyLock");
231-
match self.get() {
341+
match LazyLock::get(self) {
232342
Some(v) => d.field(v),
233343
None => d.field(&format_args!("<uninit>")),
234344
};
235345
d.finish()
236346
}
237347
}
238348

349+
#[cold]
350+
#[inline(never)]
351+
fn panic_poisoned() -> ! {
352+
panic!("LazyLock instance has previously been poisoned")
353+
}
354+
239355
// We never create a `&F` from a `&LazyLock<T, F>` so it is fine
240356
// to not impl `Sync` for `F`.
241357
#[stable(feature = "lazy_cell", since = "1.80.0")]

‎std/src/sync/lazy_lock/tests.rs

+21
Original file line numberDiff line numberDiff line change
@@ -142,3 +142,24 @@ fn is_sync_send() {
142142
fn assert_traits<T: Send + Sync>() {}
143143
assert_traits::<LazyLock<String>>();
144144
}
145+
146+
#[test]
147+
#[should_panic = "has previously been poisoned"]
148+
fn lazy_force_mut_panic() {
149+
let mut lazy = LazyLock::<String>::new(|| panic!());
150+
crate::panic::catch_unwind(crate::panic::AssertUnwindSafe(|| {
151+
let _ = LazyLock::force_mut(&mut lazy);
152+
}))
153+
.unwrap_err();
154+
let _ = &*lazy;
155+
}
156+
157+
#[test]
158+
fn lazy_force_mut() {
159+
let s = "abc".to_owned();
160+
let mut lazy = LazyLock::new(move || s);
161+
LazyLock::force_mut(&mut lazy);
162+
let p = LazyLock::force_mut(&mut lazy);
163+
p.clear();
164+
LazyLock::force_mut(&mut lazy);
165+
}

‎std/src/sync/once.rs

+10
Original file line numberDiff line numberDiff line change
@@ -314,6 +314,16 @@ impl Once {
314314
pub(crate) fn state(&mut self) -> ExclusiveState {
315315
self.inner.state()
316316
}
317+
318+
/// Sets current state of the `Once` instance.
319+
///
320+
/// Since this takes a mutable reference, no initialization can currently
321+
/// be running, so the state must be either "incomplete", "poisoned" or
322+
/// "complete".
323+
#[inline]
324+
pub(crate) fn set_state(&mut self, new_state: ExclusiveState) {
325+
self.inner.set_state(new_state);
326+
}
317327
}
318328

319329
#[stable(feature = "std_debug", since = "1.16.0")]

‎std/src/sys/sync/once/futex.rs

+9
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,15 @@ impl Once {
9191
}
9292
}
9393

94+
#[inline]
95+
pub(crate) fn set_state(&mut self, new_state: ExclusiveState) {
96+
*self.state_and_queued.get_mut() = match new_state {
97+
ExclusiveState::Incomplete => INCOMPLETE,
98+
ExclusiveState::Poisoned => POISONED,
99+
ExclusiveState::Complete => COMPLETE,
100+
};
101+
}
102+
94103
#[cold]
95104
#[track_caller]
96105
pub fn wait(&self, ignore_poisoning: bool) {

‎std/src/sys/sync/once/no_threads.rs

+9
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,15 @@ impl Once {
5555
}
5656
}
5757

58+
#[inline]
59+
pub(crate) fn set_state(&mut self, new_state: ExclusiveState) {
60+
self.state.set(match new_state {
61+
ExclusiveState::Incomplete => State::Incomplete,
62+
ExclusiveState::Poisoned => State::Poisoned,
63+
ExclusiveState::Complete => State::Complete,
64+
});
65+
}
66+
5867
#[cold]
5968
#[track_caller]
6069
pub fn wait(&self, _ignore_poisoning: bool) {

‎std/src/sys/sync/once/queue.rs

+9
Original file line numberDiff line numberDiff line change
@@ -140,6 +140,15 @@ impl Once {
140140
}
141141
}
142142

143+
#[inline]
144+
pub(crate) fn set_state(&mut self, new_state: ExclusiveState) {
145+
*self.state_and_queue.get_mut() = match new_state {
146+
ExclusiveState::Incomplete => ptr::without_provenance_mut(INCOMPLETE),
147+
ExclusiveState::Poisoned => ptr::without_provenance_mut(POISONED),
148+
ExclusiveState::Complete => ptr::without_provenance_mut(COMPLETE),
149+
};
150+
}
151+
143152
#[cold]
144153
#[track_caller]
145154
pub fn wait(&self, ignore_poisoning: bool) {

0 commit comments

Comments
 (0)
Please sign in to comment.