Allow creation of tracked associated functions (without self) (#859)

regexident · web-flow · commit e859218dea87 · 2025-05-23T10:19:12.000Z
* Rename file "setup_method_body.rs" to "setup_tracked_method_body.rs"

* Rename macro `setup_method_body!` to `setup_tracked_method_body!`

* Add support for tracked associated functions

* Rename struct `MethodArguments` to `AssociatedFunctionArguments`

(since every method is an associated function, but not every associated function is a method)
diff --git a/components/salsa-macro-rules/src/lib.rs b/components/salsa-macro-rules/src/lib.rs
@@ -19,7 +19,8 @@ mod return_mode;
 mod setup_accumulator_impl;
 mod setup_input_struct;
 mod setup_interned_struct;
-mod setup_method_body;
+mod setup_tracked_assoc_fn_body;
 mod setup_tracked_fn;
+mod setup_tracked_method_body;
 mod setup_tracked_struct;
 mod unexpected_cycle_recovery;
diff --git a/components/salsa-macro-rules/src/setup_tracked_assoc_fn_body.rs b/components/salsa-macro-rules/src/setup_tracked_assoc_fn_body.rs
@@ -0,0 +1,39 @@
+#[macro_export]
+macro_rules! setup_tracked_assoc_fn_body {
+    (
+        salsa_tracked_attr: #[$salsa_tracked_attr:meta],
+        self_ty: $self_ty:ty,
+        db_lt: $($db_lt:lifetime)?,
+        db: $db:ident,
+        db_ty: ($($db_ty:tt)*),
+        input_ids: [$($input_id:ident),*],
+        input_tys: [$($input_ty:ty),*],
+        output_ty: $output_ty:ty,
+        inner_fn_name: $inner_fn_name:ident,
+        inner_fn: $inner_fn:item,
+
+        // Annoyingly macro-rules hygiene does not extend to items defined in the macro.
+        // We have the procedural macro generate names for those items that are
+        // not used elsewhere in the user's code.
+        unused_names: [
+            $InnerTrait:ident,
+        ]
+    ) => {
+        {
+            trait $InnerTrait<$($db_lt)?> {
+                fn $inner_fn_name(db: $($db_ty)*, $($input_id: $input_ty),*) -> $output_ty;
+            }
+
+            impl<$($db_lt)?> $InnerTrait<$($db_lt)?> for $self_ty {
+                $inner_fn
+            }
+
+            #[$salsa_tracked_attr]
+            fn $inner_fn_name<$($db_lt)?>(db: $($db_ty)*, $($input_id: $input_ty),*) -> $output_ty {
+                <$self_ty as $InnerTrait>::$inner_fn_name(db, $($input_id),*)
+            }
+
+            $inner_fn_name($db, $($input_id),*)
+        }
+    };
+}
diff --git a/components/salsa-macro-rules/src/setup_tracked_method_body.rs b/components/salsa-macro-rules/src/setup_tracked_method_body.rs
@@ -1,5 +1,5 @@
 #[macro_export]
-macro_rules! setup_method_body {
+macro_rules! setup_tracked_method_body {
     (
         salsa_tracked_attr: #[$salsa_tracked_attr:meta],
         self: $self:ident,
diff --git a/components/salsa-macros/src/tracked_impl.rs b/components/salsa-macros/src/tracked_impl.rs
@@ -24,8 +24,8 @@ struct Macro {
     hygiene: Hygiene,
 }
 
-struct MethodArguments<'syn> {
-    self_token: &'syn syn::token::SelfValue,
+struct AssociatedFunctionArguments<'syn> {
+    self_token: Option<&'syn syn::token::SelfValue>,
     db_ty: &'syn syn::Type,
     db_ident: &'syn syn::Ident,
     db_lt: Option<&'syn syn::Lifetime>,
@@ -92,7 +92,7 @@ impl Macro {
         let InnerTrait = self.hygiene.ident("InnerTrait");
         let inner_fn_name = self.hygiene.ident(&fn_item.sig.ident.to_string());
 
-        let MethodArguments {
+        let AssociatedFunctionArguments {
             self_token,
             db_ty,
             db_ident,
@@ -106,30 +106,54 @@ impl Macro {
         inner_fn.vis = syn::Visibility::Inherited;
         inner_fn.sig.ident = inner_fn_name.clone();
 
-        // Construct the body of the method
-
-        let block = parse_quote!({
-            salsa::plumbing::setup_method_body! {
-                salsa_tracked_attr: #salsa_tracked_attr,
-                self: #self_token,
-                self_ty: #self_ty,
-                db_lt: #db_lt,
-                db: #db_ident,
-                db_ty: (#db_ty),
-                input_ids: [#(#input_ids),*],
-                input_tys: [#(#input_tys),*],
-                output_ty: #output_ty,
-                inner_fn_name: #inner_fn_name,
-                inner_fn: #inner_fn,
-
-                // Annoyingly macro-rules hygiene does not extend to items defined in the macro.
-                // We have the procedural macro generate names for those items that are
-                // not used elsewhere in the user's code.
-                unused_names: [
-                    #InnerTrait,
-                ]
-            }
-        });
+        // Construct the body of the method or associated function
+
+        let block = if let Some(self_token) = self_token {
+            parse_quote!({
+                salsa::plumbing::setup_tracked_method_body! {
+                    salsa_tracked_attr: #salsa_tracked_attr,
+                    self: #self_token,
+                    self_ty: #self_ty,
+                    db_lt: #db_lt,
+                    db: #db_ident,
+                    db_ty: (#db_ty),
+                    input_ids: [#(#input_ids),*],
+                    input_tys: [#(#input_tys),*],
+                    output_ty: #output_ty,
+                    inner_fn_name: #inner_fn_name,
+                    inner_fn: #inner_fn,
+
+                    // Annoyingly macro-rules hygiene does not extend to items defined in the macro.
+                    // We have the procedural macro generate names for those items that are
+                    // not used elsewhere in the user's code.
+                    unused_names: [
+                        #InnerTrait,
+                    ]
+                }
+            })
+        } else {
+            parse_quote!({
+                salsa::plumbing::setup_tracked_assoc_fn_body! {
+                    salsa_tracked_attr: #salsa_tracked_attr,
+                    self_ty: #self_ty,
+                    db_lt: #db_lt,
+                    db: #db_ident,
+                    db_ty: (#db_ty),
+                    input_ids: [#(#input_ids),*],
+                    input_tys: [#(#input_tys),*],
+                    output_ty: #output_ty,
+                    inner_fn_name: #inner_fn_name,
+                    inner_fn: #inner_fn,
+
+                    // Annoyingly macro-rules hygiene does not extend to items defined in the macro.
+                    // We have the procedural macro generate names for those items that are
+                    // not used elsewhere in the user's code.
+                    unused_names: [
+                        #InnerTrait,
+                    ]
+                }
+            })
+        };
 
         // Update the method that will actually appear in the impl to have the new body
         // and its true return type
@@ -144,18 +168,25 @@ impl Macro {
         &self,
         impl_item: &'syn syn::ItemImpl,
         fn_item: &'syn syn::ImplItemFn,
-    ) -> syn::Result<MethodArguments<'syn>> {
+    ) -> syn::Result<AssociatedFunctionArguments<'syn>> {
         let db_lt = self.extract_db_lifetime(impl_item, fn_item)?;
 
-        let self_token = self.check_self_argument(fn_item)?;
+        let is_method = matches!(&fn_item.sig.inputs[0], syn::FnArg::Receiver(_));
+
+        let (self_token, db_input_index, skipped_inputs) = if is_method {
+            (Some(self.check_self_argument(fn_item)?), 1, 2)
+        } else {
+            (None, 0, 1)
+        };
 
-        let (db_ident, db_ty) = self.check_db_argument(&fn_item.sig.inputs[1])?;
+        let (db_ident, db_ty) = self.check_db_argument(&fn_item.sig.inputs[db_input_index])?;
 
-        let input_ids: Vec<syn::Ident> = crate::fn_util::input_ids(&self.hygiene, &fn_item.sig, 2);
-        let input_tys = crate::fn_util::input_tys(&fn_item.sig, 2)?;
+        let input_ids: Vec<syn::Ident> =
+            crate::fn_util::input_ids(&self.hygiene, &fn_item.sig, skipped_inputs);
+        let input_tys = crate::fn_util::input_tys(&fn_item.sig, skipped_inputs)?;
         let output_ty = crate::fn_util::output_ty(db_lt, &fn_item.sig)?;
 
-        Ok(MethodArguments {
+        Ok(AssociatedFunctionArguments {
             self_token,
             db_ident,
             db_lt,
diff --git a/src/lib.rs b/src/lib.rs
@@ -76,8 +76,8 @@ pub mod plumbing {
     pub use salsa_macro_rules::{
         macro_if, maybe_backdate, maybe_default, maybe_default_tt, return_mode_expression,
         return_mode_ty, setup_accumulator_impl, setup_input_struct, setup_interned_struct,
-        setup_method_body, setup_tracked_fn, setup_tracked_struct, unexpected_cycle_initial,
-        unexpected_cycle_recovery,
+        setup_tracked_assoc_fn_body, setup_tracked_fn, setup_tracked_method_body,
+        setup_tracked_struct, unexpected_cycle_initial, unexpected_cycle_recovery,
     };
 
     pub use crate::__maybe_lazy_static;
diff --git a/tests/tracked_assoc_fn.rs b/tests/tracked_assoc_fn.rs
@@ -0,0 +1,85 @@
+//! Test that a `tracked` fn on a `salsa::input`
+//! compiles and executes successfully.
+#![allow(warnings)]
+
+use common::LogDatabase as _;
+use expect_test::expect;
+
+mod common;
+
+trait TrackedTrait<'db> {
+    type Output;
+
+    fn tracked_trait_fn(db: &'db dyn salsa::Database, input: MyInput) -> Self::Output;
+}
+
+#[salsa::input]
+struct MyInput {
+    field: u32,
+}
+
+#[salsa::tracked]
+struct MyOutput<'db> {
+    field: u32,
+}
+
+#[salsa::tracked]
+impl MyInput {
+    #[salsa::tracked]
+    fn tracked_fn(db: &dyn salsa::Database, input: MyInput) -> Self {
+        Self::new(db, 2 * input.field(db))
+    }
+
+    #[salsa::tracked(returns(ref))]
+    fn tracked_fn_ref(db: &dyn salsa::Database, input: MyInput) -> Self {
+        Self::new(db, 3 * input.field(db))
+    }
+}
+
+#[salsa::tracked]
+impl<'db> TrackedTrait<'db> for MyOutput<'db> {
+    type Output = Self;
+
+    #[salsa::tracked]
+    fn tracked_trait_fn(db: &'db dyn salsa::Database, input: MyInput) -> Self::Output {
+        Self::new(db, 4 * input.field(db))
+    }
+}
+
+// The self-type of a tracked impl doesn't have to be tracked itself:
+struct UntrackedHelper;
+
+#[salsa::tracked]
+impl<'db> TrackedTrait<'db> for UntrackedHelper {
+    type Output = MyOutput<'db>;
+
+    #[salsa::tracked]
+    fn tracked_trait_fn(db: &'db dyn salsa::Database, input: MyInput) -> Self::Output {
+        MyOutput::tracked_trait_fn(db, input)
+    }
+}
+
+#[test]
+fn execute() {
+    let mut db = salsa::DatabaseImpl::new();
+    let input = MyInput::new(&db, 22);
+    let output = MyOutput::tracked_trait_fn(&db, input);
+    let helper_output = UntrackedHelper::tracked_trait_fn(&db, input);
+    // assert_eq!(object.tracked_fn(&db), 44);
+    // assert_eq!(*object.tracked_fn_ref(&db), 66);
+    assert_eq!(output.field(&db), 88);
+    assert_eq!(helper_output.field(&db), 88);
+}
+
+#[test]
+fn debug_name() {
+    let mut db = common::ExecuteValidateLoggerDatabase::default();
+    let input = MyInput::new(&db, 22);
+    let output = MyOutput::tracked_trait_fn(&db, input);
+
+    assert_eq!(output.field(&db), 88);
+    db.assert_logs(expect![[r#"
+        [
+            "salsa_event(WillExecute { database_key: tracked_trait_fn_(Id(0)) })",
+        ]"#]]);
+}

Original file line number	Diff line number	Diff line change
`@@ -1,5 +1,5 @@`
`1`	`1`	`#[macro_export]`
`2`		`-macro_rules! setup_method_body {`
	`2`	`+macro_rules! setup_tracked_method_body {`
`3`	`3`	`(`
`4`	`4`	`salsa_tracked_attr: #[$salsa_tracked_attr:meta],`
`5`	`5`	`self: $self:ident,`