[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

289

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

290

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

350

this should be true.

373

the same, true.

378

should be End.

[sammccall]

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

clang-tools-extra/clangd/Selection.cpp

289

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.

290

ExpandedTokens -> MaybeSelectedExpanded

SpelledTokens -> SelectedSpelled

350

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

350

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

352

nit: remove this debug code.

378

and here.

clang-tools-extra/clangd/Selection.cpp

for (unsigned I = 0; I < Sel.size(); ++I) {
       SpelledTokens.emplace_back();
       Tok &S = SpelledTokens.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;
     }
+    ExpandedTokens = computeExpandedTokens(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())
       return NoTokens;
+    // 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 (ExpandedTokens.empty() ||
+        &ExpandedTokens.front() > &this->ExpandedTokens.back() ||
+        &ExpandedTokens.back() < &this->ExpandedTokens.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 (SpelledTokens.empty() || ExpandedTokens.empty())
       return false;
     auto B = offsetInSelFile(R.getBegin());
     auto E = offsetInSelFile(R.getEnd());
     if (B && E)
       if (*E < SpelledTokens.front().Offset || *B > SpelledTokens.back().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>
+  computeExpandedTokens(const syntax::TokenBuffer &Toks) {
+    if (SpelledTokens.empty())
+      return {};
+
+    bool StartInvalid = false;
+    const syntax::Token *Start = llvm::partition_point(
+        Toks.expandedTokens(),
+        [&, First = SpelledTokens.front().Offset](const syntax::Token &Tok) {
+          // 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 = false;
+          return false;
+        });
+    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 = SpelledTokens.back().Offset](const syntax::Token &Tok) {
+          // 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 = false;
+          return true;
+        });
+    if (EndInvalid) {
+      assert(false && "Expanded tokens could not be resolved to main file!");
+      Start = 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.
SourceLocation getExpansionStart(SourceLocation Loc) const {
     return Loc;
   }
 
   struct Tok {
     unsigned Offset;
     SelectionTree::Selection Selected;
   };
   std::vector<Tok> SpelledTokens;
+  llvm::ArrayRef<syntax::Token> ExpandedTokens;
   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",
       },