[refactor][selection] canonicalize decl ref callee to the call expr

We would like to extract the full call when just the callee function is
selected

llvm-svn: 318215

2 files changed, 74 insertions, 13 deletions

diff --git a/clang/lib/Tooling/Refactoring/ASTSelection.cpp b/clang/lib/Tooling/Refactoring/ASTSelection.cpp
index ae8a3cb7c58..7123fc32cec 100644
--- a/clang/lib/Tooling/Refactoring/ASTSelection.cpp
+++ b/clang/lib/Tooling/Refactoring/ASTSelection.cpp
@@ -259,29 +259,64 @@ struct SelectedNodeWithParents {
   /// when it makes sense to do so.
   void canonicalize();
 };
-} // end anonymous namespace
 
-void SelectedNodeWithParents::canonicalize() {
-  const Stmt *S = Node.get().Node.get<Stmt>();
-  assert(S && "non statement selection!");
-  const Stmt *Parent = Parents[Parents.size() - 1].get().Node.get<Stmt>();
-  if (!Parent)
-    return;
+enum SelectionCanonicalizationAction { KeepSelection, SelectParent };
+
+/// Returns the canonicalization action which should be applied to the
+/// selected statement.
+SelectionCanonicalizationAction
+getSelectionCanonizalizationAction(const Stmt *S, const Stmt *Parent) {
   // Select the parent expression when:
   // - The string literal in ObjC string literal is selected, e.g.:
   //     @"test"   becomes   @"test"
   //      ~~~~~~             ~~~~~~~
   if (isa<StringLiteral>(S) && isa<ObjCStringLiteral>(Parent))
-    Node = Parents.pop_back_val();
+    return SelectParent;
   // The entire call should be selected when just the member expression
-  // that refers to the method is selected.
+  // that refers to the method or the decl ref that refers to the function
+  // is selected.
   //    f.call(args)  becomes  f.call(args)
   //      ~~~~                 ~~~~~~~~~~~~
-  else if (isa<MemberExpr>(S) && isa<CXXMemberCallExpr>(Parent) &&
-           cast<CXXMemberCallExpr>(Parent)->getCallee() == S)
-    Node = Parents.pop_back_val();
+  //    func(args)  becomes  func(args)
+  //    ~~~~                 ~~~~~~~~~~
+  else if (const auto *CE = dyn_cast<CallExpr>(Parent)) {
+    if ((isa<MemberExpr>(S) || isa<DeclRefExpr>(S)) &&
+        CE->getCallee()->IgnoreImpCasts() == S)
+      return SelectParent;
+  }
   // FIXME: Syntactic form -> Entire pseudo-object expr.
-  // FIXME: Callee -> Call.
+  return KeepSelection;
+}
+
+} // end anonymous namespace
+
+void SelectedNodeWithParents::canonicalize() {
+  const Stmt *S = Node.get().Node.get<Stmt>();
+  assert(S && "non statement selection!");
+  const Stmt *Parent = Parents[Parents.size() - 1].get().Node.get<Stmt>();
+  if (!Parent)
+    return;
+
+  // Look through the implicit casts in the parents.
+  unsigned ParentIndex = 1;
+  for (; (ParentIndex + 1) <= Parents.size() && isa<ImplicitCastExpr>(Parent);
+       ++ParentIndex) {
+    const Stmt *NewParent =
+        Parents[Parents.size() - ParentIndex - 1].get().Node.get<Stmt>();
+    if (!NewParent)
+      break;
+    Parent = NewParent;
+  }
+
+  switch (getSelectionCanonizalizationAction(S, Parent)) {
+  case SelectParent:
+    Node = Parents[Parents.size() - ParentIndex];
+    for (; ParentIndex != 0; --ParentIndex)
+      Parents.pop_back();
+    break;
+  case KeepSelection:
+    break;
+  }
 }
 
 /// Finds the set of bottom-most selected AST nodes that are in the selection
diff --git a/clang/unittests/Tooling/ASTSelectionTest.cpp b/clang/unittests/Tooling/ASTSelectionTest.cpp
index fc5186ea208..2f5df8f4300 100644
--- a/clang/unittests/Tooling/ASTSelectionTest.cpp
+++ b/clang/unittests/Tooling/ASTSelectionTest.cpp
@@ -1056,4 +1056,30 @@ void dontSelectArgument(AClass &a) {
       });
 }
 
+TEST(ASTSelectionFinder, CanonicalizeFuncCalleeToCall) {
+  StringRef Source = R"(
+void function();
+
+void test() {
+  function();
+}
+     )";
+  // Just 'function':
+  findSelectedASTNodesWithRange(
+      Source, {5, 3}, FileRange{{5, 3}, {5, 11}},
+      [](SourceRange SelectionRange, Optional<SelectedASTNode> Node) {
+        EXPECT_TRUE(Node);
+        Node->dump();
+        Optional<CodeRangeASTSelection> SelectedCode =
+            CodeRangeASTSelection::create(SelectionRange, std::move(*Node));
+        EXPECT_TRUE(SelectedCode);
+        EXPECT_EQ(SelectedCode->size(), 1u);
+        EXPECT_TRUE(isa<CallExpr>((*SelectedCode)[0]));
+        EXPECT_TRUE(isa<CompoundStmt>(
+            SelectedCode->getParents()[SelectedCode->getParents().size() - 1]
+                .get()
+                .Node.get<Stmt>()));
+      });
+}
+
 } // end anonymous namespace