Merge pull request #443 from Veykril/self-span

Preserve span of receiver for query trait functions
salsa-rs · Mar 11, 2024 · 4151b09 · 4151b09
2 parents c38077f + 0038ae2
commit 4151b09
Showing 1 changed file with 9 additions and 5 deletions.
diff --git a/components/salsa-macros/src/query_group.rs b/components/salsa-macros/src/query_group.rs
@@ -109,8 +109,8 @@ pub(crate) fn query_group(args: TokenStream, input: TokenStream) -> TokenStream
 
             // Extract keys.
             let mut iter = method.sig.inputs.iter();
-            match iter.next() {
-                Some(FnArg::Receiver(sr)) if sr.mutability.is_none() => (),
+            let self_receiver = match iter.next() {
+                Some(FnArg::Receiver(sr)) if sr.mutability.is_none() => sr,
                 _ => {
                     return Error::new(
                         sig_span,
@@ -119,7 +119,7 @@ pub(crate) fn query_group(args: TokenStream, input: TokenStream) -> TokenStream
                     .to_compile_error()
                     .into();
                 }
-            }
+            };
             let mut keys: Vec<(Ident, Type)> = vec![];
             for (idx, arg) in iter.enumerate() {
                 match arg {
@@ -176,6 +176,7 @@ pub(crate) fn query_group(args: TokenStream, input: TokenStream) -> TokenStream
                     query_type: lookup_query_type,
                     query_name: format!("{}", lookup_fn_name),
                     fn_name: lookup_fn_name,
+                    receiver: self_receiver.clone(),
                     attrs: vec![], // FIXME -- some automatically generated docs on this method?
                     storage: QueryStorage::InternedLookup {
                         intern_query_type: query_type.clone(),
@@ -193,6 +194,7 @@ pub(crate) fn query_group(args: TokenStream, input: TokenStream) -> TokenStream
                 query_type,
                 query_name,
                 fn_name: method.sig.ident,
+                receiver: self_receiver.clone(),
                 attrs,
                 storage,
                 keys,
@@ -218,10 +220,11 @@ pub(crate) fn query_group(args: TokenStream, input: TokenStream) -> TokenStream
         let fn_name = &query.fn_name;
         let qt = &query.query_type;
         let attrs = &query.attrs;
+        let self_receiver = &query.receiver;
 
         query_fn_declarations.extend(quote! {
             #(#attrs)*
-            fn #fn_name(&self, #(#key_names: #keys),*) -> #value;
+            fn #fn_name(#self_receiver, #(#key_names: #keys),*) -> #value;
         });
 
         // Special case: transparent queries don't create actual storage,
@@ -277,7 +280,7 @@ pub(crate) fn query_group(args: TokenStream, input: TokenStream) -> TokenStream
                 specific durability instead of the default of
                 `Durability::LOW`. You can use `Durability::MAX`
                 to promise that its value will never change again.
- 
+
                 See `{fn_name}` for details.
 
                 *Note:* Setting values will trigger cancellation
@@ -720,6 +723,7 @@ fn filter_attrs(attrs: Vec<Attribute>) -> (Vec<Attribute>, Vec<SalsaAttr>) {
 #[derive(Debug)]
 struct Query {
     fn_name: Ident,
+    receiver: syn::Receiver,
     query_name: String,
     attrs: Vec<syn::Attribute>,
     query_type: Ident,