[clangd] Ignore semicolons, whitespace, and comments in SelectionTree.

[sammccall]

Summary

Whitespace and comments are a clear bugfix: selecting some
comments/space near a statement doesn't mean you're selecting the
surrounding block.

Semicolons are less obvious, but for similar reasons: these tokens
aren't actually claimed by any AST node (usually), so an AST-based model
like SelectionTree shouldn't take them into account.

Callers may still sometimes care about semis of course:

• when the selection is an expr with a non-expr parent, selection of the semicolon indicates intent to select the statement.
• when a statement with a trailing semi is selected, we need to know its range to ensure it can be removed.

SelectionTree may or may not play a role here, but these are separate questions
from its core function of describing which AST nodes were selected.

The mechanism here is the TokenBuffer from syntax-trees. We use it in a
fairly low-level way (just to get boundaries of raw spelled tokens). The
actual mapping of AST nodes to coordinates continues to use the (fairly
mature) SourceLocation based logic. TokenBuffer/Syntax trees
don't currently offer an alternative to getFileRange(), I think.

[sammccall]

It occurs to me that claim is O(node_tokens + log total_tokens) which is bad when the nodes are large.

Indeed for an input like namespace { namespace { namespace { ... } } } time is quadratic.

I think this is probably fine in practice. Against adversarial input clang certainly can take exponential time, and easily crash too.

If we do need to fix it my best idea is to give each "uninteresting" TokInfo (that is, !selected || claimed) a pointer to the next TokInfo with different flags. This would allow iteration to quickly skip over contiguous ranges one they had been traversed once.

[sammccall]

Improve comments in SelectedTokens.

[kadircet]

Thanks for the comments, LGTM!

[SureYeaah]

clang-tools-extra/clangd/Selection.cpp

50 ↗

