Replace IngredientCache lock with atomic primitive

src/nonce.rs

@@ -31,3 +31,13 @@ impl<T> NonceGenerator<T> {
         Nonce(NonZeroU32::new(value).unwrap(), self.phantom)
     }
 }
+
+impl<T> Nonce<T> {
+    pub(crate) fn into_u32(self) -> NonZeroU32 {
+        self.0
+    }
+
+    pub(crate) fn from_u32(u32: NonZeroU32) -> Self {
+        Self(u32, PhantomData)
+    }
+}

src/zalsa.rs

@@ -3,6 +3,8 @@ use parking_lot::{Mutex, RwLock};
 use rustc_hash::FxHashMap;
 use std::any::{Any, TypeId};
 use std::marker::PhantomData;
+use std::num::NonZeroU32;
+use std::sync::atomic::{AtomicU64, Ordering};
 use std::thread::ThreadId;
 
 use crate::cycle::CycleRecoveryStrategy;
@@ -77,7 +79,7 @@ pub struct IngredientIndex(u32);
 impl IngredientIndex {
     /// Create an ingredient index from a usize.
     pub(crate) fn from(v: usize) -> Self {
-        assert!(v < (u32::MAX as usize));
+        assert!(v <= u32::MAX as usize);
         Self(v as u32)
     }
 
@@ -107,9 +109,10 @@ pub struct MemoIngredientIndex(u32);
 
 impl MemoIngredientIndex {
     pub(crate) fn from_usize(u: usize) -> Self {
-        assert!(u < u32::MAX as usize);
+        assert!(u <= u32::MAX as usize);
         MemoIngredientIndex(u as u32)
     }
+
     pub(crate) fn as_usize(self) -> usize {
         self.0 as usize
     }
@@ -343,7 +346,10 @@ pub struct IngredientCache<I>
 where
     I: Ingredient,
 {
-    cached_data: std::sync::OnceLock<(Nonce<StorageNonce>, IngredientIndex)>,
+    // This is a packed representation of `Option<(Nonce<StorageNonce>, IngredientIndex)>`
+    // allowing us to avoid a lock in favor of an atomic load. This works thanks to `Nonce`
+    // having a niche and so the entire type can fit into a u64.
+    cached_data: AtomicU64,
     phantom: PhantomData<fn() -> I>,
 }
 
@@ -362,10 +368,12 @@ impl<I> IngredientCache<I>
 where
     I: Ingredient,
 {
+    const UNINITIALIZED: u64 = 0;
+
     /// Create a new cache
     pub const fn new() -> Self {
         Self {
-            cached_data: std::sync::OnceLock::new(),
+            cached_data: AtomicU64::new(Self::UNINITIALIZED),
             phantom: PhantomData,
         }
     }
@@ -377,11 +385,29 @@ where
         db: &'s dyn Database,
         create_index: impl Fn() -> IngredientIndex,
     ) -> &'s I {
-        let zalsa = db.zalsa();
-        let (nonce, index) = self.cached_data.get_or_init(|| {
-            let index = create_index();
-            (zalsa.nonce(), index)
+        let mut cached_data = self.cached_data.load(Ordering::Acquire);
+        if cached_data == Self::UNINITIALIZED {
+            let index = create_index().0 as u64;
+            let nonce = db.zalsa().nonce().into_u32().get() as u64;
+            let packed = (nonce << u32::BITS) | index;
+            match self.cached_data.compare_exchange(
+                Self::UNINITIALIZED,
+                packed,
+                Ordering::AcqRel,
+                Ordering::Acquire,
+            ) {
+                // raced and lost, so load the new value
+                Err(_) => cached_data = self.cached_data.load(Ordering::Acquire),
+                // successfully set the value
+                Ok(_) => cached_data = packed,
+            }
+        };
+        // unpack our u64
+        // SAFETY: We've checked against `UNINITIALIZED` (0) above and so the upper bits must be non-zero
+        let nonce = Nonce::<StorageNonce>::from_u32(unsafe {
+            NonZeroU32::new_unchecked((cached_data >> u32::BITS) as u32)
         });
+        let mut index = IngredientIndex(cached_data as u32);
 
         // FIXME: We used to cache a raw pointer to the revision but miri
         // was reporting errors because that pointer was derived from an `&`
@@ -390,12 +416,10 @@ where
         // We could fix it with orxfun/orx-concurrent-vec#18 or by "refreshing" the cache
         // when the revision changes but just caching the index is an awful lot simpler.
 
-        if db.zalsa().nonce() == *nonce {
-            zalsa.lookup_ingredient(*index).assert_type::<I>()
-        } else {
-            let index = create_index();
-            zalsa.lookup_ingredient(index).assert_type::<I>()
+        if db.zalsa().nonce() != nonce {
+            index = create_index();
         }
+        db.zalsa().lookup_ingredient(index).assert_type::<I>()
     }
 }