[clangd][ObjC] Support ObjC class rename from implementation decls

Reviewed By: kadircet

Differential Revision: https://reviews.llvm.org/D152720

5 files changed, 176 insertions, 37 deletions

diff --git a/clang-tools-extra/clangd/FindTarget.cpp b/clang-tools-extra/clangd/FindTarget.cpp
index eead9e6..630b750 100644
--- a/clang-tools-extra/clangd/FindTarget.cpp
+++ b/clang-tools-extra/clangd/FindTarget.cpp
@@ -708,8 +708,23 @@ llvm::SmallVector<ReferenceLoc> refInDecl(const Decl *D,
                                   {OCID->getClassInterface()}});
       Refs.push_back(ReferenceLoc{NestedNameSpecifierLoc(),
                                   OCID->getCategoryNameLoc(),
-                                  /*IsDecl=*/true,
+                                  /*IsDecl=*/false,
                                   {OCID->getCategoryDecl()}});
+      Refs.push_back(ReferenceLoc{NestedNameSpecifierLoc(),
+                                  OCID->getCategoryNameLoc(),
+                                  /*IsDecl=*/true,
+                                  {OCID}});
+    }
+
+    void VisitObjCImplementationDecl(const ObjCImplementationDecl *OIMD) {
+      Refs.push_back(ReferenceLoc{NestedNameSpecifierLoc(),
+                                  OIMD->getLocation(),
+                                  /*IsDecl=*/false,
+                                  {OIMD->getClassInterface()}});
+      Refs.push_back(ReferenceLoc{NestedNameSpecifierLoc(),
+                                  OIMD->getLocation(),
+                                  /*IsDecl=*/true,
+                                  {OIMD}});
     }
   };
 
diff --git a/clang-tools-extra/clangd/SemanticHighlighting.cpp b/clang-tools-extra/clangd/SemanticHighlighting.cpp
index 3826170..f6a3f7a 100644
--- a/clang-tools-extra/clangd/SemanticHighlighting.cpp
+++ b/clang-tools-extra/clangd/SemanticHighlighting.cpp
@@ -128,7 +128,7 @@ std::optional<HighlightingKind> kindForDecl(const NamedDecl *D,
     return HighlightingKind::Class;
   if (isa<ObjCProtocolDecl>(D))
     return HighlightingKind::Interface;
-  if (isa<ObjCCategoryDecl>(D))
+  if (isa<ObjCCategoryDecl, ObjCCategoryImplDecl>(D))
     return HighlightingKind::Namespace;
   if (auto *MD = dyn_cast<CXXMethodDecl>(D))
     return MD->isStatic() ? HighlightingKind::StaticMethod
diff --git a/clang-tools-extra/clangd/refactor/Rename.cpp b/clang-tools-extra/clangd/refactor/Rename.cpp
index b327053..97ea5e1 100644
--- a/clang-tools-extra/clangd/refactor/Rename.cpp
+++ b/clang-tools-extra/clangd/refactor/Rename.cpp
@@ -19,6 +19,7 @@
 #include "clang/AST/ASTTypeTraits.h"
 #include "clang/AST/Decl.h"
 #include "clang/AST/DeclCXX.h"
+#include "clang/AST/DeclObjC.h"
 #include "clang/AST/DeclTemplate.h"
 #include "clang/AST/ParentMapContext.h"
 #include "clang/AST/Stmt.h"
@@ -140,6 +141,18 @@ const NamedDecl *canonicalRenameDecl(const NamedDecl *D) {
   return dyn_cast<NamedDecl>(D->getCanonicalDecl());
 }
 
+// Some AST nodes can reference multiple declarations. We try to pick the
+// relevant one to rename here.
+const NamedDecl *pickInterestingTarget(const NamedDecl *D) {
+  // We only support renaming the class name, not the category name. This has
+  // to be done outside of canonicalization since we don't want a category name
+  // reference to be canonicalized to the class.
+  if (const auto *CD = dyn_cast<ObjCCategoryDecl>(D))
+    if (const auto CI = CD->getClassInterface())
+      return CI;
+  return D;
+}
+
 llvm::DenseSet<const NamedDecl *> locateDeclAt(ParsedAST &AST,
                                                SourceLocation TokenStartLoc) {
   unsigned Offset =
@@ -156,6 +169,7 @@ llvm::DenseSet<const NamedDecl *> locateDeclAt(ParsedAST &AST,
        targetDecl(SelectedNode->ASTNode,
                   DeclRelation::Alias | DeclRelation::TemplatePattern,
                   AST.getHeuristicResolver())) {
+    D = pickInterestingTarget(D);
     Result.insert(canonicalRenameDecl(D));
   }
   return Result;
diff --git a/clang-tools-extra/clangd/unittests/FindTargetTests.cpp b/clang-tools-extra/clangd/unittests/FindTargetTests.cpp
index 9979628..64ac524 100644
--- a/clang-tools-extra/clangd/unittests/FindTargetTests.cpp
+++ b/clang-tools-extra/clangd/unittests/FindTargetTests.cpp
@@ -1128,6 +1128,16 @@ TEST_F(TargetDeclTest, ObjC) {
   EXPECT_DECLS("ObjCCategoryImplDecl", "@interface Foo(Ext)");
 
   Code = R"cpp(
+    @interface Foo
+    @end
+    @interface Foo (Ext)
+    @end
+    @implementation Foo ([[Ext]])
+    @end
+  )cpp";
+  EXPECT_DECLS("ObjCCategoryImplDecl", "@interface Foo(Ext)");
+
+  Code = R"cpp(
     void test(id</*error-ok*/[[InvalidProtocol]]> p);
   )cpp";
   EXPECT_DECLS("ParmVarDecl", "id p");
@@ -1216,10 +1226,7 @@ protected:
     std::string DumpedReferences;
   };
 
