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.

6 files changed, 181 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")]
diff --git a/library/std/src/sync/lazy_lock/tests.rs b/library/std/src/sync/lazy_lock/tests.rs
index 8a6ab4ac4fd..94044368305 100644
--- a/library/std/src/sync/lazy_lock/tests.rs
+++ b/library/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);
+}
diff --git a/library/std/src/sync/once.rs b/library/std/src/sync/once.rs
index bf595fdea2d..5a1cd7d0b8b 100644
--- a/library/std/src/sync/once.rs
+++ b/library/std/src/sync/once.rs
@@ -314,6 +314,16 @@ impl Once {
     pub(crate) fn state(&mut self) -> ExclusiveState {
         self.inner.state()
     }
+
+    /// Sets current state of the `Once` instance.
+    ///
+    /// Since this takes a mutable reference, no initialization can currently
+    /// be running, so the state must be either "incomplete", "poisoned" or
+    /// "complete".
+    #[inline]
+    pub(crate) fn set_state(&mut self, new_state: ExclusiveState) {
+        self.inner.set_state(new_state);
+    }
 }
 
 #[stable(feature = "std_debug", since = "1.16.0")]
diff --git a/library/std/src/sys/sync/once/futex.rs b/library/std/src/sys/sync/once/futex.rs
index 2c8a054282b..25588a4217b 100644
--- a/library/std/src/sys/sync/once/futex.rs
+++ b/library/std/src/sys/sync/once/futex.rs
@@ -91,6 +91,15 @@ impl Once {
         }
     }
 
+    #[inline]
+    pub(crate) fn set_state(&mut self, new_state: ExclusiveState) {
+        *self.state_and_queued.get_mut() = match new_state {
+            ExclusiveState::Incomplete => INCOMPLETE,
+            ExclusiveState::Poisoned => POISONED,
+            ExclusiveState::Complete => COMPLETE,
+        };
+    }
+
     #[cold]
     #[track_caller]
     pub fn wait(&self, ignore_poisoning: bool) {
diff --git a/library/std/src/sys/sync/once/no_threads.rs b/library/std/src/sys/sync/once/no_threads.rs
index 12c1d9f5a6c..cdcffe790f5 100644
--- a/library/std/src/sys/sync/once/no_threads.rs
+++ b/library/std/src/sys/sync/once/no_threads.rs
@@ -55,6 +55,15 @@ impl Once {
         }
     }
 
+    #[inline]
+    pub(crate) fn set_state(&mut self, new_state: ExclusiveState) {
+        self.state.set(match new_state {
+            ExclusiveState::Incomplete => State::Incomplete,
+            ExclusiveState::Poisoned => State::Poisoned,
+            ExclusiveState::Complete => State::Complete,
+        });
+    }
+
     #[cold]
     #[track_caller]
     pub fn wait(&self, _ignore_poisoning: bool) {
diff --git a/library/std/src/sys/sync/once/queue.rs b/library/std/src/sys/sync/once/queue.rs
index 86f72c82008..17abaf0bf26 100644
--- a/library/std/src/sys/sync/once/queue.rs
+++ b/library/std/src/sys/sync/once/queue.rs
@@ -140,6 +140,15 @@ impl Once {
         }
     }
 
+    #[inline]
+    pub(crate) fn set_state(&mut self, new_state: ExclusiveState) {
+        *self.state_and_queue.get_mut() = match new_state {
+            ExclusiveState::Incomplete => ptr::without_provenance_mut(INCOMPLETE),
+            ExclusiveState::Poisoned => ptr::without_provenance_mut(POISONED),
+            ExclusiveState::Complete => ptr::without_provenance_mut(COMPLETE),
+        };
+    }
+
     #[cold]
     #[track_caller]
     pub fn wait(&self, ignore_poisoning: bool) {