SSR: Explicitly autoderef and ref placeholders as needed

Structured search replace now inserts *, & and &mut in the replacement to match any auto[de]ref in the matched code.

4 files changed, 222 insertions, 25 deletions

diff --git a/crates/ssr/src/lib.rs b/crates/ssr/src/lib.rs
index 292bd5b9a..ba669fd56 100644
--- a/crates/ssr/src/lib.rs
+++ b/crates/ssr/src/lib.rs
@@ -21,7 +21,10 @@
 // code in the `foo` module, we'll insert just `Bar`.
 //
 // Inherent method calls should generally be written in UFCS form. e.g. `foo::Bar::baz($s, $a)` will
-// match `$s.baz($a)`, provided the method call `baz` resolves to the method `foo::Bar::baz`.
+// match `$s.baz($a)`, provided the method call `baz` resolves to the method `foo::Bar::baz`. When a
+// placeholder is the receiver of a method call in the search pattern (e.g. `$s.foo()`), but not in
+// the replacement template (e.g. `bar($s)`), then *, & and &mut will be added as needed to mirror
+// whatever autoderef and autoref was happening implicitly in the matched code.
 //
 // The scope of the search / replace will be restricted to the current selection if any, otherwise
 // it will apply to the whole workspace.
diff --git a/crates/ssr/src/matching.rs b/crates/ssr/src/matching.rs
index 7f0b5061e..8bb5ced90 100644
--- a/crates/ssr/src/matching.rs
+++ b/crates/ssr/src/matching.rs
@@ -65,6 +65,10 @@ pub(crate) struct PlaceholderMatch {
     pub(crate) range: FileRange,
     /// More matches, found within `node`.
     pub(crate) inner_matches: SsrMatches,
+    /// How many times the code that the placeholder matched needed to be dereferenced. Will only be
+    /// non-zero if the placeholder matched to the receiver of a method call.
+    pub(crate) autoderef_count: usize,
+    pub(crate) autoref_kind: ast::SelfParamKind,
 }
 
 #[derive(Debug)]