-  /// Parses \p Code, finds function or namespace '::foo' and annotates its body
-  /// with results of findExplicitReferences.
-  /// See actual tests for examples of annotation format.
-  AllRefs annotateReferencesInFoo(llvm::StringRef Code) {
+  TestTU newTU(llvm::StringRef Code) {
     TestTU TU;
     TU.Code = std::string(Code);
 
@@ -1228,30 +1235,11 @@ protected:
     TU.ExtraArgs.push_back("-std=c++20");
     TU.ExtraArgs.push_back("-xobjective-c++");
 
-    auto AST = TU.build();
-    auto *TestDecl = &findDecl(AST, "foo");
-    if (auto *T = llvm::dyn_cast<FunctionTemplateDecl>(TestDecl))
-      TestDecl = T->getTemplatedDecl();
-
-    std::vector<ReferenceLoc> Refs;
-    if (const auto *Func = llvm::dyn_cast<FunctionDecl>(TestDecl))
-      findExplicitReferences(
-          Func->getBody(),
-          [&Refs](ReferenceLoc R) { Refs.push_back(std::move(R)); },
-          AST.getHeuristicResolver());
-    else if (const auto *NS = llvm::dyn_cast<NamespaceDecl>(TestDecl))
-      findExplicitReferences(
-          NS,
-          [&Refs, &NS](ReferenceLoc R) {
-            // Avoid adding the namespace foo decl to the results.
-            if (R.Targets.size() == 1 && R.Targets.front() == NS)
-              return;
-            Refs.push_back(std::move(R));
-          },
-          AST.getHeuristicResolver());
-    else
-      ADD_FAILURE() << "Failed to find ::foo decl for test";
+    return TU;
+  }
 
+  AllRefs annotatedReferences(llvm::StringRef Code, ParsedAST &AST,
+                              std::vector<ReferenceLoc> Refs) {
     auto &SM = AST.getSourceManager();
     llvm::sort(Refs, [&](const ReferenceLoc &L, const ReferenceLoc &R) {
       return SM.isBeforeInTranslationUnit(L.NameLoc, R.NameLoc);
@@ -1288,9 +1276,60 @@ protected:
 
     return AllRefs{std::move(AnnotatedCode), std::move(DumpedReferences)};
   }
+
+  /// Parses \p Code, and annotates its body with results of
+  /// findExplicitReferences on all top level decls.
+  /// See actual tests for examples of annotation format.
+  AllRefs annotateAllReferences(llvm::StringRef Code) {
+    TestTU TU = newTU(Code);
+    auto AST = TU.build();
+
+    std::vector<ReferenceLoc> Refs;
+    for (auto *TopLevel : AST.getLocalTopLevelDecls())
+      findExplicitReferences(
+          TopLevel, [&Refs](ReferenceLoc R) { Refs.push_back(std::move(R)); },
+          AST.getHeuristicResolver());
+    return annotatedReferences(Code, AST, std::move(Refs));
+  }
+
+  /// Parses \p Code, finds function or namespace '::foo' and annotates its body
+  /// with results of findExplicitReferences.
+  /// See actual tests for examples of annotation format.
+  AllRefs annotateReferencesInFoo(llvm::StringRef Code) {
+    TestTU TU = newTU(Code);
+    auto AST = TU.build();
+    auto *TestDecl = &findDecl(AST, "foo");
+    if (auto *T = llvm::dyn_cast<FunctionTemplateDecl>(TestDecl))
+      TestDecl = T->getTemplatedDecl();
+
+    std::vector<ReferenceLoc> Refs;
+    if (const auto *Func = llvm::dyn_cast<FunctionDecl>(TestDecl))
+      findExplicitReferences(
+          Func->getBody(),
+          [&Refs](ReferenceLoc R) { Refs.push_back(std::move(R)); },
+          AST.getHeuristicResolver());
+    else if (const auto *NS = llvm::dyn_cast<NamespaceDecl>(TestDecl))
+      findExplicitReferences(
+          NS,
+          [&Refs, &NS](ReferenceLoc R) {
+            // Avoid adding the namespace foo decl to the results.
+            if (R.Targets.size() == 1 && R.Targets.front() == NS)
+              return;
+            Refs.push_back(std::move(R));
+          },
+          AST.getHeuristicResolver());
+    else if (const auto *OC = llvm::dyn_cast<ObjCContainerDecl>(TestDecl))
+      findExplicitReferences(
+          OC, [&Refs](ReferenceLoc R) { Refs.push_back(std::move(R)); },
+          AST.getHeuristicResolver());
+    else
+      ADD_FAILURE() << "Failed to find ::foo decl for test";
+
+    return annotatedReferences(Code, AST, std::move(Refs));
+  }
 };
 
-TEST_F(FindExplicitReferencesTest, All) {
+TEST_F(FindExplicitReferencesTest, AllRefsInFoo) {
   std::pair</*Code*/ llvm::StringRef, /*References*/ llvm::StringRef> Cases[] =
       {// Simple expressions.
        {R"cpp(
@@ -2022,6 +2061,42 @@ TEST_F(FindExplicitReferencesTest, All) {
   }
 }
 
+TEST_F(FindExplicitReferencesTest, AllRefs) {
+  std::pair</*Code*/ llvm::StringRef, /*References*/ llvm::StringRef> Cases[] =
+      {{R"cpp(
+      @interface $0^MyClass
+      @end
+      @implementation $1^$2^MyClass
+      @end
+      )cpp",
+        "0: targets = {MyClass}, decl\n"
+        "1: targets = {MyClass}\n"
+        "2: targets = {MyClass}, decl\n"},
+       {R"cpp(
+      @interface $0^MyClass
+      @end
+      @interface $1^MyClass ($2^Category)
+      @end
+      @implementation $3^MyClass ($4^$5^Category)
+      @end
+      )cpp",
+        "0: targets = {MyClass}, decl\n"
+        "1: targets = {MyClass}\n"
+        "2: targets = {Category}, decl\n"
+        "3: targets = {MyClass}\n"
+        "4: targets = {Category}\n"
+        "5: targets = {Category}, decl\n"}};
+
+  for (const auto &C : Cases) {
+    llvm::StringRef ExpectedCode = C.first;
+    llvm::StringRef ExpectedRefs = C.second;
+
+    auto Actual = annotateAllReferences(llvm::Annotations(ExpectedCode).code());
+    EXPECT_EQ(ExpectedCode, Actual.AnnotatedCode);
+    EXPECT_EQ(ExpectedRefs, Actual.DumpedReferences) << ExpectedCode;
+  }
+}
+
 } // namespace
 } // namespace clangd
 } // namespace clang
