shard jar map

Author: ibraheemdev

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/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,35 @@ pub mod shim {
     pub use shuttle::sync::*;
     pub use shuttle::{thread, thread_local};
 
+    /// A polyfill for `dashmap::DashMap`.
+    pub struct FxDashMap<K, V>(RwLock<HashTable<K, V>>, crate::hash::FxHasher);
+
+    type HashTable<K, V> = hashbrown_14::raw::RawTable<(K, dashmap::SharedValue<V>)>;
+
+    impl<K, V> Default for FxDashMap<K, V> {
+        fn default() -> FxDashMap<K, V> {
+            FxDashMap(RwLock::default(), crate::hash::FxHasher::default())
+        }
+    }
+
+    impl<K, V> FxDashMap<K, V> {
+        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) -> &crate::hash::FxHasher {
+            &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>);
@@ -139,6 +168,8 @@ pub mod shim {
         pub use std::sync::atomic::*;
     }
 
+    pub(crate) type FxDashMap<K, V> = dashmap::DashMap<K, V, crate::hash::FxHasher>;
+
     /// A wrapper around parking-lot's `Condvar` to mirror shuttle's API.
     pub struct Condvar(parking_lot::Condvar);
 

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,18 @@
 use std::any::{Any, TypeId};
-use std::collections::hash_map;
+use std::hash::BuildHasher;
 use std::marker::PhantomData;
 use std::mem;
 use std::num::NonZeroU32;
 use std::panic::RefUnwindSafe;
 
+use dashmap::SharedValue;
 use rustc_hash::FxHashMap;
 
 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::{FxDashMap, RwLock};
 use crate::table::memo::MemoTableWithTypes;
 use crate::table::Table;
 use crate::views::Views;
@@ -146,7 +147,7 @@ pub struct Zalsa {
     /// 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: FxDashMap<TypeId, IngredientIndex>,
 
     /// A map from the `IngredientIndex` to the `TypeId` of its ID struct.
     ///
@@ -294,44 +295,81 @@ 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;
+
+        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)
+
+        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, |(_, value)| {
+                        self.jar_map.hasher().hash_one(value.get())
+                    });
+
+                    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());