@@ -169,7 +173,7 @@ impl<'db, 'sema> Matcher<'db, 'sema> {
         code: &SyntaxNode,
     ) -> Result<(), MatchFailed> {
         // Handle placeholders.
-        if let Some(placeholder) = self.get_placeholder(&SyntaxElement::Node(pattern.clone())) {
+        if let Some(placeholder) = self.get_placeholder_for_node(pattern) {
             for constraint in &placeholder.constraints {
                 self.check_constraint(constraint, code)?;
             }
@@ -178,9 +182,10 @@ impl<'db, 'sema> Matcher<'db, 'sema> {
                 // We validated the range for the node when we started the match, so the placeholder
                 // probably can't fail range validation, but just to be safe...
                 self.validate_range(&original_range)?;
-                matches_out
-                    .placeholder_values
-                    .insert(placeholder.ident.clone(), PlaceholderMatch::new(code, original_range));
+                matches_out.placeholder_values.insert(
+                    placeholder.ident.clone(),
+                    PlaceholderMatch::new(Some(code), original_range),
+                );
             }
             return Ok(());
         }
@@ -531,18 +536,40 @@ impl<'db, 'sema> Matcher<'db, 'sema> {
         if pattern_ufcs.function != code_resolved_function {
             fail_match!("Method call resolved to a different function");
         }
-        if code_resolved_function.has_self_param(self.sema.db) {
-            if let (Some(pattern_type), Some(expr)) = (&pattern_ufcs.qualifier_type, &code.expr()) {
-                self.check_expr_type(pattern_type, expr)?;
-            }
-        }
         // Check arguments.
         let mut pattern_args = pattern_ufcs
             .call_expr
             .arg_list()
             .ok_or_else(|| match_error!("Pattern function call has no args"))?
             .args();
-        self.attempt_match_opt(phase, pattern_args.next(), code.expr())?;
+        // If the function we're calling takes a self parameter, then we store additional
+        // information on the placeholder match about autoderef and autoref. This allows us to use
+        // the placeholder in a context where autoderef and autoref don't apply.
+        if code_resolved_function.has_self_param(self.sema.db) {
+            if let (Some(pattern_type), Some(expr)) = (&pattern_ufcs.qualifier_type, &code.expr()) {
+                let deref_count = self.check_expr_type(pattern_type, expr)?;
+                let pattern_receiver = pattern_args.next();
+                self.attempt_match_opt(phase, pattern_receiver.clone(), code.expr())?;
+                if let Phase::Second(match_out) = phase {
+                    if let Some(placeholder_value) = pattern_receiver
+                        .and_then(|n| self.get_placeholder_for_node(n.syntax()))
+                        .and_then(|placeholder| {
+                            match_out.placeholder_values.get_mut(&placeholder.ident)
+                        })
+                    {
+                        placeholder_value.autoderef_count = deref_count;
+                        placeholder_value.autoref_kind = self
+                            .sema
+                            .resolve_method_call_as_callable(code)
+                            .and_then(|callable| callable.receiver_param(self.sema.db))
+                            .map(|self_param| self_param.kind())
+                            .unwrap_or(ast::SelfParamKind::Owned);
+                    }
+                }
+            }
+        } else {
+            self.attempt_match_opt(phase, pattern_args.next(), code.expr())?;
+        }
         let mut code_args =
             code.arg_list().ok_or_else(|| match_error!("Code method call has no args"))?.args();
         loop {
@@ -570,26 +597,35 @@ impl<'db, 'sema> Matcher<'db, 'sema> {
         self.attempt_match_node_children(phase, pattern_ufcs.call_expr.syntax(), code.syntax())
     }
 
+    /// Verifies that `expr` matches `pattern_type`, possibly after dereferencing some number of
+    /// times. Returns the number of times it needed to be dereferenced.
     fn check_expr_type(
         &self,
         pattern_type: &hir::Type,
         expr: &ast::Expr,
-    ) -> Result<(), MatchFailed> {
+    ) -> Result<usize, MatchFailed> {
         use hir::HirDisplay;
         let code_type = self.sema.type_of_expr(&expr).ok_or_else(|| {
             match_error!("Failed to get receiver type for `{}`", expr.syntax().text())
         })?;
-        if !code_type
+        // Temporary needed to make the borrow checker happy.
+        let res = code_type
             .autoderef(self.sema.db)
-            .any(|deref_code_type| *pattern_type == deref_code_type)
-        {
-            fail_match!(
-                "Pattern type `{}` didn't match code type `{}`",
-                pattern_type.display(self.sema.db),
-                code_type.display(self.sema.db)
-            );
-        }
-        Ok(())
+            .enumerate()
+            .find(|(_, deref_code_type)| pattern_type == deref_code_type)
+            .map(|(count, _)| count)
+            .ok_or_else(|| {
+                match_error!(
+                    "Pattern type `{}` didn't match code type `{}`",
+                    pattern_type.display(self.sema.db),
+                    code_type.display(self.sema.db)
+                )
+            });
+        res
+    }
+
+    fn get_placeholder_for_node(&self, node: &SyntaxNode) -> Option<&Placeholder> {
+        self.get_placeholder(&SyntaxElement::Node(node.clone()))
     }
 
     fn get_placeholder(&self, element: &SyntaxElement) -> Option<&Placeholder> {
@@ -671,12 +707,18 @@ fn recording_match_fail_reasons() -> bool {
 }
 
 impl PlaceholderMatch {
-    fn new(node: &SyntaxNode, range: FileRange) -> Self {
-        Self { node: Some(node.clone()), range, inner_matches: SsrMatches::default() }
+    fn new(node: Option<&SyntaxNode>, range: FileRange) -> Self {
+        Self {
+            node: node.cloned(),
+            range,
+            inner_matches: SsrMatches::default(),
+            autoderef_count: 0,
+            autoref_kind: ast::SelfParamKind::Owned,
+        }
     }
 
     fn from_range(range: FileRange) -> Self {
-        Self { node: None, range, inner_matches: SsrMatches::default() }
+        Self::new(None, range)
     }
 }
 
diff --git a/crates/ssr/src/replacing.rs b/crates/ssr/src/replacing.rs
index 496a21e6e..21d0aa8a8 100644
--- a/crates/ssr/src/replacing.rs
+++ b/crates/ssr/src/replacing.rs
@@ -118,6 +118,27 @@ impl ReplacementRenderer<'_> {
                 let range = &placeholder_value.range.range;
                 let mut matched_text =
                     self.file_src[usize::from(range.start())..usize::from(range.end())].to_owned();
+                // If a method call is performed directly on the placeholder, then autoderef and
+                // autoref will apply, so we can just substitute whatever the placeholder matched to
+                // directly. If we're not applying a method call, then we need to add explicitly
+                // deref and ref in order to match whatever was being done implicitly at the match
+                // site.
+                if !token_is_method_call_receiver(token)
+                    && (placeholder_value.autoderef_count > 0
+                        || placeholder_value.autoref_kind != ast::SelfParamKind::Owned)
+                {
+                    let ref_kind = match placeholder_value.autoref_kind {
+                        ast::SelfParamKind::Owned => "",
+                        ast::SelfParamKind::Ref => "&",
+                        ast::SelfParamKind::MutRef => "&mut ",
+                    };
+                    matched_text = format!(
+                        "{}{}{}",
+                        ref_kind,
+                        "*".repeat(placeholder_value.autoderef_count),
+                        matched_text
+                    );
+                }
                 let edit = matches_to_edit_at_offset(
                     &placeholder_value.inner_matches,
                     self.file_src,
@@ -178,6 +199,31 @@ impl ReplacementRenderer<'_> {
     }
 }
 
+/// Returns whether token is the receiver of a method call. Note, being within the receiver of a
+/// method call doesn't count. e.g. if the token is `$a`, then `$a.foo()` will return true, while
+/// `($a + $b).foo()` or `x.foo($a)` will return false.
+fn token_is_method_call_receiver(token: &SyntaxToken) -> bool {
+    use syntax::ast::AstNode;
+    // Find the first method call among the ancestors of `token`, then check if the only token
+    // within the receiver is `token`.
+    if let Some(receiver) = token
+        .ancestors()
+        .find(|node| node.kind() == SyntaxKind::METHOD_CALL_EXPR)
+        .and_then(|node| ast::MethodCallExpr::cast(node).unwrap().expr())
+    {
+        let mut tokens = receiver.syntax().descendants_with_tokens().filter_map(|node_or_token| {
+            match node_or_token {
+                SyntaxElement::Token(t) => Some(t),
+                _ => None,
+            }
+        });
+        if let (Some(only_token), None) = (tokens.next(), tokens.next()) {
+            return only_token == *token;
+        }
+    }
+    false
+}
+
 fn parse_as_kind(code: &str, kind: SyntaxKind) -> Option<SyntaxNode> {
     use syntax::ast::AstNode;
     if ast::Expr::can_cast(kind) {
diff --git a/crates/ssr/src/tests.rs b/crates/ssr/src/tests.rs
index 65cd38753..0ad3512ba 100644
--- a/crates/ssr/src/tests.rs
+++ b/crates/ssr/src/tests.rs
@@ -1172,3 +1172,109 @@ fn match_trait_method_call() {
     assert_matches("Bar::foo($a, $b)", code, &["v1.foo(1)", "Bar::foo(&v1, 3)", "v1_ref.foo(5)"]);
     assert_matches("Bar2::foo($a, $b)", code, &["v2.foo(2)", "Bar2::foo(&v2, 4)", "v2_ref.foo(6)"]);
 }
+
+#[test]
+fn replace_autoref_autoderef_capture() {
+    // Here we have several calls to `$a.foo()`. In the first case autoref is applied, in the
+    // second, we already have a reference, so it isn't. When $a is used in a context where autoref
+    // doesn't apply, we need to prefix it with `&`. Finally, we have some cases where autoderef
+    // needs to be applied.
+    let code = r#"
+        struct Foo {}
+        impl Foo {
+            fn foo(&self) {}
+            fn foo2(&self) {}
+        }
+        fn bar(_: &Foo) {}
+        fn main() {
+            let f = Foo {};
+            let fr = &f;
+            let fr2 = &fr;
+            let fr3 = &fr2;
+            f.foo();
+            fr.foo();
+            fr2.foo();
+            fr3.foo();
+        }
+        "#;
+    assert_ssr_transform(
+        "Foo::foo($a) ==>> bar($a)",
+        code,
+        expect![[r#"
+            struct Foo {}
+            impl Foo {
+                fn foo(&self) {}
+                fn foo2(&self) {}
+            }
+            fn bar(_: &Foo) {}
+            fn main() {
+                let f = Foo {};
+                let fr = &f;
+                let fr2 = &fr;
+                let fr3 = &fr2;
+                bar(&f);
+                bar(&*fr);
+                bar(&**fr2);
+                bar(&***fr3);
+            }
+        "#]],
+    );
+    // If the placeholder is used as the receiver of another method call, then we don't need to
+    // explicitly autoderef or autoref.
+    assert_ssr_transform(
+        "Foo::foo($a) ==>> $a.foo2()",
+        code,
+        expect![[r#"
+            struct Foo {}
+            impl Foo {
+                fn foo(&self) {}
+                fn foo2(&self) {}
+            }
+            fn bar(_: &Foo) {}
+            fn main() {
+                let f = Foo {};
+                let fr = &f;
+                let fr2 = &fr;
+                let fr3 = &fr2;
+                f.foo2();
+                fr.foo2();
+                fr2.foo2();
+                fr3.foo2();
+            }
+        "#]],
+    );
+}
+
+#[test]
+fn replace_autoref_mut() {
+    let code = r#"
+        struct Foo {}
+        impl Foo {
+            fn foo(&mut self) {}
+        }
+        fn bar(_: &mut Foo) {}
+        fn main() {
+            let mut f = Foo {};
+            f.foo();
+            let fr = &mut f;
+            fr.foo();
+        }
+        "#;
+    assert_ssr_transform(
+        "Foo::foo($a) ==>> bar($a)",
+        code,
+        expect![[r#"
+            struct Foo {}
+            impl Foo {
+                fn foo(&mut self) {}
+            }
+            fn bar(_: &mut Foo) {}
+            fn main() {
+                let mut f = Foo {};
+                bar(&mut f);
+                let fr = &mut f;
+                bar(&mut *fr);
+            }
+        "#]],
+    );
+}