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 files changed, 123 insertions, 7 deletions

diff --git a/library/std/src/sync/lazy_lock.rs b/library/std/src/sync/lazy_lock.rs
index 953aef40e7b..afdfec43afd 100644
--- a/library/std/src/sync/lazy_lock.rs
+++ b/library/std/src/sync/lazy_lock.rs
@@ -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")]