diff --git a/clang-tools-extra/clangd/unittests/RenameTests.cpp b/clang-tools-extra/clangd/unittests/RenameTests.cpp
index 9be4a97..2414ff6 100644
--- a/clang-tools-extra/clangd/unittests/RenameTests.cpp
+++ b/clang-tools-extra/clangd/unittests/RenameTests.cpp
@@ -840,6 +840,20 @@ TEST(RenameTest, WithinFileRename) {
           foo('x');
         }
       )cpp",
+
+      // ObjC class with a category.
+      R"cpp(
+        @interface [[Fo^o]]
+        @end
+        @implementation [[F^oo]]
+        @end
+        @interface [[Fo^o]] (Category)
+        @end
+        @implementation [[F^oo]] (Category)
+        @end
+
+        void func([[Fo^o]] *f) {}
+      )cpp",
   };
   llvm::StringRef NewName = "NewName";
   for (llvm::StringRef T : Tests) {
@@ -890,6 +904,15 @@ TEST(RenameTest, Renameable) {
       )cpp",
        "not a supported kind", HeaderFile},
 
+      {R"cpp(// disallow - category rename.
+        @interface Foo
+        @end
+        @interface Foo (Cate^gory)
+        @end
+      )cpp",
+       "Cannot rename symbol: there is no symbol at the given location",
+       HeaderFile},
+
       {
           R"cpp(
          #define MACRO 1
@@ -1468,7 +1491,7 @@ TEST(CrossFileRenameTests, DeduplicateRefsFromIndex) {
 
 TEST(CrossFileRenameTests, WithUpToDateIndex) {
   MockCompilationDatabase CDB;
-  CDB.ExtraClangFlags = {"-xc++"};
+  CDB.ExtraClangFlags = {"-xobjective-c++"};
   // rename is runnning on all "^" points in FooH, and "[[]]" ranges are the
   // expected rename occurrences.
   struct Case {
@@ -1557,13 +1580,12 @@ TEST(CrossFileRenameTests, WithUpToDateIndex) {
         }
       )cpp",
       },
-      {
-          // virtual templated method
-          R"cpp(
+      {// virtual templated method
+       R"cpp(
         template <typename> class Foo { virtual void [[m]](); };
         class Bar : Foo<int> { void [[^m]]() override; };
       )cpp",
-          R"cpp(
+       R"cpp(
           #include "foo.h"
 
           template<typename T> void Foo<T>::[[m]]() {}
@@ -1571,8 +1593,7 @@ TEST(CrossFileRenameTests, WithUpToDateIndex) {
           // the canonical Foo<T>::m().
           // https://github.com/clangd/clangd/issues/1325
           class Baz : Foo<float> { void m() override; };
-        )cpp"
-      },
+        )cpp"},
       {
           // rename on constructor and destructor.
           R"cpp(
@@ -1677,6 +1698,20 @@ TEST(CrossFileRenameTests, WithUpToDateIndex) {
         }
       )cpp",
       },
+      {
+          // Objective-C classes.
+          R"cpp(
+        @interface [[Fo^o]]
+        @end
+      )cpp",
+          R"cpp(
+        #include "foo.h"
+        @implementation [[Foo]]
+        @end
+
+        void func([[Foo]] *f) {}
+      )cpp",
+      },
   };
 
   trace::TestTracer Tracer;