create a true inverse key for the lookup path

components/salsa-macros/src/query_group.rs

@@ -3,7 +3,7 @@ use heck::CamelCase;
 use proc_macro::TokenStream;
 use proc_macro2::Span;
 use quote::ToTokens;
-use syn::{parse_macro_input, FnArg, Ident, ItemTrait, ReturnType, TraitItem};
+use syn::{parse_macro_input, parse_quote, FnArg, Ident, ItemTrait, ReturnType, TraitItem, Type};
 
 /// Implementation for `[salsa::query_group]` decorator.
 pub(crate) fn query_group(args: TokenStream, input: TokenStream) -> TokenStream {
@@ -110,15 +110,56 @@ pub(crate) fn query_group(args: TokenStream, input: TokenStream) -> TokenStream
                     ),
                 };
 
+                // For `#[salsa::interned]` keys, we create a "lookup key" automatically.
+                //
+                // For a query like:
+                //
+                //     fn foo(&self, x: Key1, y: Key2) -> u32
+                //
+                // we would create
+                //
+                //     fn lookup_foo(&self, x: u32) -> (Key1, Key2)
+                let lookup_query = if let QueryStorage::Interned = storage {
+                    let lookup_query_type = Ident::new(
+                        &format!(
+                            "{}LookupQuery",
+                            method.sig.ident.to_string().to_camel_case()
+                        ),
+                        Span::call_site(),
+                    );
+                    let lookup_fn_name = Ident::new(
+                        &format!("lookup_{}", method.sig.ident.to_string()),
+                        method.sig.ident.span(),
+                    );
+                    let keys = &keys;
+                    let lookup_value: Type = parse_quote!((#(#keys),*));
+                    let lookup_keys = vec![value.clone()];
+                    Some(Query {
+                        query_type: lookup_query_type,
+                        fn_name: lookup_fn_name,
+                        attrs: vec![], // FIXME -- some automatically generated docs on this method?
+                        storage: QueryStorage::InternedLookup {
+                            intern_query_type: query_type.clone(),
+                        },
+                        keys: lookup_keys,
+                        value: lookup_value,
+                        invoke: None,
+                    })
+                } else {
+                    None
+                };
+
                 queries.push(Query {
                     query_type,
-                    fn_name: method.sig.ident.clone(),
+                    fn_name: method.sig.ident,
                     attrs,
                     storage,
                     keys,
                     value,
                     invoke,
                 });
+
+                queries.extend(lookup_query);
             }
             _ => (),
         }
@@ -199,23 +240,6 @@ pub(crate) fn query_group(args: TokenStream, input: TokenStream) -> TokenStream
             });
         }
 
