[clangd] Go-to-definition on pure virtual method decls jumps to all overrides.

[usaxena95]

LG. Thanks.

clang-tools-extra/clangd/XRefs.cpp

316

[hokein]

address comment.

clang-tools-extra/clangd/XRefs.cpp

if (const auto *DefinitionID = ID->getDefinition())
       return DefinitionID;
   if (const auto *PD = dyn_cast<ObjCProtocolDecl>(D))
     if (const auto *DefinitionID = PD->getDefinition())
       return DefinitionID;
 
   return D;
 }
 
+std::vector<LocatedSymbol> findOverrides(llvm::DenseSet<SymbolID> IDs,
+                                         const SymbolIndex *Index,
+                                         llvm::StringRef MainFilePath) {
+  if (IDs.empty())
+    return {};
+  RelationsRequest Req;
+  Req.Predicate = RelationKind::OverriddenBy;
+  Req.Subjects = std::move(IDs);
+  std::vector<LocatedSymbol> Results;
+  Index->relations(Req, [&](const SymbolID &Subject, const Symbol &Object) {
+    auto DeclLoc =
+        indexToLSPLocation(Object.CanonicalDeclaration, MainFilePath);
+    if (!DeclLoc) {
+      elog("Find overrides: {0}", DeclLoc.takeError());
+      return;
+    }
+    Results.emplace_back();
+    Results.back().Name = Object.Name.str();
+    Results.back().PreferredDeclaration = *DeclLoc;
+    auto DefLoc = indexToLSPLocation(Object.Definition, MainFilePath);
+    if (!DefLoc) {
+      elog("Failed to convert location: {0}", DefLoc.takeError());
+      return;
+    }
+    Results.back().Definition = *DefLoc;
+  });
+  return Results;
+}
+
 // Decls are more complicated.
 // The AST contains at least a declaration, maybe a definition.
 // These are up-to-date, and so generally preferred over index results.
 // We perform a single batch index lookup to find additional definitions.
 std::vector<LocatedSymbol>
 locateASTReferent(SourceLocation CurLoc, const syntax::Token *TouchedIdentifier,
                   ParsedAST &AST, llvm::StringRef MainFilePath,
                   const SymbolIndex *Index, ASTNodeKind *NodeKind) {
if (auto ID = getSymbolID(D))
       ResultIndex[ID] = Result.size() - 1;
   };
 
   // Emit all symbol locations (declaration or definition) from AST.
   DeclRelationSet Relations =
       DeclRelation::TemplatePattern | DeclRelation::Alias;
   auto Candidates =
       getDeclAtPositionWithRelations(AST, CurLoc, Relations, NodeKind);
+  llvm::DenseSet<SymbolID> VirtualMethods;
   for (const auto &E : Candidates) {
     const NamedDecl *D = E.first;
-    // Special case: void foo() ^override: jump to the overridden method.
     if (const auto *CMD = llvm::dyn_cast<CXXMethodDecl>(D)) {
+      // Special case: virtual void ^method() = 0: jump to all overrides.
+      // FIXME: extend it to ^virtual, unfortunately, virtual location is not
+      // saved in the AST.
+      if (CMD->isPure()) {
+        if (TouchedIdentifier && SM.getSpellingLoc(CMD->getLocation()) ==
+                                     TouchedIdentifier->location())
+          VirtualMethods.insert(getSymbolID(CMD));
+      }
+      // Special case: void foo() ^override: jump to the overridden method.
       const InheritableAttr *Attr = D->getAttr<OverrideAttr>();
       if (!Attr)
         Attr = D->getAttr<FinalAttr>();
       if (Attr && TouchedIdentifier &&
           SM.getSpellingLoc(Attr->getLocation()) ==
               TouchedIdentifier->location()) {
         // We may be overridding multiple methods - offer them all.
         for (const NamedDecl *ND : CMD->overridden_methods())
Index->lookup(QueryRequest, [&](const Symbol &Sym) {
                 getPreferredLocation(R.PreferredDeclaration,
                                      Sym.CanonicalDeclaration, Scratch),
                 MainFilePath))
           R.PreferredDeclaration = *Loc;
       }
     });
   }
 