(On Diff #212536)

Would this work correctly for nested templates? Or do we need to use the specialized token length function that we use for toHalfOpenFileRange?

[sammccall]

clang-tools-extra/clangd/Selection.cpp

50 ↗

(On Diff #212536)

This uses the token list as an intermediary for matching selected chars with AST nodes. TokenBuffer will indeed by default lex >> as a right shift. So we're buggy here, but I think it *mostly* doesn't matter.

If it's a double template:

• the innermost template will claim it first, if the template range touches the selection. Problem: if only the first > is selected, the inner template will only be partially selected.
• the outermost template will not get to claim it at all (if the inner template range touches the selection). Problem: if this is the *only* part of the outer template that's selected, it will be marked unselected. (This should be rare)
• If the inner template doesn't touch the selection, then the outer template will be selected but only partially, which is actually correct.

Examples:

a<b<c>>
     ~  b=partial (correct)
  ~~~~  b=partial (incorrect: b=complete) <-- this is the worst case
~~~~~~~ a=complete,b=complete (correct)
      ~ a=partial (correct)
     ~~ b=partial (incorrect: a=partial,b=partial)

I'll send a followup to fix this case tomorrow (I think we can just always split the token in half) but I don't think it's critical.

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

 //     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 "clang/Tooling/Syntax/Tokens.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.
 //
 // 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).
+//
+// Comments, directives and whitespace are completely ignored.
+// Semicolons are also ignored, as the AST generally does not model them well.
 //
 // 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);
+  SelectionTree(ASTContext &AST, const syntax::TokenBuffer &Tokens,
+                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);
+  SelectionTree(ASTContext &AST, const syntax::TokenBuffer &Tokens,
+                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

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

 #include "SourceCode.h"
 #include "clang/AST/ASTTypeTraits.h"
 #include "clang/AST/PrettyPrinter.h"
 #include "clang/AST/RecursiveASTVisitor.h"
 #include "clang/AST/TypeLoc.h"
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Basic/TokenKinds.h"
+#include "clang/Tooling/Syntax/Tokens.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/raw_ostream.h"
 #include <algorithm>
 #include <string>
 
 namespace clang {
 namespace clangd {
 namespace {
 using Node = SelectionTree::Node;
 using ast_type_traits::DynTypedNode;
 
-// Stores a collection of (possibly-overlapping) integer ranges.
-// When new ranges are added, hit-tests them against existing ones.
-class RangeSet {
+// Identifies which tokens are selected, and evaluates claims of source ranges
+// by AST nodes. Tokens may be claimed only once: first-come, first-served.
+class SelectedTokens {
 public:
-  // Returns true if any new offsets are covered.
-  // This is naive (linear in number of successful add() calls), but ok for now.
-  bool add(unsigned Begin, unsigned End) {
-    assert(std::is_sorted(Ranges.begin(), Ranges.end()));
-    assert(Begin < End);
-
-    if (covered(Begin, End))
-      return false;
-    auto Pair = std::make_pair(Begin, End);
-    Ranges.insert(llvm::upper_bound(Ranges, Pair), Pair);
-    return true;
+  SelectedTokens(llvm::ArrayRef<syntax::Token> Spelled, const SourceManager &SM,
+                 unsigned SelBegin, unsigned SelEnd)
+      : SelBegin(SelBegin), SelEnd(SelEnd) {
+    // Extract bounds and selected-ness for all tokens spelled in the file.
+    Tokens.reserve(Spelled.size());
+    for (const auto& Tok : Spelled) {
+      // As well as comments, don't count semicolons as real tokens.
+      // They're not properly claimed as expr-statement is missing from the AST.
+      if (Tok.kind() == tok::comment || Tok.kind() == tok::semi)
+        continue;
+
+      Tokens.emplace_back();
+      TokInfo &S = Tokens.back();
+      S.StartOffset = SM.getFileOffset(Tok.location());
+      S.EndOffset = S.StartOffset + Tok.length();
+      if (S.StartOffset >= SelBegin && S.EndOffset <= SelEnd)
+        S.Selected = SelectionTree::Complete;
+      else if (S.EndOffset > SelBegin && S.StartOffset < SelEnd)
+        S.Selected = SelectionTree::Partial;
+      else
+        S.Selected = SelectionTree::Unselected;
+      S.Claimed = false;
+    }
   }
 
-private:
-  bool covered(unsigned Begin, unsigned End) {
-    assert(Begin < End);
-    for (const auto &R : Ranges) {
-      if (Begin < R.first)
-        return false; // The prefix [Begin, R.first) is not covered.
-      if (Begin < R.second) {
-        Begin = R.second; // Prefix is covered, truncate the range.
-        if (Begin >= End)
-          return true;
+  // Associates any tokens overlapping [Begin, End) with an AST node.
+  // Tokens that were already claimed by another AST node are not claimed again.
+  // Returns whether the node is selected in the sense of SelectionTree.
+  SelectionTree::Selection claim(unsigned Begin, unsigned End) {
+    assert(Begin <= End);
+
+    // Fast-path for missing the selection entirely.
+    if (Begin >= SelEnd || End <= SelBegin)
+      return SelectionTree::Unselected;
+
+    // We will consider the range (at least partially) selected if it hit any
+    // selected and previously unclaimed token.
+    bool ClaimedAnyToken = false;
+    // The selection is (at most) partial if:
+    // - any claimed token is partially selected
+    // - any token in the range is unselected
+    bool PartialSelection = false;
+
+    // Find the first token that (maybe) overlaps the claimed range.
+    auto Start = llvm::partition_point(Tokens, [&](const TokInfo &Tok) {
+      return Tok.EndOffset <= Begin;
+    });
+    // Iterate over every token that overlaps the range.
+    // Claim selected tokens, and update the two result flags.
+    for (auto It = Start; It != Tokens.end() && It->StartOffset < End; ++It) {
+      if (It->Selected) {
+        if (!It->Claimed) {
+          // Token is selected, in the node's range, and unclaimed; claim it.
+          It->Claimed = true;
+          ClaimedAnyToken = true;
+          // If the token was only partially selected, so is the node.
+          PartialSelection |= (It->Selected == SelectionTree::Partial);
         }
+      } else {
+        // If the node covers an unselected token, it's not completely selected.
+        PartialSelection = true;
       }
-    return false;
     }
 
-  std::vector<std::pair<unsigned, unsigned>> Ranges; // Always sorted.
+    if (!ClaimedAnyToken)
+      return SelectionTree::Unselected;
+    return PartialSelection ? SelectionTree::Partial : SelectionTree::Complete;
+  }
+
+private:
+  struct TokInfo {
+    unsigned StartOffset;
+    unsigned EndOffset;
+    SelectionTree::Selection Selected;
+    bool Claimed;
+    bool operator<(const TokInfo &Other) const {
+      return StartOffset < Other.StartOffset;
+    }
+  };
+  std::vector<TokInfo> Tokens;
+  unsigned SelBegin, SelEnd;
 };
 
 // Show the type of a node for debugging.
 void printNodeKind(llvm::raw_ostream &OS, const DynTypedNode &N) {
   if (const TypeLoc *TL = N.get<TypeLoc>()) {
     // TypeLoc is a hierarchy, but has only a single ASTNodeKind.
     // Synthesize the name from the Type subclass (except for QualifiedTypeLoc).
     if (TL->getTypeLocClass() == TypeLoc::Qualified)

 // 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, const PrintingPolicy &PP,
-                                  unsigned Begin, unsigned End, FileID File) {
-    SelectionVisitor V(AST, PP, Begin, End, File);
+  static std::deque<Node> collect(ASTContext &AST,
+                                  const syntax::TokenBuffer &Tokens,
+                                  const PrintingPolicy &PP, unsigned Begin,
+                                  unsigned End, FileID File) {
+    SelectionVisitor V(AST, Tokens, PP, Begin, End, File);
     V.TraverseAST(AST);
     assert(V.Stack.size() == 1 && "Unpaired push/pop?");
     assert(V.Stack.top() == &V.Nodes.front());
-    // We selected TUDecl if characters were unclaimed (or the file is empty).
-    if (V.Nodes.size() == 1 || V.Claimed.add(Begin, End)) {
+    // We selected TUDecl if tokens were unclaimed (or the file is empty).
+    if (V.Nodes.size() == 1 || V.Claimed.claim(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);
   }
public:
   }
   // 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, const PrintingPolicy &PP, unsigned SelBegin,
-                   unsigned SelEnd, FileID SelFile)
+  SelectionVisitor(ASTContext &AST, const syntax::TokenBuffer &Tokens,
+                   const PrintingPolicy &PP, unsigned SelBegin, unsigned SelEnd,
+                   FileID SelFile)
       : SM(AST.getSourceManager()), LangOpts(AST.getLangOpts()),
 #ifndef NDEBUG
         PrintPolicy(PP),
 #endif
-        SelBegin(SelBegin), SelEnd(SelEnd), SelFile(SelFile),
+        Claimed(Tokens.spelledTokens(SelFile), SM, SelBegin, SelEnd),
+        SelFile(SelFile),
         SelBeginTokenStart(SM.getFileOffset(Lexer::GetBeginningOfToken(
-            SM.getComposedLoc(SelFile, SelBegin), SM, LangOpts))) {
+            SM.getComposedLoc(SelFile, SelBegin), SM, LangOpts))),
+        SelEnd(SelEnd) {
     // 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());
   }
 
SelectionTree::Selection claimRange(SourceRange S) {
     auto Range = toHalfOpenFileRange(SM, LangOpts, S);
     assert(Range && "We should be able to get the File Range");
     dlog("{1}claimRange: {0}", Range->printToString(SM), indent());
     auto B = SM.getDecomposedLoc(Range->getBegin());
     auto E = SM.getDecomposedLoc(Range->getEnd());
     // Otherwise, nodes in macro expansions can't be selected.
     if (B.first != SelFile || E.first != SelFile)
       return SelectionTree::Unselected;
-    // Is there any overlap at all between the selection and range?
-    if (B.second >= SelEnd || E.second < SelBegin)
-      return SelectionTree::Unselected;
-    // We may have hit something.
-    auto PreciseBounds = std::make_pair(B.second, E.second);
-    // Trim range using the selection, drop it if empty.
-    B.second = std::max(B.second, SelBegin);
-    E.second = std::min(E.second, SelEnd);
-    if (B.second >= E.second)
-      return SelectionTree::Unselected;
     // Attempt to claim the remaining range. If there's nothing to claim, only
     // children were selected.
-    if (!Claimed.add(B.second, E.second))
-      return SelectionTree::Unselected;
+    SelectionTree::Selection Result = Claimed.claim(B.second, E.second);
+    if (Result)
       dlog("{1}hit selection: {0}",
            SourceRange(SM.getComposedLoc(B.first, B.second),
                        SM.getComposedLoc(E.first, E.second))
                .printToString(SM),
            indent());
-    // Some of our own characters are covered, this is a true hit.
-    // Determine whether the node was completely covered.
-    return (PreciseBounds.first >= SelBegin && PreciseBounds.second <= SelEnd)
-               ? SelectionTree::Complete
-               : SelectionTree::Partial;
+    return Result;
   }
 
   std::string indent(int Offset = 0) {
     // Cast for signed arithmetic.
     int Amount = int(Stack.size()) + Offset;
     assert(Amount >= 0);
     return std::string(Amount, ' ');
   }
 
   SourceManager &SM;
   const LangOptions &LangOpts;
 #ifndef NDEBUG
   const PrintingPolicy &PrintPolicy;
 #endif
   std::stack<Node *> Stack;
-  RangeSet Claimed;
+  SelectedTokens Claimed;
   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;
+  unsigned SelEnd;
 };
 
 } // 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 ? '*'
static std::pair<unsigned, unsigned> pointBounds(unsigned Offset, FileID FID,
   // 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)
+SelectionTree::SelectionTree(ASTContext &AST, const syntax::TokenBuffer &Tokens,
+                             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.
   const SourceManager &SM = AST.getSourceManager();
   FileID FID = SM.getMainFileID();
   if (Begin == End)
     std::tie(Begin, End) = pointBounds(Begin, FID, AST);
   PrintPolicy.TerseOutput = true;
   PrintPolicy.IncludeNewlines = false;
 
   dlog("Computing selection for {0}",
        SourceRange(SM.getComposedLoc(FID, Begin), SM.getComposedLoc(FID, End))
            .printToString(SM));
-  Nodes = SelectionVisitor::collect(AST, PrintPolicy, Begin, End, FID);
+  Nodes = SelectionVisitor::collect(AST, Tokens, PrintPolicy, Begin, End, FID);
   Root = Nodes.empty() ? nullptr : &Nodes.front();
   dlog("Built selection tree\n{0}", *this);
 }
 
-SelectionTree::SelectionTree(ASTContext &AST, unsigned Offset)
-    : SelectionTree(AST, Offset, Offset) {}
+SelectionTree::SelectionTree(ASTContext &AST, const syntax::TokenBuffer &Tokens,
+                             unsigned Offset)
+    : SelectionTree(AST, Tokens, Offset, Offset) {}
 
 const Node *SelectionTree::commonAncestor() const {
   const Node *Ancestor = Root;
   while (Ancestor->Children.size() == 1 && !Ancestor->Selected)
     Ancestor = Ancestor->Children.front();
   // 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.

clang-tools-extra/trunk/clangd/refactor/Tweak.cpp

#ifndef NDEBUG
   }
 #endif
 }
 } // namespace
 
 Tweak::Selection::Selection(ParsedAST &AST, unsigned RangeBegin,
                             unsigned RangeEnd)
     : AST(AST), SelectionBegin(RangeBegin), SelectionEnd(RangeEnd),
-      ASTSelection(AST.getASTContext(), RangeBegin, RangeEnd) {
+      ASTSelection(AST.getASTContext(), AST.getTokens(), RangeBegin, RangeEnd) {
   auto &SM = AST.getSourceManager();
   Code = SM.getBufferData(SM.getMainFileID());
   Cursor = SM.getComposedLoc(SM.getMainFileID(), RangeBegin);
 }
 
 std::vector<std::unique_ptr<Tweak>>
 prepareTweaks(const Tweak::Selection &S,
               llvm::function_ref<bool(const Tweak &)> Filter) {

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

 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(),
+    return SelectionTree(AST.getASTContext(), AST.getTokens(),
                          cantFail(positionToOffset(Test.code(), Test.point())));
   case 2: // Range selection.
     return SelectionTree(
-        AST.getASTContext(),
+        AST.getASTContext(), AST.getTokens(),
         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);
+    return SelectionTree(AST.getASTContext(), AST.getTokens(), 0u, 0u);
   }
 }
 
 Range nodeRange(const SelectionTree::Node *N, ParsedAST &AST) {
   if (!N)
     return Range{};
   const SourceManager &SM = AST.getSourceManager();
   const LangOptions &LangOpts = AST.getASTContext().getLangOpts();
Case Cases[] = {
       // 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"},
 
+      // Ignores whitespace, comments, and semicolons in the selection.
+      {"void foo() { [[foo^()]]; /*comment*/^}", "CallExpr"},
+
       // Tricky case: FunctionTypeLoc in FunctionDecl has a hole in it.
       {"[[^void]] foo();", "BuiltinTypeLoc"},
       {"[[void foo^()]];", "FunctionProtoTypeLoc"},
       {"[[^void foo^()]];", "FunctionDecl"},
       {"[[void ^foo()]];", "FunctionDecl"},
       // Tricky case: two VarDecls share a specifier.
       {"[[int ^a]], b;", "VarDecl"},
       {"[[int a, ^b]];", "VarDecl"},