From 6db2587999872358e565e403504bc5ebfbfec7a5 Mon Sep 17 00:00:00 2001
From: daxpedda <daxpedda@gmail.com>
Date: Sun, 8 Oct 2023 21:10:39 +0200
Subject: [PATCH 1/2] Implement `OnceCell/Lock::try_insert()`

---
 library/core/src/cell/once.rs     | 40 ++++++++++++++++++++++++----
 library/std/src/sync/once_lock.rs | 43 ++++++++++++++++++++++++++++---
 2 files changed, 75 insertions(+), 8 deletions(-)

diff --git a/library/core/src/cell/once.rs b/library/core/src/cell/once.rs
index 2e8534f651a..f85df8a3895 100644
--- a/library/core/src/cell/once.rs
+++ b/library/core/src/cell/once.rs
@@ -87,10 +87,40 @@ impl<T> OnceCell<T> {
     #[inline]
     #[stable(feature = "once_cell", since = "1.70.0")]
     pub fn set(&self, value: T) -> Result<(), T> {
-        // SAFETY: Safe because we cannot have overlapping mutable borrows
-        let slot = unsafe { &*self.inner.get() };
-        if slot.is_some() {
-            return Err(value);
+        match self.try_insert(value) {
+            Ok(_) => Ok(()),
+            Err((_, value)) => Err(value),
+        }
+    }
+
+    /// Sets the contents of the cell to `value` if the cell was empty, then
+    /// returns a reference to it.
+    ///
+    /// # Errors
+    ///
+    /// This method returns `Ok(&value)` if the cell was empty and
+    /// `Err(&current_value, value)` if it was full.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(once_cell_try_insert)]
+    ///
+    /// use std::cell::OnceCell;
+    ///
+    /// let cell = OnceCell::new();
+    /// assert!(cell.get().is_none());
+    ///
+    /// assert_eq!(cell.try_insert(92), Ok(&92));
+    /// assert_eq!(cell.try_insert(62), Err((&92, 62)));
+    ///
+    /// assert!(cell.get().is_some());
+    /// ```
+    #[inline]
+    #[unstable(feature = "once_cell_try_insert", issue = "116693")]
+    pub fn try_insert(&self, value: T) -> Result<&T, (&T, T)> {
+        if let Some(old) = self.get() {
+            return Err((old, value));
         }
 
         // SAFETY: This is the only place where we set the slot, no races
@@ -99,7 +129,7 @@ impl<T> OnceCell<T> {
         // maintains the `inner`'s invariant.
         let slot = unsafe { &mut *self.inner.get() };
         *slot = Some(value);
-        Ok(())
+        Ok(self.get().unwrap())
     }
 
     /// Gets the contents of the cell, initializing it with `f`
diff --git a/library/std/src/sync/once_lock.rs b/library/std/src/sync/once_lock.rs
index e2b7b893cb5..f4963090795 100644
--- a/library/std/src/sync/once_lock.rs
+++ b/library/std/src/sync/once_lock.rs
@@ -126,11 +126,48 @@ impl<T> OnceLock<T> {
     #[inline]
     #[stable(feature = "once_cell", since = "1.70.0")]
     pub fn set(&self, value: T) -> Result<(), T> {
+        match self.try_insert(value) {
+            Ok(_) => Ok(()),
+            Err((_, value)) => Err(value),
+        }
+    }
+
+    /// Sets the contents of this cell to `value` if the cell was empty, then
+    /// returns a reference to it.
+    ///
+    /// May block if another thread is currently attempting to initialize the cell. The cell is
+    /// guaranteed to contain a value when set returns, though not necessarily the one provided.
+    ///
+    /// Returns `Ok(&value)` if the cell was empty and `Err(&current_value, value)` if it was full.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(once_cell_try_insert)]
+    ///
+    /// use std::sync::OnceLock;
+    ///
+    /// static CELL: OnceLock<i32> = OnceLock::new();
+    ///
+    /// fn main() {
+    ///     assert!(CELL.get().is_none());
+    ///
+    ///     std::thread::spawn(|| {
+    ///         assert_eq!(CELL.try_insert(92), Ok(&92));
+    ///     }).join().unwrap();
+    ///
+    ///     assert_eq!(CELL.try_insert(62), Err((&92, 62)));
+    ///     assert_eq!(CELL.get(), Some(&92));
+    /// }
+    /// ```
+    #[inline]
+    #[unstable(feature = "once_cell_try_insert", issue = "116693")]
+    pub fn try_insert(&self, value: T) -> Result<&T, (&T, T)> {
         let mut value = Some(value);
-        self.get_or_init(|| value.take().unwrap());
+        let res = self.get_or_init(|| value.take().unwrap());
         match value {
-            None => Ok(()),
-            Some(value) => Err(value),
+            None => Ok(res),
+            Some(value) => Err((res, value)),
         }
     }
 

From dd34d9027af3c0286aa038a365ff14387a6bf80c Mon Sep 17 00:00:00 2001
From: daxpedda <daxpedda@gmail.com>
Date: Sun, 8 Oct 2023 21:12:43 +0200
Subject: [PATCH 2/2] Add some optimizations

---
 library/core/src/cell/once.rs | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/library/core/src/cell/once.rs b/library/core/src/cell/once.rs
index f85df8a3895..3877a0c48cb 100644
--- a/library/core/src/cell/once.rs
+++ b/library/core/src/cell/once.rs
@@ -128,8 +128,7 @@ impl<T> OnceCell<T> {
         // checked that slot is currently `None`, so this write
         // maintains the `inner`'s invariant.
         let slot = unsafe { &mut *self.inner.get() };
-        *slot = Some(value);
-        Ok(self.get().unwrap())
+        Ok(slot.insert(value))
     }
 
     /// Gets the contents of the cell, initializing it with `f`
@@ -213,10 +212,9 @@ impl<T> OnceCell<T> {
         let val = outlined_call(f)?;
         // Note that *some* forms of reentrant initialization might lead to
         // UB (see `reentrant_init` test). I believe that just removing this
-        // `assert`, while keeping `set/get` would be sound, but it seems
+        // `panic`, while keeping `try_insert` would be sound, but it seems
         // better to panic, rather than to silently use an old value.
-        assert!(self.set(val).is_ok(), "reentrant init");
-        Ok(self.get().unwrap())
+        if let Ok(val) = self.try_insert(val) { Ok(val) } else { panic!("reentrant init") }
     }
 
     /// Consumes the cell, returning the wrapped value.