Optimize input queries that take no arguments

Author: Veykril

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

@@ -371,6 +371,9 @@ pub(crate) fn query_group(args: TokenStream, input: TokenStream) -> TokenStream
             QueryStorage::Dependencies => {
                 quote!(salsa::plumbing::DependencyStorage<Self>)
             }
+            QueryStorage::Input if query.keys.is_empty() => {
+                quote!(salsa::plumbing::UnitInputStorage<Self>)
+            }
             QueryStorage::Input => quote!(salsa::plumbing::InputStorage<Self>),
             QueryStorage::Interned => quote!(salsa::plumbing::InternedStorage<Self>),
             QueryStorage::InternedLookup { intern_query_type } => {

crates/salsa/src/input.rs

@@ -15,6 +15,7 @@ use crate::{DatabaseKeyIndex, QueryDb};
 use indexmap::map::Entry;
 use parking_lot::RwLock;
 use std::convert::TryFrom;
+use std::iter;
 use tracing::debug;
 
 /// Input queries store the result plus a list of the other queries
@@ -25,15 +26,12 @@ where
     Q: Query,
 {
     group_index: u16,
-    slots: RwLock<FxIndexMap<Q::Key, Slot<Q>>>,
+    slots: RwLock<FxIndexMap<Q::Key, Slot<Q::Value>>>,
 }
 
-struct Slot<Q>
-where
-    Q: Query,
-{
+struct Slot<V> {
     database_key_index: DatabaseKeyIndex,
-    stamped_value: RwLock<StampedValue<Q::Value>>,
+    stamped_value: RwLock<StampedValue<V>>,
 }
 
 impl<Q> std::panic::RefUnwindSafe for InputStorage<Q>
@@ -78,7 +76,14 @@ where
         debug_assert!(revision < db.salsa_runtime().current_revision());
         let slots = &self.slots.read();
         let slot = slots.get_index(input.key_index as usize).unwrap().1;
-        slot.maybe_changed_after(db, revision)
+
+        debug!("maybe_changed_after(slot={:?}, revision={:?})", Q::default(), revision,);
+
+        let changed_at = slot.stamped_value.read().changed_at;
+
+        debug!("maybe_changed_after: changed_at = {:?}", changed_at);
+
+        changed_at > revision
     }
 
     fn fetch(&self, db: &<Q as QueryDb<'_>>::DynDb, key: &Q::Key) -> Q::Value {
@@ -121,21 +126,6 @@ where
     }
 }
 
-impl<Q> Slot<Q>
-where
-    Q: Query,
-{
-    fn maybe_changed_after(&self, _db: &<Q as QueryDb<'_>>::DynDb, revision: Revision) -> bool {
-        debug!("maybe_changed_after(slot={:?}, revision={:?})", self, revision,);
-
-        let changed_at = self.stamped_value.read().changed_at;
-
-        debug!("maybe_changed_after: changed_at = {:?}", changed_at);
-
-        changed_at > revision
-    }
-}
-
 impl<Q> QueryStorageMassOps for InputStorage<Q>
 where
     Q: Query,
@@ -202,6 +192,167 @@ where
     }
 }
 
