[clangd] Report xref for base methods.

[usaxena95]

Summary

See: https://github.com/clangd/clangd/issues/668

struct A { virtual void foo() = 0; };
struct B : A { void foo() override; };

Find refs on B::foo() will show:

• decls of A::foo() (new)
• decls of B::foo()
• refs to A::foo() (new)
• refs to B::foo().

[hokein]

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

1908

ah, this is a good case! I think we should include the base's base method as part of the refs -- because BB->func() can actually call D::func if BB points to a Derived object.

the fix is to add ids from overridden_methods recursively.

[usaxena95]

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

1908

Yeah we can do that. Although I had intentionally left those out considering the number of references would increase significantly if we do it for the complete type hierarchy. Also since LSP doesn't provide a way to annotate these kind of references, it would be hard for users to find exact refs (to Derived::foo()).
WDYT ?

[usaxena95]

Addressed comments.

[hokein]

Thanks!

clang-tools-extra/clangd/XRefs.cpp

1294

The recursive will be infinite if we have a cycle, but I guess clang will never generate a broken AST like that.

clang-tools-extra/clangd/XRefs.cpp

for (const NamedDecl *ND : getDeclAtPosition(AST, *CurLoc, Relations)) {
     } else if (const auto *RD = dyn_cast<CXXRecordDecl>(ND)) {
       IDs.insert(getSymbolID(RD));
       QueryKind = RelationKind::BaseOf;
     }
   }
   return findImplementors(std::move(IDs), QueryKind, Index, *MainFilePath);
 }
 
+namespace {
+// Recursively finds all the overridden methods of `CMD` in complete type
+// hierarchy.
+void getOverriddenMethods(const CXXMethodDecl *CMD,
+                          llvm::DenseSet<SymbolID> &OverriddenMethods) {
+  if (!CMD)
+    return;
+  for (const CXXMethodDecl *Base : CMD->overridden_methods()) {
+    if (auto ID = getSymbolID(Base))
+      OverriddenMethods.insert(ID);
+    getOverriddenMethods(Base, OverriddenMethods);
+  }
+}
+} // namespace
+
 ReferencesResult findReferences(ParsedAST &AST, Position Pos, uint32_t Limit,
                                 const SymbolIndex *Index) {
   if (!Limit)
     Limit = std::numeric_limits<uint32_t>::max();
   ReferencesResult Results;
   const SourceManager &SM = AST.getSourceManager();
   auto MainFilePath =
       getCanonicalPath(SM.getFileEntryForID(SM.getMainFileID()), SM);
   if (!MainFilePath) {
     elog("Failed to get a path for the main file, so no references");
     return Results;
   }
   auto URIMainFile = URIForFile::canonicalize(*MainFilePath, *MainFilePath);
   auto CurLoc = sourceLocationInMainFile(SM, Pos);
   if (!CurLoc) {
     llvm::consumeError(CurLoc.takeError());
     return {};
   }
 
-  llvm::DenseSet<SymbolID> IDs;
+  llvm::DenseSet<SymbolID> IDs, OverriddenMethods;
 
   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 (Macro) {
     for (const NamedDecl *D : Decls)
       if (auto ID = getSymbolID(D))
         Targets.insert(ID);
 
     RelationsRequest OverriddenBy;
     if (Index) {
       OverriddenBy.Predicate = RelationKind::OverriddenBy;
       for (const NamedDecl *ND : Decls) {
-        // Special case: Inlcude declaration of overridding methods.
+        // Special case: For virtual methods, report decl/def of overrides and
+        // references to all overridden methods in complete type hierarchy.
         if (const auto *CMD = llvm::dyn_cast<CXXMethodDecl>(ND)) {
           if (CMD->isVirtual())
             if (IdentifierAtCursor && SM.getSpellingLoc(CMD->getLocation()) ==
-                                          IdentifierAtCursor->location())
+                                          IdentifierAtCursor->location()) {
               if (auto ID = getSymbolID(CMD))
                 OverriddenBy.Subjects.insert(ID);
+              getOverriddenMethods(CMD, OverriddenMethods);
+            }
         }
       }
     }
 
     // We traverse the AST to find references in the main file.
     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
if (Index && Results.References.size() <= Limit) {
         if (D->getParentFunctionOrMethod())
           continue;
         if (auto ID = getSymbolID(D))
           IDs.insert(ID);
       }
     }
   }
   // Now query the index for references from other files.
