[clangd] Selection: Prune gtest TEST()s earlier

[sammccall]

Summary

When searching for AST nodes that may overlap the selection, mayHit() was only
attempting to prune nodes whose begin/end are both in the main file.

While failing to prune never gives wrong results, it hurts performance.
In GTest unit-tests, TEST() macros at the top level declare classes.
These were never pruned and we traversed *every* such class for any selection.

We fix this by reasoning about what tokens such a node might claim.
They must lie within its ultimate macro expansion range, so if this doesn't
overlap with the selection, we can prune the node.

[sammccall]

Use the other end too

[sammccall]

Rebase (offsetInSelFile)

[hokein]

Thanks, this looks good.

clang-tools-extra/clangd/Selection.cpp

306

sorry, I don't quite understand this sentence...

315

I think the current implementation is fine, just notice there is an alternative

• SourceManager::getExpansionRange(SourceRange) method does similar things, we could use that to save some bits of code, but we will have some extra cost for unnecessary getImmediateExpansionRange/getFileID.

318

as discussed offline, offsetInSelFile is doing smart and efficient thing which relies on the underlying offset implementation in source manger, would be nice to have some comments and docs there and SourceManager comparison operator.

[sammccall]

clang-tools-extra/clangd/Selection.cpp

306

Rephrased.

315

Yeah, it does a fair amount of unneccesary work in order for the endpoint to be exact (handling whether intermediate exp ranges for endpoint are char ranges or token ranges).
For our case we don't care if we're a token too far to the right, and we don't want to pay for all this.

318

Oops, I forgot to document that function! Done.

clang-tools-extra/clangd/Selection.cpp

 //===--- Selection.cpp ----------------------------------------------------===//
 //
 // 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 "AST.h"
-#include "SourceCode.h"
 #include "support/Logger.h"
 #include "support/Trace.h"
 #include "clang/AST/ASTTypeTraits.h"
 #include "clang/AST/Decl.h"
 #include "clang/AST/DeclCXX.h"
 #include "clang/AST/Expr.h"
 #include "clang/AST/ExprCXX.h"
 #include "clang/AST/PrettyPrinter.h"
public:
   }
 
   // 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())
       return false;
-    auto B = offsetInSelFile(R.getBegin());
-    auto E = offsetInSelFile(R.getEnd());
-    if (B && E)
-      if (*E < SpelledTokens.front().Offset || *B > SpelledTokens.back().Offset)
+    // 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)
+        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)
         return false;
     return true;
   }
 
 private:
   // Hit-test a consecutive range of tokens from a single file ID.
   SelectionTree::Selection
   testChunk(FileID FID, llvm::ArrayRef<syntax::Token> Batch) const {
SelectionTree::Selection testToken(unsigned Offset) const {
     // 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)
       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
+    // can use cheap integer operations instead.
     if (Loc < SelFileBounds.getBegin() || Loc >= SelFileBounds.getEnd())
       return llvm::None;
+    // FIXME: subtracting getRawEncoding() is dubious, move this logic into SM.
     return Loc.getRawEncoding() - SelFileBounds.getBegin().getRawEncoding();
   }
 
   SourceLocation getExpansionStart(SourceLocation Loc) const {
     while (Loc.isMacroID())
       Loc = SM.getImmediateExpansionRange(Loc).getBegin();
     return Loc;
   }