+/// Same as `InputStorage`, but optimized for queries that take no inputs.
+pub struct UnitInputStorage<Q>
+where
+    Q: Query<Key = ()>,
+{
+    group_index: u16,
+    slot: UnitSlot<Q::Value>,
+}
+
+struct UnitSlot<V> {
+    database_key_index: DatabaseKeyIndex,
+    stamped_value: RwLock<Option<StampedValue<V>>>,
+}
+
+impl<Q> std::panic::RefUnwindSafe for UnitInputStorage<Q>
+where
+    Q: Query<Key = ()>,
+    Q::Key: std::panic::RefUnwindSafe,
+    Q::Value: std::panic::RefUnwindSafe,
+{
+}
+
+impl<Q> QueryStorageOps<Q> for UnitInputStorage<Q>
+where
+    Q: Query<Key = ()>,
+{
+    const CYCLE_STRATEGY: crate::plumbing::CycleRecoveryStrategy = CycleRecoveryStrategy::Panic;
+
+    fn new(group_index: u16) -> Self {
+        let database_key_index =
+            DatabaseKeyIndex { group_index, query_index: Q::QUERY_INDEX, key_index: 0 };
+        UnitInputStorage {
+            group_index,
+            slot: UnitSlot { database_key_index, stamped_value: RwLock::new(None) },
+        }
+    }
+
+    fn fmt_index(
+        &self,
+        _db: &<Q as QueryDb<'_>>::DynDb,
+        index: DatabaseKeyIndex,
+        fmt: &mut std::fmt::Formatter<'_>,
+    ) -> std::fmt::Result {
+        assert_eq!(index.group_index, self.group_index);
+        assert_eq!(index.query_index, Q::QUERY_INDEX);
+        write!(fmt, "{}", Q::QUERY_NAME)
+    }
+
+    fn maybe_changed_after(
+        &self,
+        db: &<Q as QueryDb<'_>>::DynDb,
+        input: DatabaseKeyIndex,
+        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());
+
+        debug!("maybe_changed_after(slot={:?}, revision={:?})", Q::default(), revision,);
+
+        let changed_at = self.slot.stamped_value.read().as_ref().unwrap().changed_at;
+
+        debug!("maybe_changed_after: changed_at = {:?}", changed_at);
+
+        changed_at > revision
+    }
+
+    fn fetch(&self, db: &<Q as QueryDb<'_>>::DynDb, &(): &Q::Key) -> Q::Value {
+        db.unwind_if_cancelled();
+
+        let StampedValue { value, durability, changed_at } = self
+            .slot
+            .stamped_value
+            .read()
+            .clone()
+            .unwrap_or_else(|| panic!("no value set for {:?}", Q::default()));
+
+        db.salsa_runtime().report_query_read_and_unwind_if_cycle_resulted(
+            self.slot.database_key_index,
+            durability,
+            changed_at,
+        );
+
+        value
+    }
+
+    fn durability(&self, _db: &<Q as QueryDb<'_>>::DynDb, &(): &Q::Key) -> Durability {
+        match &*self.slot.stamped_value.read() {
+            Some(stamped_value) => stamped_value.durability,
+            None => panic!("no value set for {:?}", Q::default(),),
+        }
+    }
+
+    fn entries<C>(&self, _db: &<Q as QueryDb<'_>>::DynDb) -> C
+    where
+        C: std::iter::FromIterator<TableEntry<Q::Key, Q::Value>>,
+    {
+        iter::once(TableEntry::new(
+            (),
+            self.slot.stamped_value.read().as_ref().map(|it| it.value.clone()),
+        ))
+        .collect()
+    }
+}
+
+impl<Q> QueryStorageMassOps for UnitInputStorage<Q>
+where
+    Q: Query<Key = ()>,
+{
+    fn purge(&self) {
+        *self.slot.stamped_value.write() = Default::default();
+    }
+}
+
+impl<Q> InputQueryStorageOps<Q> for UnitInputStorage<Q>
+where
+    Q: Query<Key = ()>,
+{
+    fn set(&self, runtime: &mut Runtime, (): &Q::Key, value: Q::Value, durability: Durability) {
+        tracing::debug!("{:?} = {:?} ({:?})", Q::default(), value, durability);
+
+        // The value is changing, so we need a new revision (*). We also
+        // need to update the 'last changed' revision by invoking
+        // `guard.mark_durability_as_changed`.
+        //
+        // CAREFUL: This will block until the global revision lock can
+        // be acquired. If there are still queries executing, they may
+        // need to read from this input. Therefore, we wait to acquire
+        // the lock on `map` until we also hold the global query write
+        // lock.
+        //
+        // (*) Technically, since you can't presently access an input
+        // for a non-existent key, and you can't enumerate the set of
+        // keys, we only need a new revision if the key used to
+        // exist. But we may add such methods in the future and this
+        // case doesn't generally seem worth optimizing for.
+        runtime.with_incremented_revision(|next_revision| {
+            let mut stamped_value_slot = self.slot.stamped_value.write();
+
+            // Do this *after* we acquire the lock, so that we are not
+            // racing with somebody else to modify this same cell.
+            // (Otherwise, someone else might write a *newer* revision
+            // into the same cell while we block on the lock.)
+            let stamped_value = StampedValue { value, durability, changed_at: next_revision };
+
+            match &mut *stamped_value_slot {
+                Some(slot_stamped_value) => {
+                    let old_durability = slot_stamped_value.durability;
+                    *slot_stamped_value = stamped_value;
+                    Some(old_durability)
+                }
+
+                stamped_value_slot @ None => {
+                    *stamped_value_slot = Some(stamped_value);
+                    None
+                }
+            }
+        });
+    }
+}
+
 /// Check that `Slot<Q, MP>: Send + Sync` as long as
 /// `DB::DatabaseData: Send + Sync`, which in turn implies that
 /// `Q::Key: Send + Sync`, `Q::Value: Send + Sync`.
@@ -213,7 +364,8 @@ where
     Q::Value: Send + Sync,
 {
     fn is_send_sync<T: Send + Sync>() {}
-    is_send_sync::<Slot<Q>>();
+    is_send_sync::<Slot<Q::Value>>();
+    is_send_sync::<UnitSlot<Q::Value>>();
 }
 
 /// Check that `Slot<Q, MP>: 'static` as long as
@@ -227,14 +379,6 @@ where
     Q::Value: 'static,
 {
     fn is_static<T: 'static>() {}
-    is_static::<Slot<Q>>();
-}
-
-impl<Q> std::fmt::Debug for Slot<Q>
-where
-    Q: Query,
-{
-    fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(fmt, "{:?}", Q::default())
-    }
+    is_static::<Slot<Q::Value>>();
+    is_static::<UnitSlot<Q::Value>>();
 }

crates/salsa/src/plumbing.rs

@@ -15,7 +15,7 @@ use triomphe::Arc;
 
 pub use crate::derived::DependencyStorage;
 pub use crate::derived::MemoizedStorage;
-pub use crate::input::InputStorage;
+pub use crate::input::{InputStorage, UnitInputStorage};
 pub use crate::interned::InternedStorage;
 pub use crate::interned::LookupInternedStorage;
 pub use crate::{revision::Revision, DatabaseKeyIndex, QueryDb, Runtime};