[refactor] selection: new CodeRangeASTSelection represents a set of selected consecutive statements

[arphaman]

Summary

This patch adds a CodeRangeASTSelection value to the refactoring library. This value represents a set of selected statements in one body of code.

The refactoring rules will be able to use this value with the CodeRangeSelectionRequirement that will be introduced in follow up patches. This selection will be used to implement the "extract" refactoring.

[hokein]

include/clang/Tooling/Refactoring/ASTSelection.h

74

I see all your tests are for function body, does it support other body, i.e. "global namespace", "compound statement"?

if yes, how about adding more test cases to cover it.

// variable in global namespace
int a; // select int.

void f() {
   {
       int a;  // select a.
   }
}

95

nit: instead of a structure, this feels more like a class.

lib/Tooling/Refactoring/ASTSelection.cpp

247

nit: remove the empty line.

unittests/Tooling/ASTSelectionTest.cpp

685

I think "No selection range" is more precise -- you are passing None parameter to the function.

For empty range, it should be something like FileRange{{2, 2}, {2, 2}}, we can also add a test for it.

688

I'm a bit confused here, if the selection range is none/empty, shouldn't SelectedASTNode be empty?

If this behavior is intended, I'd suggest documenting it in findSelectedASTNodesWithRange.

718

nit: would be clearer to add a comment indicating the corresponding code, the same below.

like

EXPECT_TRUE(...); // function f definition
EXPECT_TRUE(...); // function body of function f

[arphaman]

Address review comments

include/clang/Tooling/Refactoring/ASTSelection.h

74

Yes, I added a couple that you suggested.

unittests/Tooling/ASTSelectionTest.cpp

688

No, the AST selection will work even with a cursor location, hence it won't be empty.

However, the CodeRangeASTSelection requires a non-empty selection range, so it won't work with just cursors.

[hokein]

LGTM.

unittests/Tooling/ASTSelectionTest.cpp

688

I see, thanks.

the AST selection will work even with a cursor location, hence it won't be empty.

Is this documented somewhere in the code? I couldn't find any comments in the definition of findSelectedASTNodesWithRange or findSelectedASTNodes in this file. Would be nice to document it although this is a test API, so that it won't confuse future code readers.

[arphaman]

unittests/Tooling/ASTSelectionTest.cpp

688

Not quite, I'll add a comment in a follow-up commit.

include/clang/Tooling/Refactoring/ASTSelection.h

struct SelectedASTNode {
 
   SelectedASTNode(const ast_type_traits::DynTypedNode &Node,
                   SourceSelectionKind SelectionKind)
       : Node(Node), SelectionKind(SelectionKind) {}
   SelectedASTNode(SelectedASTNode &&) = default;
   SelectedASTNode &operator=(SelectedASTNode &&) = default;
 
   void dump(llvm::raw_ostream &OS = llvm::errs()) const;
+
+  using ReferenceType = std::reference_wrapper<const SelectedASTNode>;
 };
 
 /// Traverses the given ASTContext and creates a tree of selected AST nodes.
 ///
 /// \returns None if no nodes are selected in the AST, or a selected AST node
 /// that corresponds to the TranslationUnitDecl otherwise.
 Optional<SelectedASTNode> findSelectedASTNodes(const ASTContext &Context,
                                                SourceRange SelectionRange);
 
