[clangd] SelectionTree treats TranslationUnitDecl consistently with other containers.

[sammccall]

Summary

Previously TranslationUnitDecl would never be selected.
This means root() and commonAncestor() are now never null, and therefore changed
them to references.

[hokein]

clang-tools-extra/clangd/Selection.h

100 ↗

(On Diff #211138)

I think it would be nice to mention "what means no selection" (TUDecl with partial)?

And it seems that we are missing some documentation describing the selection tree behavior on the preprocessor elements e.g. #incl^ude, header guard.

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

47 ↗

(On Diff #211138)

we'll get a TUDecl if the whole file is selected, however this tweak doesn't want to traversal the whole TU (otherwise we will generate highlighting tokens for #included files), we only interested in the top level decls of the main file.

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

38 ↗

(On Diff #211138)

Is the N->Parent intended? this seems like we'd exclude the TUDecl?

[sammccall]

Reverted commonAncestor() back to a pointer, return null if ancestor is TUDecl.

[sammccall]

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

47 ↗

(On Diff #211138)

Right! We discussed this further offline, and came to the conclusions that:

• traversing the whole TU is undesirable almost always
• it's already possible to get the whole TU from commonAncestor today (by selecting multiple top-level decls), and there are resulting bugs
• returning null instead of TU from commonAncestor() is a safer default

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

38 ↗

(On Diff #211138)

This is obsolete now, but was useful to flag: it's exactly the type of question that tweaks shouldn't have to worry about by default.

[sammccall]

doc selectiontree's interactions with preprocessor

[hokein]

looks good!

clang-tools-extra/trunk/clangd/Selection.h

 //
 // Usually commonAncestor() is the place to start:
 //  - it's the simplest answer to "what node is under the cursor"
 //  - the selected Expr (for example) can be found by walking up the parent
 //    chain and checking Node->ASTNode.
 //  - 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 {
 public:
   // Creates a selection tree at the given byte offset in the main file.
   // This is approximately equivalent to a range of one character.
   // (Usually, the character to the right of Offset, sometimes to the left).
   SelectionTree(ASTContext &AST, unsigned Offset);
struct Node {
     Selection Selected;
     // Walk up the AST to get the DeclContext of this Node,
     // which is not the node itself.
     const DeclContext& getDeclContext() const;
     // Printable node kind, like "CXXRecordDecl" or "AutoTypeLoc".
     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.
   const Node *Root;
   clang::PrintingPolicy PrintPolicy;
 
   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;
   }
 };
 
 } // namespace clangd
 } // namespace clang
 #endif

clang-tools-extra/trunk/clangd/Selection.cpp

public:
   // Runs the visitor to gather selected nodes and their ancestors.
   // If there is any selection, the root (TUDecl) is the first node.
   static std::deque<Node> collect(ASTContext &AST, const PrintingPolicy &PP,
                                   unsigned Begin, unsigned End, FileID File) {
     SelectionVisitor V(AST, PP, Begin, End, File);
     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);
   }
 
   // We traverse all "well-behaved" nodes the same way:
   //  - push the node onto the stack
   //  - traverse its children recursively
   //  - pop it from the stack
   //  - hit testing: is intersection(node, selection) - union(children) empty?
SelectionTree::SelectionTree(ASTContext &AST, unsigned Begin, unsigned End)
   Root = Nodes.empty() ? nullptr : &Nodes.front();
   dlog("Built selection tree\n{0}", *this);
 }
 
 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 {
   for (const Node* CurrentNode = this; CurrentNode != nullptr;
        CurrentNode = CurrentNode->Parent) {
     if (const Decl* Current = CurrentNode->ASTNode.get<Decl>()) {
       if (CurrentNode != this)
         if (auto *DC = dyn_cast<DeclContext>(Current))

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

 ///        ImplicitCastExpr fcall
 ///         .DeclRefExpr fcall
 ///        BinaryOperator 2 + 2
 ///          *IntegerLiteral 2
 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));
   }
   std::string title() const override { return "Show selection tree"; }
   Intent intent() const override { return Info; }
   bool hidden() const override { return true; }
 };
 REGISTER_TWEAK(ShowSelectionTree)

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

 //===-- SelectionTests.cpp - ----------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 #include "Annotations.h"
 #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"
 
 namespace clang {
 namespace clangd {
 namespace {
 using ::testing::UnorderedElementsAreArray;
 

 }
 
 Range nodeRange(const SelectionTree::Node *N, ParsedAST &AST) {
   if (!N)
     return Range{};
   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");
   return Range{
       offsetToPosition(Buffer, SM.getFileOffset(FileRange->getBegin())),
       offsetToPosition(Buffer, SM.getFileOffset(FileRange->getEnd()))};
 }
 
 std::string nodeKind(const SelectionTree::Node *N) {
   return N ? N->kind() : "<null>";
 }
 
 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());
   }
   return Result;
 }
 
 // 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;
     }
   return Seen;
 }
 
 TEST(SelectionTest, CommonAncestor) {
Case Cases[] = {
       {"^", nullptr},
       {"void foo() { [[foo^^]] (); }", "DeclRefExpr"},
 
       // FIXME: Ideally we'd get a declstmt or the VarDecl itself here.
       // This doesn't happen now; the RAV doesn't traverse a node containing ;.
       {"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"},
 
       // No crash
       {
           R"cpp(
             template <class T> struct Foo {};
             template <[[template<class> class /*cursor here*/^U]]>
              struct Foo<U<int>*> {};
           )cpp",
           "TemplateTemplateParmDecl"},
   };
   for (const Case &C : Cases) {
     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"
           << T;
 
       // Check that common ancestor is reachable on exactly one path from root,
       // and no nodes outside it are selected.
       EXPECT_TRUE(verifyCommonAncestor(T.root(), T.commonAncestor(), C.Code))
const char *Cases[] = {
         }
       )cpp",
       R"cpp(
           template <class T>
           struct unique_ptr {};
           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);
     auto AST = TestTU::withCode(Test.code()).build();
     auto T = makeSelectionTree(C, AST);
 
     std::vector<Range> Complete, Partial;
     for (const SelectionTree::Node *N : allNodes(T))

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

EXPECT_THAT(Result,
               ::testing::HasSubstr("<col:9> 'int' 2\n`-IntegerLiteral"));
   EXPECT_THAT(Result, ::testing::HasSubstr("<col:13> 'int' 2"));
 }
 
 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 
+  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));
 }
 
 TEST(TweakTest, DumpRecordLayout) {
   llvm::StringLiteral ID = "DumpRecordLayout";
   checkAvailable(ID, "^s^truct ^X ^{ int x; ^};");
   checkNotAvailable(ID, "struct X { int ^a; };");