From fae9f73da65a92812ff96aadb18edb98af4a3d43 Mon Sep 17 00:00:00 2001
From: Gavin Li <gavinli@thegavinli.com>
Date: Sat, 24 Dec 2022 18:33:19 -0800
Subject: [PATCH] Add once_cell::race::OnceRef

Like OnceBox, but stores a shared reference instead of a Box.
---
 CHANGELOG.md |  4 +++
 Cargo.toml   |  2 +-
 src/race.rs  | 92 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 97 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 5cf6ea6..3f3a862 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,6 +4,10 @@
 
 -
 
+## 1.17.0
+
+- Add `race::OnceRef` for storing a `&'a T`.
+
 ## 1.16.0
 
 - Add `no_std` implementation based on `critical-section`,
diff --git a/Cargo.toml b/Cargo.toml
index 9a9a6df..5ceec5e 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "once_cell"
-version = "1.16.0"
+version = "1.17.0"
 authors = ["Aleksey Kladov <aleksey.kladov@gmail.com>"]
 license = "MIT OR Apache-2.0"
 edition = "2021"
diff --git a/src/race.rs b/src/race.rs
index fd255c4..dff5847 100644
--- a/src/race.rs
+++ b/src/race.rs
@@ -25,6 +25,8 @@ use atomic_polyfill as atomic;
 use core::sync::atomic;
 
 use atomic::{AtomicUsize, Ordering};
+use core::cell::UnsafeCell;
+use core::marker::PhantomData;
 use core::num::NonZeroUsize;
 
 /// A thread-safe cell which can be written to only once.
@@ -172,6 +174,96 @@ impl OnceBool {
     }
 }
 
+/// A thread-safe cell which can be written to only once.
+pub struct OnceRef<'a, T> {
+    inner: OnceNonZeroUsize,
+    ghost: PhantomData<UnsafeCell<&'a T>>,
+}
+
+// TODO: Replace UnsafeCell with SyncUnsafeCell once stabilized
+unsafe impl<'a, T: Sync> Sync for OnceRef<'a, T> {}
+
+impl<'a, T> core::fmt::Debug for OnceRef<'a, T> {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        write!(f, "OnceRef({:?})", self.inner)
+    }
+}
+
+impl<'a, T> Default for OnceRef<'a, T> {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl<'a, T> OnceRef<'a, T> {
+    /// Creates a new empty cell.
+    pub const fn new() -> OnceRef<'a, T> {
+        OnceRef { inner: OnceNonZeroUsize::new(), ghost: PhantomData }
+    }
+
+    /// Gets a reference to the underlying value.
+    pub fn get(&self) -> Option<&'a T> {
+        self.inner.get().map(|ptr| unsafe { &*(ptr.get() as *const T) })
+    }
+
+    /// Sets the contents of this cell to `value`.
+    ///
+    /// Returns `Ok(())` if the cell was empty and `Err(value)` if it was
+    /// full.
+    pub fn set(&self, value: &'a T) -> Result<(), ()> {
+        let ptr = NonZeroUsize::new(value as *const T as usize).unwrap();
+        self.inner.set(ptr)
+    }
+
+    /// Gets the contents of the cell, initializing it with `f` if the cell was
+    /// empty.
+    ///
+    /// If several threads concurrently run `get_or_init`, more than one `f` can
+    /// be called. However, all threads will return the same value, produced by
+    /// some `f`.
+    pub fn get_or_init<F>(&self, f: F) -> &'a T
+    where
+        F: FnOnce() -> &'a T,
+    {
+        let f = || NonZeroUsize::new(f() as *const T as usize).unwrap();
+        let ptr = self.inner.get_or_init(f);
+        unsafe { &*(ptr.get() as *const T) }
+    }
+
+    /// Gets the contents of the cell, initializing it with `f` if
+    /// the cell was empty. If the cell was empty and `f` failed, an
+    /// error is returned.
+    ///
+    /// If several threads concurrently run `get_or_init`, more than one `f` can
+    /// be called. However, all threads will return the same value, produced by
+    /// some `f`.
+    pub fn get_or_try_init<F, E>(&self, f: F) -> Result<&'a T, E>
+    where
+        F: FnOnce() -> Result<&'a T, E>,
+    {
+        let f = || f().map(|value| NonZeroUsize::new(value as *const T as usize).unwrap());
+        let ptr = self.inner.get_or_try_init(f)?;
+        unsafe { Ok(&*(ptr.get() as *const T)) }
+    }
+
+    /// ```compile_fail
+    /// use once_cell::race::OnceRef;
+    ///
+    /// let mut l = OnceRef::new();
+    ///
+    /// {
+    ///     let y = 2;
+    ///     let mut r = OnceRef::new();
+    ///     r.set(&y).unwrap();
+    ///     core::mem::swap(&mut l, &mut r);
+    /// }
+    ///
+    /// // l now contains a dangling reference to y
+    /// eprintln!("uaf: {}", l.get().unwrap());
+    /// ```
+    fn _dummy() {}
+}
+
 #[cfg(feature = "alloc")]
 pub use self::once_box::OnceBox;