Repository: salsa-rs/salsa (rust)

Issue:
Support tracked associated functions
In https://github.com/salsa-rs/salsa/issues/319 @nikomatsakis motivated the introduction of support for tracked methods (i.e. associated functions with a `self` argument).

I would like to propose extending this to type-level associated functions (i.e. without `self` argument) as well.

Tracked associated functions would have to follow the same function signature rules as their equivalent free functions.

As such the following:

```rust
#[salsa::input]
struct MyInput { ... }

#[salsa::tracked]
struct MyOutput { ... }

#[salsa::tracked]
fn output_from_input(db: &dyn crate::Db, input: MyInput) -> MyOutput { ... }

// ...

let input = MyInput::new(db, …);
let output: MyOutput = output_from_input(db, input);
```

… could be rewritten as such:

```rust
#[salsa::input]
struct MyInput { ... }

#[salsa::tracked]
struct MyOutput { ... }

#[salsa::tracked]
impl MyOutput {
    #[salsa::tracked]
    fn from_input(db: &dyn crate::Db, input: MyInput) -> Self { ... }
}

// ...

let input = MyInput::new(db, …);
let output = MyOutput::from_input(db, input);
```

While I'm using a "constructor" function here as a motivational example, they are obviously merely one of may use cases for tracked associated functions.

As such support for tracked associated functions would allow for neatly grouping tracked functions into (even untracked, and otherwise unrelated) types:

```rust
#[salsa::input]
struct MyInput { ... }

#[salsa::tracked]
struct MyOutput { ... }

struct OutputGenerator;

#[salsa::tracked]
impl OutputGenerator {
    #[salsa::tracked]
    fn output_from_input(db: &dyn crate::Db, input: MyInput) -> MyOutput { ... }
}

// ...

let input = MyInput::new(db, …);
let output: MyOutput = OutputGenerator::output(db, input);
```

The `OutputGenerator` type kind of like acts like a "namespace" type here, but with additional benefits (for those who need id), such as providing the tracked functions access to its associated types (a major limitation of free functions in a "namespace" module).

---

I'd argue that tracked associated functions would not only fill a gap in salsa's current expressivity, they also tend to feel more idiomatic and and support for semantic grouping (beyond merely showing functions in a common module) makes the code look more cleaned up overall.

Pull request description:
Allow creation of tracked associated functions (without `self`)
Fixes #858 & #585.

This allows users to annotate impl blocks with `#[salsa::tracked]` and then create tracked associated type-level functions by marking individual functions with `#[salsa::tracked]`.

Note this requires your code that looks like:

```rust
#[salsa::tracked]
impl MyOutput {
    #[salsa::tracked]
    fn tracked_fn(db: &dyn Db, input: MyInput) -> MyOutput {
        MyOutput::new(db, input.field(db) * 2)
    }
}
```

You get an error if you annotate an associated function with `#[salsa::tracked]` but forget to mark the impl block.


Patch:
diff --git a/components/salsa-macro-rules/src/lib.rs b/components/salsa-macro-rules/src/lib.rs
index 45f75355..897ff4cc 100644
--- 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
new file mode 100644
index 00000000..5b3788bf
--- /dev/null
+++ 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_method_body.rs b/components/salsa-macro-rules/src/setup_tracked_method_body.rs
similarity index 96%
rename from components/salsa-macro-rules/src/setup_method_body.rs
rename to components/salsa-macro-rules/src/setup_tracked_method_body.rs
index 3d900b2f..39a23167 100644
--- a/components/salsa-macro-rules/src/setup_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
index 256eae3f..fbbf93c3 100644
--- 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
index f12cf3fc..7009f360 100644
--- 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;


Test patch:
diff --git a/tests/tracked_assoc_fn.rs b/tests/tracked_assoc_fn.rs
new file mode 100644
index 00000000..c1f1ca67
--- /dev/null
+++ 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)) })",
+        ]"#]]);
+}