+/// An AST selection value that corresponds to a selection of a set of
+/// statements that belong to one body of code (like one function).
+///
+/// For example, the following selection in the source.
+///
+/// \code
+/// void function() {
+///  // selection begin:
+///  int x = 0;
+///  {
+///     // selection end
+///     x = 1;
+///  }
+///  x = 2;
+/// }
+/// \endcode
+///
+/// Would correspond to a code range selection of statements "int x = 0"
+/// and the entire compound statement that follows it.
+///
+/// A \c CodeRangeASTSelection value stores references to the full
+/// \c SelectedASTNode tree and should not outlive it.
+struct CodeRangeASTSelection {
+  CodeRangeASTSelection(CodeRangeASTSelection &&) = default;
+  CodeRangeASTSelection &operator=(CodeRangeASTSelection &&) = default;
+
+  /// Returns the parent hierarchy (top to bottom) for the selected nodes.
+  ArrayRef<SelectedASTNode::ReferenceType> getParents() { return Parents; }
+
+  /// Returns the number of selected statements.
+  size_t size() const {
+    if (!AreChildrenSelected)
+      return 1;
+    return SelectedNode.get().Children.size();
+  }
+
+  const Stmt *operator[](size_t I) const {
+    if (!AreChildrenSelected) {
+      assert(I == 0 && "Invalid index");
+      return SelectedNode.get().Node.get<Stmt>();
+    }
+    return SelectedNode.get().Children[I].Node.get<Stmt>();
+  }
+
+  static Optional<CodeRangeASTSelection>
+  create(SourceRange SelectionRange, const SelectedASTNode &ASTSelection);
+
+private:
+  CodeRangeASTSelection(SelectedASTNode::ReferenceType SelectedNode,
+                        ArrayRef<SelectedASTNode::ReferenceType> Parents,
+                        bool AreChildrenSelected)
+      : SelectedNode(SelectedNode), Parents(Parents.begin(), Parents.end()),
+        AreChildrenSelected(AreChildrenSelected) {}
+
+  /// The reference to the selected node (or reference to the selected
+  /// child nodes).
+  SelectedASTNode::ReferenceType SelectedNode;
+  /// The parent hierarchy (top to bottom) for the selected noe.
+  llvm::SmallVector<SelectedASTNode::ReferenceType, 8> Parents;
+  /// True only when the children of the selected node are actually selected.
+  bool AreChildrenSelected;
+};
+
 } // end namespace tooling
 } // end namespace clang
 
 #endif // LLVM_CLANG_TOOLING_REFACTOR_AST_SELECTION_H

lib/Tooling/Refactoring/ASTSelection.cpp

if (const Decl *D = Node.Node.get<Decl>()) {
     OS << S->getStmtClassName();
   }
   OS << ' ' << selectionKindToString(Node.SelectionKind) << "\n";
   for (const auto &Child : Node.Children)
     dump(Child, OS, Indent + 1);
 }
 
 void SelectedASTNode::dump(llvm::raw_ostream &OS) const { ::dump(*this, OS); }
