[clangd] Start using SyntaxTrees for folding ranges feature

[kbobyrev]

Summary

This is an initial attempt to start using Syntax Trees in clangd while improving state of folding ranges feature and experimenting with Syntax Tree capabilities.

[kbobyrev]

Add support for compound statements.

[kbobyrev]

Add one more test and remove unused include.

[kbobyrev]

@sammccall I've reduced the patch to the bare minimum (compound statements) as I've had some issues with couple of other kinds of folding ranges and I'll be adding support for the node kinds one by one.

[eduucaldas]

clang-tools-extra/clangd/SemanticSelection.cpp

60–62

63–66

Take a look at clang/include/clang/Tooling/Syntax/Nodes.h, syntax constructs usually have nice classes with accessors.

For instance CompoundStatement has the accessors getLbrace and getRbrace that seem to be exactly what you want.

However these might not give exactly the first leaf and last leaf in the case of syntactically incorrect code.

[gribozavr2]

clang-tools-extra/clangd/SemanticSelection.cpp

40–48

WDYT about "makeFoldingRange" or even "toFoldingRange" to emphasize it is a factory / conversion function (no actual computation)?

58

"collectFoldingRanges" was a better name I think.

59

Why not return the vector?

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

247

"Will" makes it sound like it is not intentional.

[sammccall]

clang-tools-extra/clangd/SemanticSelection.cpp

50–100

Have you considered how you want macro-expanded code to work?

