Optimize salsa some more

Author: Veykril

Cargo.lock

@@ -21,6 +21,7 @@ rustc-hash = "1.0"
 smallvec = "1.0.0"
 oorandom = "11"
 triomphe = "0.1.11"
+itertools.workspace = true
 
 salsa-macros = { version = "0.0.0", path = "salsa-macros" }
 

crates/salsa/Cargo.toml

@@ -526,7 +526,7 @@ pub(crate) fn query_group(args: TokenStream, input: TokenStream) -> TokenStream
         fmt_ops.extend(quote! {
             #query_index => {
                 salsa::plumbing::QueryStorageOps::fmt_index(
-                    &*self.#fn_name, db, input, fmt,
+                    &*self.#fn_name, db, input.key_index(), fmt,
                 )
             }
         });
@@ -537,7 +537,7 @@ pub(crate) fn query_group(args: TokenStream, input: TokenStream) -> TokenStream
         maybe_changed_ops.extend(quote! {
             #query_index => {
                 salsa::plumbing::QueryStorageOps::maybe_changed_after(
-                    &*self.#fn_name, db, input, revision
+                    &*self.#fn_name, db, input.key_index(), revision
                 )
             }
         });

crates/salsa/salsa-macros/src/query_group.rs

@@ -102,13 +102,13 @@ where
 
         let mut write = self.slot_map.write();
         let entry = write.entry(key.clone());
-        let key_index = u32::try_from(entry.index()).unwrap();
+        let key_index = entry.index() as u32;
         let database_key_index = DatabaseKeyIndex {
             group_index: self.group_index,
             query_index: Q::QUERY_INDEX,
             key_index,
         };
-        entry.or_insert_with(|| Arc::new(Slot::new(key.clone(), database_key_index))).clone()
+        entry.or_insert_with(|| Arc::new(Slot::new(database_key_index))).clone()
     }
 }
 
@@ -131,34 +131,33 @@ where
     fn fmt_index(
         &self,
         _db: &<Q as QueryDb<'_>>::DynDb,
-        index: DatabaseKeyIndex,
+        index: u32,
         fmt: &mut std::fmt::Formatter<'_>,
     ) -> std::fmt::Result {
-        assert_eq!(index.group_index, self.group_index);
-        assert_eq!(index.query_index, Q::QUERY_INDEX);
         let slot_map = self.slot_map.read();
-        let key = slot_map.get_index(index.key_index as usize).unwrap().0;
+        let key = slot_map.get_index(index as usize).unwrap().0;
         write!(fmt, "{}({:?})", Q::QUERY_NAME, key)
     }
 
     fn maybe_changed_after(
         &self,
         db: &<Q as QueryDb<'_>>::DynDb,
-        input: DatabaseKeyIndex,
+        index: u32,
         revision: Revision,
     ) -> bool {
-        assert_eq!(input.group_index, self.group_index);
-        assert_eq!(input.query_index, Q::QUERY_INDEX);
         debug_assert!(revision < db.salsa_runtime().current_revision());
-        let slot = self.slot_map.read().get_index(input.key_index as usize).unwrap().1.clone();
-        slot.maybe_changed_after(db, revision)
+        let read = &self.slot_map.read();
+        let Some((key, slot)) = read.get_index(index as usize) else {
+            return false;
+        };
+        slot.maybe_changed_after(db, revision, key)
     }
 
     fn fetch(&self, db: &<Q as QueryDb<'_>>::DynDb, key: &Q::Key) -> Q::Value {
         db.unwind_if_cancelled();
 
         let slot = self.slot(key);
-        let StampedValue { value, durability, changed_at } = slot.read(db);
+        let StampedValue { value, durability, changed_at } = slot.read(db, key);
 
         if let Some(evicted) = self.lru_list.record_use(&slot) {
             evicted.evict();
@@ -182,7 +181,7 @@ where
         C: std::iter::FromIterator<TableEntry<Q::Key, Q::Value>>,
     {
         let slot_map = self.slot_map.read();
-        slot_map.values().filter_map(|slot| slot.as_table_entry()).collect()
+        slot_map.iter().filter_map(|(key, slot)| slot.as_table_entry(key)).collect()
     }
 }