+
+/// Returns true if the given node has any direct children with the following
+/// selection kind.
+///
+/// Note: The direct children also include children of direct children with the
+/// "None" selection kind.
+static bool hasAnyDirectChildrenWithKind(const SelectedASTNode &Node,
+                                         SourceSelectionKind Kind) {
+  assert(Kind != SourceSelectionKind::None && "invalid predicate!");
+  for (const auto &Child : Node.Children) {
+    if (Child.SelectionKind == Kind)
+      return true;
+    if (Child.SelectionKind == SourceSelectionKind::None)
+      return hasAnyDirectChildrenWithKind(Child, Kind);
+  }
+  return false;
+}
+
+namespace {
+
+struct SelectedNodeWithParents {
+  SelectedNodeWithParents(SelectedNodeWithParents &&) = default;
+  SelectedNodeWithParents &operator=(SelectedNodeWithParents &&) = default;
+  SelectedASTNode::ReferenceType Node;
+  llvm::SmallVector<SelectedASTNode::ReferenceType, 8> Parents;
+};
+}
+
+/// Finds the set of bottom-most selected AST nodes that are in the selection
+/// tree with the specified selection kind.
+///
+/// For example, given the following selection tree:
+///
+/// FunctionDecl "f" contains-selection
+///   CompoundStmt contains-selection [#1]
+///     CallExpr inside
+///     ImplicitCastExpr inside
+///       DeclRefExpr inside
+///     IntegerLiteral inside
+///     IntegerLiteral inside
+/// FunctionDecl "f2" contains-selection
+///   CompoundStmt contains-selection [#2]
+///     CallExpr inside
+///     ImplicitCastExpr inside
+///       DeclRefExpr inside
+///     IntegerLiteral inside
+///     IntegerLiteral inside
+///
+/// This function will find references to nodes #1 and #2 when searching for the
+/// \c ContainsSelection kind.
+static void findDeepestWithKind(
+    const SelectedASTNode &ASTSelection,
+    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;
+  }
+  // Search in the children.
+  ParentStack.push_back(std::cref(ASTSelection));
+  for (const auto &Child : ASTSelection.Children)
+    findDeepestWithKind(Child, MatchingNodes, Kind, ParentStack);
+  ParentStack.pop_back();
+}
+
+static void findDeepestWithKind(
+    const SelectedASTNode &ASTSelection,
+    llvm::SmallVectorImpl<SelectedNodeWithParents> &MatchingNodes,
+    SourceSelectionKind Kind) {
+  llvm::SmallVector<SelectedASTNode::ReferenceType, 16> ParentStack;
+  findDeepestWithKind(ASTSelection, MatchingNodes, Kind, ParentStack);
+}
+
+Optional<CodeRangeASTSelection>
+CodeRangeASTSelection::create(SourceRange SelectionRange,
+                              const SelectedASTNode &ASTSelection) {
+  // Code range is selected when the selection range is not empty.
+  if (SelectionRange.getBegin() == SelectionRange.getEnd())
+    return None;
+  llvm::SmallVector<SelectedNodeWithParents, 4> ContainSelection;
+  findDeepestWithKind(ASTSelection, ContainSelection,
+                      SourceSelectionKind::ContainsSelection);
+  // We are looking for a selection in one body of code, so let's focus on
+  // one matching result.
+  if (ContainSelection.size() != 1)
+    return None;
+  SelectedNodeWithParents &Selected = ContainSelection[0];
+  if (!Selected.Node.get().Node.get<Stmt>())
+    return None;
+  const Stmt *CodeRangeStmt = Selected.Node.get().Node.get<Stmt>();
+  if (!isa<CompoundStmt>(CodeRangeStmt)) {
+    // FIXME (Alex L): Canonicalize.
+    return CodeRangeASTSelection(Selected.Node, Selected.Parents,
+                                 /*AreChildrenSelected=*/false);
+  }
+  // FIXME (Alex L): Tweak selection rules for compound statements, see:
+  // https://github.com/apple/swift-clang/blob/swift-4.1-branch/lib/Tooling/
+  // Refactor/ASTSlice.cpp#L513
+  // The user selected multiple statements in a compound statement.
+  Selected.Parents.push_back(Selected.Node);
+  return CodeRangeASTSelection(Selected.Node, Selected.Parents,
+                               /*AreChildrenSelected=*/true);
+}

unittests/Tooling/ASTSelectionTest.cpp

struct FileLocation {
   }
 };
 
 using FileRange = std::pair<FileLocation, FileLocation>;
 
 class SelectionFinderVisitor : public TestVisitor<SelectionFinderVisitor> {
   FileLocation Location;
   Optional<FileRange> SelectionRange;
-  llvm::function_ref<void(Optional<SelectedASTNode>)> Consumer;
+  llvm::function_ref<void(SourceRange SelectionRange,
+                          Optional<SelectedASTNode>)>
+      Consumer;
 
 public:
-  SelectionFinderVisitor(
-      FileLocation Location, Optional<FileRange> SelectionRange,
-      llvm::function_ref<void(Optional<SelectedASTNode>)> Consumer)
+  SelectionFinderVisitor(FileLocation Location,
+                         Optional<FileRange> SelectionRange,
+                         llvm::function_ref<void(SourceRange SelectionRange,
+                                                 Optional<SelectedASTNode>)>
+                             Consumer)
       : Location(Location), SelectionRange(SelectionRange), Consumer(Consumer) {
   }
 
   bool VisitTranslationUnitDecl(const TranslationUnitDecl *TU) {
     const ASTContext &Context = TU->getASTContext();
     const SourceManager &SM = Context.getSourceManager();
 
     SourceRange SelRange;
     if (SelectionRange) {
       SelRange = SourceRange(SelectionRange->first.translate(SM),
                              SelectionRange->second.translate(SM));
     } else {
       SourceLocation Loc = Location.translate(SM);
       SelRange = SourceRange(Loc, Loc);
     }
-    Consumer(findSelectedASTNodes(Context, SelRange));
+    Consumer(SelRange, findSelectedASTNodes(Context, SelRange));
     return false;
   }
 };
 