-  auto QueryIndex = [&](llvm::DenseSet<SymbolID> IDs, bool AllowAttributes) {
+  auto QueryIndex = [&](llvm::DenseSet<SymbolID> IDs, bool AllowAttributes,
+                        bool AllowMainFileSymbols) {
     RefsRequest Req;
     Req.IDs = std::move(IDs);
     Req.Limit = Limit;
     if (Req.IDs.empty() || !Index || Results.References.size() > Limit)
       return;
     Results.HasMore |= Index->refs(Req, [&](const Ref &R) {
       // No need to continue process if we reach the limit.
       if (Results.References.size() > Limit)
         return;
       auto LSPLoc = toLSPLocation(R.Location, *MainFilePath);
       // Avoid indexed results for the main file - the AST is authoritative.
-      if (!LSPLoc || LSPLoc->uri.file() == *MainFilePath)
+      if (!LSPLoc ||
+          (!AllowMainFileSymbols && LSPLoc->uri.file() == *MainFilePath))
         return;
       ReferencesResult::Reference Result;
       Result.Loc = std::move(*LSPLoc);
       if (AllowAttributes) {
         if ((R.Kind & RefKind::Declaration) == RefKind::Declaration)
           Result.Attributes |= ReferencesResult::Declaration;
         // FIXME: our index should definitely store def | decl separately!
         if ((R.Kind & RefKind::Definition) == RefKind::Definition)
           Result.Attributes |=
               ReferencesResult::Declaration | ReferencesResult::Definition;
       }
       Results.References.push_back(std::move(Result));
     });
   };
-  QueryIndex(std::move(IDs), /*AllowAttributes=*/true);
+  QueryIndex(std::move(IDs), /*AllowAttributes=*/true,
+             /*AllowMainFileSymbols=*/false);
+  // For a virtual method: Occurrences of BaseMethod should be treated as refs
+  // and not as decl/def. Allow symbols from main file since AST does not report
+  // these.
+  QueryIndex(std::move(OverriddenMethods), /*AllowAttributes=*/false,
+             /*AllowMainFileSymbols=*/true);
   if (Results.References.size() > Limit) {
     Results.HasMore = true;
     Results.References.resize(Limit);
   }
-  // FIXME: Report refs of base methods.
   return Results;
 }
 
 std::vector<SymbolDetails> getSymbolInfo(ParsedAST &AST, Position Pos) {
   const SourceManager &SM = AST.getSourceManager();
   auto CurLoc = sourceLocationInMainFile(SM, Pos);
   if (!CurLoc) {
     llvm::consumeError(CurLoc.takeError());

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

llvm::StringRef Test =
         };
         void Derived::$overridedef[[func]]() {}
         void test(Derived* D) {
           D->func();  // No references to the overrides.
         })cpp";
   checkFindRefs(Test, /*UseIndex=*/true);
 }
 
+TEST(FindReferences, RefsToBaseMethod) {
+  llvm::StringRef Test =
+      R"cpp(
+        class BaseBase {
+        public:
+          virtual void [[func]]();
+        };
+        class Base : public BaseBase {
+        public:
+          void [[func]]() override;
+        };
+        class Derived : public Base {
+        public:
+          void $decl[[fu^nc]]() override;
+        };
+        void test(BaseBase* BB, Base* B, Derived* D) {
+          // refs to overridden methods in complete type hierarchy are reported.
+          BB->[[func]]();
+          B->[[func]]();
+          D->[[func]]();
+        })cpp";
+  checkFindRefs(Test, /*UseIndex=*/true);
+}
+
 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";