Optimise findRefs for XRefs and docHighlights

[usaxena95]

Summary

Reduces time spent in findRef by ~66%.

Ran crude measurement clangd --check=clang/lib/Serialization/ASTReader.cpp
Previous time: 52 min
New time: 12 min
This include time for much more than just docHighlight (hover, definition, tweaks,..). So in this case the speedup is at least 77%.
Moreover running docHighlight on each token previously was O(N^2) and is now O(N) in size of main file. Therefore speedup % would be O(N) instead of a constant. In practice for long-tail large files, I have seen 60-80% speedups.

[kadircet]

66% win sounds great, it would be nice to have some detailed numbers (but this is clearly a huge win, so no need to reperform the experiments if numbers are gone)

clang-tools-extra/clangd/XRefs.cpp

895

nit: no need for braces

you can directly build the set with canonicaldecls and consume it here

953

we're still calling getsymbolid here lots of times, for probably the same symbol.

as discussed, what about building a lookup table on construction from canonicaldecl to symbolid and use that lookup table instead of calling getsymbolid over and over here?

1257

i mean directly inserting canonicaldecls here.

1258–1259

no need for this and the following loop anymore.

[usaxena95]

Address comments.

clang-tools-extra/clangd/XRefs.cpp

953

I don't mind adding a lookup table for this but this is not huge, only O(#ref) as compared to previous O(size of TU).

1257

This would make it messy to duplicate this on each end of the call.
+ This has potential of introducing bugs in future findRef calls if we don't pass the canonical decls. Or we would have to assert that we only get canonical decls.
I think this is an implementation detail of findRefs and should not be visible outside its API for simplicity.

[usaxena95]

Format.

[kadircet]

thanks for noticing this, LGTM!

clang-tools-extra/clangd/XRefs.cpp

953

