From ed8a16203bd96d50875a9ce3dff6ada8daa50cff Mon Sep 17 00:00:00 2001
From: Alex Lorenz <arphaman@gmail.com>
Date: Tue, 31 Oct 2017 01:28:17 +0000
Subject: [PATCH] [refactor] select the entire DeclStmt if one ifs decls is
 selected

git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@316971 91177308-0d34-0410-b5e6-96231b3b80d8
---
 lib/Tooling/Refactoring/ASTSelection.cpp | 22 +++++++---
 unittests/Tooling/ASTSelectionTest.cpp   | 56 ++++++++++++++++++++++++
 2 files changed, 73 insertions(+), 5 deletions(-)

diff --git a/lib/Tooling/Refactoring/ASTSelection.cpp b/lib/Tooling/Refactoring/ASTSelection.cpp
index 9d0683a285..6ac432622c 100644
--- a/lib/Tooling/Refactoring/ASTSelection.cpp
+++ b/lib/Tooling/Refactoring/ASTSelection.cpp
@@ -279,11 +279,23 @@ static void findDeepestWithKind(
     llvm::SmallVectorImpl<SelectedNodeWithParents> &MatchingNodes,
     SourceSelectionKind Kind,
     llvm::SmallVectorImpl<SelectedASTNode::ReferenceType> &ParentStack) {
-  if (!hasAnyDirectChildrenWithKind(ASTSelection, Kind)) {
-    // This node is the bottom-most.
-    MatchingNodes.push_back(SelectedNodeWithParents{
-        std::cref(ASTSelection), {ParentStack.begin(), ParentStack.end()}});
-    return;
+  if (ASTSelection.Node.get<DeclStmt>()) {
+    // Select the entire decl stmt when any of its child declarations is the
+    // bottom-most.
+    for (const auto &Child : ASTSelection.Children) {
+      if (!hasAnyDirectChildrenWithKind(Child, Kind)) {
+        MatchingNodes.push_back(SelectedNodeWithParents{
+            std::cref(ASTSelection), {ParentStack.begin(), ParentStack.end()}});
+        return;
+      }
+    }
+  } else {
+    if (!hasAnyDirectChildrenWithKind(ASTSelection, Kind)) {
+      // This node is the bottom-most.
+      MatchingNodes.push_back(SelectedNodeWithParents{
+          std::cref(ASTSelection), {ParentStack.begin(), ParentStack.end()}});
+      return;
+    }
   }
   // Search in the children.
   ParentStack.push_back(std::cref(ASTSelection));
diff --git a/unittests/Tooling/ASTSelectionTest.cpp b/unittests/Tooling/ASTSelectionTest.cpp
index 79e89f90f4..94435d49a8 100644
--- a/unittests/Tooling/ASTSelectionTest.cpp
+++ b/unittests/Tooling/ASTSelectionTest.cpp
@@ -840,4 +840,60 @@ void f() {
       });
 }
 
+TEST(ASTSelectionFinder, SelectEntireDeclStmtRange) {
+  StringRef Source = R"(
+void f(int x, int y) {
+   int a = x * y;
+}
+)";
+  // 'int a = x * y'
+  findSelectedASTNodesWithRange(
+      Source, {3, 4}, FileRange{{3, 4}, {3, 17}},
+      [](SourceRange SelectionRange, Optional<SelectedASTNode> Node) {
+        EXPECT_TRUE(Node);
+        Optional<CodeRangeASTSelection> SelectedCode =
+            CodeRangeASTSelection::create(SelectionRange, std::move(*Node));
+        EXPECT_TRUE(SelectedCode);
+        EXPECT_EQ(SelectedCode->size(), 1u);
+        EXPECT_TRUE(isa<DeclStmt>((*SelectedCode)[0]));
+        ArrayRef<SelectedASTNode::ReferenceType> Parents =
+            SelectedCode->getParents();
+        EXPECT_EQ(Parents.size(), 3u);
+        EXPECT_TRUE(
+            isa<TranslationUnitDecl>(Parents[0].get().Node.get<Decl>()));
+        // Function 'f' definition.
+        EXPECT_TRUE(isa<FunctionDecl>(Parents[1].get().Node.get<Decl>()));
+        // Function body of function 'F'.
+        EXPECT_TRUE(isa<CompoundStmt>(Parents[2].get().Node.get<Stmt>()));
+      });
+}
+
+TEST(ASTSelectionFinder, SelectEntireDeclStmtRangeWithMultipleDecls) {
+  StringRef Source = R"(
+void f(int x, int y) {
+   int a = x * y, b = x - y;
+}
+)";
+  // 'b = x - y'
+  findSelectedASTNodesWithRange(
+      Source, {3, 19}, FileRange{{3, 19}, {3, 28}},
+      [](SourceRange SelectionRange, Optional<SelectedASTNode> Node) {
+        EXPECT_TRUE(Node);
+        Optional<CodeRangeASTSelection> SelectedCode =
+            CodeRangeASTSelection::create(SelectionRange, std::move(*Node));
+        EXPECT_TRUE(SelectedCode);
+        EXPECT_EQ(SelectedCode->size(), 1u);
+        EXPECT_TRUE(isa<DeclStmt>((*SelectedCode)[0]));
+        ArrayRef<SelectedASTNode::ReferenceType> Parents =
+            SelectedCode->getParents();
+        EXPECT_EQ(Parents.size(), 3u);
+        EXPECT_TRUE(
+            isa<TranslationUnitDecl>(Parents[0].get().Node.get<Decl>()));
+        // Function 'f' definition.
+        EXPECT_TRUE(isa<FunctionDecl>(Parents[1].get().Node.get<Decl>()));
+        // Function body of function 'F'.
+        EXPECT_TRUE(isa<CompoundStmt>(Parents[2].get().Node.get<Stmt>()));
+      });
+}
+
 } // end anonymous namespace
-- 
2.40.0