Rollup merge of #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`: https://github.com/rust-lang/rust/issues/129333

3 files changed, 142 insertions, 6 deletions

diff --git a/library/std/src/sync/lazy_lock.rs b/library/std/src/sync/lazy_lock.rs
index 1dfe1dc538b..b05615035d7 100644
--- a/library/std/src/sync/lazy_lock.rs
+++ b/library/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,7 +325,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>")),
         };
@@ -234,6 +333,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")]