[clangd] Update XRefs to support overriden ObjC methods

clang-tools-extra/clangd/XRefs.cpp

@@ -372,6 +372,15 @@ void enhanceLocatedSymbolsFromIndex(llvm::MutableArrayRef<LocatedSymbol> Result,
   });
 }
 
+bool objcMethodIsTouched(const SourceManager &SM, const ObjCMethodDecl *OMD,
+                         SourceLocation Loc) {
+  unsigned NumSels = OMD->getNumSelectorLocs();
+  for (unsigned I = 0; I < NumSels; ++I)
+    if (SM.getSpellingLoc(OMD->getSelectorLoc(I)) == Loc)
+      return true;
+  return false;
+}
+
 // 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.
@@ -430,6 +439,26 @@ locateASTReferent(SourceLocation CurLoc, const syntax::Token *TouchedIdentifier,
         continue;
       }
     }
+    // Special case: - (void)^method {} should jump to overrides, but the decl
+    // shouldn't, only the definition. Note that an Objective-C method can
+    // override a parent class or protocol.
+    //
+    // FIXME: Support jumping from a protocol decl to overrides on go-to
+    // definition.
+    if (const auto *OMD = llvm::dyn_cast<ObjCMethodDecl>(D)) {
+      if (OMD->isThisDeclarationADefinition() && TouchedIdentifier &&
+          objcMethodIsTouched(SM, OMD, TouchedIdentifier->location())) {
+        llvm::SmallVector<const ObjCMethodDecl *, 4> Overrides;
+        OMD->getOverriddenMethods(Overrides);
+        if (!Overrides.empty()) {
+          for (const auto *Override : Overrides)
+            AddResultDecl(Override);
+          LocateASTReferentMetric.record(1, "objc-overriden-method");
+        }
+        AddResultDecl(OMD);
+        continue;
+      }
+    }
 
     // Special case: the cursor is on an alias, prefer other results.
     // This targets "using ns::^Foo", where the target is more interesting.
@@ -1283,6 +1312,12 @@ std::vector<LocatedSymbol> findImplementations(ParsedAST &AST, Position Pos,
     } else if (const auto *RD = dyn_cast<CXXRecordDecl>(ND)) {
       IDs.insert(getSymbolID(RD));
       QueryKind = RelationKind::BaseOf;
+    } else if (const auto *OMD = dyn_cast<ObjCMethodDecl>(ND)) {
+      IDs.insert(getSymbolID(OMD));
+      QueryKind = RelationKind::OverriddenBy;
+    } else if (const auto *ID = dyn_cast<ObjCInterfaceDecl>(ND)) {
+      IDs.insert(getSymbolID(ID));
+      QueryKind = RelationKind::BaseOf;
     }
   }
   return findImplementors(std::move(IDs), QueryKind, Index, AST.tuPath());
@@ -1302,6 +1337,21 @@ void getOverriddenMethods(const CXXMethodDecl *CMD,
   }
 }
 
+// Recursively finds all the overridden methods of `OMD` in complete type
+// hierarchy.
+void getOverriddenMethods(const ObjCMethodDecl *OMD,
+                          llvm::DenseSet<SymbolID> &OverriddenMethods) {
+  if (!OMD)
+    return;
+  llvm::SmallVector<const ObjCMethodDecl *, 4> Overrides;
+  OMD->getOverriddenMethods(Overrides);
+  for (const ObjCMethodDecl *Base : Overrides) {
+    if (auto ID = getSymbolID(Base))
+      OverriddenMethods.insert(ID);
+    getOverriddenMethods(Base, OverriddenMethods);
+  }
+}
+
 std::optional<std::string>
 stringifyContainerForMainFileRef(const Decl *Container) {
   // FIXME We might also want to display the signature here
@@ -1438,6 +1488,12 @@ ReferencesResult findReferences(ParsedAST &AST, Position Pos, uint32_t Limit,
             getOverriddenMethods(CMD, OverriddenMethods);
           }
         }
+        // Special case: Objective-C methods can override a parent class or
+        // protocol, we should be sure to report references to those.
+        if (const auto *OMD = llvm::dyn_cast<ObjCMethodDecl>(ND)) {
+          OverriddenBy.Subjects.insert(getSymbolID(OMD));
+          getOverriddenMethods(OMD, OverriddenMethods);
+        }
       }
     }
 

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

@@ -1335,6 +1335,42 @@ TEST_F(SymbolCollectorTest, OverrideRelationsMultipleInheritance) {
                           OverriddenBy(CBar, DBar), OverriddenBy(CBaz, DBaz)));
 }
 
