[clangd] SelectionTree treats TranslationUnitDecl (mostly) consistently with other containers.

Summary:
Previously TranslationUnitDecl would never be selected.
This means root() is never null, and returns a reference.

commonAncestor() is in principle never null also, but returning TUDecl
here requires tweaks to be careful not to traverse it (this was already
possible when selecting multiple top-level decls, and there are associated bugs!)
Instead, never allow commonAncestor() to return TUDecl, return null instead.

Reviewers: hokein

Subscribers: ilya-biryukov, MaskRay, jkorous, arphaman, kadircet, cfe-commits

Tags: #clang

Differential Revision: https://reviews.llvm.org/D65101

llvm-svn: 366893

Author: sam-mccall

clang-tools-extra/clangd/Selection.cpp

@@ -103,8 +103,14 @@ class SelectionVisitor : public RecursiveASTVisitor<SelectionVisitor> {
     V.TraverseAST(AST);
     assert(V.Stack.size() == 1 && "Unpaired push/pop?");
     assert(V.Stack.top() == &V.Nodes.front());
-    if (V.Nodes.size() == 1) // TUDecl, but no nodes under it.
-      V.Nodes.clear();
+    // We selected TUDecl if characters were unclaimed (or the file is empty).
+    if (V.Nodes.size() == 1 || V.Claimed.add(Begin, End)) {
+      StringRef FileContent = AST.getSourceManager().getBufferData(File);
+      // Don't require the trailing newlines to be selected.
+      bool SelectedAll = Begin == 0 && End >= FileContent.rtrim().size();
+      V.Stack.top()->Selected =
+          SelectedAll ? SelectionTree::Complete : SelectionTree::Partial;
+    }
     return std::move(V.Nodes);
   }
 
