[clangd] Extend find-refs to include overrides.

[hokein]

Summary

find-references on virtual void meth^od() = 0 will include override references.

[usaxena95]

LG. Thanks!

clang-tools-extra/clangd/XRefs.cpp

1335

nit: s/OverrideSymbols/Overrides

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

1848

nit: May be include tests for CRTP too.

[njames93]

Should we be providing virtual method overrides when finding references, surely that's what query implementations is for?

[hokein]

address review comment.

[hokein]

In D94390#2491046, @njames93 wrote:

Should we be providing virtual method overrides when finding references, surely that's what query implementations is for?

yeah, go-to-implementation is designed for this purpose. However, not every LSP-client would expose a UI for go-to-implementation, we want to surface these results for these clients as well via other common features like go-to-def, find-references, hover etc.

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

1848

I think CRTP is about class inheritance, while in this patch we just handle virtual methods which seems not related to CRTP.

clang-tools-extra/clangd/XRefs.cpp

struct Reference {
     index::SymbolRoleSet Role;
 
     Range range(const SourceManager &SM) const {
       return halfOpenToRange(SM, SpelledTok.range(SM).toCharRange(SM));
     }
   };
 
   ReferenceFinder(const ParsedAST &AST,
-                  const std::vector<const NamedDecl *> &TargetDecls)
-      : AST(AST) {
-    for (const NamedDecl *D : TargetDecls)
-      CanonicalTargets.insert(D->getCanonicalDecl());
-  }
+                  const llvm::DenseSet<SymbolID> &TargetIDs)
+      : AST(AST), TargetIDs(TargetIDs) {}
 
   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 {
-    assert(D->isCanonicalDecl() && "expect D to be a canonical declaration");
     const SourceManager &SM = AST.getSourceManager();
-    if (!CanonicalTargets.count(D) || !isInsideMainFile(Loc, SM))
+    if (!isInsideMainFile(Loc, SM) ||
+        TargetIDs.find(getSymbolID(D)) == TargetIDs.end())
       return true;
     const auto &TB = AST.getTokens();
     Loc = SM.getFileLoc(Loc);
     if (const auto *Tok = TB.spelledTokenAt(Loc))
       References.push_back({*Tok, Roles});
     return true;
   }
 
 private:
-  llvm::SmallSet<const Decl *, 4> CanonicalTargets;
   std::vector<Reference> References;
   const ParsedAST &AST;
+  const llvm::DenseSet<SymbolID> &TargetIDs;
 };
 
 std::vector<ReferenceFinder::Reference>
