From bac4c668cfb20ad2971e244d6fe5337c651f0f17 Mon Sep 17 00:00:00 2001
From: Jack Rickard <jack.rickard@outlook.com>
Date: Sat, 3 Sep 2022 15:44:56 +0100
Subject: [PATCH] Add more tests for tracked methods

---
 .../salsa-2022-macros/src/tracked_fn.rs       | 35 ++++++++++++-------
 .../tracked_method_on_untracked_impl.rs       | 18 ++++++++++
 .../tracked_method_on_untracked_impl.stderr   | 11 ++++++
 salsa-2022-tests/tests/tracked_method.rs      | 20 ++++++++++-
 4 files changed, 71 insertions(+), 13 deletions(-)
 create mode 100644 salsa-2022-tests/tests/compile-fail/tracked_method_on_untracked_impl.rs
 create mode 100644 salsa-2022-tests/tests/compile-fail/tracked_method_on_untracked_impl.stderr

diff --git a/components/salsa-2022-macros/src/tracked_fn.rs b/components/salsa-2022-macros/src/tracked_fn.rs
index 3285bf84..274a673a 100644
--- a/components/salsa-2022-macros/src/tracked_fn.rs
+++ b/components/salsa-2022-macros/src/tracked_fn.rs
@@ -94,6 +94,15 @@ pub(crate) fn tracked_impl(
             ))
         }
     };
+    let self_type_name = &self_type.path.segments.last().unwrap().ident;
+    let name_prefix = match &item_impl.trait_ {
+        Some((_, trait_name, _)) => format!(
+            "{}_{}",
+            self_type_name,
+            trait_name.segments.last().unwrap().ident
+        ),
+        None => format!("{}", self_type_name),
+    };
     let extra_impls = item_impl
         .items
         .iter_mut()
@@ -116,10 +125,18 @@ pub(crate) fn tracked_impl(
             } else {
                 Ok(FnArgs::default())
             };
-            Some(match inner_args {
-                Ok(inner_args) => tracked_method(&args, inner_args, item_method, self_type),
-                Err(e) => Err(e),
-            })
+            let inner_args = match inner_args {
+                Ok(inner_args) => inner_args,
+                Err(err) => return Some(Err(err)),
+            };
+            let name = format!("{}_{}", name_prefix, item_method.sig.ident);
+            Some(tracked_method(
+                &args,
+                inner_args,
+                item_method,
+                self_type,
+                &name,
+            ))
         })
         // Collate all the errors so we can display them all at once
         .fold(Ok(Vec::new()), |mut acc, res| {
@@ -166,6 +183,7 @@ fn tracked_method(
     mut args: FnArgs,
     item_method: &mut syn::ImplItemMethod,
     self_type: &syn::TypePath,
+    name: &str,
 ) -> syn::Result<TokenStream> {
     args.jar_ty = args.jar_ty.or_else(|| outer_args.jar_ty.clone());
 
@@ -182,14 +200,7 @@ fn tracked_method(
         sig: item_method.sig.clone(),
         block: Box::new(rename_self_in_block(item_method.block.clone())?),
     };
-    item_fn.sig.ident = syn::Ident::new(
-        &format!(
-            "{}_{}",
-            self_type.path.segments.last().unwrap().ident,
-            item_fn.sig.ident
-        ),
-        item_fn.sig.ident.span(),
-    );
+    item_fn.sig.ident = syn::Ident::new(name, item_fn.sig.ident.span());
     // Flip the first and second arguments as the rest of the code expects the
     // database to come first and the struct to come second. We also need to
     // change the self argument to a normal typed argument called __salsa_self.
diff --git a/salsa-2022-tests/tests/compile-fail/tracked_method_on_untracked_impl.rs b/salsa-2022-tests/tests/compile-fail/tracked_method_on_untracked_impl.rs
new file mode 100644
index 00000000..ecda59ed
--- /dev/null
+++ b/salsa-2022-tests/tests/compile-fail/tracked_method_on_untracked_impl.rs
@@ -0,0 +1,18 @@
+#[salsa::jar(db = Db)]
+struct Jar(MyInput, tracked_method_on_untracked_impl);
+
+trait Db: salsa::DbWithJar<Jar> {}
+
+#[salsa::input]
+struct MyInput {
+    field: u32,
+}
+
+impl MyInput {
+    #[salsa::tracked]
+    fn tracked_method_on_untracked_impl(self, db: &dyn Db) -> u32 {
+        input.field(db)
+    }
+}
+
+fn main() {}
diff --git a/salsa-2022-tests/tests/compile-fail/tracked_method_on_untracked_impl.stderr b/salsa-2022-tests/tests/compile-fail/tracked_method_on_untracked_impl.stderr
new file mode 100644
index 00000000..03bffa3d
--- /dev/null
+++ b/salsa-2022-tests/tests/compile-fail/tracked_method_on_untracked_impl.stderr
@@ -0,0 +1,11 @@
+error: #[salsa::tracked] must also be applied to the impl block for tracked methods
+  --> tests/compile-fail/tracked_method_on_untracked_impl.rs:13:41
+   |
+13 |     fn tracked_method_on_untracked_impl(self, db: &dyn Db) -> u32 {
+   |                                         ^^^^
+
+error[E0412]: cannot find type `tracked_method_on_untracked_impl` in this scope
+ --> tests/compile-fail/tracked_method_on_untracked_impl.rs:2:21
+  |
+2 | struct Jar(MyInput, tracked_method_on_untracked_impl);
+  |                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ not found in this scope
diff --git a/salsa-2022-tests/tests/tracked_method.rs b/salsa-2022-tests/tests/tracked_method.rs
index 35d2cabf..ed1449f9 100644
--- a/salsa-2022-tests/tests/tracked_method.rs
+++ b/salsa-2022-tests/tests/tracked_method.rs
@@ -3,10 +3,19 @@
 #![allow(warnings)]
 
 #[salsa::jar(db = Db)]
-struct Jar(MyInput, MyInput_tracked_fn, MyInput_tracked_fn_ref);
+struct Jar(
+    MyInput,
+    MyInput_tracked_fn,
+    MyInput_tracked_fn_ref,
+    MyInput_TrackedTrait_tracked_trait_fn,
+);
 
 trait Db: salsa::DbWithJar<Jar> {}
 
+trait TrackedTrait {
+    fn tracked_trait_fn(self, db: &dyn Db) -> u32;
+}
+
 #[salsa::input]
 struct MyInput {
     field: u32,
@@ -25,6 +34,14 @@ impl MyInput {
     }
 }
 
+#[salsa::tracked]
+impl TrackedTrait for MyInput {
+    #[salsa::tracked]
+    fn tracked_trait_fn(self, db: &dyn Db) -> u32 {
+        self.field(db) * 4
+    }
+}
+
 #[test]
 fn execute() {
     #[salsa::db(Jar)]
@@ -41,4 +58,5 @@ fn execute() {
     let object = MyInput::new(&mut db, 22);
     assert_eq!(object.tracked_fn(&db), 44);
     assert_eq!(*object.tracked_fn_ref(&db), 66);
+    assert_eq!(object.tracked_trait_fn(&db), 88);
 }