+TEST_F(SymbolCollectorTest, ObjCOverrideRelationsSimpleInheritance) {
+  std::string Header = R"cpp(
+    @interface A
+    - (void)foo;
+    @end
+    @interface B : A
+    - (void)foo;  // A::foo
+    - (void)bar;
+    @end
+    @interface C : B
+    - (void)bar;  // B::bar
+    @end
+    @interface D : C
+    - (void)foo;  // B::foo
+    - (void)bar;  // C::bar
+    @end
+  )cpp";
+  runSymbolCollector(Header, /*Main=*/"",
+                     {"-xobjective-c++", "-Wno-objc-root-class"});
+  const Symbol &AFoo = findSymbol(Symbols, "A::foo");
+  const Symbol &BFoo = findSymbol(Symbols, "B::foo");
+  const Symbol &DFoo = findSymbol(Symbols, "D::foo");
+
+  const Symbol &BBar = findSymbol(Symbols, "B::bar");
+  const Symbol &CBar = findSymbol(Symbols, "C::bar");
+  const Symbol &DBar = findSymbol(Symbols, "D::bar");
+
+  std::vector<Relation> Result;
+  for (const Relation &R : Relations)
+    if (R.Predicate == RelationKind::OverriddenBy)
+      Result.push_back(R);
+  EXPECT_THAT(Result, UnorderedElementsAre(
+                          OverriddenBy(AFoo, BFoo), OverriddenBy(BBar, CBar),
+                          OverriddenBy(BFoo, DFoo), OverriddenBy(CBar, DBar)));
+}
+
 TEST_F(SymbolCollectorTest, CountReferences) {
   const std::string Header = R"(
     class W;

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

@@ -411,6 +411,85 @@ TEST(LocateSymbol, FindOverrides) {
                                    sym("foo", Code.range("2"), std::nullopt)));
 }
 
