[clangd] Lib to compute and represent selection under cursor.

[sammccall]

Summary

The primary problem this solves is to expose the codeAction selection to
AST-based refactorings in a way that makes it easy and efficient for them to
bind to the right parts of the AST.

It should also allow us to make XRefs based features (textDocument/definition)
more robust, more easily implement textDocument/typeDefinition etc.
As an example, template parameter references can be identified without special
handling.
There should be slight speedup too: we can prune most of the AST traversal
in most cases.

Elephant in the room: this is similar-but-different to Tooling/Refactoring/ASTSelection.
That captures a smaller set of AST nodes, has a slightly different way of
representing selections, and generally has mare features and does more work.
The overall shape is pretty similar, and yet I can't quite get to behave as I
expect.

[sammccall]

eliminate a gratuitous template

[ilya-biryukov]

Awesome job. Any data on how much time this takes to build? Do you think this would ever be a bottleneck?

clangd/Selection.cpp

30 ↗

(On Diff #184672)

Any reason to limit this to the main file and not take a FileID as a parameter? This looks like an artificial limitation, even if that's the only way we can use it now in clangd.

51 ↗

(On Diff #184672)

NIT: maybe make the ctor private and create a static function? Alternatively, maybe create a free function that creates returns deque<Node>?
Would simplify the callsite a bit, making it a bit more readable.

113 ↗

(On Diff #184672)

NIT: maybe specify auto twice?

auto B = ...;
auto E = ...;

158 ↗

(On Diff #184672)

NIT: specify auto twice?

223 ↗

(On Diff #184672)

Why not vector? We seem to only push_back/pop_back

224 ↗

(On Diff #184672)

NIT: maybe split into three declarations?

[sammccall]

In D57562#1380170, @ilya-biryukov wrote:

Awesome job. Any data on how much time this takes to build? Do you think this would ever be a bottleneck?

I don't have data yet, I need to add some tracing to D57570 and try it out with big files in VSCode.
Unless I'm *seriously* miscalculating though, it should be:

• almost free on code that's not mostly macro-expanded where the selection is smallish (canSafelySkipNode optimization means not traversing much)
• a fairly cheap RecursiveASTVisitor if the selection is fairly small (we just do push/pop and integer compares on SourceLocations, lexing only happens for nodes that actually are partially selected)
• maybe quite expensive if you select-all on a huge file - this is what I should test. Could put some caps on size if needed.

clangd/Selection.cpp

30 ↗

(On Diff #184672)

I can't think of anything that specificallly would break (unless someone passed a FileID corresponding to a macro...)

On the other hand, I can't think of any reason to allow it: this is very geared to user-driven selections, and users only ever see code in the main file. Do you have something in mind?

(In the absence of a need for it, I find plumbing an extra parameter and losing a little bit of concreteness in the mental model to be the main reasons to resist...)

113 ↗

(On Diff #184672)

hmm, when these are logically pairs I prefer to omit the second auto, to emphasize that they're the same type...
is there a gotcha here?

223 ↗

(On Diff #184672)

deque guarantees pointer stability, if we used vectors then push/pop would break the parent/child links between nodes.

224 ↗

(On Diff #184672)

as above, I think SelBegin/SelEnd belong together.
Split out SelBeginTokenStart and gave it a lengthy comment, as it's pretty subtle.

[sammccall]

Organization in Selection.cpp

[ilya-biryukov]

The "select everything in a very large file" is exactly the problematic case that came to mind before. I'm sure it's possible to build the selection tree in linear time, but the traversals might also become a problem.

For the code tweaks, we definitely don't want each tweak traversing the whole AST. There are multiple potential options there, e.g. cutting the tree at a reasonable depth.

clangd/Selection.cpp

30 ↗

(On Diff #184672)

Nothing in particular, it was just a code style thing, I sometimes find more generic functions easier to read, e.g. in this case I won't need to keep the fact that we're calling getMainFileID() because selections are currently only used in the main file (even though in principle they're also applicable to any file).

Not a big deal, though, I'm happy either way.

113 ↗

(On Diff #184672)

Only if one starts using pointers(int *a, *b), but that alone is enough of a reason for me to keep away from multiple declarators in a single declaration.

I think in most of our code in clangd we tend to use separate declarations, so we may also want to keep away from them for consistency.

Not a big deal, though, keeping them is also totally fine.

223 ↗

(On Diff #184672)

Got you! And thanks for adding a comment.

[sammccall]

In D57562#1380335, @ilya-biryukov wrote:

The "select everything in a very large file" is exactly the problematic case that came to mind before. I'm sure it's possible to build the selection tree in linear time, but the traversals might also become a problem.

(Just so we're on the same page: this is good discussion regardless, but do you see this as an issue we should try and resolve in this patch?)

So recursive traversal downward is definitely the case that could be slow. My goal was to make it hard to be slow and I didn't entirely succeed (yet?).

I don't expect that to be a really common pattern. More common is traversing *upward* from the LCA (cheap because no branching), and possibly traversing downward but only looking for fairly specific nodes.

For the code tweaks, we definitely don't want each tweak traversing the whole AST. There are multiple potential options there, e.g. cutting the tree at a reasonable depth.

Another fairly natural-seeming idea is not to store children of Completely covered nodes - as soon as you hit something that's completely covered you should take it or leave it.
This should keep the number of nodes in the tree low, because it's hard to fractally select half of all parts of the tree - in practice you're selecting big subtrees.
However it's slightly tricky from an implementation point of view: we should still require a node own at least one character to be complete, which requires knowing the children, which we only know after traversal, which means it's too late to prune traversal of this subtree. And I'm also not sure if this might be too limiting for some tweaks we want to write.

clangd/Selection.cpp

30 ↗

(On Diff #184672)

Fair enough. I've put a FileID member in SelectionVisitor, and have removed all the hardcoded main-file references, it's clearer and the plumbing is only a little noisy. IIUC the code-style thing mainly applies to the implementation?

I'll still keep this out of the *public* API for now though, until we have some need for it.
The SelectionTree constructors document that offsets are with respect to the main file, and I don't think that interface is particularly prone to misinterpretation.

113 ↗

(On Diff #184672)

Expanded them all to separate decls, this is indeed more common.

[sammccall]

Review comments

[ilya-biryukov]

LGTM and sorry for delaying this with discussions.
The change looks solid, regardless of potential problems we may still run into.

clangd/Selection.cpp

30 ↗

(On Diff #184672)

Totally agree, public API is a totally different beast, the comment only applied to the implementation.

[lebedev.ri]

(And another review that omitted lists..)

clang-tools-extra/trunk/clangd/CMakeLists.txt

add_clang_library(clangDaemon
   GlobalCompilationDatabase.cpp
   Headers.cpp
   IncludeFixer.cpp
   JSONTransport.cpp
   Logger.cpp
   Protocol.cpp
   Quality.cpp
   RIFF.cpp
+  Selection.cpp
   SourceCode.cpp
   Threading.cpp
   Trace.cpp
   TUScheduler.cpp
   URI.cpp
   XRefs.cpp
 
   index/Background.cpp

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

+//===--- Selection.h - What's under the cursor? -------------------*-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
+//
+//===----------------------------------------------------------------------===//
+// Many features are triggered at locations/ranges and operate on AST nodes.
+// (e.g. go-to-definition or code tweaks).
+// At a high level, such features need to work out which node is the correct
+// target.
+//
+// There are a few levels of ambiguity here:
+//
+// Which tokens are included:
+//   int x = one + two;  // what should "go to definition" do?
+//            ^^^^^^
+//
+// Same token means multiple things:
+//   string("foo")       // class string, or a constructor?
+//   ^
+//
+// Which level of the AST is interesting?
+//   if (err) {          // reference to 'err', or operator bool(),
+//       ^               // or the if statement itself?
+//
+// Here we build and expose a data structure that allows features to resolve
+// these ambiguities in an appropriate way:
+//   - we determine which low-level nodes are partly or completely covered
+//     by the selection.
+//   - we expose a tree of the selected nodes and their lexical parents.
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_SELECTION_H
+#define LLVM_CLANG_TOOLS_EXTRA_CLANGD_SELECTION_H
+#include "clang/AST/ASTTypeTraits.h"
+#include "clang/AST/PrettyPrinter.h"
+#include "llvm/ADT/SmallVector.h"
+
+namespace clang {
+namespace clangd {
+class ParsedAST;
+
+// A selection can partially or completely cover several AST nodes.
+// The SelectionTree contains nodes that are covered, and their parents.
+// SelectionTree does not contain all AST nodes, rather only:
+//   Decl, Stmt, TypeLoc, NestedNamespaceSpecifierLoc, CXXCtorInitializer.
+// (These are the nodes with source ranges that fit in DynTypedNode).
+//
+// 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.
+//
+// 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);
+  // Creates a selection tree for the given range in the main file.
+  // The range includes bytes [Start, End).
+  // If Start == End, uses the same heuristics as SelectionTree(AST, Start).
+  SelectionTree(ASTContext &AST, unsigned Start, unsigned End);
+
+  // Describes to what extent an AST node is covered by the selection.
+  enum Selection {
+    // The AST node owns no characters covered by the selection.
+    // Note that characters owned by children don't count:
+    //   if (x == 0) scream();
+    //       ^^^^^^
+    // The IfStmt would be Unselected because all the selected characters are
+    // associated with its children.
+    // (Invisible nodes like ImplicitCastExpr are always unselected).
+    Unselected,
+    // The AST node owns selected characters, but is not completely covered.
+    Partial,
+    // The AST node owns characters, and is covered by the selection.
+    Complete,
+  };
+  // An AST node that is implicated in the selection.
+  // (Either selected directly, or some descendant is selected).
+  struct Node {
+    // The parent within the selection tree. nullptr for TranslationUnitDecl.
+    Node *Parent;
+    // Direct children within the selection tree.
+    llvm::SmallVector<const Node *, 8> Children;
+    // The corresponding node from the full AST.
+    ast_type_traits::DynTypedNode ASTNode;
+    // The extent to which this node is covered by the selection.
+    Selection Selected;
+  };
+
+  // The most specific common ancestor of all the selected nodes.
+  // If there is no selection, this is nullptr.
+  const Node *commonAncestor() const;
+  // The selection node corresponding to TranslationUnitDecl.
+  // If there is no selection, this is nullptr.
+  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";
+    return OS;
+  }
+};
+
+} // namespace clangd
+} // namespace clang
+#endif

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

+//===--- Selection.h ------------------------------------------------------===//
+//
+// 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 "Selection.h"
+#include "ClangdUnit.h"
+#include "clang/AST/RecursiveASTVisitor.h"
+
+namespace clang {
+namespace clangd {
+namespace {
+using Node = SelectionTree::Node;
+using ast_type_traits::DynTypedNode;
+
+// We find the selection by visiting written nodes in the AST, looking for nodes
+// that intersect with the selected character range.
+//
+// While traversing, we maintain a parent stack. As nodes pop off the stack,
+// we decide whether to keep them or not. To be kept, they must either be
+// selected or contain some nodes that are.
+//
+// For simple cases (not inside macros) we prune subtrees that don't intersect.
+class SelectionVisitor : public RecursiveASTVisitor<SelectionVisitor> {
+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, unsigned Begin,
+                                  unsigned End, FileID File) {
+    SelectionVisitor V(AST, 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();
+    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?
+  //  - attach it to the tree if it or any children hit the selection
+  //
+  // Two categories of nodes are not "well-behaved":
+  //  - those without source range information, we don't record those
+  //  - those that can't be stored in DynTypedNode.
+  // We're missing some interesting things like Attr due to the latter.
+  bool TraverseDecl(Decl *X) {
+    if (isa<TranslationUnitDecl>(X))
+      return Base::TraverseDecl(X); // Already pushed by constructor.
+    return traverseNode(X, [&] { return Base::TraverseDecl(X); });
+  }
+  bool TraverseTypeLoc(TypeLoc X) {
+    return traverseNode(&X, [&] { return Base::TraverseTypeLoc(X); });
+  }
+  bool TraverseTypeNestedNameSpecifierLoc(NestedNameSpecifierLoc X) {
+    return traverseNode(
+        &X, [&] { return Base::TraverseNestedNameSpecifierLoc(X); });
+  }
+  bool TraverseConstructorInitializer(CXXCtorInitializer *X) {
+    return traverseNode(
+        X, [&] { return Base::TraverseConstructorInitializer(X); });
+  }
+  // Stmt is the same, but this form allows the data recursion optimization.
+  bool dataTraverseStmtPre(Stmt *X) {
+    if (!X || canSafelySkipNode(X->getSourceRange()))
+      return false;
+    push(DynTypedNode::create(*X));
+    return true;
+  }
+  bool dataTraverseStmtPost(Stmt *X) {
+    pop();
+    return true;
+  }
+  // Uninteresting parts of the AST that don't have locations within them.
+  bool TraverseNestedNameSpecifier(NestedNameSpecifier *) { return true; }
+  bool TraverseType(QualType) { return true; }
+
+private:
+  using Base = RecursiveASTVisitor<SelectionVisitor>;
+  SelectionVisitor(ASTContext &AST, unsigned SelBegin, unsigned SelEnd,
+                   FileID SelFile)
+      : SM(AST.getSourceManager()), LangOpts(AST.getLangOpts()),
+        SelBegin(SelBegin), SelEnd(SelEnd), SelFile(SelFile),
+        SelBeginTokenStart(SM.getFileOffset(Lexer::GetBeginningOfToken(
+            SM.getComposedLoc(SelFile, SelBegin), SM, LangOpts))) {
+    // Ensure we have a node for the TU decl, regardless of traversal scope.
+    Nodes.emplace_back();
+    Nodes.back().ASTNode = DynTypedNode::create(*AST.getTranslationUnitDecl());
+    Nodes.back().Parent = nullptr;
+    Nodes.back().Selected = SelectionTree::Unselected;
+    Stack.push(&Nodes.back());
+  }
+
+  // Generic case of TraverseFoo. Func should be the call to Base::TraverseFoo.
+  // Node is always a pointer so the generic code can handle any null checks.
+  template <typename T, typename Func>
+  bool traverseNode(T *Node, const Func &Body) {
+    if (Node == nullptr || canSafelySkipNode(Node->getSourceRange()))
+      return true;
+    push(DynTypedNode::create(*Node));
+    bool Ret = Body();
+    pop();
+    return Ret;
+  }
+
+  // An optimization for a common case: nodes outside macro expansions that
+  // don't intersect the selection may be recursively skipped.
+  bool canSafelySkipNode(SourceRange S) {
+    auto B = SM.getDecomposedLoc(S.getBegin());
+    auto E = SM.getDecomposedLoc(S.getEnd());
+    if (B.first != SelFile || E.first != SelFile)
+      return false;
+    return B.second >= SelEnd || E.second < SelBeginTokenStart;
+  }
+
+  // Pushes a node onto the ancestor stack. Pairs with pop().
+  void push(DynTypedNode Node) {
+    Nodes.emplace_back();
+    Nodes.back().ASTNode = std::move(Node);
+    Nodes.back().Parent = Stack.top();
+    Nodes.back().Selected = SelectionTree::Unselected;
+    Stack.push(&Nodes.back());
+  }
+
+  // Pops a node off the ancestor stack, and finalizes it. Pairs with push().
+  void pop() {
+    Node &N = *Stack.top();
+    N.Selected = computeSelection(N);
+    if (N.Selected || !N.Children.empty()) {
+      // Attach to the tree.
+      N.Parent->Children.push_back(&N);
+    } else {
+      // Neither N any children are selected, it doesn't belong in the tree.
+      assert(&N == &Nodes.back());
+      Nodes.pop_back();
+    }
+    Stack.pop();
+  }
+
+  // Perform hit-testing of a complete Node against the selection.
+  // This runs for every node in the AST, and must be fast in common cases.
+  // This is called from pop(), so we can take children into account.
+  SelectionTree::Selection computeSelection(const Node &N) {
+    SourceRange S = N.ASTNode.getSourceRange();
+    if (!S.isValid())
+      return SelectionTree::Unselected;
+    // getTopMacroCallerLoc() allows selection of constructs in macro args. e.g:
+    //   #define LOOP_FOREVER(Body) for(;;) { Body }
+    //   void IncrementLots(int &x) {
+    //     LOOP_FOREVER( ++x; )
+    //   }
+    // Selecting "++x" or "x" will do the right thing.
+    auto B = SM.getDecomposedLoc(SM.getTopMacroCallerLoc(S.getBegin()));
+    auto E = SM.getDecomposedLoc(SM.getTopMacroCallerLoc(S.getEnd()));
+    // Otherwise, nodes in macro expansions can't be selected.
+    if (B.first != SelFile || E.first != SelFile)
+      return SelectionTree::Unselected;
+    // Cheap test: is there any overlap at all between the selection and range?
+    // Note that E.second is the *start* of the last token, which is why we
+    // compare against the "rounded-down" SelBegin.
+    if (B.second >= SelEnd || E.second < SelBeginTokenStart)
+      return SelectionTree::Unselected;
+
+    // We hit something, need some more precise checks.
+    // Adjust [B, E) to be a half-open character range.
+    E.second += Lexer::MeasureTokenLength(S.getEnd(), SM, LangOpts);
+    // This node's own selected text is (this range ^ selection) - child ranges.
+    // If that's empty, then we've only collided with children.
+    if (nodesCoverRange(N.Children, std::max(SelBegin, B.second),
+                        std::min(SelEnd, E.second)))
+      return SelectionTree::Unselected; // Hit children only.
+    // Some of our own characters are covered, this is a true hit.
+    return (B.second >= SelBegin && E.second <= SelEnd)
+               ? SelectionTree::Complete
+               : SelectionTree::Partial;
+  }
+
+  // Is the range [Begin, End) entirely covered by the union of the Nodes?
+  // (The range is a parent node's extent, and the covering nodes are children).
+  bool nodesCoverRange(llvm::ArrayRef<const Node *> Nodes, unsigned Begin,
+                       unsigned End) {
+    if (Begin >= End)
+      return true;
+    if (Nodes.empty())
+      return false;
+
+    // Collect all the expansion ranges, as offsets.
+    SmallVector<std::pair<unsigned, unsigned>, 8> ChildRanges;
+    for (const Node *N : Nodes) {
+      CharSourceRange R = SM.getExpansionRange(N->ASTNode.getSourceRange());
+      auto B = SM.getDecomposedLoc(R.getBegin());
+      auto E = SM.getDecomposedLoc(R.getEnd());
+      if (B.first != SelFile || E.first != SelFile)
+        continue;
+      assert(R.isTokenRange());
+      // Try to cover up to the next token, spaces between children don't count.
+      if (auto Tok = Lexer::findNextToken(R.getEnd(), SM, LangOpts))
+        E.second = SM.getFileOffset(Tok->getLocation());
+      else
+        E.second += Lexer::MeasureTokenLength(R.getEnd(), SM, LangOpts);
+      ChildRanges.push_back({B.second, E.second});
+    }
+    llvm::sort(ChildRanges);
+
+    // Scan through the child ranges, removing as we go.
+    for (const auto R : ChildRanges) {
+      if (R.first > Begin)
+        return false;   // [Begin, R.first) is not covered.
+      Begin = R.second; // Eliminate [R.first, R.second).
+      if (Begin >= End)
+        return true; // Remaining range is empty.
+    }
+    return false; // Went through all children, trailing characters remain.
+  }
+
+  SourceManager &SM;
+  const LangOptions &LangOpts;
+  std::stack<Node *> Stack;
+  std::deque<Node> Nodes; // Stable pointers as we add more nodes.
+  // Half-open selection range.
+  unsigned SelBegin;
+  unsigned SelEnd;
+  FileID SelFile;
+  // If the selection start slices a token in half, the beginning of that token.
+  // This is useful for checking whether the end of a token range overlaps
+  // the selection: range.end < SelBeginTokenStart is equivalent to
+  // range.end + measureToken(range.end) < SelBegin (assuming range.end points
+  // to a token), and it saves a lex every time.
+  unsigned SelBeginTokenStart;
+};
+
+} // namespace
+
+void SelectionTree::print(llvm::raw_ostream &OS, const SelectionTree::Node &N,
+                          int Indent) const {
+  if (N.Selected)
+    OS.indent(Indent - 1) << (N.Selected == SelectionTree::Complete ? '*'
+                                                                    : '.');
+  else
+    OS.indent(Indent);
+  OS << N.ASTNode.getNodeKind().asStringRef() << " ";
+  N.ASTNode.print(OS, PrintPolicy);
+  OS << "\n";
+  for (const Node *Child : N.Children)
+    print(OS, *Child, Indent + 2);
+}
+
+// Decide which selection emulates a "point" query in between characters.
+static std::pair<unsigned, unsigned> pointBounds(unsigned Offset, FileID FID,
+                                                 ASTContext &AST) {
+  StringRef Buf = AST.getSourceManager().getBufferData(FID);
+  // Edge-cases where the choice is forced.
+  if (Buf.size() == 0)
+    return {0, 0};
+  if (Offset == 0)
+    return {0, 1};
+  if (Offset == Buf.size())
+    return {Offset - 1, Offset};
+  // We could choose either this byte or the previous. Usually we prefer the
+  // character on the right of the cursor (or under a block cursor).
+  // But if that's whitespace, we likely want the token on the left.
+  if (isWhitespace(Buf[Offset]) && !isWhitespace(Buf[Offset - 1]))
+    return {Offset - 1, Offset};
+  return {Offset, Offset + 1};
+}
+
+SelectionTree::SelectionTree(ASTContext &AST, unsigned Begin, unsigned End)
+    : PrintPolicy(AST.getLangOpts()) {
+  // No fundamental reason the selection needs to be in the main file,
+  // but that's all clangd has needed so far.
+  FileID FID = AST.getSourceManager().getMainFileID();
+  if (Begin == End)
+    std::tie(Begin, End) = pointBounds(Begin, FID, AST);
+  PrintPolicy.TerseOutput = true;
+
+  Nodes = SelectionVisitor::collect(AST, Begin, End, FID);
+  Root = Nodes.empty() ? nullptr : &Nodes.front();
+}
+
+SelectionTree::SelectionTree(ASTContext &AST, unsigned Offset)
+    : SelectionTree(AST, Offset, Offset) {}
+
+const Node *SelectionTree::commonAncestor() const {
+  if (!Root)
+    return nullptr;
+  for (const Node *Ancestor = Root;; Ancestor = Ancestor->Children.front()) {
+    if (Ancestor->Selected || Ancestor->Children.size() > 1)
+      return Ancestor;
+    // The tree only contains ancestors of the interesting nodes.
+    assert(!Ancestor->Children.empty() && "bad node in selection tree");
+  }
+}
+
+} // namespace clangd
+} // namespace clang

clang-tools-extra/trunk/unittests/clangd/CMakeLists.txt

add_extra_unittest(ClangdTests
   FuzzyMatchTests.cpp
   GlobalCompilationDatabaseTests.cpp
   HeadersTests.cpp
   IndexActionTests.cpp
   IndexTests.cpp
   JSONTransportTests.cpp
   QualityTests.cpp
   RIFFTests.cpp
+  SelectionTests.cpp
   SerializationTests.cpp
   SourceCodeTests.cpp
   SymbolCollectorTests.cpp
   SymbolInfoTests.cpp
   SyncAPI.cpp
   TUSchedulerTests.cpp
   TestFS.cpp
   TestIndex.cpp

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

+//===-- RIFFTests.cpp - Binary container unit tests -----------------------===//
+//
+// 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 "gmock/gmock.h"
+#include "gtest/gtest.h"
+
+namespace clang {
+namespace clangd {
+namespace {
+using ::testing::UnorderedElementsAreArray;
+
+SelectionTree makeSelectionTree(const StringRef MarkedCode, ParsedAST &AST) {
+  Annotations Test(MarkedCode);
+  switch (Test.points().size()) {
+  case 1: // Point selection.
+    return SelectionTree(AST.getASTContext(),
+                         cantFail(positionToOffset(Test.code(), Test.point())));
+  case 2: // Range selection.
+    return SelectionTree(
+        AST.getASTContext(),
+        cantFail(positionToOffset(Test.code(), Test.points()[0])),
+        cantFail(positionToOffset(Test.code(), Test.points()[1])));
+  default:
+    ADD_FAILURE() << "Expected 1-2 points for selection.\n" << MarkedCode;
+    return SelectionTree(AST.getASTContext(), 0u, 0u);
+  }
+}
+
+Range nodeRange(const SelectionTree::Node *N, ParsedAST &AST) {
+  if (!N)
+    return Range{};
+  SourceManager &SM = AST.getASTContext().getSourceManager();
+  StringRef Buffer = SM.getBufferData(SM.getMainFileID());
+  SourceRange SR = N->ASTNode.getSourceRange();
+  SR.setBegin(SM.getFileLoc(SR.getBegin()));
+  SR.setEnd(SM.getFileLoc(SR.getEnd()));
+  CharSourceRange R =
+      Lexer::getAsCharRange(SR, SM, AST.getASTContext().getLangOpts());
+  return Range{offsetToPosition(Buffer, SM.getFileOffset(R.getBegin())),
+               offsetToPosition(Buffer, SM.getFileOffset(R.getEnd()) + 1)};
+}
+
+std::string nodeKind(const SelectionTree::Node *N) {
+  if (!N)
+    return "<null>";
+  return N->ASTNode.getNodeKind().asStringRef().str();
+}
+
+std::vector<const SelectionTree::Node *> allNodes(const SelectionTree &T) {
+  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,
+                          const SelectionTree::Node *Common,
+                          StringRef MarkedCode) {
+  if (Root == Common)
+    return true;
+  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)) {
+      if (Seen)
+        ADD_FAILURE() << "Saw common ancestor twice\n" << MarkedCode;
+      Seen = true;
+    }
+  return Seen;
+}
+
+TEST(SelectionTest, CommonAncestor) {
+  struct Case {
+    // Selection is between ^marks^.
+    // common ancestor marked with a [[range]].
+    const char *Code;
+    const char *CommonAncestorKind;
+  };
+  Case Cases[] = {
+      {
+          R"cpp(
+            struct AAA { struct BBB { static int ccc(); };};
+            int x = AAA::[[B^B^B]]::ccc();
+          )cpp",
+          "TypeLoc",
+      },
+      {
+          R"cpp(
+            struct AAA { struct BBB { static int ccc(); };};
+            int x = AAA::[[B^BB^]]::ccc();
+          )cpp",
+          "TypeLoc",
+      },
+      {
+          R"cpp(
+            struct AAA { struct BBB { static int ccc(); };};
+            int x = [[AAA::BBB::c^c^c]]();
+          )cpp",
+          "DeclRefExpr",
+      },
+      {
+          R"cpp(
+            struct AAA { struct BBB { static int ccc(); };};
+            int x = [[AAA::BBB::cc^c(^)]];
+          )cpp",
+          "CallExpr",
+      },
+
+      {
+          R"cpp(
+            void foo() { [[if (1^11) { return; } else {^ }]] }
+          )cpp",
+          "IfStmt",
+      },
+      {
+          R"cpp(
+            void foo();
+            #define CALL_FUNCTION(X) X()
+            void bar() { CALL_FUNCTION([[f^o^o]]); }
+          )cpp",
+          "DeclRefExpr",
+      },
+      {
+          R"cpp(
+            void foo();
+            #define CALL_FUNCTION(X) X()
+            void bar() { CALL_FUNC^TION([[fo^o]]); }
+          )cpp",
+          "DeclRefExpr",
+      },
+      {
+          R"cpp(
+            void foo();
+            #define CALL_FUNCTION(X) X()
+            void bar() [[{ C^ALL_FUNC^TION(foo); }]]
+          )cpp",
+          "CompoundStmt",
+      },
+      {
+          R"cpp(
+            void foo();
+            #define CALL_FUNCTION(X) X^()^
+            void bar() { CALL_FUNCTION(foo); }
+          )cpp",
+          nullptr,
+      },
+
+      // Point selections.
+      {"void foo() { [[^foo]](); }", "DeclRefExpr"},
+      {"void foo() { [[f^oo]](); }", "DeclRefExpr"},
+      {"void foo() { [[fo^o]](); }", "DeclRefExpr"},
+      {"void foo() { [[foo^()]]; }", "CallExpr"},
+      {"void foo() { [[foo^]] (); }", "DeclRefExpr"},
+      {"int bar; void foo() [[{ foo (); }]]^", "CompoundStmt"},
+      {"[[^void]] foo();", "TypeLoc"},
+      {"^", 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},
+
+      // Node types that have caused problems in the past.
+      {"template <typename T> void foo() { [[^T]] t; }", "TypeLoc"},
+  };
+  for (const Case &C : Cases) {
+    Annotations Test(C.Code);
+    auto AST = TestTU::withCode(Test.code()).build();
+    auto T = makeSelectionTree(C.Code, AST);
+
+    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.
+      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))
+          << C.Code;
+    }
+  }
+}
+
+TEST(SelectionTest, Selected) {
+  // Selection with ^marks^.
+  // Partially selected nodes marked with a [[range]].
+  // Completely selected nodes marked with a $C[[range]].
+  const char *Cases[] = {
+      R"cpp( int abc, xyz = [[^ab^c]]; )cpp",
+      R"cpp( int abc, xyz = [[a^bc^]]; )cpp",
+      R"cpp( int abc, xyz = $C[[^abc^]]; )cpp",
+      R"cpp(
+        void foo() {
+          [[if ([[1^11]]) $C[[{
+            $C[[return]];
+          }]] else [[{^
+          }]]]]
+        }
+      )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))
+      if (N->Selected == SelectionTree::Complete)
+        Complete.push_back(nodeRange(N, AST));
+      else if (N->Selected == SelectionTree::Partial)
+        Partial.push_back(nodeRange(N, AST));
+    EXPECT_THAT(Complete, UnorderedElementsAreArray(Test.ranges("C"))) << C;
+    EXPECT_THAT(Partial, UnorderedElementsAreArray(Test.ranges())) << C;
+  }
+}
+
+} // namespace
+} // namespace clangd
+} // namespace clang