Replace ingredient cache with faster ingredient map (#921)

* replace ingredient cache with faster ingredient map

* avoid checking for downcasters in ingredient cache slow-path

* pre-size ingredient map

* avoid double lookup in ingredient creation slow-path

Cargo.toml

@@ -19,6 +19,7 @@ hashbrown = "0.15"
 hashlink = "0.10"
 indexmap = "2"
 intrusive-collections = "0.9.7"
+papaya = "0.2.2"
 parking_lot = "0.12"
 portable-atomic = "1"
 rustc-hash = "2"

components/salsa-macro-rules/src/setup_accumulator_impl.rs

@@ -26,7 +26,9 @@ macro_rules! setup_accumulator_impl {
                     $zalsa::IngredientCache::new();
 
                 $CACHE.get_or_create(zalsa, || {
-                    zalsa.add_or_lookup_jar_by_type::<$zalsa_struct::JarImpl<$Struct>>()
+                    zalsa
+                        .lookup_jar_by_type::<$zalsa_struct::JarImpl<$Struct>>()
+                        .get_or_create()
                 })
             }
 

components/salsa-macro-rules/src/setup_input_struct.rs

@@ -101,14 +101,14 @@ macro_rules! setup_input_struct {
                         $zalsa::IngredientCache::new();
 
                     CACHE.get_or_create(zalsa, || {
-                        zalsa.add_or_lookup_jar_by_type::<$zalsa_struct::JarImpl<$Configuration>>()
+                        zalsa.lookup_jar_by_type::<$zalsa_struct::JarImpl<$Configuration>>().get_or_create()
                     })
                 }
 
                 pub fn ingredient_mut(db: &mut dyn $zalsa::Database) -> (&mut $zalsa_struct::IngredientImpl<Self>, &mut $zalsa::Runtime) {
                     let zalsa_mut = db.zalsa_mut();
                     zalsa_mut.new_revision();
-                    let index = zalsa_mut.add_or_lookup_jar_by_type::<$zalsa_struct::JarImpl<$Configuration>>();
+                    let index = zalsa_mut.lookup_jar_by_type::<$zalsa_struct::JarImpl<$Configuration>>().get_or_create();
                     let (ingredient, runtime) = zalsa_mut.lookup_ingredient_mut(index);
                     let ingredient = ingredient.assert_type_mut::<$zalsa_struct::IngredientImpl<Self>>();
                     (ingredient, runtime)
@@ -150,7 +150,7 @@ macro_rules! setup_input_struct {
                 type MemoIngredientMap = $zalsa::MemoIngredientSingletonIndex;
 
                 fn lookup_or_create_ingredient_index(aux: &$zalsa::Zalsa) -> $zalsa::IngredientIndices {
-                    aux.add_or_lookup_jar_by_type::<$zalsa_struct::JarImpl<$Configuration>>().into()
+                    aux.lookup_jar_by_type::<$zalsa_struct::JarImpl<$Configuration>>().get_or_create().into()
                 }
 
                 #[inline]

components/salsa-macro-rules/src/setup_interned_struct.rs

@@ -149,7 +149,7 @@ macro_rules! setup_interned_struct {
 
                     let zalsa = db.zalsa();
                     CACHE.get_or_create(zalsa, || {
-                        zalsa.add_or_lookup_jar_by_type::<$zalsa_struct::JarImpl<$Configuration>>()
+                        zalsa.lookup_jar_by_type::<$zalsa_struct::JarImpl<$Configuration>>().get_or_create()
                     })
                 }
             }
@@ -182,7 +182,7 @@ macro_rules! setup_interned_struct {
                 type MemoIngredientMap = $zalsa::MemoIngredientSingletonIndex;
 
                 fn lookup_or_create_ingredient_index(aux: &$zalsa::Zalsa) -> $zalsa::IngredientIndices {
-                    aux.add_or_lookup_jar_by_type::<$zalsa_struct::JarImpl<$Configuration>>().into()
+                    aux.lookup_jar_by_type::<$zalsa_struct::JarImpl<$Configuration>>().get_or_create().into()
                 }
 
                 #[inline]

components/salsa-macro-rules/src/setup_tracked_fn.rs

@@ -154,15 +154,24 @@ macro_rules! setup_tracked_fn {
                 fn fn_ingredient(db: &dyn $Db) -> &$zalsa::function::IngredientImpl<$Configuration> {
                     let zalsa = db.zalsa();
                     $FN_CACHE.get_or_create(zalsa, || {
+                        let jar_entry = zalsa.lookup_jar_by_type::<$Configuration>();
+
+                        // If the ingredient has already been inserted, we know that the downcaster
+                        // has also been registered. This is a fast-path for multi-database use cases
+                        // that bypass the ingredient cache and will always execute this closure.
+                        if let Some(index) = jar_entry.get() {
+                            return index;
+                        }
+
                         <dyn $Db as $Db>::zalsa_register_downcaster(db);
-                        zalsa.add_or_lookup_jar_by_type::<$Configuration>()
+                        jar_entry.get_or_create()
                     })
                 }
 
                 pub fn fn_ingredient_mut(db: &mut dyn $Db) -> &mut $zalsa::function::IngredientImpl<Self> {
                     <dyn $Db as $Db>::zalsa_register_downcaster(db);
                     let zalsa_mut = db.zalsa_mut();
-                    let index = zalsa_mut.add_or_lookup_jar_by_type::<$Configuration>();
+                    let index = zalsa_mut.lookup_jar_by_type::<$Configuration>().get_or_create();
                     let (ingredient, _) = zalsa_mut.lookup_ingredient_mut(index);
                     ingredient.assert_type_mut::<$zalsa::function::IngredientImpl<Self>>()
                 }
@@ -174,7 +183,7 @@ macro_rules! setup_tracked_fn {
                         let zalsa = db.zalsa();
                         $INTERN_CACHE.get_or_create(zalsa, || {
                             <dyn $Db as $Db>::zalsa_register_downcaster(db);
-                            zalsa.add_or_lookup_jar_by_type::<$Configuration>().successor(0)
+                            zalsa.lookup_jar_by_type::<$Configuration>().get_or_create().successor(0)
                         })
                     }
                 }

components/salsa-macro-rules/src/setup_tracked_struct.rs

@@ -188,7 +188,7 @@ macro_rules! setup_tracked_struct {
                         $zalsa::IngredientCache::new();
 
                     CACHE.get_or_create(zalsa, || {
-                        zalsa.add_or_lookup_jar_by_type::<$zalsa_struct::JarImpl<$Configuration>>()
+                        zalsa.lookup_jar_by_type::<$zalsa_struct::JarImpl<$Configuration>>().get_or_create()
                     })
                 }
             }
@@ -211,7 +211,7 @@ macro_rules! setup_tracked_struct {
                 type MemoIngredientMap = $zalsa::MemoIngredientSingletonIndex;
 
                 fn lookup_or_create_ingredient_index(aux: &$zalsa::Zalsa) -> $zalsa::IngredientIndices {
-                    aux.add_or_lookup_jar_by_type::<$zalsa_struct::JarImpl<$Configuration>>().into()
+                    aux.lookup_jar_by_type::<$zalsa_struct::JarImpl<$Configuration>>().get_or_create().into()
                 }
 
                 #[inline]

src/accumulator.rs

@@ -64,7 +64,7 @@ pub struct IngredientImpl<A: Accumulator> {
 impl<A: Accumulator> IngredientImpl<A> {
     /// Find the accumulator ingredient for `A` in the database, if any.
     pub fn from_zalsa(zalsa: &Zalsa) -> Option<&Self> {
-        let index = zalsa.add_or_lookup_jar_by_type::<JarImpl<A>>();
+        let index = zalsa.lookup_jar_by_type::<JarImpl<A>>().get_or_create();
         let ingredient = zalsa.lookup_ingredient(index).assert_type::<Self>();
         Some(ingredient)
     }

src/hash.rs

@@ -1,4 +1,4 @@
-use std::hash::{BuildHasher, Hash};
+use std::hash::{BuildHasher, Hash, Hasher};
 
 pub(crate) type FxHasher = std::hash::BuildHasherDefault<rustc_hash::FxHasher>;
 pub(crate) type FxIndexSet<K> = indexmap::IndexSet<K, FxHasher>;
@@ -8,3 +8,24 @@ pub(crate) type FxHashSet<K> = std::collections::HashSet<K, FxHasher>;
 pub(crate) fn hash<T: Hash>(t: &T) -> u64 {
     FxHasher::default().hash_one(t)
 }
+
+// `TypeId` is a 128-bit hash internally, and it's `Hash` implementation
+// writes the lower 64-bits. Hashing it again would be unnecessary.
+#[derive(Default)]
+pub(crate) struct TypeIdHasher(u64);
+
+impl Hasher for TypeIdHasher {
+    fn write(&mut self, _: &[u8]) {
+        unreachable!("`TypeId` calls `write_u64`");
+    }
+
+    #[inline]
+    fn write_u64(&mut self, id: u64) {
+        self.0 = id;
+    }
+
+    #[inline]
+    fn finish(&self) -> u64 {
+        self.0
+    }
+}

src/salsa_struct.rs

@@ -10,12 +10,12 @@ pub trait SalsaStructInDb: Sized {
     /// Lookup or create ingredient indices.
     ///
     /// Note that this method does *not* create the ingredients themselves, this is handled by
-    /// [`Zalsa::add_or_lookup_jar_by_type()`]. This method only creates
+    /// [`crate::zalsa::JarEntry::get_or_create`]. This method only creates
     /// or looks up the indices corresponding to the ingredients.
     ///
-    /// While implementors of this trait may call [`Zalsa::add_or_lookup_jar_by_type()`]
+    /// While implementors of this trait may call [`crate::zalsa::JarEntry::get_or_create`]
     /// to create the ingredient, they aren't required to. For example, supertypes recursively
-    /// call [`Zalsa::add_or_lookup_jar_by_type()`] for their variants and combine them.
+    /// call [`crate::zalsa::JarEntry::get_or_create`] for their variants and combine them.
     fn lookup_or_create_ingredient_index(zalsa: &Zalsa) -> IngredientIndices;
 
     /// Plumbing to support nested salsa supertypes.

src/sync.rs

@@ -5,6 +5,40 @@ pub mod shim {
     pub use shuttle::sync::*;
     pub use shuttle::{thread, thread_local};
 
+    pub mod papaya {
+        use std::hash::{BuildHasher, Hash};
+        use std::marker::PhantomData;
+
+        pub struct HashMap<K, V, S>(super::Mutex<std::collections::HashMap<K, V, S>>);
+
+        impl<K, V, S: Default> Default for HashMap<K, V, S> {
+            fn default() -> Self {
+                Self(super::Mutex::default())
+            }
+        }
+
+        pub struct LocalGuard<'a>(PhantomData<&'a ()>);
+
+        impl<K, V, S> HashMap<K, V, S>
+        where
+            K: Eq + Hash,
+            V: Clone,
+            S: BuildHasher,
+        {
+            pub fn guard(&self) -> LocalGuard<'_> {
+                LocalGuard(PhantomData)
+            }
+
+            pub fn get(&self, key: &K, _guard: &LocalGuard<'_>) -> Option<V> {
+                self.0.lock().get(key).cloned()
+            }
+
+            pub fn insert(&self, key: K, value: V, _guard: &LocalGuard<'_>) {
+                self.0.lock().insert(key, value);
+            }
+        }
+    }
+
     /// A wrapper around shuttle's `Mutex` to mirror parking-lot's API.
     #[derive(Default, Debug)]
     pub struct Mutex<T>(shuttle::sync::Mutex<T>);
@@ -139,6 +173,48 @@ pub mod shim {
         pub use std::sync::atomic::*;
     }
 
+    pub mod papaya {
+        use std::hash::{BuildHasher, Hash};
+
+        pub use papaya::LocalGuard;
+
+        pub struct HashMap<K, V, S>(papaya::HashMap<K, V, S>);
+
+        impl<K, V, S: Default> Default for HashMap<K, V, S> {
+            fn default() -> Self {
+                Self(
+                    papaya::HashMap::builder()
+                        .capacity(256) // A relatively large capacity to hopefully avoid resizing.
+                        .resize_mode(papaya::ResizeMode::Blocking)
+                        .hasher(S::default())
+                        .build(),
+                )
+            }
+        }
+
+        impl<K, V, S> HashMap<K, V, S>
+        where
+            K: Eq + Hash,
+            V: Clone,
+            S: BuildHasher,
+        {
+            #[inline]
+            pub fn guard(&self) -> LocalGuard<'_> {
+                self.0.guard()
+            }
+
+            #[inline]
+            pub fn get(&self, key: &K, guard: &LocalGuard<'_>) -> Option<V> {
+                self.0.get(key, guard).cloned()
+            }
+
+            #[inline]
+            pub fn insert(&self, key: K, value: V, guard: &LocalGuard<'_>) {
+                self.0.insert(key, value, guard);
+            }
+        }
+    }
+
     /// A wrapper around parking-lot's `Condvar` to mirror shuttle's API.
     pub struct Condvar(parking_lot::Condvar);
 

src/zalsa.rs

@@ -1,5 +1,5 @@
 use std::any::{Any, TypeId};
-use std::collections::hash_map;
+use std::hash::BuildHasherDefault;
 use std::marker::PhantomData;
 use std::mem;
 use std::num::NonZeroU32;
@@ -7,11 +7,12 @@ use std::panic::RefUnwindSafe;
 
 use rustc_hash::FxHashMap;
 
+use crate::hash::TypeIdHasher;
 use crate::ingredient::{Ingredient, Jar};
 use crate::nonce::{Nonce, NonceGenerator};
 use crate::runtime::Runtime;
 use crate::sync::atomic::{AtomicU64, Ordering};
-use crate::sync::{Mutex, RwLock};
+use crate::sync::{papaya, Mutex, RwLock};
 use crate::table::memo::MemoTableWithTypes;
 use crate::table::Table;
 use crate::views::Views;
@@ -141,12 +142,10 @@ pub struct Zalsa {
     memo_ingredient_indices: RwLock<Vec<Vec<IngredientIndex>>>,
 
     /// Map from the type-id of an `impl Jar` to the index of its first ingredient.
-    /// This is using a `Mutex<FxHashMap>` (versus, say, a `FxDashMap`)
-    /// so that we can protect `ingredients_vec` as well and predict what the
-    /// first ingredient index will be. This allows ingredients to store their own indices.
-    /// This may be worth refactoring in the future because it naturally adds more overhead to
-    /// adding new kinds of ingredients.
-    jar_map: Mutex<FxHashMap<TypeId, IngredientIndex>>,
+    jar_map: papaya::HashMap<TypeId, IngredientIndex, BuildHasherDefault<TypeIdHasher>>,
+
+    /// The write-lock for `jar_map`.
+    jar_map_lock: Mutex<()>,
 
     /// A map from the `IngredientIndex` to the `TypeId` of its ID struct.
     ///
@@ -182,7 +181,8 @@ impl Zalsa {
         Self {
             views_of: Views::new::<Db>(),
             nonce: NONCE.nonce(),
-            jar_map: Default::default(),
+            jar_map: papaya::HashMap::default(),
+            jar_map_lock: Mutex::default(),
             ingredient_to_id_struct_type_id_map: Default::default(),
             ingredients_vec: boxcar::Vec::new(),
             ingredients_requiring_reset: boxcar::Vec::new(),
@@ -299,26 +299,35 @@ impl Zalsa {
     /// **NOT SEMVER STABLE**
     #[doc(hidden)]
     #[inline]
-    pub fn add_or_lookup_jar_by_type<J: Jar>(&self) -> IngredientIndex {
+    pub fn lookup_jar_by_type<J: Jar>(&self) -> JarEntry<'_, J> {
         let jar_type_id = TypeId::of::<J>();
-        if let Some(index) = self.jar_map.lock().get(&jar_type_id) {
-            return *index;
-        };
-        self.add_or_lookup_jar_by_type_slow::<J>(jar_type_id)
+        let guard = self.jar_map.guard();
+
+        match self.jar_map.get(&jar_type_id, &guard) {
+            Some(index) => JarEntry::Occupied(index),
+            None => JarEntry::Vacant {
+                guard,
+                zalsa: self,
+                _jar: PhantomData,
+            },
+        }
     }
 
+    #[cold]
     #[inline(never)]
-    fn add_or_lookup_jar_by_type_slow<J: Jar>(&self, jar_type_id: TypeId) -> IngredientIndex {
+    fn add_or_lookup_jar_by_type<J: Jar>(&self, guard: &papaya::LocalGuard<'_>) -> IngredientIndex {
+        let jar_type_id = TypeId::of::<J>();
         let dependencies = J::create_dependencies(self);
-        let mut jar_map = self.jar_map.lock();
+
+        let jar_map_lock = self.jar_map_lock.lock();
+
         let index = IngredientIndex::from(self.ingredients_vec.count());
-        match jar_map.entry(jar_type_id) {
-            hash_map::Entry::Occupied(entry) => {
-                // Someone made it earlier than us.
-                return *entry.get();
-            }
-            hash_map::Entry::Vacant(entry) => entry.insert(index),
+
+        // Someone made it earlier than us.
+        if let Some(index) = self.jar_map.get(&jar_type_id, guard) {
+            return index;
         };
+
         let ingredients = J::create_ingredients(self, index, dependencies);
         for ingredient in ingredients {
             let expected_index = ingredient.ingredient_index();
@@ -338,7 +347,12 @@ impl Zalsa {
             );
         }
 
-        drop(jar_map);
+        // Insert the index after all ingredients are inserted to avoid exposing
+        // partially initialized jars to readers.
+        self.jar_map.insert(jar_type_id, index, guard);
+
+        drop(jar_map_lock);
+
         self.ingredient_to_id_struct_type_id_map
             .write()
             .insert(index, J::id_struct_type_id());
@@ -420,6 +434,36 @@ impl Zalsa {
     }
 }
 
+pub enum JarEntry<'a, J> {
+    Occupied(IngredientIndex),
+    Vacant {
+        zalsa: &'a Zalsa,
+        guard: papaya::LocalGuard<'a>,
+        _jar: PhantomData<J>,
+    },
+}
+
+impl<J> JarEntry<'_, J>
+where
+    J: Jar,
+{
+    #[inline]
+    pub fn get(&self) -> Option<IngredientIndex> {
+        match *self {
+            JarEntry::Occupied(index) => Some(index),
+            JarEntry::Vacant { .. } => None,
+        }
+    }
+
+    #[inline]
+    pub fn get_or_create(&self) -> IngredientIndex {
+        match self {
+            JarEntry::Occupied(index) => *index,
+            JarEntry::Vacant { zalsa, guard, _jar } => zalsa.add_or_lookup_jar_by_type::<J>(guard),
+        }
+    }
+}
+
 /// Caches a pointer to an ingredient in a database.
 /// Optimized for the case of a single database.
 pub struct IngredientCache<I>
@@ -482,6 +526,7 @@ where
         );
         let cached_data = self.cached_data.load(Ordering::Acquire);
         if cached_data == Self::UNINITIALIZED {
+            #[cold]
             #[inline(never)]
             fn get_or_create_index_slow<I: Ingredient>(
                 this: &IngredientCache<I>,

tests/interned-structs_self_ref.rs

@@ -87,7 +87,9 @@ const _: () = {
 
             let zalsa = db.zalsa();
             CACHE.get_or_create(zalsa, || {
-                zalsa.add_or_lookup_jar_by_type::<zalsa_struct_::JarImpl<Configuration_>>()
+                zalsa
+                    .lookup_jar_by_type::<zalsa_struct_::JarImpl<Configuration_>>()
+                    .get_or_create()
             })
         }
     }
@@ -114,7 +116,8 @@ const _: () = {
         type MemoIngredientMap = zalsa_::MemoIngredientSingletonIndex;
 
         fn lookup_or_create_ingredient_index(aux: &Zalsa) -> salsa::plumbing::IngredientIndices {
-            aux.add_or_lookup_jar_by_type::<zalsa_struct_::JarImpl<Configuration_>>()
+            aux.lookup_jar_by_type::<zalsa_struct_::JarImpl<Configuration_>>()
+                .get_or_create()
                 .into()
         }