As written, this code can produce ranges inside macro definitions, invalid ranges (if { and } aren't from the same macro expansion) or even crash (FirstToken->endLocation() may be an invalid one-past-the-end location).

My suggestion would be to have this return optional<FoldingRange> and simply bail out if the location are not file locations for now. (Later we can deal with macro args in some cases)

63–66

I think we should treat all bracket-like things generically. Today this is just CompoundStmt, but we want to handle init lists, function calls, parens around for-loop conditions, template parameter and arg lists etc in the same way.

This sort of use is why the OpenParen/CloseParen NodeRoles are generic - we can have one set of logic to handle all of these. (No opinion on whether that should live here or in the syntax trees library, but putting it here for now seems fine).

So in the end I think checking the class name and then grabbing the braces by role (not kind) is the right thing here.
We definitely want to avoid asserting that the code looks the way we expect though.

67

this is worthy of a comment (fold the entire range inside the brackets, including whitespace)

71

until we support FoldingRangeClientCapabilities, should we just have the line check?

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

206–208

This is too many test cases for straightforward compoundstmt, I think. We have only once type of node handled for now, but we can't afford to write and maintain this many tests once we have lots.

I think a single if{}-elseif(nobraces)-else{} along with the existing function examples is probably enough. (Function examples show missing range for empty body, though it could use a comment)

We definitely need cases for macros:

• entire {} within one macro arg
• entire {} in the macro body
• some combination (e.g. { in macro body, } outside the macro)

(we won't need these for every node type, assuming we have some consistent handling)

[kbobyrev]

Resolve review comments.

clang-tools-extra/clangd/SemanticSelection.cpp

63–66

So in the end I think checking the class name and then grabbing the braces by role (not kind) is the right thing here.
We definitely want to avoid asserting that the code looks the way we expect though.

Can you elaborate a bit on how this would work? Is your proposal to iterate through CompoundStatement first-level children and grab the OpenParen + CloseParen roles?

[sammccall]

clang-tools-extra/clangd/SemanticSelection.cpp

63–66

Exactly. And bail out if both don't exist.

And this can be done on Tree, so it's trivial to add support for function calls etc (but feel free to keep the scope small)

clang/include/clang/Tooling/Syntax/Nodes.h

771 ↗

(On Diff #299438)

This doesn't look right: now a const method grants mutable access to the child. I think we need both overloads :-(

(Fortunately this is to be tablegen'd one day...)

[kbobyrev]

Resolve comments.

[eduucaldas]

clang-tools-extra/clangd/SemanticSelection.cpp

63–66

I very much agree on all the points you made Sam.

I think the logic to get ranges from node roles should eventually live in the syntax trees library. We might want to hide some of these lower-level functions in the future, so it would be better to not rely on them. But it's no harm for now :).

[eduucaldas]

clang/include/clang/Tooling/Syntax/Tree.h

172–177

I think that makes sense, since all the functions used by findChild are public anyways.

WDYT Dmitri? Now that the API is being used I realize its surface ^^. Perhaps we should pour some thought into that in the future :)

clang/lib/Tooling/Syntax/Tree.cpp

304–308

Similarly to the const version of findFirstLeaf. I think this should work :)

Also you could put the definition in clang/Tooling/Syntax/Tree.h.

[kbobyrev]

Prevent code duplication for const-nonconst versions of the same function.

[kbobyrev]

Fix a typo and remove invalid const-ness.

[gribozavr2]

clang/include/clang/Tooling/Syntax/Tree.h

172–177

I think it makes sense. Const-correctness often creates this type of API duplication.

[sammccall]

clang-tools-extra/clangd/SemanticSelection.cpp

40–48

Here you're looking at positions but ignoring file ID, which is always a reason to be cautious...

We've dropped macro IDs, and we're traversing only the syntax tree for this file, but we could still end up in the wrong file:

void foo() {
  #include "SomeTablegen.inc"
}

We need to verify both start/end are in the main file.

I'd suggest passing in the main FileID as a param here, and decomposing the start/end locations into (FileID, Offset), and bailing out if either FID doesn't match.
This subsumes the macro check (since macro expansions have their own FileIDs) though you might want to keep it to be explicit, as we likely want to support some cases later.

As a bonus, then you get to use SourceManager::get{Line,Column}Number(FileID, Offset) instead of the "spelling" variants that are slightly confusing as we've already established we have a file location already.

63

strictly this should probably be findLastChild both semantically and for performance... but that doesn't exist, because it's a single-linked list for now.

Not a big deal in practice, but we should fix this (and add STL iterators!)

166

this actually does the right thing (traverses the ASTContext's registered roots rather than the TUDecl itself), but... what a misleading API. We should fix that to take ASTContext instead...

(Not in this patch, maybe a followup?)

clang/include/clang/Tooling/Syntax/Tree.h

175

I think it's more conventional to implement the const variant and use const_cast for the non-const variant, so the compiler will verify that the function implementation is const, and the cast is just adding "I know if I'm not const then the returned object isn't either".

By implementing the *non-const* variant, the compiler doesn't help you verify any part of the const contract.

[kbobyrev]

Decompose locations and add checks for FileID == MainFileID.

Also, rebase on top of master.

[kbobyrev]

clang-tools-extra/clangd/SemanticSelection.cpp

63

Added a FIXME.

[sammccall]

clang-tools-extra/clangd/SemanticSelection.cpp

42

just use SourceLocations or SourceRange here? We have the SourceManager to decompose them anyway.

[kbobyrev]

Address post-LGTM comments.

[kbobyrev]

Resolve merge conflict.

clang-tools-extra/clangd/SemanticSelection.cpp

 //===--- SemanticSelection.cpp -----------------------------------*- C++-*-===//
 //
 // 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 "SemanticSelection.h"
 #include "FindSymbols.h"
 #include "ParsedAST.h"
 #include "Protocol.h"
 #include "Selection.h"
 #include "SourceCode.h"
 #include "clang/AST/DeclBase.h"
 #include "clang/Basic/SourceLocation.h"
+#include "clang/Basic/SourceManager.h"
+#include "clang/Basic/TokenKinds.h"
+#include "clang/Tooling/Syntax/BuildTree.h"
+#include "clang/Tooling/Syntax/Nodes.h"
+#include "clang/Tooling/Syntax/Tree.h"
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/Support/Casting.h"
 #include "llvm/Support/Error.h"
+#include <queue>
+#include <vector>
 
 namespace clang {
 namespace clangd {
 namespace {
 
 // Adds Range \p R to the Result if it is distinct from the last added Range.
 // Assumes that only consecutive ranges can coincide.
 void addIfDistinct(const Range &R, std::vector<Range> &Result) {
   if (Result.empty() || Result.back() != R) {
     Result.push_back(R);
   }
 }
 
-// Recursively collects FoldingRange from a symbol and its children.
-void collectFoldingRanges(DocumentSymbol Symbol,
-                          std::vector<FoldingRange> &Result) {
+typedef std::pair<FileID, unsigned> LocInfo;
+
+llvm::Optional<FoldingRange> toFoldingRange(LocInfo Begin, LocInfo End,
+                                            const SourceManager &SM) {
+  // Do not produce folding ranges if either range ends is not within the main
+  // file. Macros have their own FileID so this also checks if locations are not
+  // within the macros.
+  if ((Begin.first != SM.getMainFileID()) || (End.first != SM.getMainFileID()))
+    return llvm::None;
   FoldingRange Range;
-  Range.startLine = Symbol.range.start.line;
-  Range.startCharacter = Symbol.range.start.character;
-  Range.endLine = Symbol.range.end.line;
-  Range.endCharacter = Symbol.range.end.character;
-  Result.push_back(Range);
-  for (const auto &Child : Symbol.children)
-    collectFoldingRanges(Child, Result);
+  Range.startCharacter = SM.getColumnNumber(Begin.first, Begin.second) - 1;
+  Range.startLine = SM.getLineNumber(Begin.first, Begin.second) - 1;
+  Range.endCharacter = SM.getColumnNumber(End.first, End.second) - 1;
+  Range.endLine = SM.getLineNumber(End.first, End.second) - 1;
+  return Range;
+}
+
+llvm::Optional<FoldingRange> extractFoldingRange(const syntax::Node *Node,
+                                                 const SourceManager &SM) {
+  if (const auto *Stmt = dyn_cast<syntax::CompoundStatement>(Node)) {
+    const auto *LBrace = cast_or_null<syntax::Leaf>(
+        Stmt->findChild(syntax::NodeRole::OpenParen));
+    // FIXME(kirillbobyrev): This should find the last child. Compound
+    // statements have only one pair of braces so this is valid but for other
+    // node kinds it might not be correct.
+    const auto *RBrace = cast_or_null<syntax::Leaf>(
+        Stmt->findChild(syntax::NodeRole::CloseParen));
+    if (!LBrace || !RBrace)
+      return llvm::None;
+    // Fold the entire range within braces, including whitespace.
+    const LocInfo LBraceLocInfo =
+                      SM.getDecomposedLoc(LBrace->getToken()->endLocation()),
+                  RBraceLocInfo =
+                      SM.getDecomposedLoc(RBrace->getToken()->location());
+    auto Range = toFoldingRange(LBraceLocInfo, RBraceLocInfo, SM);
+    // Do not generate folding range for compound statements without any
+    // nodes and newlines.
+    if (Range && Range->startLine != Range->endLine)
+      return Range;
+  }
+  return llvm::None;
+}
+
+// Traverse the tree and collect folding ranges along the way.
+std::vector<FoldingRange> collectFoldingRanges(const syntax::Node *Root,
+                                               const SourceManager &SM) {
+  std::queue<const syntax::Node *> Nodes;
+  Nodes.push(Root);
+  std::vector<FoldingRange> Result;
+  while (!Nodes.empty()) {
+    const syntax::Node *Node = Nodes.front();
+    Nodes.pop();
+    const auto Range = extractFoldingRange(Node, SM);
+    if (Range)
+      Result.push_back(*Range);
+    if (const auto *T = dyn_cast<syntax::Tree>(Node))
+      for (const auto *NextNode = T->getFirstChild(); NextNode;
+           NextNode = NextNode->getNextSibling())
+        Nodes.push(NextNode);
+  }
+  return Result;
 }
 
 } // namespace
 
 llvm::Expected<SelectionRange> getSemanticRanges(ParsedAST &AST, Position Pos) {
   std::vector<Range> Ranges;
   const auto &SM = AST.getSourceManager();
   const auto &LangOpts = AST.getLangOpts();
llvm::Expected<SelectionRange> getSemanticRanges(ParsedAST &AST, Position Pos) {
   }
 
   return std::move(Head);
 }
 
 // FIXME(kirillbobyrev): Collect comments, PP conditional regions, includes and
 // other code regions (e.g. public/private/protected sections of classes,
 // control flow statement bodies).
-// Related issue:
-// https://github.com/clangd/clangd/issues/310
+// Related issue: https://github.com/clangd/clangd/issues/310
 llvm::Expected<std::vector<FoldingRange>> getFoldingRanges(ParsedAST &AST) {
-  // FIXME(kirillbobyrev): getDocumentSymbols() is conveniently available but
-  // limited (e.g. doesn't yield blocks inside functions and provides ranges for
-  // nodes themselves instead of their contents which is less useful). Replace
-  // this with a more general RecursiveASTVisitor implementation instead.
-  auto DocumentSymbols = getDocumentSymbols(AST);
-  if (!DocumentSymbols)
-    return DocumentSymbols.takeError();
-  std::vector<FoldingRange> Result;
-  for (const auto &Symbol : *DocumentSymbols)
-    collectFoldingRanges(Symbol, Result);
-  return Result;
+  syntax::Arena A(AST.getSourceManager(), AST.getLangOpts(), AST.getTokens());
+  const auto *SyntaxTree =
+      syntax::buildSyntaxTree(A, *AST.getASTContext().getTranslationUnitDecl());
+  return collectFoldingRanges(SyntaxTree, AST.getSourceManager());
 }
 
 } // namespace clangd
 } // namespace clang

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

EXPECT_THAT(gatherRanges(Ranges->front()),
               ElementsAreArray(SourceAnnotations.ranges()));
   EXPECT_THAT(gatherRanges(Ranges->back()),
               ElementsAre(SourceAnnotations.range("empty")));
 }
 
 TEST(FoldingRanges, All) {
   const char *Tests[] = {
       R"cpp(
-        [[int global_variable]];
+        #define FOO int foo() {\
+          int Variable = 42; \
+        }
 
-        [[void func() {
-          int v = 100;
-        }]]
+        // Do not generate folding range for braces within macro expansion.
+        FOO
+
+        // Do not generate folding range within macro arguments.
+        #define FUNCTOR(functor) functor
+        void func() {[[
+          FUNCTOR([](){});
+        ]]}
+
+        // Do not generate folding range with a brace coming from macro.
+        #define LBRACE {
+        void bar() LBRACE
+          int X = 42;
+        }
+      )cpp",
+      R"cpp(
+        void func() {[[
+          int Variable = 100;
+
+          if (Variable > 5) {[[
+            Variable += 42;
+          ]]} else if (Variable++)
+            ++Variable;
+          else {[[
+            Variable--;
+          ]]}
+
+          // Do not generate FoldingRange for empty CompoundStmts.
+          for (;;) {}
+
+          // If there are newlines between {}, we should generate one.
+          for (;;) {[[
+
+          ]]}
+        ]]}
       )cpp",
       R"cpp(
-        [[class Foo {
+        class Foo {
         public:
-          [[Foo() {
+          Foo() {[[
             int X = 1;
-          }]]
+          ]]}
 
         private:
-          [[int getBar() {
+          int getBar() {[[
             return 42;
-          }]]
+          ]]}
 
-          [[void getFooBar() { }]]
-        }]];
+          // Braces are located at the same line: no folding range here.
+          void getFooBar() { }
+        };
       )cpp",
   };
   for (const char *Test : Tests) {
     auto T = Annotations(Test);
     auto AST = TestTU::withCode(T.code()).build();
     EXPECT_THAT(gatherFoldingRanges(llvm::cantFail(getFoldingRanges(AST))),
                 UnorderedElementsAreArray(T.ranges()))
         << Test;
   }
 }
 
 } // namespace
 } // namespace clangd
 } // namespace clang

