Rollup merge of #130476 - workingjubilee:more-lazy-methods-take-2, r=… · qinheping/verify-rust-std@f63c0c1
11use super::UnsafeCell;
2+use crate::hint::unreachable_unchecked;
23use crate::ops::Deref;
34use crate::{fmt, mem};
45@@ -82,7 +83,7 @@ impl<T, F: FnOnce() -> T> LazyCell<T, F> {
8283match this.state.into_inner() {
8384State::Init(data) => Ok(data),
8485State::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> {
114115State::Init(data) => data,
115116// SAFETY: The state is uninitialized.
116117State::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}
153219154220impl<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() };
162270match state {
163271State::Init(data) => Some(data),
164272 _ => None,
@@ -188,10 +296,16 @@ impl<T: Default> Default for LazyCell<T> {
188296impl<T: fmt::Debug, F> fmt::Debug for LazyCell<T, F> {
189297fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
190298let mut d = f.debug_tuple("LazyCell");
191-match self.get() {
299+match LazyCell::get(self) {
192300Some(data) => d.field(data),
193301None => 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+}