@@ -424,12 +430,13 @@ SelectionTree::SelectionTree(ASTContext &AST, unsigned Offset)
     : SelectionTree(AST, Offset, Offset) {}
 
 const Node *SelectionTree::commonAncestor() const {
-  if (!Root)
-    return nullptr;
   const Node *Ancestor = Root;
   while (Ancestor->Children.size() == 1 && !Ancestor->Selected)
     Ancestor = Ancestor->Children.front();
-  return Ancestor;
+  // Returning nullptr here is a bit unprincipled, but it makes the API safer:
+  // the TranslationUnitDecl contains all of the preamble, so traversing it is a
+  // performance cliff. Callers can check for null and use root() if they want.
+  return Ancestor != Root ? Ancestor : nullptr;
 }
 
 const DeclContext& SelectionTree::Node::getDeclContext() const {

clang-tools-extra/clangd/Selection.h

@@ -54,6 +54,11 @@ class ParsedAST;
 //  - if you want to traverse the selected nodes, they are all under
 //    commonAncestor() in the tree.
 //
+// SelectionTree tries to behave sensibly in the presence of macros, but does
+// not model any preprocessor concepts: the output is a subset of the AST.
+// Currently comments, directives etc are treated as part of the lexically
+// containing AST node, (though we may want to change this in future).
+//
 // The SelectionTree owns the Node structures, but the ASTNode attributes
 // point back into the AST it was constructed with.
 class SelectionTree {
@@ -100,11 +105,11 @@ class SelectionTree {
     std::string kind() const;
   };
   // The most specific common ancestor of all the selected nodes.
-  // If there is no selection, this is nullptr.
+  // Returns nullptr if the common ancestor is the root.
+  // (This is to avoid accidentally traversing the TUDecl and thus preamble).
   const Node *commonAncestor() const;
   // The selection node corresponding to TranslationUnitDecl.
-  // If there is no selection, this is nullptr.
-  const Node *root() const { return Root; }
+  const Node &root() const { return *Root; }
 
 private:
   std::deque<Node> Nodes; // Stable-pointer storage.
@@ -114,10 +119,7 @@ class SelectionTree {
   void print(llvm::raw_ostream &OS, const Node &N, int Indent) const;
   friend llvm::raw_ostream &operator<<(llvm::raw_ostream &OS,
                                        const SelectionTree &T) {
-    if (auto R = T.root())
-      T.print(OS, *R, 0);
-    else
-      OS << "(empty selection)\n";
+    T.print(OS, T.root(), 1);
     return OS;
   }
 };

clang-tools-extra/clangd/refactor/tweaks/DumpAST.cpp

@@ -84,9 +84,7 @@ class ShowSelectionTree : public Tweak {
 public:
   const char *id() const override final;
 
-  bool prepare(const Selection &Inputs) override {
-    return Inputs.ASTSelection.root() != nullptr;
-  }
+  bool prepare(const Selection &Inputs) override { return true; }
   Expected<Effect> apply(const Selection &Inputs) override {
     return Effect::showMessage(llvm::to_string(Inputs.ASTSelection));
   }

clang-tools-extra/clangd/unittests/SelectionTests.cpp

@@ -9,6 +9,8 @@
 #include "Selection.h"
 #include "SourceCode.h"
 #include "TestTU.h"
+#include "clang/AST/Decl.h"
+#include "llvm/Support/Casting.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 
@@ -40,6 +42,8 @@ Range nodeRange(const SelectionTree::Node *N, ParsedAST &AST) {
   const SourceManager &SM = AST.getSourceManager();
   const LangOptions &LangOpts = AST.getASTContext().getLangOpts();
   StringRef Buffer = SM.getBufferData(SM.getMainFileID());
+  if (llvm::isa_and_nonnull<TranslationUnitDecl>(N->ASTNode.get<Decl>()))
+    return Range{Position{}, offsetToPosition(Buffer, Buffer.size())};
   auto FileRange =
       toHalfOpenFileRange(SM, LangOpts, N->ASTNode.getSourceRange());
   assert(FileRange && "We should be able to get the File Range");
@@ -53,7 +57,7 @@ std::string nodeKind(const SelectionTree::Node *N) {
 }
 
 std::vector<const SelectionTree::Node *> allNodes(const SelectionTree &T) {
-  std::vector<const SelectionTree::Node *> Result = {T.root()};
+  std::vector<const SelectionTree::Node *> Result = {&T.root()};
   for (unsigned I = 0; I < Result.size(); ++I) {
     const SelectionTree::Node *N = Result[I];
     Result.insert(Result.end(), N->Children.begin(), N->Children.end());
@@ -63,16 +67,16 @@ std::vector<const SelectionTree::Node *> allNodes(const SelectionTree &T) {
 
 // Returns true if Common is a descendent of Root.
 // Verifies nothing is selected above Common.
-bool verifyCommonAncestor(const SelectionTree::Node *Root,
+bool verifyCommonAncestor(const SelectionTree::Node &Root,
                           const SelectionTree::Node *Common,
                           StringRef MarkedCode) {
-  if (Root == Common)
+  if (&Root == Common)
     return true;
-  if (Root->Selected)
+  if (Root.Selected)
     ADD_FAILURE() << "Selected nodes outside common ancestor\n" << MarkedCode;
   bool Seen = false;
-  for (const SelectionTree::Node *Child : Root->Children)
-    if (verifyCommonAncestor(Child, Common, MarkedCode)) {
+  for (const SelectionTree::Node *Child : Root.Children)
+    if (verifyCommonAncestor(*Child, Common, MarkedCode)) {
       if (Seen)
         ADD_FAILURE() << "Saw common ancestor twice\n" << MarkedCode;
       Seen = true;
@@ -239,6 +243,9 @@ TEST(SelectionTest, CommonAncestor) {
       {"int x = 42;^", nullptr},
       {"int x = 42^;", nullptr},
 
+      // Common ancestor is logically TUDecl, but we never return that.
+      {"^int x; int y;^", nullptr},
+
       // Node types that have caused problems in the past.
       {"template <typename T> void foo() { [[^T]] t; }",
        "TemplateTypeParmTypeLoc"},
@@ -256,18 +263,17 @@ TEST(SelectionTest, CommonAncestor) {
     Annotations Test(C.Code);
     auto AST = TestTU::withCode(Test.code()).build();
     auto T = makeSelectionTree(C.Code, AST);
+    EXPECT_EQ("TranslationUnitDecl", nodeKind(&T.root())) << C.Code;
 
     if (Test.ranges().empty()) {
       // If no [[range]] is marked in the example, there should be no selection.
       EXPECT_FALSE(T.commonAncestor()) << C.Code << "\n" << T;
-      EXPECT_FALSE(T.root()) << C.Code << "\n" << T;
     } else {
-      // If there is an expected selection, both common ancestor and root
-      // should exist with the appropriate node types in them.
+      // If there is an expected selection, common ancestor should exist
+      // with the appropriate node type.
       EXPECT_EQ(C.CommonAncestorKind, nodeKind(T.commonAncestor()))
           << C.Code << "\n"
           << T;
-      EXPECT_EQ("TranslationUnitDecl", nodeKind(T.root())) << C.Code;
       // Convert the reported common ancestor to a range and verify it.
       EXPECT_EQ(nodeRange(T.commonAncestor(), AST), Test.range())
           << C.Code << "\n"
@@ -316,10 +322,10 @@ TEST(SelectionTest, Selected) {
           void foo(^$C[[unique_ptr<$C[[unique_ptr<$C[[int]]>]]>]]^ a) {}
       )cpp",
       R"cpp(int a = [[5 >^> 1]];)cpp",
-      R"cpp(
+      R"cpp([[
         #define ECHO(X) X
         ECHO(EC^HO([[$C[[int]]) EC^HO(a]]));
-      )cpp",
+      ]])cpp",
   };
   for (const char *C : Cases) {
     Annotations Test(C);

clang-tools-extra/clangd/unittests/TweakTests.cpp

@@ -266,16 +266,16 @@ TEST(TweakTest, ShowSelectionTree) {
   llvm::StringLiteral ID = "ShowSelectionTree";
 
   checkAvailable(ID, "^int f^oo() { re^turn 2 ^+ 2; }");
-  checkNotAvailable(ID, "/*c^omment*/ int foo() return 2 ^ + 2; }");
+  checkAvailable(ID, "/*c^omment*/ int foo() return 2 ^ + 2; }");
 
   const char *Input = "int fcall(int); int x = fca[[ll(2 +]]2);";
-  const char *Output = R"(TranslationUnitDecl 
-  VarDecl int x = fcall(2 + 2)
-   .CallExpr fcall(2 + 2)
-      ImplicitCastExpr fcall
-       .DeclRefExpr fcall
-     .BinaryOperator 2 + 2
-       *IntegerLiteral 2
+  const char *Output = R"( TranslationUnitDecl 
+   VarDecl int x = fcall(2 + 2)
+    .CallExpr fcall(2 + 2)
+       ImplicitCastExpr fcall
+        .DeclRefExpr fcall
+      .BinaryOperator 2 + 2
+        *IntegerLiteral 2
 )";
   EXPECT_EQ(Output, getMessage(ID, Input));
 }