thanks. previous one was also size of main file, not the whole TU (if you see indications of the latter in a different application let's take a look at that separately).

but there are big enough files in which you can have thousands of references to stringref and what not. so i think this actually matters.

962

s/TargetDeclAndID/TargetDeclToID

966–968

i'd actually update the interface here to just take an ArrayRef<Decl*> as targets, as we're currently building those extra dense sets just to iterate over them. both allTargetDecls and getDeclAtPosition (which uses allTargetDecls underneath) are guaranteed to not return duplicates. So we can pass their result directly here. Up to you though.

1257

SG, it's an implementation detail of this file already so not a big deal. But moreover we're building a separate lookup table already now, so obsolete.

1917

can you revert this change?

[usaxena95]

Addressed comments and added performance measurements.

clang-tools-extra/clangd/XRefs.cpp

 #include "Protocol.h"
 #include "Quality.h"
 #include "Selection.h"
 #include "SourceCode.h"
 #include "URI.h"
 #include "index/Index.h"
 #include "index/Merge.h"
 #include "index/Relation.h"
+#include "index/SymbolID.h"
 #include "index/SymbolLocation.h"
 #include "support/Logger.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/ASTTypeTraits.h"
 #include "clang/AST/Attr.h"
 #include "clang/AST/Attrs.inc"
 #include "clang/AST/Decl.h"
 #include "clang/AST/DeclCXX.h"

 #include "clang/Basic/TokenKinds.h"
 #include "clang/Index/IndexDataConsumer.h"
 #include "clang/Index/IndexSymbol.h"
 #include "clang/Index/IndexingAction.h"
 #include "clang/Index/IndexingOptions.h"
 #include "clang/Index/USRGeneration.h"
 #include "clang/Tooling/Syntax/Tokens.h"
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/None.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/SmallSet.h"
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/raw_ostream.h"
 
 namespace clang {
 namespace clangd {
if (const auto *CMD = llvm::dyn_cast<CXXMethodDecl>(D)) {
       if (CMD->isPure()) {
         if (TouchedIdentifier && SM.getSpellingLoc(CMD->getLocation()) ==
                                      TouchedIdentifier->location()) {
           VirtualMethods.insert(getSymbolID(CMD));
           LocateASTReferentMetric.record(1, "method-to-override");
         }
       }
       // Special case: void foo() ^override: jump to the overridden method.
       if (NodeKind.isSame(ASTNodeKind::getFromNodeKind<OverrideAttr>()) ||
           NodeKind.isSame(ASTNodeKind::getFromNodeKind<FinalAttr>())) {
         // We may be overridding multiple methods - offer them all.
         for (const NamedDecl *ND : CMD->overridden_methods())
           AddResultDecl(ND);
         continue;
       }
     }
 
     // Special case: the cursor is on an alias, prefer other results.
struct Reference {
     SymbolID Target;
 
     Range range(const SourceManager &SM) const {
       return halfOpenToRange(SM, SpelledTok.range(SM).toCharRange(SM));
     }
   };
 
   ReferenceFinder(const ParsedAST &AST,
-                  const llvm::DenseSet<SymbolID> &TargetIDs, bool PerToken)
-      : PerToken(PerToken), AST(AST), TargetIDs(TargetIDs) {}
+                  const llvm::ArrayRef<const NamedDecl *> Targets, bool PerToken)
+      : PerToken(PerToken), AST(AST) {
+    for (const NamedDecl *ND : Targets) {
+      const Decl *CD = ND->getCanonicalDecl();
+      TargetDeclToID[CD] = getSymbolID(CD);
+    }
+  }
 
   std::vector<Reference> take() && {
     llvm::sort(References, [](const Reference &L, const Reference &R) {
       auto LTok = L.SpelledTok.location();
       auto RTok = R.SpelledTok.location();
       return std::tie(LTok, L.Role) < std::tie(RTok, R.Role);
     });
     // We sometimes see duplicates when parts of the AST get traversed twice.
     References.erase(std::unique(References.begin(), References.end(),
                                  [](const Reference &L, const Reference &R) {
                                    auto LTok = L.SpelledTok.location();
                                    auto RTok = R.SpelledTok.location();
                                    return std::tie(LTok, L.Role) ==
                                           std::tie(RTok, R.Role);
                                  }),
                      References.end());
     return std::move(References);
   }
 
   bool
   handleDeclOccurrence(const Decl *D, index::SymbolRoleSet Roles,
                        llvm::ArrayRef<index::SymbolRelation> Relations,
                        SourceLocation Loc,
                        index::IndexDataConsumer::ASTNodeInfo ASTNode) override {
+    auto DeclID = TargetDeclToID.find(D->getCanonicalDecl());
+    if (DeclID == TargetDeclToID.end())
+      return true;
     const SourceManager &SM = AST.getSourceManager();
     if (!isInsideMainFile(Loc, SM))
       return true;
-    SymbolID ID = getSymbolID(D);
-    if (!TargetIDs.contains(ID))
-      return true;
     const auto &TB = AST.getTokens();
 
     llvm::SmallVector<SourceLocation, 1> Locs;
     if (PerToken) {
       // Check whether this is one of the few constructs where the reference
       // can be split over several tokens.
       if (auto *OME = llvm::dyn_cast_or_null<ObjCMessageExpr>(ASTNode.OrigE)) {
         OME->getSelectorLocs(Locs);
       } else if (auto *OMD =
                      llvm::dyn_cast_or_null<ObjCMethodDecl>(ASTNode.OrigD)) {
         OMD->getSelectorLocs(Locs);
       }
       // Sanity check: we expect the *first* token to match the reported loc.
       // Otherwise, maybe it was e.g. some other kind of reference to a Decl.
       if (!Locs.empty() && Locs.front() != Loc)
         Locs.clear(); // First token doesn't match, assume our guess was wrong.
     }
     if (Locs.empty())
       Locs.push_back(Loc);
 
     for (SourceLocation L : Locs) {
       L = SM.getFileLoc(L);
       if (const auto *Tok = TB.spelledTokenAt(L))
-        References.push_back({*Tok, Roles, ID});
+        References.push_back({*Tok, Roles, DeclID->getSecond()});
     }
     return true;
   }
 
 private:
   bool PerToken; // If true, report 3 references for split ObjC selector names.
   std::vector<Reference> References;
   const ParsedAST &AST;
-  const llvm::DenseSet<SymbolID> &TargetIDs;
+  llvm::DenseMap<const Decl *, SymbolID> TargetDeclToID;
 };
 
 std::vector<ReferenceFinder::Reference>
-findRefs(const llvm::DenseSet<SymbolID> &IDs, ParsedAST &AST, bool PerToken) {
-  ReferenceFinder RefFinder(AST, IDs, PerToken);
+findRefs(const llvm::ArrayRef<const NamedDecl*> TargetDecls, ParsedAST &AST,
+         bool PerToken) {
+  ReferenceFinder RefFinder(AST, TargetDecls, PerToken);
   index::IndexingOptions IndexOpts;
   IndexOpts.SystemSymbolFilter =
       index::IndexingOptions::SystemSymbolFilterKind::All;
   IndexOpts.IndexFunctionLocals = true;
   IndexOpts.IndexParametersInDeclarations = true;
   IndexOpts.IndexTemplateParameters = true;
   indexTopLevelDecls(AST.getASTContext(), AST.getPreprocessor(),
                      AST.getLocalTopLevelDecls(), RefFinder, IndexOpts);
if (!CurLoc) {
     llvm::consumeError(CurLoc.takeError());
     return {};
   }
   std::vector<DocumentHighlight> Result;
   auto TryTree = [&](SelectionTree ST) {
     if (const SelectionTree::Node *N = ST.commonAncestor()) {
       DeclRelationSet Relations =
           DeclRelation::TemplatePattern | DeclRelation::Alias;
-      auto Decls =
+      auto TargetDecls=
            targetDecl(N->ASTNode, Relations, AST.getHeuristicResolver());
-      if (!Decls.empty()) {
+      if (!TargetDecls.empty()) {
         // FIXME: we may get multiple DocumentHighlights with the same location
         // and different kinds, deduplicate them.
-        llvm::DenseSet<SymbolID> Targets;
-        for (const NamedDecl *ND : Decls)
-          if (auto ID = getSymbolID(ND))
-            Targets.insert(ID);
-        for (const auto &Ref : findRefs(Targets, AST, /*PerToken=*/true))
+        for (const auto &Ref : findRefs(TargetDecls, AST, /*PerToken=*/true))
           Result.push_back(toHighlight(Ref, SM));
         return true;
       }
       auto ControlFlow = relatedControlFlow(*N);
       if (!ControlFlow.empty()) {
         for (SourceLocation Loc : ControlFlow)
           if (auto Highlight = toHighlight(Loc, AST.getTokens()))
             Result.push_back(std::move(*Highlight));
if (Macro) {
     }
   } else {
     // Handle references to Decls.
 
     DeclRelationSet Relations =
         DeclRelation::TemplatePattern | DeclRelation::Alias;
     std::vector<const NamedDecl *> Decls =
         getDeclAtPosition(AST, *CurLoc, Relations);
-    llvm::DenseSet<SymbolID> TargetsInMainFile;
+    llvm::SmallVector<const NamedDecl *> TargetsInMainFile;
     for (const NamedDecl *D : Decls) {
       auto ID = getSymbolID(D);
       if (!ID)
         continue;
-      TargetsInMainFile.insert(ID);
+      TargetsInMainFile.push_back(D);
       // Not all symbols can be referenced from outside (e.g. function-locals).
       // TODO: we could skip TU-scoped symbols here (e.g. static functions) if
       // we know this file isn't a header. The details might be tricky.
       if (D->getParentFunctionOrMethod())
         continue;
       IDsToQuery.insert(ID);
     }
 

 
 // Given a type targeted by the cursor, return a type that's more interesting
 // to target.
 static QualType unwrapFindType(QualType T) {
   if (T.isNull())
     return T;
 
   // If there's a specific type alias, point at that rather than unwrapping.
   if (const auto* TDT = T->getAs<TypedefType>())
     return QualType(TDT, 0);
 
   // Pointers etc => pointee type.
   if (const auto *PT = T->getAs<PointerType>())
     return unwrapFindType(PT->getPointeeType());
   if (const auto *RT = T->getAs<ReferenceType>())
     return unwrapFindType(RT->getPointeeType());
   if (const auto *AT = T->getAsArrayTypeUnsafe())