-void findSelectedASTNodes(
+void findSelectedASTNodesWithRange(
     StringRef Source, FileLocation Location, Optional<FileRange> SelectionRange,
-    llvm::function_ref<void(Optional<SelectedASTNode>)> Consumer,
+    llvm::function_ref<void(SourceRange SelectionRange,
+                            Optional<SelectedASTNode>)>
+        Consumer,
     SelectionFinderVisitor::Language Language =
         SelectionFinderVisitor::Lang_CXX11) {
   SelectionFinderVisitor Visitor(Location, SelectionRange, Consumer);
   EXPECT_TRUE(Visitor.runOver(Source, Language));
 }
 
+void findSelectedASTNodes(
+    StringRef Source, FileLocation Location, Optional<FileRange> SelectionRange,
+    llvm::function_ref<void(Optional<SelectedASTNode>)> Consumer,
+    SelectionFinderVisitor::Language Language =
+        SelectionFinderVisitor::Lang_CXX11) {
+  findSelectedASTNodesWithRange(
+      Source, Location, SelectionRange,
+      [&](SourceRange, Optional<SelectedASTNode> Selection) {
+        Consumer(std::move(Selection));
+      },
+      Language);
+}
+
 void checkNodeImpl(bool IsTypeMatched, const SelectedASTNode &Node,
                    SourceSelectionKind SelectionKind, unsigned NumChildren) {
   ASSERT_TRUE(IsTypeMatched);
   EXPECT_EQ(Node.Children.size(), NumChildren);
   ASSERT_EQ(Node.SelectionKind, SelectionKind);
 }
 
 void checkDeclName(const SelectedASTNode &Node, StringRef Name) {
findSelectedASTNodes(
         checkNode<DeclRefExpr>(Cast2.Children[0],
                                SourceSelectionKind::InsideSelection);
         checkNode<DeclRefExpr>(BinOp.Children[1],
                                SourceSelectionKind::InsideSelection);
       },
       SelectionFinderVisitor::Lang_OBJC);
 }
 
