Skip to content

Commit f63c0c1

Browse files
workingjubileeworkingjubilee
authored
Rollup merge of rust-lang#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`: rust-lang#129333
2 parents 09e36ad + 262a08b commit f63c0c1

File tree

11 files changed

+312
-11
lines changed

11 files changed

+312
-11
lines changed

‎core/src/cell/lazy.rs

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

@@ -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,72 @@ 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+
/// # Examples
126+
///
127+
/// ```
128+
/// #![feature(lazy_get)]
129+
/// use std::cell::LazyCell;
130+
///
131+
/// let mut lazy = LazyCell::new(|| 92);
132+
///
133+
/// let p = LazyCell::force_mut(&mut lazy);
134+
/// assert_eq!(*p, 92);
135+
/// *p = 44;
136+
/// assert_eq!(*lazy, 44);
137+
/// ```
138+
#[inline]
139+
#[unstable(feature = "lazy_get", issue = "129333")]
140+
pub fn force_mut(this: &mut LazyCell<T, F>) -> &mut T {
141+
#[cold]
142+
/// # Safety
143+
/// May only be called when the state is `Uninit`.
144+
unsafe fn really_init_mut<T, F: FnOnce() -> T>(state: &mut State<T, F>) -> &mut T {
145+
// INVARIANT: Always valid, but the value may not be dropped.
146+
struct PoisonOnPanic<T, F>(*mut State<T, F>);
147+
impl<T, F> Drop for PoisonOnPanic<T, F> {
148+
#[inline]
149+
fn drop(&mut self) {
150+
// SAFETY: Invariant states it is valid, and we don't drop the old value.
151+
unsafe {
152+
self.0.write(State::Poisoned);
153+
}
154+
}
155+
}
156+
157+
let State::Uninit(f) = state else {
158+
// `unreachable!()` here won't optimize out because the function is cold.
159+
// SAFETY: Precondition.
160+
unsafe { unreachable_unchecked() };
161+
};
162+
// SAFETY: We never drop the state after we read `f`, and we write a valid value back
163+
// in any case, panic or success. `f` can't access the `LazyCell` because it is mutably
164+
// borrowed.
165+
let f = unsafe { core::ptr::read(f) };
166+
// INVARIANT: Initiated from mutable reference, don't drop because we read it.
167+
let guard = PoisonOnPanic(state);
168+
let data = f();
169+
// SAFETY: `PoisonOnPanic` invariant, and we don't drop the old value.
170+
unsafe {
171+
core::ptr::write(guard.0, State::Init(data));
172+
}
173+
core::mem::forget(guard);
174+
let State::Init(data) = state else { unreachable!() };
175+
data
176+
}
177+
178+
let state = this.state.get_mut();
179+
match state {
180+
State::Init(data) => data,
181+
// SAFETY: `state` is `Uninit`.
182+
State::Uninit(_) => unsafe { really_init_mut(state) },
183+
State::Poisoned => panic_poisoned(),
118184
}
119185
}
120186

@@ -152,13 +218,55 @@ impl<T, F: FnOnce() -> T> LazyCell<T, F> {
152218
}
153219

