Shard jar map

Cargo.toml

@@ -12,10 +12,13 @@ description = "A generic framework for on-demand, incrementalized computation (e
 salsa-macro-rules = { version = "0.22.0", path = "components/salsa-macro-rules" }
 salsa-macros = { version = "0.22.0", path = "components/salsa-macros", optional = true }
 
-boxcar = "0.2.13"
+boxcar = { git = "https://github.com/ibraheemdev/boxcar", rev = "574c893f0a6d8b2cde17dd9933fd9e3d74ea8e0f" }
 crossbeam-queue = "0.3.11"
 crossbeam-utils = "0.8.21"
+dashmap = { version = "6", features = ["raw-api"] }
 hashbrown = "0.15"
+# The version of hashbrown used by dashmap.
+hashbrown_14 = { package = "hashbrown", version = "0.14" }
 hashlink = "0.10"
 indexmap = "2"
 intrusive-collections = "0.9.7"
@@ -51,7 +54,6 @@ salsa-macros = { version = "=0.22.0", path = "components/salsa-macros" }
 [dev-dependencies]
 # examples
 crossbeam-channel = "0.5.14"
-dashmap = { version = "6", features = ["raw-api"] }
 eyre = "0.6.8"
 notify-debouncer-mini = "0.4.1"
 ordered-float = "4.2.1"

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

@@ -242,6 +242,16 @@ macro_rules! setup_tracked_fn {
                     }
                 }
 
