[clangd] Elide even more checks in SelectionTree.

[sammccall]

Summary

During pop() we convert nodes into spans of expanded syntax::Tokens.
If we precompute a range of plausible (expanded) tokens, then we can do an
extremely cheap approximate hit-test against it, because syntax::Tokens are
ordered by pointer.

This would seem not to buy anything (we don't enter nodes unless they overlap
the selection), but in fact the spans we have are for *newly* claimed ranges
(i.e. those unclaimed by any child node).

So if you have:

{ { [[2+2]]; } }

then all of the CompoundStmts pass the hit test and are pushed, but we skip
full hit-testing of the brackets during pop() as they lie outside the range.

This is ~10x average speedup for selectiontree on a bad case I've seen
(large gtest file).

[sammccall]

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

206

This test isn't strictly related, but it failed with a buggy version of this patch.
(It's extracted from a DefineOutline test which also failed but in a less direct way)

[hokein]

this looks good in general.

clang-tools-extra/clangd/Selection.cpp

300

if it is empty, I think it is more appropriate to return NoTokens (the old behavior)

301

I would suggest rename the member ExpandedTokens, and SpellingTokens to something like AffectedSpelledTokens, AffectedExpandedTokens, I think it is important to express "these are affected-tokens by the selection" in their name

374

this should be true.

397

the same, true.

402

should be End.

[sammccall]

Thanks! You caught a serious bug, I'm digging into it.

clang-tools-extra/clangd/Selection.cpp

300

Whoops, this was meant to be this->ExpandedTokens.

You're right that in principle ExpandedTokens.empty() should return NoTokens. (We would never actually get here).

Similarly in principle !ExpandedTokens.empty() && SpelledTokens.empty() should return Unselected instead of NoTokens. This could happen but didn't make any difference to observable behavior.

I've fixed all these in any case.

301

ExpandedTokens -> MaybeSelectedExpanded

SpelledTokens -> SelectedSpelled

374

Oh shoot, now the assert is firing, I must have missed some cases :-(
Need to fix this.

[sammccall]

Handle EOF

[sammccall]

PTAL

clang-tools-extra/clangd/Selection.cpp

374

OK, this was just not handling EOF token location (which offsetInSelFile rejects, correctly I think)

[hokein]

Thanks, this looks great!

clang-tools-extra/clangd/Selection.cpp

376

nit: remove this debug code.

402

and here.

clang-tools-extra/clangd/Selection.cpp

for (const syntax::TokenBuffer::Expansion &X :
             PPIgnored[&Tok - SelFirst] = true;
         }
       }
     }
     // Precompute selectedness and offset for selected spelled tokens.
     for (unsigned I = 0; I < Sel.size(); ++I) {
       if (shouldIgnore(Sel[I]) || PPIgnored[I])
         continue;
-      SpelledTokens.emplace_back();
-      Tok &S = SpelledTokens.back();
+      SelectedSpelled.emplace_back();
+      Tok &S = SelectedSpelled.back();
       S.Offset = SM.getFileOffset(Sel[I].location());
       if (S.Offset >= SelBegin && S.Offset + Sel[I].length() <= SelEnd)
         S.Selected = SelectionTree::Complete;
       else
         S.Selected = SelectionTree::Partial;
     }