154220
impl<T, F> LazyCell<T, F> {
221+
/// Returns a reference to the value if initialized, or `None` if not.
222+
///
223+
/// # Examples
224+
///
225+
/// ```
226+
/// #![feature(lazy_get)]
227+
///
228+
/// use std::cell::LazyCell;
229+
///
230+
/// let mut lazy = LazyCell::new(|| 92);
231+
///
232+
/// assert_eq!(LazyCell::get_mut(&mut lazy), None);
233+
/// let _ = LazyCell::force(&lazy);
234+
/// *LazyCell::get_mut(&mut lazy).unwrap() = 44;
235+
/// assert_eq!(*lazy, 44);
236+
/// ```
155237
#[inline]
156-
fn get(&self) -> Option<&T> {
238+
#[unstable(feature = "lazy_get", issue = "129333")]
239+
pub fn get_mut(this: &mut LazyCell<T, F>) -> Option<&mut T> {
240+
let state = this.state.get_mut();
241+
match state {
242+
State::Init(data) => Some(data),
243+
_ => None,
244+
}
245+
}
246+
247+
/// Returns a mutable reference to the value if initialized, or `None` if not.
248+
///
249+
/// # Examples
250+
///
251+
/// ```
252+
/// #![feature(lazy_get)]
253+
///
254+
/// use std::cell::LazyCell;
255+
///
256+
/// let lazy = LazyCell::new(|| 92);
257+
///
258+
/// assert_eq!(LazyCell::get(&lazy), None);
259+
/// let _ = LazyCell::force(&lazy);
260+
/// assert_eq!(LazyCell::get(&lazy), Some(&92));
261+
/// ```
262+
#[inline]
263+
#[unstable(feature = "lazy_get", issue = "129333")]
264+
pub fn get(this: &LazyCell<T, F>) -> Option<&T> {
157265
// SAFETY:
158266
// This is sound for the same reason as in `force`: once the state is
159267
// initialized, it will not be mutably accessed again, so this reference
160268
// will stay valid for the duration of the borrow to `self`.
161-
let state = unsafe { &*self.state.get() };
269+
let state = unsafe { &*this.state.get() };
162270
match state {
163271
State::Init(data) => Some(data),
164272
_ => None,
@@ -188,10 +296,16 @@ impl<T: Default> Default for LazyCell<T> {
188296
impl<T: fmt::Debug, F> fmt::Debug for LazyCell<T, F> {
189297
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
190298
let mut d = f.debug_tuple("LazyCell");
191-
match self.get() {
299+
match LazyCell::get(self) {
192300
Some(data) => d.field(data),
193301
None => d.field(&format_args!("<uninit>")),
194302
};
195303
d.finish()
196304
}
197305
}
306+
307+
#[cold]
308+
#[inline(never)]
309+
fn panic_poisoned() -> ! {
310+
panic!("LazyCell instance has previously been poisoned")
311+
}

‎core/src/lib.rs

+1
Original file line numberDiff line numberDiff line change
@@ -170,6 +170,7 @@
170170
#![feature(is_ascii_octdigit)]
171171
#![feature(is_val_statically_known)]
172172
#![feature(isqrt)]
173+
#![feature(lazy_get)]
173174
#![feature(link_cfg)]
174175
#![feature(offset_of_enum)]
175176
#![feature(panic_internals)]

‎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();

‎core/tests/lib.rs

+1
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,7 @@
7676
#![feature(iterator_try_collect)]
7777
#![feature(iterator_try_reduce)]
7878
#![feature(layout_for_ptr)]
79+
#![feature(lazy_get)]
7980
#![feature(maybe_uninit_fill)]
8081
#![feature(maybe_uninit_uninit_array_transpose)]
8182
#![feature(maybe_uninit_write_slice)]

‎std/src/lib.rs

+1
Original file line numberDiff line numberDiff line change
@@ -337,6 +337,7 @@
337337
#![feature(hasher_prefixfree_extras)]
338338
#![feature(hashmap_internals)]
339339
#![feature(ip)]
340+
#![feature(lazy_get)]
340341
#![feature(maybe_uninit_slice)]
341342
#![feature(maybe_uninit_write_slice)]
342343
#![feature(panic_can_unwind)]

‎std/src/sync/lazy_lock.rs

+111-6
Original file line numberDiff line numberDiff line change
@@ -119,7 +119,7 @@ impl<T, F: FnOnce() -> T> LazyLock<T, F> {
119119
pub fn into_inner(mut this: Self) -> Result<T, F> {
120120
let state = this.once.state();
121121
match state {
122-
ExclusiveState::Poisoned => panic!("LazyLock instance has previously been poisoned"),
122+
ExclusiveState::Poisoned => panic_poisoned(),
123123
state => {
124124
let this = ManuallyDrop::new(this);
125125
let data = unsafe { ptr::read(&this.data) }.into_inner();
@@ -132,6 +132,60 @@ impl<T, F: FnOnce() -> T> LazyLock<T, F> {
132132
}
133133
}
134134

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

174228
impl<T, F> LazyLock<T, F> {
175-
/// Gets the inner value if it has already been initialized.
176-
fn get(&self) -> Option<&T> {
177-
if self.once.is_completed() {
229+
/// Returns a reference to the value if initialized, or `None` if not.
230+
///
231+
/// # Examples
232+
///
233+
/// ```
234+
/// #![feature(lazy_get)]
235+
///
236+
/// use std::sync::LazyLock;
237+
///
238+
/// let mut lazy = LazyLock::new(|| 92);
239+
///
240+
/// assert_eq!(LazyLock::get_mut(&mut lazy), None);
241+
/// let _ = LazyLock::force(&lazy);
242+
/// *LazyLock::get_mut(&mut lazy).unwrap() = 44;
243+
/// assert_eq!(*lazy, 44);
244+
/// ```
245+
#[inline]
246+
#[unstable(feature = "lazy_get", issue = "129333")]
247+
pub fn get_mut(this: &mut LazyLock<T, F>) -> Option<&mut T> {
248+
// `state()` does not perform an atomic load, so prefer it over `is_complete()`.
249+
let state = this.once.state();
250+
match state {
251+
// SAFETY:
252+
// The closure has been run successfully, so `value` has been initialized.
253+
ExclusiveState::Complete => Some(unsafe { &mut this.data.get_mut().value }),
254+
_ => None,
255+
}
256+
}
257+
258+
/// Returns a mutable reference to the value if initialized, or `None` if not.
259+
///
260+
/// # Examples
261+
///
262+
/// ```
263+
/// #![feature(lazy_get)]
264+
///
265+
/// use std::sync::LazyLock;
266+
///
267+
/// let lazy = LazyLock::new(|| 92);
268+
///
269+
/// assert_eq!(LazyLock::get(&lazy), None);
270+
/// let _ = LazyLock::force(&lazy);
271+
/// assert_eq!(LazyLock::get(&lazy), Some(&92));
272+
/// ```
273+
#[inline]
274+
#[unstable(feature = "lazy_get", issue = "129333")]
275+
pub fn get(this: &LazyLock<T, F>) -> Option<&T> {
276+
if this.once.is_completed() {
178277
// SAFETY:
179278
// The closure has been run successfully, so `value` has been initialized
180279
// and will not be modified again.
181-
Some(unsafe { &*(*self.data.get()).value })
280+
Some(unsafe { &(*this.data.get()).value })
182281
} else {
183282
None
184283
}
@@ -226,14 +325,20 @@ impl<T: Default> Default for LazyLock<T> {
226325
impl<T: fmt::Debug, F> fmt::Debug for LazyLock<T, F> {
227326
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
228327
let mut d = f.debug_tuple("LazyLock");
229-
match self.get() {
328+
match LazyLock::get(self) {
230329
Some(v) => d.field(v),
231330
None => d.field(&format_args!("<uninit>")),
232331
};
233332
d.finish()
234333
}
235334
}
236335

336+
#[cold]
337+
#[inline(never)]
338+
fn panic_poisoned() -> ! {
339+
panic!("LazyLock instance has previously been poisoned")
340+
}
341+
237342
// We never create a `&F` from a `&LazyLock<T, F>` so it is fine
238343
// to not impl `Sync` for `F`.
239344
#[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+
}

0 commit comments

Comments
 (0)