From 4d005e529b460cb05152000f7b109fc0e657e053 Mon Sep 17 00:00:00 2001
From: Daiki Ihara <sasurau4@gmail.com>
Date: Fri, 25 Jun 2021 22:14:42 +0900
Subject: [PATCH] Fix extract_function with macro arg

---
 .../src/handlers/extract_function.rs          | 92 +++++++++++++------
 crates/syntax/src/ast/expr_ext.rs             |  5 +
 2 files changed, 69 insertions(+), 28 deletions(-)

diff --git a/crates/ide_assists/src/handlers/extract_function.rs b/crates/ide_assists/src/handlers/extract_function.rs
index 1bf9c073a4..192403b15d 100644
--- a/crates/ide_assists/src/handlers/extract_function.rs
+++ b/crates/ide_assists/src/handlers/extract_function.rs
@@ -2,7 +2,7 @@ use std::{hash::BuildHasherDefault, iter};
 
 use ast::make;
 use either::Either;
-use hir::{HirDisplay, Local, Semantics, TypeInfo};
+use hir::{HirDisplay, InFile, Local, Semantics, TypeInfo};
 use ide_db::{
     defs::{Definition, NameRefClass},
     search::{FileReference, ReferenceAccess, SearchScope},
@@ -17,7 +17,7 @@ use syntax::{
         edit::{AstNodeEdit, IndentLevel},
         AstNode,
     },
-    match_ast, ted,
+    match_ast, ted, SyntaxElement,
     SyntaxKind::{self, COMMENT},
     SyntaxNode, SyntaxToken, TextRange, TextSize, TokenAtOffset, WalkEvent, T,
 };
@@ -582,38 +582,45 @@ impl FunctionBody {
         &self,
         sema: &Semantics<RootDatabase>,
     ) -> (FxIndexSet<Local>, Option<ast::SelfParam>) {
-        // FIXME: currently usages inside macros are not found
         let mut self_param = None;
         let mut res = FxIndexSet::default();
-        self.walk_expr(&mut |expr| {
-            let name_ref = match expr {
-                ast::Expr::PathExpr(path_expr) => {
-                    path_expr.path().and_then(|it| it.as_single_name_ref())
-                }
-                _ => return,
-            };
-            if let Some(name_ref) = name_ref {
-                if let Some(
-                    NameRefClass::Definition(Definition::Local(local_ref))
-                    | NameRefClass::FieldShorthand { local_ref, field_ref: _ },
-                ) = NameRefClass::classify(sema, &name_ref)
-                {
-                    if local_ref.is_self(sema.db) {
-                        match local_ref.source(sema.db).value {
-                            Either::Right(it) => {
-                                self_param.replace(it);
-                            }
-                            Either::Left(_) => {
-                                stdx::never!(
-                                    "Local::is_self returned true, but source is IdentPat"
-                                );
-                            }
-                        }
-                    } else {
+        let mut cb = |name_ref: Option<_>| {
+            let local_ref =
+                match name_ref.and_then(|name_ref| NameRefClass::classify(sema, &name_ref)) {
+                    Some(
+                        NameRefClass::Definition(Definition::Local(local_ref))
+                        | NameRefClass::FieldShorthand { local_ref, field_ref: _ },
+                    ) => local_ref,
+                    _ => return,
+                };
+            let InFile { file_id, value } = local_ref.source(sema.db);
+            // locals defined inside macros are not relevant to us
+            if !file_id.is_macro() {
+                match value {
+                    Either::Right(it) => {
+                        self_param.replace(it);
+                    }
+                    Either::Left(_) => {
                         res.insert(local_ref);
                     }
                 }
             }
+        };
+        self.walk_expr(&mut |expr| match expr {
+            ast::Expr::PathExpr(path_expr) => {
+                cb(path_expr.path().and_then(|it| it.as_single_name_ref()))
+            }
+            ast::Expr::MacroCall(call) => {
+                if let Some(tt) = call.token_tree() {
+                    tt.syntax()
+                        .children_with_tokens()
+                        .flat_map(SyntaxElement::into_token)
+                        .filter(|it| it.kind() == SyntaxKind::IDENT)
+                        .flat_map(|t| sema.descend_into_macros_many(t))
+                        .for_each(|t| cb(t.parent().and_then(ast::NameRef::cast)));
+                }
+            }
+            _ => (),
         });
         (res, self_param)
     }
@@ -4155,6 +4162,35 @@ fn foo() {
 fn $0fun_name(y: &mut Foo) {
     y.foo();
 }
+"#,
+        );
+    }
+
+    #[test]
+    fn extract_with_macro_arg() {
+        check_assist(
+            extract_function,
+            r#"
+macro_rules! m {
+    ($val:expr) => { $val };
+}
+fn main() {
+    let bar = "bar";
+    $0m!(bar);$0
+}
+"#,
+            r#"
+macro_rules! m {
+    ($val:expr) => { $val };
+}
+fn main() {
+    let bar = "bar";
+    fun_name(bar);
+}
+
+fn $0fun_name(bar: &str) {
+    m!(bar);
+}
 "#,
         );
     }
diff --git a/crates/syntax/src/ast/expr_ext.rs b/crates/syntax/src/ast/expr_ext.rs
index 4598066bc9..1c40de1e60 100644
--- a/crates/syntax/src/ast/expr_ext.rs
+++ b/crates/syntax/src/ast/expr_ext.rs
@@ -77,6 +77,11 @@ impl ast::Expr {
                 }
                 // Don't skip subtree since we want to process the expression child next
                 Some(ast::Stmt::ExprStmt(_)) => (),
+                // This might be an expression
+                Some(ast::Stmt::Item(ast::Item::MacroCall(mcall))) => {
+                    cb(WalkEvent::Enter(ast::Expr::MacroCall(mcall)));
+                    preorder.skip_subtree();
+                }
                 // skip inner items which might have their own expressions
                 Some(ast::Stmt::Item(_)) => preorder.skip_subtree(),
                 None => {