-findRefs(const std::vector<const NamedDecl *> &Decls, ParsedAST &AST) {
-  ReferenceFinder RefFinder(AST, Decls);
+findRefs(const llvm::DenseSet<SymbolID> &IDs, ParsedAST &AST) {
+  ReferenceFinder RefFinder(AST, IDs);
   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);
std::vector<DocumentHighlight> findDocumentHighlights(ParsedAST &AST,
   auto TryTree = [&](SelectionTree ST) {
     if (const SelectionTree::Node *N = ST.commonAncestor()) {
       DeclRelationSet Relations =
           DeclRelation::TemplatePattern | DeclRelation::Alias;
       auto Decls = targetDecl(N->ASTNode, Relations);
       if (!Decls.empty()) {
         // FIXME: we may get multiple DocumentHighlights with the same location
         // and different kinds, deduplicate them.
-        for (const auto &Ref : findRefs({Decls.begin(), Decls.end()}, AST))
+        llvm::DenseSet<SymbolID> Targets;
+        for (const NamedDecl *ND : Decls)
+          if (auto ID = getSymbolID(ND))
+            Targets.insert(ID);
+        for (const auto &Ref : findRefs(Targets, AST))
           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 (!MainFilePath) {
     return Results;
   }
   auto URIMainFile = URIForFile::canonicalize(*MainFilePath, *MainFilePath);
   auto CurLoc = sourceLocationInMainFile(SM, Pos);
   if (!CurLoc) {
     llvm::consumeError(CurLoc.takeError());
     return {};
   }
-  llvm::Optional<DefinedMacro> Macro;
-  if (const auto *IdentifierAtCursor =
-          syntax::spelledIdentifierTouching(*CurLoc, AST.getTokens())) {
-    Macro = locateMacroAt(*IdentifierAtCursor, AST.getPreprocessor());
-  }
 
   RefsRequest Req;
+
+  const auto *IdentifierAtCursor =
+      syntax::spelledIdentifierTouching(*CurLoc, AST.getTokens());
+  llvm::Optional<DefinedMacro> Macro;
+  if (IdentifierAtCursor)
+    Macro = locateMacroAt(*IdentifierAtCursor, AST.getPreprocessor());
   if (Macro) {
     // Handle references to macro.
     if (auto MacroSID = getSymbolID(Macro->Name, Macro->Info, SM)) {
       // Collect macro references from main file.
       const auto &IDToRefs = AST.getMacros().MacroRefs;
       auto Refs = IDToRefs.find(MacroSID);
       if (Refs != IDToRefs.end()) {
         for (const auto &Ref : Refs->second) {
           Location Result;
           Result.range = Ref.Rng;
           Result.uri = URIMainFile;
           Results.References.push_back(std::move(Result));
         }
       }
       Req.IDs.insert(MacroSID);
     }
   } else {
     // Handle references to Decls.
 
     DeclRelationSet Relations =
         DeclRelation::TemplatePattern | DeclRelation::Alias;
     std::vector<const NamedDecl *> Decls =
         getDeclAtPosition(AST, *CurLoc, Relations);
+    llvm::DenseSet<SymbolID> Targets;
+    for (const NamedDecl *D : Decls)
+      if (auto ID = getSymbolID(D))
+        Targets.insert(ID);
+
+    llvm::DenseSet<SymbolID> Overrides;
+    if (Index) {
+      RelationsRequest FindOverrides;
+      FindOverrides.Predicate = RelationKind::OverriddenBy;
+      for (const NamedDecl *ND : Decls) {
+        // Special case: virtual void meth^od() = 0 includes refs of overrides.
+        if (const auto *CMD = llvm::dyn_cast<CXXMethodDecl>(ND)) {
+          if (CMD->isPure())
+            if (IdentifierAtCursor && SM.getSpellingLoc(CMD->getLocation()) ==
+                                          IdentifierAtCursor->location())
+              if (auto ID = getSymbolID(CMD))
+                FindOverrides.Subjects.insert(ID);
+        }
+      }
+      if (!FindOverrides.Subjects.empty())
+        Index->relations(FindOverrides,
+                         [&](const SymbolID &Subject, const Symbol &Object) {
+                           Overrides.insert(Object.ID);
+                         });
+      Targets.insert(Overrides.begin(), Overrides.end());
+    }
 
     // We traverse the AST to find references in the main file.
-    auto MainFileRefs = findRefs(Decls, AST);
+    auto MainFileRefs = findRefs(Targets, AST);
     // We may get multiple refs with the same location and different Roles, as
     // cross-reference is only interested in locations, we deduplicate them
     // by the location to avoid emitting duplicated locations.
     MainFileRefs.erase(std::unique(MainFileRefs.begin(), MainFileRefs.end(),
                                    [](const ReferenceFinder::Reference &L,
                                       const ReferenceFinder::Reference &R) {
                                      return L.SpelledTok.location() ==
                                             R.SpelledTok.location();
                                    }),
                        MainFileRefs.end());
     for (const auto &Ref : MainFileRefs) {
       Location Result;
       Result.range = Ref.range(SM);
       Result.uri = URIMainFile;
       Results.References.push_back(std::move(Result));
     }
     if (Index && Results.References.size() <= Limit) {
+      Req.IDs = std::move(Overrides);
       for (const Decl *D : Decls) {
         // 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;
         if (auto ID = getSymbolID(D))

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

for (const char *Test : Tests) {
     for (const auto &R : T.ranges())
       ExpectedLocations.push_back(RangeIs(R));
     EXPECT_THAT(findReferences(AST, T.point(), 0).References,
                 ElementsAreArray(ExpectedLocations))
         << Test;
   }
 }
 
+TEST(FindReferences, IncludeOverrides) {
+  llvm::StringRef Test =
+      R"cpp(
+        class Base {
+        public:
+          virtual void [[f^unc]]() = 0;
+        };
+        class Derived : public Base {
+        public:
+          void [[func]]() override;
+        };
+        void test(Derived* D) {
+          D->[[func]]();
+        })cpp";
+  Annotations T(Test);
+  auto TU = TestTU::withCode(T.code());
+  auto AST = TU.build();
+  std::vector<Matcher<Location>> ExpectedLocations;
+  for (const auto &R : T.ranges())
+    ExpectedLocations.push_back(RangeIs(R));
+  EXPECT_THAT(findReferences(AST, T.point(), 0, TU.index().get()).References,
+              ElementsAreArray(ExpectedLocations))
+      << Test;
+}
+
 TEST(FindReferences, MainFileReferencesOnly) {
   llvm::StringRef Test =
       R"cpp(
         void test() {
           int [[fo^o]] = 1;
           // refs not from main file should not be included.
           #include "foo.inc"
         })cpp";