+TEST(ASTSelectionFinder, SimpleCodeRangeASTSelection) {
+  StringRef Source = R"(void f(int x, int y) {
+  int z = x;
+  f(2, 3);
+  if (x == 0) {
+    return;
+  }
+  x = 1;
+  return;
+}
+void f2() {
+  int m = 0;
+}
+)";
+  // Empty range is an invalid code range.
+  findSelectedASTNodesWithRange(
+      Source, {2, 2}, None,
+      [](SourceRange SelectionRange, Optional<SelectedASTNode> Node) {
+        EXPECT_TRUE(Node);
+        Optional<CodeRangeASTSelection> SelectedCode =
+            CodeRangeASTSelection::create(SelectionRange, std::move(*Node));
+        EXPECT_FALSE(SelectedCode);
+      });
+  // Range that spans multiple functions is an invalid code range.
+  findSelectedASTNodesWithRange(
+      Source, {2, 2}, FileRange{{7, 2}, {12, 1}},
+      [](SourceRange SelectionRange, Optional<SelectedASTNode> Node) {
+        EXPECT_TRUE(Node);
+        Optional<CodeRangeASTSelection> SelectedCode =
+            CodeRangeASTSelection::create(SelectionRange, std::move(*Node));
+        EXPECT_FALSE(SelectedCode);
+      });
+  // Just 'z = x;':
+  findSelectedASTNodesWithRange(
+      Source, {2, 2}, FileRange{{2, 2}, {2, 13}},
+      [](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>()));
+        EXPECT_TRUE(isa<FunctionDecl>(Parents[1].get().Node.get<Decl>()));
+        EXPECT_TRUE(isa<CompoundStmt>(Parents[2].get().Node.get<Stmt>()));
+      });
+  // From 'f(2,3)' until just before 'x = 1;':
+  findSelectedASTNodesWithRange(
+      Source, {3, 2}, FileRange{{3, 2}, {7, 1}},
+      [](SourceRange SelectionRange, Optional<SelectedASTNode> Node) {
+        EXPECT_TRUE(Node);
+        Optional<CodeRangeASTSelection> SelectedCode =
+            CodeRangeASTSelection::create(SelectionRange, std::move(*Node));
+        EXPECT_TRUE(SelectedCode);
+        EXPECT_EQ(SelectedCode->size(), 2u);
+        EXPECT_TRUE(isa<CallExpr>((*SelectedCode)[0]));
+        EXPECT_TRUE(isa<IfStmt>((*SelectedCode)[1]));
+        ArrayRef<SelectedASTNode::ReferenceType> Parents =
+            SelectedCode->getParents();
+        EXPECT_EQ(Parents.size(), 3u);
+        EXPECT_TRUE(
+            isa<TranslationUnitDecl>(Parents[0].get().Node.get<Decl>()));
+        EXPECT_TRUE(isa<FunctionDecl>(Parents[1].get().Node.get<Decl>()));
+        EXPECT_TRUE(isa<CompoundStmt>(Parents[2].get().Node.get<Stmt>()));
+      });
+  // From 'f(2,3)' until just before ';' in 'x = 1;':
+  findSelectedASTNodesWithRange(
+      Source, {3, 2}, FileRange{{3, 2}, {7, 8}},
+      [](SourceRange SelectionRange, Optional<SelectedASTNode> Node) {
+        EXPECT_TRUE(Node);
+        Optional<CodeRangeASTSelection> SelectedCode =
+            CodeRangeASTSelection::create(SelectionRange, std::move(*Node));
+        EXPECT_TRUE(SelectedCode);
+        EXPECT_EQ(SelectedCode->size(), 3u);
+        EXPECT_TRUE(isa<CallExpr>((*SelectedCode)[0]));
+        EXPECT_TRUE(isa<IfStmt>((*SelectedCode)[1]));
+        EXPECT_TRUE(isa<BinaryOperator>((*SelectedCode)[2]));
+      });
+  // From the middle of 'int z = 3' until the middle of 'x = 1;':
+  findSelectedASTNodesWithRange(
+      Source, {2, 10}, FileRange{{2, 10}, {7, 5}},
+      [](SourceRange SelectionRange, Optional<SelectedASTNode> Node) {
+        EXPECT_TRUE(Node);
+        EXPECT_TRUE(Node);
+        Optional<CodeRangeASTSelection> SelectedCode =
+            CodeRangeASTSelection::create(SelectionRange, std::move(*Node));
+        EXPECT_TRUE(SelectedCode);
+        EXPECT_EQ(SelectedCode->size(), 4u);
+        EXPECT_TRUE(isa<DeclStmt>((*SelectedCode)[0]));
+        EXPECT_TRUE(isa<CallExpr>((*SelectedCode)[1]));
+        EXPECT_TRUE(isa<IfStmt>((*SelectedCode)[2]));
+        EXPECT_TRUE(isa<BinaryOperator>((*SelectedCode)[3]));
+      });
+}
+
 } // end anonymous namespace