Auto merge of #107329 - joboet:optimize_lazylock, r=m-ou-se

Optimize `LazyLock` size

The initialization function was unnecessarily stored separately from the data to be initialized. Since both cannot exist at the same time, a `union` can be used, with the `Once` acting as discriminant. This unfortunately requires some extra methods on `Once` so that `Drop` can be implemented correctly and efficiently.

`@rustbot` label +T-libs +A-atomic

6 files changed, 120 insertions, 17 deletions

diff --git a/library/std/src/sync/lazy_lock.rs b/library/std/src/sync/lazy_lock.rs
index 4a15305301d..7e85d6a063a 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;
diff --git a/library/std/src/sync/mod.rs b/library/std/src/sync/mod.rs
index ba20bab87a4..4edc956173b 100644
--- a/library/std/src/sync/mod.rs
+++ b/library/std/src/sync/mod.rs
@@ -186,7 +186,7 @@ mod condvar;
 mod lazy_lock;
 mod mpmc;
 mod mutex;
-mod once;
+pub(crate) mod once;
 mod once_lock;
 mod poison;
 mod remutex;
diff --git a/library/std/src/sync/once.rs b/library/std/src/sync/once.rs
index 0f25417d6b5..1b17c31089f 100644
--- a/library/std/src/sync/once.rs
+++ b/library/std/src/sync/once.rs
@@ -43,6 +43,12 @@ pub struct OnceState {
     pub(crate) inner: sys::OnceState,
 }
 
+pub(crate) enum ExclusiveState {
+    Incomplete,
+    Poisoned,
+    Complete,
+}
+
 /// Initialization value for static [`Once`] values.
 ///
 /// # Examples
@@ -248,6 +254,16 @@ impl Once {
     pub fn is_completed(&self) -> bool {
         self.inner.is_completed()
     }
+
+    /// Returns the 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 state(&mut self) -> ExclusiveState {
+        self.inner.state()
+    }
 }
 
 #[stable(feature = "std_debug", since = "1.16.0")]
diff --git a/library/std/src/sys/unsupported/once.rs b/library/std/src/sys/unsupported/once.rs
index b4bb4975f41..11fde1888ba 100644
--- a/library/std/src/sys/unsupported/once.rs
+++ b/library/std/src/sys/unsupported/once.rs
@@ -1,5 +1,6 @@
 use crate::cell::Cell;
 use crate::sync as public;
+use crate::sync::once::ExclusiveState;
 
 pub struct Once {
     state: Cell<State>,
@@ -44,6 +45,16 @@ impl Once {
         self.state.get() == State::Complete
     }
 
+    #[inline]
+    pub(crate) fn state(&mut self) -> ExclusiveState {
+        match self.state.get() {
+            State::Incomplete => ExclusiveState::Incomplete,
+            State::Poisoned => ExclusiveState::Poisoned,
+            State::Complete => ExclusiveState::Complete,
+            _ => unreachable!("invalid Once state"),
+        }
+    }
+
     #[cold]
     #[track_caller]
     pub fn call(&self, ignore_poisoning: bool, f: &mut impl FnMut(&public::OnceState)) {
diff --git a/library/std/src/sys_common/once/futex.rs b/library/std/src/sys_common/once/futex.rs
index 5c7e6c01371..42db5fad4b4 100644
--- a/library/std/src/sys_common/once/futex.rs
+++ b/library/std/src/sys_common/once/futex.rs
@@ -4,6 +4,7 @@ use crate::sync::atomic::{
     AtomicU32,
     Ordering::{Acquire, Relaxed, Release},
 };
+use crate::sync::once::ExclusiveState;
 use crate::sys::futex::{futex_wait, futex_wake_all};
 
 // On some platforms, the OS is very nice and handles the waiter queue for us.
@@ -78,6 +79,16 @@ impl Once {
         self.state.load(Acquire) == COMPLETE
     }
 
+    #[inline]
+    pub(crate) fn state(&mut self) -> ExclusiveState {
+        match *self.state.get_mut() {
+            INCOMPLETE => ExclusiveState::Incomplete,
+            POISONED => ExclusiveState::Poisoned,
+            COMPLETE => ExclusiveState::Complete,
+            _ => unreachable!("invalid Once state"),
+        }
+    }
+
     // This uses FnMut to match the API of the generic implementation. As this
     // implementation is quite light-weight, it is generic over the closure and
     // so avoids the cost of dynamic dispatch.
diff --git a/library/std/src/sys_common/once/queue.rs b/library/std/src/sys_common/once/queue.rs
index d953a674592..def0bcd6fac 100644
--- a/library/std/src/sys_common/once/queue.rs
+++ b/library/std/src/sys_common/once/queue.rs
@@ -60,6 +60,7 @@ use crate::fmt;
 use crate::ptr;
 use crate::sync as public;
 use crate::sync::atomic::{AtomicBool, AtomicPtr, Ordering};
+use crate::sync::once::ExclusiveState;
 use crate::thread::{self, Thread};
 
 type Masked = ();
@@ -121,6 +122,16 @@ impl Once {
         self.state_and_queue.load(Ordering::Acquire).addr() == COMPLETE
     }
 
+    #[inline]
+    pub(crate) fn state(&mut self) -> ExclusiveState {
+        match self.state_and_queue.get_mut().addr() {
+            INCOMPLETE => ExclusiveState::Incomplete,
+            POISONED => ExclusiveState::Poisoned,
+            COMPLETE => ExclusiveState::Complete,
+            _ => unreachable!("invalid Once state"),
+        }
+    }
+
     // This is a non-generic function to reduce the monomorphization cost of
     // using `call_once` (this isn't exactly a trivial or small implementation).
     //