Fix Disambiguator- and IdentityMap hashing

Author: Veykril

src/tracked_struct.rs

@@ -187,7 +187,7 @@ impl Identity {
 }
 
 /// Stores the data that (almost) uniquely identifies a tracked struct.
-/// This includes the ingredient index of that struct type plus the hash of its id fields.
+/// This includes the ingredient index of that struct type plus the hash of its untracked fields.
 /// This is mapped to a disambiguator -- a value that starts as 0 but increments each round,
 /// allowing for multiple tracked structs with the same hash and ingredient_index
 /// created within the query to each have a unique id.
@@ -222,10 +222,7 @@ impl IdentityMap {
     pub(crate) fn insert(&mut self, key: Identity, id: Id) -> Option<Id> {
         use hashbrown::hash_map::RawEntryMut;
 
-        let eq_modulo_hash = |k: &Identity| {
-            k.ingredient_index == key.ingredient_index && k.disambiguator == key.disambiguator
-        };
-        let entry = self.map.raw_entry_mut().from_hash(key.hash, eq_modulo_hash);
+        let entry = self.map.raw_entry_mut().from_hash(key.hash, |k| *k == key);
         match entry {
             RawEntryMut::Occupied(mut occupied) => Some(occupied.insert(id)),
             RawEntryMut::Vacant(vacant) => {
@@ -236,12 +233,9 @@ impl IdentityMap {
     }
 
     pub(crate) fn get(&self, key: &Identity) -> Option<Id> {
-        let eq_modulo_hash = |k: &Identity| {
-            k.ingredient_index == key.ingredient_index && k.disambiguator == key.disambiguator
-        };
         self.map
             .raw_entry()
-            .from_hash(key.hash, eq_modulo_hash)
+            .from_hash(key.hash, |k| *k == *key)
             .map(|(_, &v)| v)
     }
 
@@ -318,8 +312,7 @@ impl DisambiguatorMap {
     pub(crate) fn disambiguate(&mut self, key: IdentityHash) -> Disambiguator {
         use hashbrown::hash_map::RawEntryMut;
 
-        let eq_modulo_hash = |k: &IdentityHash| k.ingredient_index == key.ingredient_index;
-        let entry = self.map.raw_entry_mut().from_hash(key.hash, eq_modulo_hash);
+        let entry = self.map.raw_entry_mut().from_hash(key.hash, |k| *k == key);
         let disambiguator = match entry {
             RawEntryMut::Occupied(occupied) => occupied.into_mut(),
             RawEntryMut::Vacant(vacant) => {
@@ -388,7 +381,7 @@ where
 
         let identity = Identity {
             hash: identity_hash.hash,
-            ingredient_index: self.ingredient_index,
+            ingredient_index: identity_hash.ingredient_index,
             disambiguator,
         };
 

src/update.rs

@@ -5,6 +5,8 @@ use std::{
     sync::Arc,
 };
 
+use rayon::iter::Either;
+
 use crate::Revision;
 
 /// This is used by the macro generated code.
@@ -348,6 +350,24 @@ where
     }
 }
 
+unsafe impl<L, R> Update for Either<L, R>
+where
+    L: Update,
+    R: Update,
+{
+    unsafe fn maybe_update(old_pointer: *mut Self, new_value: Self) -> bool {
+        let old_value = unsafe { &mut *old_pointer };
+        match (old_value, new_value) {
+            (Either::Left(old), Either::Left(new)) => L::maybe_update(old, new),
+            (Either::Right(old), Either::Right(new)) => R::maybe_update(old, new),
+            (old_value, new_value) => {
+                *old_value = new_value;
+                true
+            }
+        }
+    }
+}
+
 macro_rules! fallback_impl {
     ($($t:ty,)*) => {
         $(

tests/tracked_struct_disambiguates.rs

@@ -0,0 +1,107 @@
+//! Test that disambiguation works, that is when we have a revision where we track multiple structs
+//! that have the same hash, we can still differentiate between them.
+#![allow(warnings)]
+
+use std::hash::Hash;
+
+use rayon::iter::Either;
+use salsa::Setter;
+
+#[salsa::input]
+struct MyInput {
+    field: u32,
+}
+
+#[salsa::input]
+struct MyInputs {
+    field: Vec<MyInput>,
+}
+
+#[salsa::tracked]
+struct TrackedStruct<'db> {
+    field: DumbHashable,
+}
+
+#[salsa::tracked]
+struct TrackedStruct2<'db> {
+    field: DumbHashable,
+}
+
+#[derive(Debug, Clone)]
+pub struct DumbHashable {
+    field: u32,
+}
+
+impl Eq for DumbHashable {}
+impl PartialEq for DumbHashable {
+    fn eq(&self, other: &Self) -> bool {
+        self.field == other.field
+    }
+}
+
+// Force collisions, note that this is still a correct implementation wrt. PartialEq / Eq above
+// as keep the property that k1 == k2 -> hash(k1) == hash(k2)
+impl Hash for DumbHashable {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        (self.field % 3).hash(state);
+    }
+}
+
+fn alternate(
+    db: &dyn salsa::Database,
+    input: MyInput,
+) -> Either<TrackedStruct<'_>, TrackedStruct2<'_>> {
+    if input.field(db) % 2 == 0 {
+        Either::Left(TrackedStruct::new(
+            db,
+            DumbHashable {
+                field: input.field(db),
+            },
+        ))
+    } else {
+        Either::Right(TrackedStruct2::new(
+            db,
+            DumbHashable {
+                field: input.field(db),
+            },
+        ))
+    }
+}
+
+#[salsa::tracked]
+fn batch(
+    db: &dyn salsa::Database,
+    inputs: MyInputs,
+) -> Vec<Either<TrackedStruct<'_>, TrackedStruct2<'_>>> {
+    inputs
+        .field(db)
+        .iter()
+        .map(|input| alternate(db, input.clone()))
+        .collect()
+}
+
+#[test]
+fn execute() {
+    let mut db = salsa::DatabaseImpl::new();
+    let inputs = MyInputs::new(
+        &db,
+        (0..1024)
+            .into_iter()
+            .map(|i| MyInput::new(&db, i))
+            .collect(),
+    );
+    let trackeds = batch(&db, inputs);
+    for (id, tracked) in trackeds.into_iter().enumerate() {
+        assert_eq!(id % 2 == 0, tracked.is_left());
+        assert_eq!(id % 2 != 0, tracked.is_right());
+    }
+    for input in inputs.field(&db) {
+        let prev = input.field(&db);
+        input.set_field(&mut db).to(prev + 1);
+    }
+    let trackeds = batch(&db, inputs);
+    for (id, tracked) in trackeds.into_iter().enumerate() {
+        assert_eq!(id % 2 == 0, tracked.is_right());
+        assert_eq!(id % 2 != 0, tracked.is_left());
+    }
+}