From afb385fba44ee395db6065168a3d54cf11afe95f Mon Sep 17 00:00:00 2001
From: Luke Parker <lukeparker5132@gmail.com>
Date: Wed, 26 Jul 2023 02:59:24 -0400
Subject: [PATCH] Use spin's Once for OnceLock

---
 common/std-shims/Cargo.toml  |  2 +-
 common/std-shims/src/sync.rs | 32 ++++++--------------------------
 2 files changed, 7 insertions(+), 27 deletions(-)

diff --git a/common/std-shims/Cargo.toml b/common/std-shims/Cargo.toml
index 0f20c333..b7d3abbe 100644
--- a/common/std-shims/Cargo.toml
+++ b/common/std-shims/Cargo.toml
@@ -13,7 +13,7 @@ all-features = true
 rustdoc-args = ["--cfg", "docsrs"]
 
 [dependencies]
-spin = "0.9"
+spin = { version = "0.9", features = ["mutex", "once"] }
 hashbrown = "0.14"
 
 [features]
diff --git a/common/std-shims/src/sync.rs b/common/std-shims/src/sync.rs
index 75cfe39a..6b8070bc 100644
--- a/common/std-shims/src/sync.rs
+++ b/common/std-shims/src/sync.rs
@@ -29,42 +29,22 @@ pub use mutex_shim::{ShimMutex as Mutex, MutexGuard};
 pub use std::sync::OnceLock;
 #[cfg(not(feature = "std"))]
 mod oncelock_shim {
-  use super::Mutex;
+  use spin::Once;
 
-  pub struct OnceLock<T>(Mutex<bool>, Option<T>);
+  pub struct OnceLock<T>(Once<T>);
   impl<T> OnceLock<T> {
     pub const fn new() -> OnceLock<T> {
-      OnceLock(Mutex::new(false), None)
+      OnceLock(Once::new())
     }
-
-    // These return a distinct Option in case of None so another caller using get_or_init doesn't
-    // transform it from None to Some
     pub fn get(&self) -> Option<&T> {
-      if !*self.0.lock() {
-        None
-      } else {
-        self.1.as_ref()
-      }
+      self.0.poll()
     }
     pub fn get_mut(&mut self) -> Option<&mut T> {
-      if !*self.0.lock() {
-        None
-      } else {
-        self.1.as_mut()
-      }
+      self.0.get_mut()
     }
 
     pub fn get_or_init<F: FnOnce() -> T>(&self, f: F) -> &T {
-      let mut lock = self.0.lock();
-      if !*lock {
-        unsafe {
-          core::ptr::addr_of!(self.1).cast_mut().write_unaligned(Some(f()));
-        }
-      }
-      *lock = true;
-      drop(lock);
-
-      self.get().unwrap()
+      self.0.call_once(f)
     }
   }
 }