+    MaybeSelectedExpanded = computeMaybeSelectedExpandedTokens(Buf);
   }
 
   // Test whether a consecutive range of tokens is selected.
   // The tokens are taken from the expanded token stream.
   SelectionTree::Selection
   test(llvm::ArrayRef<syntax::Token> ExpandedTokens) const {
-    if (SpelledTokens.empty())
+    if (ExpandedTokens.empty())
       return NoTokens;
+    if (SelectedSpelled.empty())
+      return SelectionTree::Unselected;
+    // Cheap (pointer) check whether any of the tokens could touch selection.
+    // In most cases, the node's overall source range touches ExpandedTokens,
+    // or we would have failed mayHit(). However now we're only considering
+    // the *unclaimed* spans of expanded tokens.
+    // This is a significant performance improvement when a lot of nodes
+    // surround the selection, including when generated by macros.
+    if (MaybeSelectedExpanded.empty() ||
+        &ExpandedTokens.front() > &MaybeSelectedExpanded.back() ||
+        &ExpandedTokens.back() < &MaybeSelectedExpanded.front()) {
+      return SelectionTree::Unselected;
+    }
+
     SelectionTree::Selection Result = NoTokens;
     while (!ExpandedTokens.empty()) {
       // Take consecutive tokens from the same context together for efficiency.
       SourceLocation Start = ExpandedTokens.front().location();
       FileID FID = SM.getFileID(Start);
       // Comparing SourceLocations against bounds is cheaper than getFileID().
       SourceLocation Limit = SM.getComposedLoc(FID, SM.getFileIDSize(FID));
       auto Batch = ExpandedTokens.take_while([&](const syntax::Token &T) {
         return T.location() >= Start && T.location() < Limit;
       });
       assert(!Batch.empty());
       ExpandedTokens = ExpandedTokens.drop_front(Batch.size());
 
       update(Result, testChunk(FID, Batch));
     }
     return Result;
   }
 
   // Cheap check whether any of the tokens in R might be selected.
   // If it returns false, test() will return NoTokens or Unselected.
   // If it returns true, test() may return any value.
   bool mayHit(SourceRange R) const {
-    if (SpelledTokens.empty())
+    if (SelectedSpelled.empty() || MaybeSelectedExpanded.empty())
       return false;
     // If the node starts after the selection ends, it is not selected.
     // Tokens a macro location might claim are >= its expansion start.
     // So if the expansion start > last selected token, we can prune it.
     // (This is particularly helpful for GTest's TEST macro).
     if (auto B = offsetInSelFile(getExpansionStart(R.getBegin())))
-      if (*B > SpelledTokens.back().Offset)
+      if (*B > SelectedSpelled.back().Offset)
         return false;
     // If the node ends before the selection begins, it is not selected.
     SourceLocation EndLoc = R.getEnd();
     while (EndLoc.isMacroID())
       EndLoc = SM.getImmediateExpansionRange(EndLoc).getEnd();
     // In the rare case that the expansion range is a char range, EndLoc is
     // ~one token too far to the right. We may fail to prune, that's OK.
     if (auto E = offsetInSelFile(EndLoc))
-      if (*E < SpelledTokens.front().Offset)
+      if (*E < SelectedSpelled.front().Offset)
         return false;
     return true;
   }
 
 private:
+  // Plausible expanded tokens that might be affected by the selection.
+  // This is an overestimate, it may contain tokens that are not selected.
+  // The point is to allow cheap pruning in test()
+  llvm::ArrayRef<syntax::Token>
+  computeMaybeSelectedExpandedTokens(const syntax::TokenBuffer &Toks) {
+    if (SelectedSpelled.empty())
+      return {};
+
+    bool StartInvalid = false;
+    const syntax::Token *Start = llvm::partition_point(
+        Toks.expandedTokens(),
+        [&, First = SelectedSpelled.front().Offset](const syntax::Token &Tok) {
+          if (Tok.kind() == tok::eof)
+            return false;
+          // Implausible if upperbound(Tok) < First.
+          SourceLocation Loc = Tok.location();
+          auto Offset = offsetInSelFile(Loc);
+          while (Loc.isValid() && !Offset) {
+            Loc = Loc.isMacroID() ? SM.getImmediateExpansionRange(Loc).getEnd()
+                                  : SM.getIncludeLoc(SM.getFileID(Loc));
+            Offset = offsetInSelFile(Loc);
+          }
+          if (Offset)
+            return *Offset < First;
+          StartInvalid = true;
+          return false; // conservatively assume this token can overlap
+        });
+    if (StartInvalid) {
+      assert(false && "Expanded tokens could not be resolved to main file!");
+      Start = Toks.expandedTokens().begin();
+    }
+
+    bool EndInvalid = false;
+    const syntax::Token *End = llvm::partition_point(
+        Toks.expandedTokens(),
+        [&, Last = SelectedSpelled.back().Offset](const syntax::Token &Tok) {
+          if (Tok.kind() == tok::eof)
+            return false;
+          // Plausible if lowerbound(Tok) <= Last.
+          SourceLocation Loc = Tok.location();
+          auto Offset = offsetInSelFile(Loc);
+          while (Loc.isValid() && !Offset) {
+            Loc = Loc.isMacroID()
+                      ? SM.getImmediateExpansionRange(Loc).getBegin()
+                      : SM.getIncludeLoc(SM.getFileID(Loc));
+            Offset = offsetInSelFile(Loc);
+          }
+          if (Offset)
+            return *Offset <= Last;
+          EndInvalid = true;
+          return true; // conservatively assume this token can overlap
+        });
+    if (EndInvalid) {
+      assert(false && "Expanded tokens could not be resolved to main file!");
+      End = Toks.expandedTokens().end();
+    }
+
+    return llvm::makeArrayRef(Start, End);
+  }
+
   // Hit-test a consecutive range of tokens from a single file ID.
   SelectionTree::Selection
   testChunk(FileID FID, llvm::ArrayRef<syntax::Token> Batch) const {
     assert(!Batch.empty());
     SourceLocation StartLoc = Batch.front().location();
     // There are several possible categories of FileID depending on how the
     // preprocessor was used to generate these tokens:
     //   main file, #included file, macro args, macro bodies.
if (auto ExpansionOffset = offsetInSelFile(getExpansionStart(StartLoc)))
       return testToken(*ExpansionOffset);
     return NoTokens;
   }
 
   // Is the closed token range [Begin, End] selected?
   SelectionTree::Selection testTokenRange(unsigned Begin, unsigned End) const {
     assert(Begin <= End);
     // Outside the selection entirely?
-    if (End < SpelledTokens.front().Offset ||
-        Begin > SpelledTokens.back().Offset)
+    if (End < SelectedSpelled.front().Offset ||
+        Begin > SelectedSpelled.back().Offset)
       return SelectionTree::Unselected;
 
     // Compute range of tokens.
     auto B = llvm::partition_point(
-        SpelledTokens, [&](const Tok &T) { return T.Offset < Begin; });
-    auto E = std::partition_point(
-        B, SpelledTokens.end(), [&](const Tok &T) { return T.Offset <= End; });
+        SelectedSpelled, [&](const Tok &T) { return T.Offset < Begin; });
+    auto E = std::partition_point(B, SelectedSpelled.end(), [&](const Tok &T) {
+      return T.Offset <= End;
+    });
 
     // Aggregate selectedness of tokens in range.