+  auto Overrides = findOverrides(VirtualMethods, Index, MainFilePath);
+  Result.insert(Result.end(), Overrides.begin(), Overrides.end());
   return Result;
 }
 
 bool tokenSpelledAt(SourceLocation SpellingLoc, const syntax::TokenBuffer &TB) {
   auto ExpandedTokens = TB.expandedTokens(
       TB.sourceManager().getMacroArgExpandedLocation(SpellingLoc));
   return !ExpandedTokens.empty();
 }
if (!MainFilePath) {
     return {};
   }
   auto CurLoc = sourceLocationInMainFile(SM, Pos);
   if (!CurLoc) {
     elog("Failed to convert position to source location: {0}",
          CurLoc.takeError());
     return {};
   }
-  std::vector<LocatedSymbol> Results;
   DeclRelationSet Relations =
       DeclRelation::TemplatePattern | DeclRelation::Alias;
-  RelationsRequest Req;
-  Req.Predicate = RelationKind::OverriddenBy;
+  llvm::DenseSet<SymbolID> VirtualMethods;
   for (const NamedDecl *ND : getDeclAtPosition(AST, *CurLoc, Relations))
     if (const CXXMethodDecl *CXXMD = llvm::dyn_cast<CXXMethodDecl>(ND))
       if (CXXMD->isVirtual())
-        Req.Subjects.insert(getSymbolID(ND));
-
-  if (Req.Subjects.empty())
-    return Results;
-  Index->relations(Req, [&](const SymbolID &Subject, const Symbol &Object) {
-    auto DeclLoc =
-        indexToLSPLocation(Object.CanonicalDeclaration, *MainFilePath);
-    if (!DeclLoc) {
-      elog("Find implementation: {0}", DeclLoc.takeError());
-      return;
-    }
-    LocatedSymbol Loc;
-    Loc.Name = Object.Name.str();
-    Loc.PreferredDeclaration = *DeclLoc;
-    auto DefLoc = indexToLSPLocation(Object.Definition, *MainFilePath);
-    if (DefLoc)
-      Loc.Definition = *DefLoc;
-    else
-      llvm::consumeError(DefLoc.takeError());
-    Results.push_back(Loc);
-  });
-  return Results;
+        VirtualMethods.insert(getSymbolID(ND));
+  return findOverrides(std::move(VirtualMethods), Index, *MainFilePath);
 }
 
 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();

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

Test = Annotations(R"cpp(// definition in AST.
         class [[Forward]] {};
         F^orward create();
       )cpp");
   EXPECT_THAT(
       LocateWithIndex(Test),
       ElementsAre(Sym("Forward", SymbolHeader.range("forward"), Test.range())));
 }
 
+TEST(LocateSymbol, FindOverrides) {
+  auto Code = Annotations(R"cpp(
+    class Foo {
+      virtual void $1[[fo^o]]() = 0;
+    };
+    class Bar : public Foo {
+      void $2[[foo]]() override;
+    };
+  )cpp");
+  TestTU TU = TestTU::withCode(Code.code());
+  auto AST = TU.build();
+  EXPECT_THAT(locateSymbolAt(AST, Code.point(), TU.index().get()),
+              UnorderedElementsAre(Sym("foo", Code.range("1"), llvm::None),
+                                   Sym("foo", Code.range("2"), llvm::None)));
+}
+
 TEST(LocateSymbol, WithIndexPreferredLocation) {
   Annotations SymbolHeader(R"cpp(
         class $p[[Proto]] {};
         void $f[[func]]() {};
       )cpp");
   TestTU TU;
   TU.HeaderCode = std::string(SymbolHeader.code());
   TU.HeaderFilename = "x.proto"; // Prefer locations in codegen files.