[clangd][ObjC] Improve xrefs for protocols and classes

Summary:
Previously clangd would jump to forward declarations for protocols
and classes instead of their definition/implementation.

Reviewers: sammccall

Subscribers: ilya-biryukov, MaskRay, jkorous, arphaman, kadircet, usaxena95, cfe-commits

Tags: #clang

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

clang-tools-extra/clangd/FindTarget.cpp

@@ -355,7 +355,21 @@ public:
       Flags |= Rel::Underlying; // continue with the underlying decl.
     } else if (const auto *DG = dyn_cast<CXXDeductionGuideDecl>(D)) {
       D = DG->getDeducedTemplate();
+    } else if (const ObjCImplementationDecl *IID =
+                   dyn_cast<ObjCImplementationDecl>(D)) {
+      // Treat ObjC{Interface,Implementation}Decl as if they were a decl/def
+      // pair as long as the interface isn't implicit.
+      if (const auto *CID = IID->getClassInterface())
+        if (const auto *DD = CID->getDefinition())
+          if (!DD->isImplicitInterfaceDecl())
+            D = DD;
+    } else if (const ObjCCategoryImplDecl *CID =
+                   dyn_cast<ObjCCategoryImplDecl>(D)) {
+      // Treat ObjC{Category,CategoryImpl}Decl as if they were a decl/def pair.
+      D = CID->getCategoryDecl();
     }