-    bool ExtendsOutsideSelection = Begin < SpelledTokens.front().Offset ||
-                                   End > SpelledTokens.back().Offset;
+    bool ExtendsOutsideSelection = Begin < SelectedSpelled.front().Offset ||
+                                   End > SelectedSpelled.back().Offset;
     SelectionTree::Selection Result =
         ExtendsOutsideSelection ? SelectionTree::Unselected : NoTokens;
     for (auto It = B; It != E; ++It)
       update(Result, It->Selected);
     return Result;
   }
 
   // Is the token at `Offset` selected?
   SelectionTree::Selection testToken(unsigned Offset) const {
     // Outside the selection entirely?
-    if (Offset < SpelledTokens.front().Offset ||
-        Offset > SpelledTokens.back().Offset)
+    if (Offset < SelectedSpelled.front().Offset ||
+        Offset > SelectedSpelled.back().Offset)
       return SelectionTree::Unselected;
     // Find the token, if it exists.
     auto It = llvm::partition_point(
-        SpelledTokens, [&](const Tok &T) { return T.Offset < Offset; });
-    if (It != SpelledTokens.end() && It->Offset == Offset)
+        SelectedSpelled, [&](const Tok &T) { return T.Offset < Offset; });
+    if (It != SelectedSpelled.end() && It->Offset == Offset)
       return It->Selected;
     return NoTokens;
   }
 
   // Decomposes Loc and returns the offset if the file ID is SelFile.
   llvm::Optional<unsigned> offsetInSelFile(SourceLocation Loc) const {
     // Decoding Loc with SM.getDecomposedLoc is relatively expensive.
     // But SourceLocations for a file are numerically contiguous, so we
while (Loc.isMacroID())
       Loc = SM.getImmediateExpansionRange(Loc).getBegin();
     return Loc;
   }
 
   struct Tok {
     unsigned Offset;
     SelectionTree::Selection Selected;
   };
-  std::vector<Tok> SpelledTokens;
+  std::vector<Tok> SelectedSpelled;
+  llvm::ArrayRef<syntax::Token> MaybeSelectedExpanded;
   FileID SelFile;
   SourceRange SelFileBounds;
   const SourceManager &SM;
 };
 
 // Show the type of a node for debugging.
 void printNodeKind(llvm::raw_ostream &OS, const DynTypedNode &N) {
   if (const TypeLoc *TL = N.get<TypeLoc>()) {

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

Case Cases[] = {
             void fo^o() {}
             #endif
             }
           )cpp",
           nullptr,
       },
       {
           R"cpp(
+            #define TARGET void foo()
+            [[TAR^GET{ return; }]]
+          )cpp",
+          "FunctionDecl",
+      },
+      {
+          R"cpp(
             struct S { S(const char*); };
             [[S s ^= "foo"]];
           )cpp",
           // The AST says a CXXConstructExpr covers the = sign in C++14.
           // But we consider CXXConstructExpr to only own brackets.
           // (It's not the interesting constructor anyway, just S(&&)).
           "VarDecl",
       },