1 files changed, 70 insertions, 16 deletions

diff --git a/library/std/src/sync/lazy_lock.rs b/library/std/src/sync/lazy_lock.rs
index 4a1530530..7e85d6a06 100644
--- a/library/std/src/sync/lazy_lock.rs
+++ b/library/std/src/sync/lazy_lock.rs
@@ -1,8 +1,21 @@
-use crate::cell::Cell;
+use crate::cell::UnsafeCell;
 use crate::fmt;
+use crate::mem::ManuallyDrop;
 use crate::ops::Deref;
 use crate::panic::{RefUnwindSafe, UnwindSafe};
-use crate::sync::OnceLock;
+use crate::sync::Once;
+
+use super::once::ExclusiveState;
+
+// We use the state of a Once as discriminant value. Upon creation, the state is
+// "incomplete" and `f` contains the initialization closure. In the first call to
+// `call_once`, `f` is taken and run. If it succeeds, `value` is set and the state
+// is changed to "complete". If it panics, the Once is poisoned, so none of the
+// two fields is initialized.
+union Data<T, F> {
+    value: ManuallyDrop<T>,
+    f: ManuallyDrop<F>,
+}
 
 /// A value which is initialized on the first access.
 ///
@@ -43,16 +56,17 @@ use crate::sync::OnceLock;
 /// ```
 #[unstable(feature = "once_cell", issue = "74465")]
 pub struct LazyLock<T, F = fn() -> T> {
-    cell: OnceLock<T>,
-    init: Cell<Option<F>>,
+    once: Once,
+    data: UnsafeCell<Data<T, F>>,
 }
+
 impl<T, F: FnOnce() -> T> LazyLock<T, F> {
     /// Creates a new lazy value with the given initializing
     /// function.
     #[inline]
     #[unstable(feature = "once_cell", issue = "74465")]
     pub const fn new(f: F) -> LazyLock<T, F> {
-        LazyLock { cell: OnceLock::new(), init: Cell::new(Some(f)) }
+        LazyLock { once: Once::new(), data: UnsafeCell::new(Data { f: ManuallyDrop::new(f) }) }
     }
 
     /// Forces the evaluation of this lazy value and
@@ -74,10 +88,50 @@ impl<T, F: FnOnce() -> T> LazyLock<T, F> {
     #[inline]
     #[unstable(feature = "once_cell", issue = "74465")]
     pub fn force(this: &LazyLock<T, F>) -> &T {
-        this.cell.get_or_init(|| match this.init.take() {
-            Some(f) => f(),
-            None => panic!("Lazy instance has previously been poisoned"),
-        })
+        this.once.call_once(|| {
+            // SAFETY: `call_once` only runs this closure once, ever.
+            let data = unsafe { &mut *this.data.get() };
+            let f = unsafe { ManuallyDrop::take(&mut data.f) };
+            let value = f();
+            data.value = ManuallyDrop::new(value);
+        });
+
+        // SAFETY:
+        // There are four possible scenarios:
+        // * the closure was called and initialized `value`.
+        // * the closure was called and panicked, so this point is never reached.
+        // * the closure was not called, but a previous call initialized `value`.
+        // * the closure was not called because the Once is poisoned, so this point
+        //   is never reached.
+        // So `value` has definitely been initialized and will not be modified again.
+        unsafe { &*(*this.data.get()).value }
+    }
+}
+
+impl<T, F> LazyLock<T, F> {
+    /// Get the inner value if it has already been initialized.
+    fn get(&self) -> Option<&T> {
+        if self.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 })
+        } else {
+            None
+        }
+    }
+}
+
+#[unstable(feature = "once_cell", issue = "74465")]
+impl<T, F> Drop for LazyLock<T, F> {
+    fn drop(&mut self) {
+        match self.once.state() {
+            ExclusiveState::Incomplete => unsafe { ManuallyDrop::drop(&mut self.data.get_mut().f) },
+            ExclusiveState::Complete => unsafe {
+                ManuallyDrop::drop(&mut self.data.get_mut().value)
+            },
+            ExclusiveState::Poisoned => {}
+        }
     }
 }
 
@@ -103,23 +157,23 @@ impl<T: Default> Default for LazyLock<T> {
 #[unstable(feature = "once_cell", issue = "74465")]
 impl<T: fmt::Debug, F> fmt::Debug for LazyLock<T, F> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        f.debug_struct("Lazy").field("cell", &self.cell).finish_non_exhaustive()
+        match self.get() {
+            Some(v) => f.debug_tuple("LazyLock").field(v).finish(),
+            None => f.write_str("LazyLock(Uninit)"),
+        }
     }
 }
 
 // We never create a `&F` from a `&LazyLock<T, F>` so it is fine
 // to not impl `Sync` for `F`
-// we do create a `&mut Option<F>` in `force`, but this is
-// properly synchronized, so it only happens once
-// so it also does not contribute to this impl.
 #[unstable(feature = "once_cell", issue = "74465")]
-unsafe impl<T, F: Send> Sync for LazyLock<T, F> where OnceLock<T>: Sync {}
+unsafe impl<T: Sync + Send, F: Send> Sync for LazyLock<T, F> {}
 // auto-derived `Send` impl is OK.
 
 #[unstable(feature = "once_cell", issue = "74465")]
-impl<T, F: UnwindSafe> RefUnwindSafe for LazyLock<T, F> where OnceLock<T>: RefUnwindSafe {}
+impl<T: RefUnwindSafe + UnwindSafe, F: UnwindSafe> RefUnwindSafe for LazyLock<T, F> {}
 #[unstable(feature = "once_cell", issue = "74465")]
-impl<T, F: UnwindSafe> UnwindSafe for LazyLock<T, F> where OnceLock<T>: UnwindSafe {}
+impl<T: UnwindSafe, F: UnwindSafe> UnwindSafe for LazyLock<T, F> {}
 
 #[cfg(test)]
 mod tests;