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

Author: workingjubilee

core/src/cell/lazy.rs

@@ -1,4 +1,5 @@
 use super::UnsafeCell;
+use crate::hint::unreachable_unchecked;
 use crate::ops::Deref;
 use crate::{fmt, mem};
 
@@ -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,72 @@ 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.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(lazy_get)]
+    /// 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);
+    /// ```
+    #[inline]
+    #[unstable(feature = "lazy_get", issue = "129333")]
+    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_mut<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_mut(state) },
+            State::Poisoned => panic_poisoned(),
         }
     }
 
@@ -152,13 +218,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,
@@ -188,10 +296,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")
+}

core/src/lib.rs

@@ -170,6 +170,7 @@
 #![feature(is_ascii_octdigit)]
 #![feature(is_val_statically_known)]
 #![feature(isqrt)]
+#![feature(lazy_get)]
 #![feature(link_cfg)]
 #![feature(offset_of_enum)]
 #![feature(panic_internals)]

core/tests/lazy.rs

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

core/tests/lib.rs

@@ -76,6 +76,7 @@
 #![feature(iterator_try_collect)]
 #![feature(iterator_try_reduce)]
 #![feature(layout_for_ptr)]
+#![feature(lazy_get)]
 #![feature(maybe_uninit_fill)]
 #![feature(maybe_uninit_uninit_array_transpose)]
 #![feature(maybe_uninit_write_slice)]

std/src/lib.rs

@@ -337,6 +337,7 @@
 #![feature(hasher_prefixfree_extras)]
 #![feature(hashmap_internals)]
 #![feature(ip)]
+#![feature(lazy_get)]
 #![feature(maybe_uninit_slice)]
 #![feature(maybe_uninit_write_slice)]
 #![feature(panic_can_unwind)]

std/src/sync/lazy_lock.rs

@@ -119,7 +119,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();
@@ -132,6 +132,60 @@ impl<T, F: FnOnce() -> T> LazyLock<T, F> {
         }
     }
 
+    /// Forces the evaluation of this lazy value and returns a mutable reference to
+    /// the result.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(lazy_get)]
+    /// 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);
+    /// ```
+    #[inline]
+    #[unstable(feature = "lazy_get", issue = "129333")]
+    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_mut<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_mut(this) },
+        }
+    }
+
     /// Forces the evaluation of this lazy value and returns a reference to
     /// result. This is equivalent to the `Deref` impl, but is explicit.
     ///
@@ -172,13 +226,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
         }
@@ -226,14 +325,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")]

std/src/sync/lazy_lock/tests.rs

@@ -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);
+}