+                fn ingredients_count() -> usize {
+                    $zalsa::macro_if! {
+                        if $needs_interner {
+                            2
+                        } else {
+                            1
+                        }
+                    }
+                }
+
                 fn create_ingredients(
                     zalsa: &$zalsa::Zalsa,
                     first_index: $zalsa::IngredientIndex,

src/accumulator.rs

@@ -43,6 +43,10 @@ impl<A: Accumulator> Default for JarImpl<A> {
 }
 
 impl<A: Accumulator> Jar for JarImpl<A> {
+    fn ingredients_count() -> usize {
+        1
+    }
+
     fn create_ingredients(
         _zalsa: &Zalsa,
         first_index: IngredientIndex,

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/ingredient.rs

@@ -30,8 +30,15 @@ pub trait Jar: Any {
         IngredientIndices::empty()
     }
 
+    /// Returns the number of ingredients that `create_ingredients` will return.
+    fn ingredients_count() -> usize
+    where
+        Self: Sized;
+
     /// Create the ingredients given the index of the first one.
     /// All subsequent ingredients will be assigned contiguous indices.
+    ///
+    /// Note that the vector returned must be of length `ingredients_count`.
     fn create_ingredients(
         zalsa: &Zalsa,
         first_index: IngredientIndex,

src/input.rs

@@ -55,6 +55,10 @@ impl<C: Configuration> Default for JarImpl<C> {
 }
 
 impl<C: Configuration> Jar for JarImpl<C> {
+    fn ingredients_count() -> usize {
+        1 + C::FIELD_DEBUG_NAMES.len()
+    }
+
     fn create_ingredients(
         _zalsa: &Zalsa,
         struct_index: crate::zalsa::IngredientIndex,

src/interned.rs

@@ -201,6 +201,10 @@ impl<C: Configuration> Default for JarImpl<C> {
 }
 
 impl<C: Configuration> Jar for JarImpl<C> {
+    fn ingredients_count() -> usize {
+        1
+    }
+
     fn create_ingredients(
         _zalsa: &Zalsa,
         first_index: IngredientIndex,

src/memo_ingredient_indices.rs

@@ -9,7 +9,7 @@ use crate::{Id, IngredientIndex};
 /// be viewed as a *set* of [`IngredientIndex`], where each instance of the enum can belong
 /// to one, potentially different, index. This is what this type represents: a set of
 /// `IngredientIndex`.
-#[derive(Clone)]
+#[derive(Clone, Default)]
 pub struct IngredientIndices {
     indices: Box<[IngredientIndex]>,
 }

src/sync.rs

@@ -5,6 +5,38 @@ pub mod shim {
     pub use shuttle::sync::*;
     pub use shuttle::{thread, thread_local};
 
+    /// A polyfill for `dashmap::DashMap`.
+    pub struct DashMap<K, V, H>(RwLock<HashTable<K, V>>, H);
+
+    type HashTable<K, V> = hashbrown_14::raw::RawTable<(K, dashmap::SharedValue<V>)>;
+
+    impl<K, V, H> Default for DashMap<K, V, H>
+    where
+        H: Default,
+    {
+        fn default() -> DashMap<K, V, H> {
+            DashMap(RwLock::default(), H::default())
+        }
+    }
+
+    impl<K, V, H> DashMap<K, V, H> {
+        pub fn shards(&self) -> &[RwLock<HashTable<K, V>>] {
+            std::slice::from_ref(&self.0)
+        }
+
+        pub fn determine_shard(&self, _hash: usize) -> usize {
+            0
+        }
+
+        pub fn hasher(&self) -> &H {
+            &self.1
+        }
+
+        pub fn clear(&self) {
+            self.0.write().clear();
+        }
+    }
+
     /// A wrapper around shuttle's `Mutex` to mirror parking-lot's API.
     #[derive(Default, Debug)]
     pub struct Mutex<T>(shuttle::sync::Mutex<T>);
@@ -130,6 +162,7 @@ pub mod shim {
 
 #[cfg(not(feature = "shuttle"))]
 pub mod shim {
+    pub use dashmap::DashMap;
     pub use parking_lot::{Mutex, MutexGuard, RwLock};
     pub use std::sync::*;
     pub use std::{thread, thread_local};

src/tracked_struct.rs

@@ -109,6 +109,10 @@ impl<C: Configuration> Default for JarImpl<C> {
 }
 
 impl<C: Configuration> Jar for JarImpl<C> {
+    fn ingredients_count() -> usize {
+        1 + C::TRACKED_FIELD_INDICES.len()
+    }
+
     fn create_ingredients(
         _zalsa: &Zalsa,
         struct_index: crate::zalsa::IngredientIndex,

src/zalsa.rs

@@ -1,17 +1,19 @@
 use std::any::{Any, TypeId};
-use std::collections::hash_map;
+use std::hash::{BuildHasher, BuildHasherDefault};
 use std::marker::PhantomData;
 use std::mem;
 use std::num::NonZeroU32;
 use std::panic::RefUnwindSafe;
 
+use dashmap::SharedValue;
 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::{DashMap, RwLock};
 use crate::table::memo::MemoTableWithTypes;
 use crate::table::Table;
 use crate::views::Views;
@@ -141,12 +143,7 @@ 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: DashMap<TypeId, IngredientIndex, BuildHasherDefault<TypeIdHasher>>,
 
     /// A map from the `IngredientIndex` to the `TypeId` of its ID struct.
     ///
@@ -294,44 +291,83 @@ impl Zalsa {
     #[inline]
     pub fn add_or_lookup_jar_by_type<J: Jar>(&self) -> IngredientIndex {
         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 jar_hash = self.jar_map.hasher().hash_one(jar_type_id);
+        let shard = self.jar_map.determine_shard(jar_hash as usize);
+
+        {
+            let jar_map = self.jar_map.shards()[shard].read();
+            if let Some((_, index)) = jar_map.get(jar_hash, |&(key, _)| key == jar_type_id) {
+                return *index.get();
+            };
+        }
+
+        self.add_or_lookup_jar_by_type_slow::<J>(jar_type_id, jar_hash, shard)
     }
 
+    #[cold]
     #[inline(never)]
-    fn add_or_lookup_jar_by_type_slow<J: Jar>(&self, jar_type_id: TypeId) -> IngredientIndex {
-        let dependencies = J::create_dependencies(self);
-        let mut jar_map = self.jar_map.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),
+    fn add_or_lookup_jar_by_type_slow<J: Jar>(
+        &self,
+        jar_type_id: TypeId,
+        jar_hash: u64,
+        shard: usize,
+    ) -> IngredientIndex {
+        let mut dependencies = J::create_dependencies(self);
+
+        let mut jar_map = self.jar_map.shards()[shard].write();
+
+        // Someone made it earlier than us.
+        if let Some((_, index)) = jar_map.get(jar_hash, |&(key, _)| key == jar_type_id) {
+            return *index.get();
         };
-        let ingredients = J::create_ingredients(self, index, dependencies);
-        for ingredient in ingredients {
-            let expected_index = ingredient.ingredient_index();
 
-            if ingredient.requires_reset_for_new_revision() {
-                self.ingredients_requiring_reset.push(expected_index);
-            }
+        let mut ingredients = None;
+        let index = self.ingredients_vec
+            // Claim the ingredient slots eagerly, which guarantees that the indices will be sequential.
+            .push_many(J::ingredients_count(), |index| {
+                let ingredients = ingredients.get_or_insert_with(|| {
+                    // Create the ingredients list with the first index.
+                    let ingredient_index = IngredientIndex::from(index);
+                    let ingredients = J::create_ingredients(self, ingredient_index, mem::take(&mut dependencies));
+
+                    // Create the jar.
+                    let value = (jar_type_id, SharedValue::new(ingredient_index));
+                    jar_map.insert(jar_hash, value, |(key, _)| {
+                        self.jar_map.hasher().hash_one(key)
+                    });
+
+                    ingredients.into_iter()
+                });
+
+                // Get the next ingredient from the list.
+                let ingredient = ingredients.next().expect(
+                    "`Jar::ingredients` returned less ingredients than specified by `Jar::ingredients_count`",
+                );
 
-            let actual_index = self.ingredients_vec.push(ingredient);
-            assert_eq!(
-                expected_index.as_u32() as usize,
-                actual_index,
-                "ingredient `{:?}` was predicted to have index `{:?}` but actually has index `{:?}`",
-                self.ingredients_vec[actual_index],
-                expected_index.as_u32(),
-                actual_index,
-            );
-        }
+                let expected_index = ingredient.ingredient_index();
+                if ingredient.requires_reset_for_new_revision() {
+                    self.ingredients_requiring_reset.push(expected_index);
+                }
+
+                // Ensure we are assigning the ingredient to the correct index in the vector.
+                assert_eq!(
+                    expected_index.as_u32() as usize,
+                    index,
+                    "ingredient `{:?}` was predicted to have index `{:?}` but actually has index `{:?}`",
+                    ingredient,
+                    expected_index.as_u32(),
+                    index,
+                );
 
+                ingredient
+            });
+
+        // We need to hold the shard lock until all ingredients have been initialized, to avoid
+        // returning partially-initialized jars.
         drop(jar_map);
+
+        let index = IngredientIndex::from(index);
         self.ingredient_to_id_struct_type_id_map
             .write()
             .insert(index, J::id_struct_type_id());