+TEST(LocateSymbol, FindOverridesFromDefObjC) {
+  auto Code = Annotations(R"objc(
+    @protocol Fooey
+    - (void)foo;
+    @end
+    @interface Base
+    - (void)foo;
+    @end
+    @interface Foo : Base<Fooey>
+    - (void)$1[[foo]];
+    @end
+
+    @interface Bar : Foo
+    - (void)$2[[foo]];
+    @end
+    @implementation Bar
+    - (void)$3[[fo^o]] {}
+    @end
+  )objc");
+  TestTU TU = TestTU::withCode(Code.code());
+  TU.ExtraArgs.push_back("-xobjective-c++");
+  auto AST = TU.build();
+  EXPECT_THAT(
+      locateSymbolAt(AST, Code.point(), TU.index().get()),
+      UnorderedElementsAre(sym("foo", Code.range("1"), std::nullopt),
+                           sym("foo", Code.range("2"), Code.range("3"))));
+}
+
+TEST(LocateSymbol, NoOverridesFromDeclObjC) {
+  auto Code = Annotations(R"objc(
+    @protocol Fooey
+    - (void)foo;
+    @end
+    @interface Base
+    - (void)foo;
+    @end
+    @interface Foo : Base<Fooey>
+    - (void)foo;
+    @end
+
+    @interface Bar : Foo
+    - (void)$2[[fo^o]];
+    @end
+    @implementation Bar
+    - (void)$3[[foo]] {}
+    @end
+  )objc");
+  TestTU TU = TestTU::withCode(Code.code());
+  TU.ExtraArgs.push_back("-xobjective-c++");
+  auto AST = TU.build();
+  EXPECT_THAT(
+      locateSymbolAt(AST, Code.point(), TU.index().get()),
+      UnorderedElementsAre(sym("foo", Code.range("2"), Code.range("3"))));
+}
+
+TEST(LocateSymbol, ObjCNoOverridesOnUsage) {
+  auto Code = Annotations(R"objc(
+    @interface Foo
+    - (void)foo;
+    @end
+
+    @interface Bar : Foo
+    - (void)$1[[foo]];
+    @end
+    @implementation Bar
+    - (void)$2[[foo]] {}
+    @end
+    void doSomething(Bar *bar) {
+      [bar fo^o];
+    }
+  )objc");
+  TestTU TU = TestTU::withCode(Code.code());
+  TU.ExtraArgs.push_back("-xobjective-c++");
+  auto AST = TU.build();
+  EXPECT_THAT(
+      locateSymbolAt(AST, Code.point(), TU.index().get()),
+      UnorderedElementsAre(sym("foo", Code.range("1"), Code.range("2"))));
+}
+
 TEST(LocateSymbol, WithIndexPreferredLocation) {
   Annotations SymbolHeader(R"cpp(
         class $p[[Proto]] {};
@@ -1834,6 +1913,41 @@ TEST(FindImplementations, Inheritance) {
   }
 }
 
+TEST(FindImplementations, InheritanceObjC) {
+  llvm::StringRef Test = R"objc(
+    @interface $base^Base
+    - (void)fo$foo^o;
+    @end
+    @protocol Protocol
+    - (void)$protocol^protocol;
+    @end
+    @interface $ChildDecl[[Child]] : Base <Protocol>
+    - (void)concrete;
+    - (void)$fooDecl[[foo]];
+    @end
+    @implementation $ChildDef[[Child]]
+    - (void)concrete {}
+    - (void)$fooDef[[foo]] {}
+    - (void)$protocolDef[[protocol]] {}
+    @end
+  )objc";
+
+  Annotations Code(Test);
+  auto TU = TestTU::withCode(Code.code());
+  TU.ExtraArgs.push_back("-xobjective-c++");
+  auto AST = TU.build();
+  auto Index = TU.index();
+  EXPECT_THAT(findImplementations(AST, Code.point("base"), Index.get()),
+              UnorderedElementsAre(sym("Child", Code.range("ChildDecl"),
+                                       Code.range("ChildDef"))));
+  EXPECT_THAT(findImplementations(AST, Code.point("foo"), Index.get()),
+              UnorderedElementsAre(
+                  sym("foo", Code.range("fooDecl"), Code.range("fooDef"))));
+  EXPECT_THAT(findImplementations(AST, Code.point("protocol"), Index.get()),
+              UnorderedElementsAre(sym("protocol", Code.range("protocolDef"),
+                                       Code.range("protocolDef"))));
+}
+
 TEST(FindImplementations, CaptureDefinition) {
   llvm::StringRef Test = R"cpp(
     struct Base {
@@ -1963,6 +2077,7 @@ void checkFindRefs(llvm::StringRef Test, bool UseIndex = false) {
   Annotations T(Test);
   auto TU = TestTU::withCode(T.code());
   TU.ExtraArgs.push_back("-std=c++20");
+  TU.ExtraArgs.push_back("-xobjective-c++");
 
   auto AST = TU.build();
   std::vector<Matcher<ReferencesResult::Reference>> ExpectedLocations;
@@ -2260,6 +2375,25 @@ TEST(FindReferences, IncludeOverrides) {
   checkFindRefs(Test, /*UseIndex=*/true);
 }
 
+TEST(FindReferences, IncludeOverridesObjC) {
+  llvm::StringRef Test =
+      R"objc(
+        @interface Base
+        - (void)$decl(Base)[[f^unc]];
+        @end
+        @interface Derived : Base
+        - (void)$overridedecl(Derived::func)[[func]];
+        @end
+        @implementation Derived
+        - (void)$overridedef[[func]] {}
+        @end
+        void test(Derived *derived, Base *base) {
+          [derived func];  // No references to the overrides.
+          [base $(test)[[func]]];
+        })objc";
+  checkFindRefs(Test, /*UseIndex=*/true);
+}
+
 TEST(FindReferences, RefsToBaseMethod) {
   llvm::StringRef Test =
       R"cpp(
@@ -2284,6 +2418,27 @@ TEST(FindReferences, RefsToBaseMethod) {
   checkFindRefs(Test, /*UseIndex=*/true);
 }
 
+TEST(FindReferences, RefsToBaseMethodObjC) {
+  llvm::StringRef Test =
+      R"objc(
+        @interface BaseBase
+        - (void)$(BaseBase)[[func]];
+        @end
+        @interface Base : BaseBase
+        - (void)$(Base)[[func]];
+        @end
+        @interface Derived : Base
+        - (void)$decl(Derived)[[fu^nc]];
+        @end
+        void test(BaseBase *bb, Base *b, Derived *d) {
+          // refs to overridden methods in complete type hierarchy are reported.
+          [bb $(test)[[func]]];
+          [b $(test)[[func]]];
+          [d $(test)[[fu^nc]]];
+        })objc";
+  checkFindRefs(Test, /*UseIndex=*/true);
+}
+
 TEST(FindReferences, MainFileReferencesOnly) {
   llvm::StringRef Test =
       R"cpp(