clang/include/clang/Tooling/Syntax/Tree.h

 class Tree : public Node {
 public:
   using Node::Node;
   static bool classof(const Node *N);
 
   Node *getFirstChild() { return FirstChild; }
   const Node *getFirstChild() const { return FirstChild; }
 
-  Leaf *findFirstLeaf();
-  const Leaf *findFirstLeaf() const {
-    return const_cast<Tree *>(this)->findFirstLeaf();
+  const Leaf *findFirstLeaf() const;
+  Leaf *findFirstLeaf() {
+    return const_cast<Leaf *>(const_cast<const Tree *>(this)->findFirstLeaf());
   }
 
-  Leaf *findLastLeaf();
-  const Leaf *findLastLeaf() const {
-    return const_cast<Tree *>(this)->findLastLeaf();
+  const Leaf *findLastLeaf() const;
+  Leaf *findLastLeaf() {
+    return const_cast<Leaf *>(const_cast<const Tree *>(this)->findLastLeaf());
   }
 
-protected:
   /// Find the first node with a corresponding role.
-  Node *findChild(NodeRole R);
+  const Node *findChild(NodeRole R) const;
+  Node *findChild(NodeRole R) {
+    return const_cast<Node *>(const_cast<const Tree *>(this)->findChild(R));
+  }
 
 private:
   /// Prepend \p Child to the list of children and and sets the parent pointer.
   /// A very low-level operation that does not check any invariants, only used
   /// by TreeBuilder and FactoryImpl.
   /// EXPECTS: Role != Detached.
   void prependChildLowLevel(Node *Child, NodeRole Role);
   /// Like the previous overload, but does not set role for \p Child.
   /// EXPECTS: Child->Role != Detached

clang/lib/Tooling/Syntax/Tree.cpp

 }
 
 void syntax::Node::assertInvariantsRecursive() const {
 #ifndef NDEBUG
   traverse(this, [&](const syntax::Node *N) { N->assertInvariants(); });
 #endif
 }
 
-syntax::Leaf *syntax::Tree::findFirstLeaf() {
-  for (auto *C = getFirstChild(); C; C = C->getNextSibling()) {
-    if (auto *L = dyn_cast<syntax::Leaf>(C))
+const syntax::Leaf *syntax::Tree::findFirstLeaf() const {
+  for (const auto *C = getFirstChild(); C; C = C->getNextSibling()) {
+    if (const auto *L = dyn_cast<syntax::Leaf>(C))
       return L;
-    if (auto *L = cast<syntax::Tree>(C)->findFirstLeaf())
+    if (const auto *L = cast<syntax::Tree>(C)->findFirstLeaf())
       return L;
   }
   return nullptr;
 }
 
-syntax::Leaf *syntax::Tree::findLastLeaf() {
-  syntax::Leaf *Last = nullptr;
-  for (auto *C = getFirstChild(); C; C = C->getNextSibling()) {
-    if (auto *L = dyn_cast<syntax::Leaf>(C))
+const syntax::Leaf *syntax::Tree::findLastLeaf() const {
+  const syntax::Leaf *Last = nullptr;
+  for (const auto *C = getFirstChild(); C; C = C->getNextSibling()) {
+    if (const auto *L = dyn_cast<syntax::Leaf>(C))
       Last = L;
-    else if (auto *L = cast<syntax::Tree>(C)->findLastLeaf())
+    else if (const auto *L = cast<syntax::Tree>(C)->findLastLeaf())
       Last = L;
   }
   return Last;
 }
 
-syntax::Node *syntax::Tree::findChild(NodeRole R) {
-  for (auto *C = FirstChild; C; C = C->getNextSibling()) {
+const syntax::Node *syntax::Tree::findChild(NodeRole R) const {
+  for (const auto *C = FirstChild; C; C = C->getNextSibling()) {
     if (C->getRole() == R)
       return C;
   }
   return nullptr;
 }
 
 bool syntax::List::classof(const syntax::Node *N) {
   switch (N->getKind()) {
   case syntax::NodeKind::NestedNameSpecifier:
   case syntax::NodeKind::CallArguments:
   case syntax::NodeKind::ParameterDeclarationList:
     return true;
   default:
     return false;
   }
 }
 
 std::vector<syntax::List::ElementAndDelimiter<syntax::Node>>
 syntax::List::getElementsAsNodesAndDelimiters() {
   if (!getFirstChild())