-        // For interned queries, we need `lookup_foo`
-        if let QueryStorage::Interned = query.storage {
-            let lookup_fn_name = Ident::new(&format!("lookup_{}", fn_name), fn_name.span());
-
-            query_fn_declarations.extend(quote! {
-                /// Lookup the value(s) interned with a specific key.
-                fn #lookup_fn_name(&mut self, value: #value) -> (#(#keys),*);
-            });
-
-            query_fn_definitions.extend(quote! {
-                fn #lookup_fn_name(&mut self, value: #value) -> (#(#keys),*) {
-                    <Self as salsa::plumbing::GetQueryTable<#qt>>::get_query_table(self)
-                        .lookup(value)
-                }
-            });
-        }
-
         // A variant for the group descriptor below
         query_descriptor_variants.extend(quote! {
             #fn_name((#(#keys),*)),
@@ -266,6 +290,7 @@ pub(crate) fn query_group(args: TokenStream, input: TokenStream) -> TokenStream
         impl<DB__> salsa::plumbing::QueryGroup<DB__> for #group_struct
         where
             DB__: #trait_name,
+            DB__: salsa::plumbing::HasQueryGroup<#group_struct>,
             DB__: salsa::Database,
         {
             type GroupStorage = #group_storage<DB__>;
@@ -291,16 +316,19 @@ pub(crate) fn query_group(args: TokenStream, input: TokenStream) -> TokenStream
     for query in &queries {
         let fn_name = &query.fn_name;
         let qt = &query.query_type;
-        let storage = Ident::new(
+
-            match query.storage {
+        let db = quote! {DB};
-                QueryStorage::Memoized => "MemoizedStorage",
+
-                QueryStorage::Volatile => "VolatileStorage",
+        let storage = match &query.storage {
-                QueryStorage::Dependencies => "DependencyStorage",
+            QueryStorage::Memoized => quote!(salsa::plumbing::MemoizedStorage<#db, Self>),
-                QueryStorage::Input => "InputStorage",
+            QueryStorage::Volatile => quote!(salsa::plumbing::VolatileStorage<#db, Self>),
-                QueryStorage::Interned => "InternedStorage",
+            QueryStorage::Dependencies => quote!(salsa::plumbing::DependencyStorage<#db, Self>),
-            },
+            QueryStorage::Input => quote!(salsa::plumbing::InputStorage<#db, Self>),
-            Span::call_site(),
+            QueryStorage::Interned => quote!(salsa::plumbing::InternedStorage<#db, Self>),
-        );
+            QueryStorage::InternedLookup { intern_query_type } => {
+                quote!(salsa::plumbing::LookupInternedStorage<#db, Self, #intern_query_type>)
+            }
+        };
         let keys = &query.keys;
         let value = &query.value;
 
@@ -309,16 +337,17 @@ pub(crate) fn query_group(args: TokenStream, input: TokenStream) -> TokenStream
             #[derive(Default, Debug)]
             #trait_vis struct #qt;
 
-            impl<DB> salsa::Query<DB> for #qt
+            impl<#db> salsa::Query<#db> for #qt
             where
                 DB: #trait_name,
+                DB: salsa::plumbing::HasQueryGroup<#group_struct>,
                 DB: salsa::Database,
             {
                 type Key = (#(#keys),*);
                 type Value = #value;
-                type Storage = salsa::plumbing::#storage<DB, Self>;
+                type Storage = #storage;
                 type Group = #group_struct;
-                type GroupStorage = #group_storage<DB>;
+                type GroupStorage = #group_storage<#db>;
                 type GroupKey = #group_key;
 
                 fn query_storage(group_storage: &Self::GroupStorage) -> &Self::Storage {
@@ -331,7 +360,7 @@ pub(crate) fn query_group(args: TokenStream, input: TokenStream) -> TokenStream
             }
         });
 
-        // Implement the QueryFunction trait for all queries except inputs.
+        // Implement the QueryFunction trait for queries which need it.
         if query.storage.needs_query_function() {
             let span = query.fn_name.span();
             let key_names: &Vec<_> = &(0..query.keys.len())
@@ -401,6 +430,7 @@ pub(crate) fn query_group(args: TokenStream, input: TokenStream) -> TokenStream
         #trait_vis struct #group_storage<DB__>
         where
             DB__: #trait_name,
+            DB__: salsa::plumbing::HasQueryGroup<#group_struct>,
             DB__: salsa::Database,
         {
             #storage_fields
@@ -409,6 +439,7 @@ pub(crate) fn query_group(args: TokenStream, input: TokenStream) -> TokenStream
         impl<DB__> Default for #group_storage<DB__>
         where
             DB__: #trait_name,
+            DB__: salsa::plumbing::HasQueryGroup<#group_struct>,
             DB__: salsa::Database,
         {
             #[inline]
@@ -462,19 +493,22 @@ struct Query {
     invoke: Option<syn::Path>,
 }
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+#[derive(Debug, Clone, PartialEq, Eq)]
 enum QueryStorage {
     Memoized,
     Volatile,
     Dependencies,
     Input,
     Interned,
+    InternedLookup { intern_query_type: Ident },
 }
 
 impl QueryStorage {
-    fn needs_query_function(self) -> bool {
+    fn needs_query_function(&self) -> bool {
         match self {
-            QueryStorage::Input | QueryStorage::Interned => false,
+            QueryStorage::Input | QueryStorage::Interned | QueryStorage::InternedLookup { .. } => {
+                false
+            }
             QueryStorage::Memoized | QueryStorage::Volatile | QueryStorage::Dependencies => true,
         }
     }

src/interned.rs

@@ -1,6 +1,6 @@
 use crate::debug::TableEntry;
 use crate::plumbing::CycleDetected;
-use crate::plumbing::InternedQueryStorageOps;
+use crate::plumbing::HasQueryGroup;
 use crate::plumbing::QueryStorageMassOps;
 use crate::plumbing::QueryStorageOps;
 use crate::runtime::ChangedAt;
@@ -25,6 +25,25 @@ where
     tables: RwLock<InternTables<Q::Key>>,
 }
 
+/// Storage for the looking up interned things.
+pub struct LookupInternedStorage<DB, Q, IQ>
+where
+    Q: Query<DB>,
+    Q::Key: InternKey,
+    Q::Value: Eq + Hash,
+    IQ: Query<
+        DB,
+        Key = Q::Value,
+        Value = Q::Key,
+        Group = Q::Group,
+        GroupStorage = Q::GroupStorage,
+        GroupKey = Q::GroupKey,
+    >,
+    DB: Database,
+{
+    phantom: std::marker::PhantomData<(DB, Q, IQ)>,
+}
+
 struct InternTables<K> {
     /// Map from the key to the corresponding intern-index.
     map: FxHashMap<K, InternIndex>,
@@ -133,6 +152,28 @@ where
     }
 }
 
+impl<DB, Q, IQ> Default for LookupInternedStorage<DB, Q, IQ>
+where
+    Q: Query<DB>,
+    Q::Key: InternKey,
+    Q::Value: Eq + Hash,
+    IQ: Query<
+        DB,
+        Key = Q::Value,
+        Value = Q::Key,
+        Group = Q::Group,
+        GroupStorage = Q::GroupStorage,
+        GroupKey = Q::GroupKey,
+    >,
+    DB: Database,
+{
+    fn default() -> Self {
+        LookupInternedStorage {
+            phantom: std::marker::PhantomData,
+        }
+    }
+}
+
 impl<K> Default for InternTables<K>
 where
     K: Eq + Hash,
@@ -301,7 +342,12 @@ where
 
     /// Given an index, lookup and clone its value, updating the
     /// `accessed_at` time if necessary.
-    fn lookup_value(&self, db: &DB, index: u32) -> StampedValue<Q::Key> {
+    fn lookup_value<R>(
+        &self,
+        db: &DB,
+        index: u32,
+        op: impl FnOnce(&Q::Key) -> R,
+    ) -> StampedValue<R> {
         let index = index as usize;
         let revision_now = db.salsa_runtime().current_revision();
 
@@ -315,7 +361,7 @@ where
                 } => {
                     if *accessed_at == revision_now {
                         return StampedValue {
-                            value: value.clone(),
+                            value: op(value),
                             changed_at: ChangedAt {
                                 is_constant: false,
                                 revision: *interned_at,
@@ -338,7 +384,7 @@ where
                 *accessed_at = revision_now;
 
                 return StampedValue {
-                    value: value.clone(),
+                    value: op(value),
                     changed_at: ChangedAt {
                         is_constant: false,
                         revision: *interned_at,
@@ -406,29 +452,6 @@ where
     }
 }
 
-impl<DB, Q> InternedQueryStorageOps<DB, Q> for InternedStorage<DB, Q>
-where
-    Q: Query<DB>,
-    Q::Value: InternKey,
-    DB: Database,
-{
-    fn lookup(&self, db: &DB, value: Q::Value) -> Q::Key {
-        let index: u32 = value.as_u32();
-        let StampedValue {
-            value,
-            changed_at: _,
-        } = self.lookup_value(db, index);
-
-        // XXX -- this setup is wrong, we can't report the read. We
-        // should create a *second* query that is linked to this query
-        // somehow. Or, at least, we need a distinct entry in the
-        // group key so that we can implement the "maybe changed" and
-        // all that stuff.
-
-        value
-    }
-}
-
 impl<DB, Q> QueryStorageMassOps<DB> for InternedStorage<DB, Q>
 where
     Q: Query<DB>,
@@ -468,3 +491,98 @@ where
         });
     }
 }
+
+impl<DB, Q, IQ> QueryStorageOps<DB, Q> for LookupInternedStorage<DB, Q, IQ>
+where
+    Q: Query<DB>,
+    Q::Key: InternKey,
+    Q::Value: Eq + Hash,
+    IQ: Query<
+        DB,
+        Key = Q::Value,
+        Value = Q::Key,
+        Storage = InternedStorage<DB, IQ>,
+        Group = Q::Group,
+        GroupStorage = Q::GroupStorage,
+        GroupKey = Q::GroupKey,
+    >,
+    DB: Database + HasQueryGroup<Q::Group>,
+{
+    fn try_fetch(
+        &self,
+        db: &DB,
+        key: &Q::Key,
+        database_key: &DB::DatabaseKey,
+    ) -> Result<Q::Value, CycleDetected> {
+        let index: u32 = key.as_u32();
+
+        let group_storage = <DB as HasQueryGroup<Q::Group>>::group_storage(db);
+        let interned_storage = IQ::query_storage(group_storage);
+        let StampedValue { value, changed_at } =
+            interned_storage.lookup_value(db, index, Clone::clone);
+
+        db.salsa_runtime()
+            .report_query_read(database_key, changed_at);
+
+        Ok(value)
+    }
+
+    fn maybe_changed_since(
+        &self,
+        db: &DB,
+        revision: Revision,
+        key: &Q::Key,
+        _database_key: &DB::DatabaseKey,
+    ) -> bool {
+        let index: u32 = key.as_u32();
+
+        // FIXME -- This seems maybe not quite right, as it will panic
+        // if `key` has been removed from the map since, but it should
+        // return true in that event.
+
+        let group_storage = <DB as HasQueryGroup<Q::Group>>::group_storage(db);
+        let interned_storage = IQ::query_storage(group_storage);
+        let StampedValue {
+            value: (),
+            changed_at,
+        } = interned_storage.lookup_value(db, index, |_| ());
+
+        changed_at.changed_since(revision)
+    }
+
+    fn is_constant(&self, _db: &DB, _key: &Q::Key) -> bool {
+        false
+    }
+
+    fn entries<C>(&self, db: &DB) -> C
+    where
+        C: std::iter::FromIterator<TableEntry<Q::Key, Q::Value>>,
+    {
+        let group_storage = <DB as HasQueryGroup<Q::Group>>::group_storage(db);
+        let interned_storage = IQ::query_storage(group_storage);
+        let tables = interned_storage.tables.read();
+        tables
+            .map
+            .iter()
+            .map(|(key, index)| TableEntry::new(<Q::Key>::from_u32(index.index), Some(key.clone())))
+            .collect()
+    }
+}
+
+impl<DB, Q, IQ> QueryStorageMassOps<DB> for LookupInternedStorage<DB, Q, IQ>
+where
+    Q: Query<DB>,
+    Q::Key: InternKey,
+    Q::Value: Eq + Hash,
+    IQ: Query<
+        DB,
+        Key = Q::Value,
+        Value = Q::Key,
+        Group = Q::Group,
+        GroupStorage = Q::GroupStorage,
+        GroupKey = Q::GroupKey,
+    >,
+    DB: Database,
+{
+    fn sweep(&self, _db: &DB, _strategy: SweepStrategy) {}
+}

src/lib.rs

@@ -21,7 +21,6 @@ pub mod plumbing;
 
 use crate::plumbing::CycleDetected;
 use crate::plumbing::InputQueryStorageOps;
-use crate::plumbing::InternedQueryStorageOps;
 use crate::plumbing::QueryStorageMassOps;
 use crate::plumbing::QueryStorageOps;
 use derive_new::new;
@@ -465,17 +464,6 @@ where
             })
     }
 
-    /// For `#[salsa::interned]` queries only, does a reverse lookup
-    /// from the interned value (which must be some newtype'd integer)
-    /// to get the key that was interned.
-    pub fn lookup(&self, value: Q::Value) -> Q::Key
-    where
-        Q::Storage: plumbing::InternedQueryStorageOps<DB, Q>,
-        Q::Value: InternKey,
-    {
-        self.storage.lookup(self.db, value)
-    }
-
     /// Remove all values for this query that have not been used in
     /// the most recent revision.
     pub fn sweep(&self, strategy: SweepStrategy)

src/plumbing.rs

@@ -14,6 +14,7 @@ pub use crate::derived::MemoizedStorage;
 pub use crate::derived::VolatileStorage;
 pub use crate::input::InputStorage;
 pub use crate::interned::InternedStorage;
+pub use crate::interned::LookupInternedStorage;
 pub use crate::runtime::Revision;
 
 pub struct CycleDetected;
@@ -184,15 +185,6 @@ where
         C: std::iter::FromIterator<TableEntry<Q::Key, Q::Value>>;
 }
 
-/// An optional trait that is implemented for "interned" storage.
-pub trait InternedQueryStorageOps<DB, Q>: Default
-where
-    DB: Database,
-    Q: Query<DB>,
-{
-    fn lookup(&self, db: &DB, value: Q::Value) -> Q::Key;
-}
-
 /// An optional trait that is implemented for "user mutable" storage:
 /// that is, storage whose value is not derived from other storage but
 /// is set independently.

tests/interned.rs

@@ -39,7 +39,7 @@ impl salsa::InternKey for InternKey {
 
 #[test]
 fn test_intern1() {
-    let mut db = Database::default();
+    let db = Database::default();
     let foo0 = db.intern1(format!("foo"));
     let bar0 = db.intern1(format!("bar"));
     let foo1 = db.intern1(format!("foo"));
@@ -55,7 +55,7 @@ fn test_intern1() {
 
 #[test]
 fn test_intern2() {
-    let mut db = Database::default();
+    let db = Database::default();
     let foo0 = db.intern2(format!("x"), format!("foo"));
     let bar0 = db.intern2(format!("x"), format!("bar"));
     let foo1 = db.intern2(format!("x"), format!("foo"));
@@ -71,7 +71,7 @@ fn test_intern2() {
 
 #[test]
 fn test_intern_key() {
-    let mut db = Database::default();
+    let db = Database::default();
     let foo0 = db.intern_key(format!("foo"));
     let bar0 = db.intern_key(format!("bar"));
     let foo1 = db.intern_key(format!("foo"));