+    if (!D)
+      return;
 
     if (const Decl *Pat = getTemplatePattern(D)) {
       assert(Pat != D);

clang-tools-extra/clangd/XRefs.cpp

@@ -78,6 +78,32 @@ const NamedDecl *getDefinition(const NamedDecl *D) {
     return VD->getDefinition();
   if (const auto *FD = dyn_cast<FunctionDecl>(D))
     return FD->getDefinition();
+  // Objective-C classes can have three types of declarations:
+  //
+  // - forward declaration: @class MyClass;
+  // - true declaration (interface definition): @interface MyClass ... @end
+  // - true definition (implementation): @implementation MyClass ... @end
+  //
+  // Objective-C categories are extensions are on classes:
+  //
+  // - declaration: @interface MyClass (Ext) ... @end
+  // - definition: @implementation MyClass (Ext) ... @end
+  //
+  // With one special case, a class extension, which is normally used to keep
+  // some declarations internal to a file without exposing them in a header.
+  //
+  // - class extension declaration: @interface MyClass () ... @end
+  // - which really links to class definition: @implementation MyClass ... @end
+  if (const auto *ID = dyn_cast<ObjCInterfaceDecl>(D))
+    return ID->getImplementation();
+  if (const auto *CD = dyn_cast<ObjCCategoryDecl>(D)) {
+    if (CD->IsClassExtension()) {
+      if (const auto *ID = CD->getClassInterface())
+        return ID->getImplementation();
+      return nullptr;
+    }
+    return CD->getImplementation();
+  }
   // Only a single declaration is allowed.
   if (isa<ValueDecl>(D) || isa<TemplateTypeParmDecl>(D) ||
       isa<TemplateTemplateParmDecl>(D)) // except cases above
@@ -223,6 +249,37 @@ locateMacroReferent(const syntax::Token &TouchedIdentifier, ParsedAST &AST,
   return llvm::None;
 }
 
+// A wrapper around `Decl::getCanonicalDecl` to support cases where Clang's
+// definition of a canonical declaration doesn't match up to what a programmer
+// would expect. For example, Objective-C classes can have three types of
+// declarations:
+//
+// - forward declaration(s): @class MyClass;
+// - true declaration (interface definition): @interface MyClass ... @end
+// - true definition (implementation): @implementation MyClass ... @end
+//
+// Clang will consider the forward declaration to be the canonical declaration
+// because it is first. We actually want the class definition if it is
+// available since that is what a programmer would consider the primary
+// declaration to be.
+const NamedDecl *getPreferredDecl(const NamedDecl *D) {
+  // FIXME: Canonical declarations of some symbols might refer to built-in
+  // decls with possibly-invalid source locations (e.g. global new operator).
+  // In such cases we should pick up a redecl with valid source location
+  // instead of failing.
+  D = llvm::cast<NamedDecl>(D->getCanonicalDecl());
+
+  // Prefer Objective-C class/protocol definitions over the forward declaration.
+  if (const auto *ID = dyn_cast<ObjCInterfaceDecl>(D))
+    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;
+}
+
 // 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.
@@ -238,11 +295,7 @@ locateASTReferent(SourceLocation CurLoc, const syntax::Token *TouchedIdentifier,
   llvm::DenseMap<SymbolID, size_t> ResultIndex;
 
   auto AddResultDecl = [&](const NamedDecl *D) {
-    // FIXME: Canonical declarations of some symbols might refer to built-in
-    // decls with possibly-invalid source locations (e.g. global new operator).
-    // In such cases we should pick up a redecl with valid source location
-    // instead of failing.
-    D = llvm::cast<NamedDecl>(D->getCanonicalDecl());
+    D = getPreferredDecl(D);
     auto Loc =
         makeLocation(AST.getASTContext(), nameLocation(*D, SM), MainFilePath);
     if (!Loc)
@@ -302,6 +355,19 @@ locateASTReferent(SourceLocation CurLoc, const syntax::Token *TouchedIdentifier,
       continue;
     }
 
+    // Special case: if the class name is selected, also map Objective-C
+    // categories and category implementations back to their class interface.
+    //
+    // Since `TouchedIdentifier` might refer to the `ObjCCategoryImplDecl`
+    // instead of the `ObjCCategoryDecl` we intentionally check the contents
+    // of the locs when checking for class name equivalence.
+    if (const auto *CD = dyn_cast<ObjCCategoryDecl>(D))
+      if (const auto *ID = CD->getClassInterface())
+        if (TouchedIdentifier &&
+            (CD->getLocation() == TouchedIdentifier->location() ||
+             ID->getName() == TouchedIdentifier->text(SM)))
+          AddResultDecl(ID);
+
     // Otherwise the target declaration is the right one.
     AddResultDecl(D);
   }

clang-tools-extra/clangd/index/SymbolCollector.cpp

@@ -20,6 +20,7 @@
 #include "clang/AST/Decl.h"
 #include "clang/AST/DeclBase.h"
 #include "clang/AST/DeclCXX.h"
+#include "clang/AST/DeclObjC.h"
 #include "clang/AST/DeclTemplate.h"
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/SourceManager.h"
@@ -156,16 +157,21 @@ getTokenLocation(SourceLocation TokLoc, const SourceManager &SM,
   return Result;
 }
 
-// Checks whether \p ND is a definition of a TagDecl (class/struct/enum/union)
-// in a header file, in which case clangd would prefer to use ND as a canonical
-// declaration.
-// FIXME: handle symbol types that are not TagDecl (e.g. functions), if using
-// the first seen declaration as canonical declaration is not a good enough
-// heuristic.
+// Checks whether \p ND is a good candidate to be the *canonical* declaration of
+// its symbol (e.g. a go-to-declaration target). This overrides the default of
+// using Clang's canonical declaration, which is the first in the TU.
+//
+// Example: preferring a class declaration over its forward declaration.
 bool isPreferredDeclaration(const NamedDecl &ND, index::SymbolRoleSet Roles) {
   const auto &SM = ND.getASTContext().getSourceManager();
-  return (Roles & static_cast<unsigned>(index::SymbolRole::Definition)) &&
-         isa<TagDecl>(&ND) && !isInsideMainFile(ND.getLocation(), SM);
+  if (isa<TagDecl>(ND))
+    return (Roles & static_cast<unsigned>(index::SymbolRole::Definition)) &&
+           !isInsideMainFile(ND.getLocation(), SM);
+  if (const auto *ID = dyn_cast<ObjCInterfaceDecl>(&ND))
+    return ID->isThisDeclarationADefinition();
+  if (const auto *PD = dyn_cast<ObjCProtocolDecl>(&ND))
+    return PD->isThisDeclarationADefinition();
+  return false;
 }
 
 RefKind toRefKind(index::SymbolRoleSet Roles, bool Spelled = false) {
@@ -267,6 +273,25 @@ bool SymbolCollector::handleDeclOccurrence(
   // picked a replacement for D
   if (D->getFriendObjectKind() != Decl::FriendObjectKind::FOK_None)
     D = CanonicalDecls.try_emplace(D, ASTNode.OrigD).first->second;
+  // Flag to mark that D should be considered canonical meaning its declaration
+  // will override any previous declaration for the Symbol.
+  bool DeclIsCanonical = false;
+  // Avoid treating ObjCImplementationDecl as a canonical declaration if it has
+  // a corresponding non-implicit and non-forward declared ObjcInterfaceDecl.
+  if (const auto *IID = dyn_cast<ObjCImplementationDecl>(D)) {
+    DeclIsCanonical = true;
+    if (const auto *CID = IID->getClassInterface())
+      if (const auto *DD = CID->getDefinition())
+        if (!DD->isImplicitInterfaceDecl())
+          D = DD;
+  }
+  // Avoid treating ObjCCategoryImplDecl as a canonical declaration in favor of
+  // its ObjCCategoryDecl if it has one.
+  if (const auto *CID = dyn_cast<ObjCCategoryImplDecl>(D)) {
+    DeclIsCanonical = true;
+    if (const auto *CD = CID->getCategoryDecl())
+      D = CD;
+  }
   const NamedDecl *ND = dyn_cast<NamedDecl>(D);
   if (!ND)
     return true;
@@ -331,14 +356,14 @@ bool SymbolCollector::handleDeclOccurrence(
     return true;
 
   const Symbol *BasicSymbol = Symbols.find(*ID);
-  if (!BasicSymbol) // Regardless of role, ND is the canonical declaration.
-    BasicSymbol = addDeclaration(*ND, std::move(*ID), IsMainFileOnly);
-  else if (isPreferredDeclaration(*OriginalDecl, Roles))
-    // If OriginalDecl is preferred, replace the existing canonical
+  if (isPreferredDeclaration(*OriginalDecl, Roles))
+    // If OriginalDecl is preferred, replace/create the existing canonical
     // declaration (e.g. a class forward declaration). There should be at most
     // one duplicate as we expect to see only one preferred declaration per
     // TU, because in practice they are definitions.
     BasicSymbol = addDeclaration(*OriginalDecl, std::move(*ID), IsMainFileOnly);
+  else if (!BasicSymbol || DeclIsCanonical)
+    BasicSymbol = addDeclaration(*ND, std::move(*ID), IsMainFileOnly);
 
   if (Roles & static_cast<unsigned>(index::SymbolRole::Definition))
     addDefinition(*OriginalDecl, *BasicSymbol);

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

@@ -756,6 +756,30 @@ TEST_F(TargetDeclTest, ObjC) {
   )cpp";
   EXPECT_DECLS("ObjCInterfaceTypeLoc", "@interface Foo");
 
+  Code = R"cpp(// Don't consider implicit interface as the target.
+    @implementation [[Implicit]]
+    @end
+  )cpp";
+  EXPECT_DECLS("ObjCImplementationDecl", "@implementation Implicit");
+
+  Code = R"cpp(
+    @interface Foo
+    @end
+    @implementation [[Foo]]
+    @end
+  )cpp";
+  EXPECT_DECLS("ObjCImplementationDecl", "@interface Foo");
+
+  Code = R"cpp(
+    @interface Foo
+    @end
+    @interface Foo (Ext)
+    @end
+    @implementation [[Foo]] (Ext)
+    @end
+  )cpp";
+  EXPECT_DECLS("ObjCCategoryImplDecl", "@interface Foo(Ext)");
+
   Code = R"cpp(
     @protocol Foo
     @end

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

@@ -542,6 +542,96 @@ TEST_F(SymbolCollectorTest, ObjCPropertyImpl) {
   //        Figure out why it's platform-dependent.
 }
 
+TEST_F(SymbolCollectorTest, ObjCLocations) {
+  Annotations Header(R"(
+    // Declared in header, defined in main.
+    @interface $dogdecl[[Dog]]
+    @end
+    @interface $fluffydecl[[Dog]] (Fluffy)
+    @end
+  )");
+  Annotations Main(R"(
+    @interface Dog ()
+    @end
+    @implementation $dogdef[[Dog]]
+    @end
+    @implementation $fluffydef[[Dog]] (Fluffy)
+    @end
+    // Category with no declaration (only implementation).
+    @implementation $ruff[[Dog]] (Ruff)
+    @end
+    // Implicitly defined interface.
+    @implementation $catdog[[CatDog]]
+    @end
+  )");
+  runSymbolCollector(Header.code(), Main.code(),
+                     {"-xobjective-c++", "-Wno-objc-root-class"});
+  EXPECT_THAT(Symbols,
+              UnorderedElementsAre(
+                  AllOf(QName("Dog"), DeclRange(Header.range("dogdecl")),
+                        DefRange(Main.range("dogdef"))),
+                  AllOf(QName("Fluffy"), DeclRange(Header.range("fluffydecl")),
+                        DefRange(Main.range("fluffydef"))),
+                  AllOf(QName("CatDog"), DeclRange(Main.range("catdog")),
+                        DefRange(Main.range("catdog"))),
+                  AllOf(QName("Ruff"), DeclRange(Main.range("ruff")),
+                        DefRange(Main.range("ruff")))));
+}
+
+TEST_F(SymbolCollectorTest, ObjCForwardDecls) {
+  Annotations Header(R"(
+    // Forward declared in header, declared and defined in main.
+    @protocol Barker;
+    @class Dog;
+    // Never fully declared so Clang latches onto this decl.
+    @class $catdogdecl[[CatDog]];
+  )");
+  Annotations Main(R"(
+    @protocol $barkerdecl[[Barker]]
+    - (void)woof;
+    @end
+    @interface $dogdecl[[Dog]]<Barker>
+    - (void)woof;
+    @end
+    @implementation $dogdef[[Dog]]
+    - (void)woof {}
+    @end
+    @implementation $catdogdef[[CatDog]]
+    @end
+  )");
+  runSymbolCollector(Header.code(), Main.code(),
+                     {"-xobjective-c++", "-Wno-objc-root-class"});
+  EXPECT_THAT(Symbols,
+              UnorderedElementsAre(
+                  AllOf(QName("CatDog"), DeclRange(Header.range("catdogdecl")),
+                        DefRange(Main.range("catdogdef"))),
+                  AllOf(QName("Dog"), DeclRange(Main.range("dogdecl")),
+                        DefRange(Main.range("dogdef"))),
+                  AllOf(QName("Barker"), DeclRange(Main.range("barkerdecl"))),
+                  QName("Barker::woof"), QName("Dog::woof")));
+}
+
+TEST_F(SymbolCollectorTest, ObjCClassExtensions) {
+  Annotations Header(R"(
+    @interface $catdecl[[Cat]]
+    @end
+  )");
+  Annotations Main(R"(
+    @interface Cat ()
+    - (void)meow;
+    @end
+    @interface Cat ()
+    - (void)pur;
+    @end
+  )");
+  runSymbolCollector(Header.code(), Main.code(),
+                     {"-xobjective-c++", "-Wno-objc-root-class"});
+  EXPECT_THAT(Symbols,
+              UnorderedElementsAre(
+                  AllOf(QName("Cat"), DeclRange(Header.range("catdecl"))),
+                  QName("Cat::meow"), QName("Cat::pur")));
+}
+
 TEST_F(SymbolCollectorTest, Locations) {
   Annotations Header(R"cpp(
     // Declared in header, defined in main.

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

@@ -684,7 +684,57 @@ TEST(LocateSymbol, All) {
           enum class E { [[A]], B };
           E e = E::A^;
         };
-      )cpp"};
+      )cpp",
+
+      R"objc(
+        @protocol Dog;
+        @protocol $decl[[Dog]]
+        - (void)bark;
+        @end
+        id<Do^g> getDoggo() {
+          return 0;
+        }
+      )objc",
+
+      R"objc(
+        @interface Cat
+        @end
+        @implementation Cat
+        @end
+        @interface $decl[[Cat]] (Exte^nsion)
+        - (void)meow;
+        @end
+        @implementation $def[[Cat]] (Extension)
+        - (void)meow {}
+        @end
+      )objc",
+
+      R"objc(
+        @class $decl[[Foo]];
+        Fo^o * getFoo() {
+          return 0;
+        }
+      )objc",
+
+      R"objc(// Prefer interface definition over forward declaration
+        @class Foo;
+        @interface $decl[[Foo]]
+        @end
+        Fo^o * getFoo() {
+          return 0;
+        }
+      )objc",
+
+      R"objc(
+        @class Foo;
+        @interface $decl[[Foo]]
+        @end
+        @implementation $def[[Foo]]
+        @end
+        Fo^o * getFoo() {
+          return 0;
+        }
+      )objc"};
   for (const char *Test : Tests) {
     Annotations T(Test);
     llvm::Optional<Range> WantDecl;
@@ -702,6 +752,7 @@ TEST(LocateSymbol, All) {
     // FIXME: Auto-completion in a template requires disabling delayed template
     // parsing.
     TU.ExtraArgs.push_back("-fno-delayed-template-parsing");
+    TU.ExtraArgs.push_back("-xobjective-c++");
 
     auto AST = TU.build();
     auto Results = locateSymbolAt(AST, T.point());
@@ -719,6 +770,90 @@ TEST(LocateSymbol, All) {
   }
 }
 
+TEST(LocateSymbol, AllMulti) {
+  // Ranges in tests:
+  //   $declN is the declaration location
+  //   $defN is the definition location (if absent, symbol has no definition)
+  //
+  // NOTE:
+  //   N starts at 0.
+  struct ExpectedRanges {
+    Range WantDecl;
+    llvm::Optional<Range> WantDef;
+  };
+  const char *Tests[] = {
+      R"objc(
+        @interface $decl0[[Cat]]
+        @end
+        @implementation $def0[[Cat]]
+        @end
+        @interface $decl1[[Ca^t]] (Extension)
+        - (void)meow;
+        @end
+        @implementation $def1[[Cat]] (Extension)
+        - (void)meow {}
+        @end
+      )objc",
+
+      R"objc(
+        @interface $decl0[[Cat]]
+        @end
+        @implementation $def0[[Cat]]
+        @end
+        @interface $decl1[[Cat]] (Extension)
+        - (void)meow;
+        @end
+        @implementation $def1[[Ca^t]] (Extension)
+        - (void)meow {}
+        @end
+      )objc",
+
+      R"objc(
+        @interface $decl0[[Cat]]
+        @end
+        @interface $decl1[[Ca^t]] ()
+        - (void)meow;
+        @end
+        @implementation $def0[[$def1[[Cat]]]]
+        - (void)meow {}
+        @end
+      )objc",
+  };
+  for (const char *Test : Tests) {
+    Annotations T(Test);
+    std::vector<ExpectedRanges> Ranges;
+    for (int Idx = 0; true; Idx++) {
+      bool HasDecl = !T.ranges("decl" + std::to_string(Idx)).empty();
+      bool HasDef = !T.ranges("def" + std::to_string(Idx)).empty();
+      if (!HasDecl && !HasDef)
+        break;
+      ExpectedRanges Range;
+      if (HasDecl)
+        Range.WantDecl = T.range("decl" + std::to_string(Idx));
+      if (HasDef)
+        Range.WantDef = T.range("def" + std::to_string(Idx));
+      Ranges.push_back(Range);
+    }
+
+    TestTU TU;
+    TU.Code = std::string(T.code());
+    TU.ExtraArgs.push_back("-xobjective-c++");
+
+    auto AST = TU.build();
+    auto Results = locateSymbolAt(AST, T.point());
+
+    ASSERT_THAT(Results, ::testing::SizeIs(Ranges.size())) << Test;
+    for (size_t Idx = 0; Idx < Ranges.size(); Idx++) {
+      EXPECT_EQ(Results[Idx].PreferredDeclaration.range, Ranges[Idx].WantDecl)
+          << "($decl" << Idx << ")" << Test;
+      llvm::Optional<Range> GotDef;
+      if (Results[Idx].Definition)
+        GotDef = Results[Idx].Definition->range;
+      EXPECT_EQ(GotDef, Ranges[Idx].WantDef) << "($def" << Idx << ")" << Test;
+    }
+  }
+}
+
 // LocateSymbol test cases that produce warnings.
 // These are separated out from All so that in All we can assert
 // that there are no diagnostics.