[clangd] Use `ObjCProtocolLoc` for generalized ObjC protocol support

This removes clangd's existing workaround in favor of proper support
via the newly added `ObjCProtocolLoc`. This improves support by
allowing clangd to properly identify which protocol is selected
now that `ObjCProtocolLoc` gets its own ASTNode.

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

4 files changed, 46 insertions, 50 deletions

diff --git a/clang-tools-extra/clangd/FindTarget.cpp b/clang-tools-extra/clangd/FindTarget.cpp
index e96aa25fd780..1b7b7de4f904 100644
--- a/clang-tools-extra/clangd/FindTarget.cpp
+++ b/clang-tools-extra/clangd/FindTarget.cpp
@@ -453,15 +453,6 @@ public:
       void VisitObjCInterfaceType(const ObjCInterfaceType *OIT) {
         Outer.add(OIT->getDecl(), Flags);
       }
-      void VisitObjCObjectType(const ObjCObjectType *OOT) {
-        // Make all of the protocols targets since there's no child nodes for
-        // protocols. This isn't needed for the base type, which *does* have a
-        // child `ObjCInterfaceTypeLoc`. This structure is a hack, but it works
-        // well for go-to-definition.
-        unsigned NumProtocols = OOT->getNumProtocols();
-        for (unsigned I = 0; I < NumProtocols; I++)
-          Outer.add(OOT->getProtocol(I), Flags);
-      }
     };
     Visitor(*this, Flags).Visit(T.getTypePtr());
   }
@@ -547,6 +538,8 @@ allTargetDecls(const DynTypedNode &N, const HeuristicResolver *Resolver) {
     Finder.add(TAL->getArgument(), Flags);
   else if (const CXXBaseSpecifier *CBS = N.get<CXXBaseSpecifier>())
     Finder.add(CBS->getTypeSourceInfo()->getType(), Flags);
+  else if (const ObjCProtocolLoc *PL = N.get<ObjCProtocolLoc>())
+    Finder.add(PL->getProtocol(), Flags);
   return Finder.takeDecls();
 }
 
@@ -669,25 +662,7 @@ llvm::SmallVector<ReferenceLoc> refInDecl(const Decl *D,
                                   {OMD}});
     }
 
-    void visitProtocolList(
-        llvm::iterator_range<ObjCProtocolList::iterator> Protocols,
-        llvm::iterator_range<const SourceLocation *> Locations) {
-      for (const auto &P : llvm::zip(Protocols, Locations)) {
-        Refs.push_back(ReferenceLoc{NestedNameSpecifierLoc(),
-                                    std::get<1>(P),
-                                    /*IsDecl=*/false,
-                                    {std::get<0>(P)}});
-      }
-    }
-
-    void VisitObjCInterfaceDecl(const ObjCInterfaceDecl *OID) {
-      if (OID->isThisDeclarationADefinition())
-        visitProtocolList(OID->protocols(), OID->protocol_locs());
-      Base::VisitObjCInterfaceDecl(OID); // Visit the interface's name.
-    }
-
     void VisitObjCCategoryDecl(const ObjCCategoryDecl *OCD) {
-      visitProtocolList(OCD->protocols(), OCD->protocol_locs());
       // getLocation is the extended class's location, not the category's.
       Refs.push_back(ReferenceLoc{NestedNameSpecifierLoc(),
                                   OCD->getLocation(),
@@ -709,12 +684,6 @@ llvm::SmallVector<ReferenceLoc> refInDecl(const Decl *D,
                                   /*IsDecl=*/true,
                                   {OCID->getCategoryDecl()}});
     }
-
-    void VisitObjCProtocolDecl(const ObjCProtocolDecl *OPD) {
-      if (OPD->isThisDeclarationADefinition())
-        visitProtocolList(OPD->protocols(), OPD->protocol_locs());
-      Base::VisitObjCProtocolDecl(OPD); // Visit the protocol's name.
-    }
   };
 
   Visitor V{Resolver};
@@ -944,16 +913,6 @@ refInTypeLoc(TypeLoc L, const HeuristicResolver *Resolver) {
                                   /*IsDecl=*/false,
                                   {L.getIFaceDecl()}});
     }
-
-    void VisitObjCObjectTypeLoc(ObjCObjectTypeLoc L) {
-      unsigned NumProtocols = L.getNumProtocols();
-      for (unsigned I = 0; I < NumProtocols; I++) {
-        Refs.push_back(ReferenceLoc{NestedNameSpecifierLoc(),
-                                    L.getProtocolLoc(I),
-                                    /*IsDecl=*/false,
-                                    {L.getProtocol(I)}});
-      }
-    }
   };
 
   Visitor V{Resolver};
@@ -1049,6 +1008,11 @@ public:
     return RecursiveASTVisitor::TraverseNestedNameSpecifierLoc(L);
   }
 
+  bool TraverseObjCProtocolLoc(ObjCProtocolLoc ProtocolLoc) {
+    visitNode(DynTypedNode::create(ProtocolLoc));
+    return true;
+  }
+
   bool TraverseConstructorInitializer(CXXCtorInitializer *Init) {
     visitNode(DynTypedNode::create(*Init));
     return RecursiveASTVisitor::TraverseConstructorInitializer(Init);
@@ -1094,6 +1058,12 @@ private:
                              {CCI->getAnyMember()}}};
       }
     }
+    if (const ObjCProtocolLoc *PL = N.get<ObjCProtocolLoc>())
+      return {ReferenceLoc{NestedNameSpecifierLoc(),
+                           PL->getLocation(),
+                           /*IsDecl=*/false,
+                           {PL->getProtocol()}}};
+
     // We do not have location information for other nodes (QualType, etc)
     return {};
   }
diff --git a/clang-tools-extra/clangd/Selection.cpp b/clang-tools-extra/clangd/Selection.cpp
index dbfe05c8e8b5..ba2f253eb075 100644
--- a/clang-tools-extra/clangd/Selection.cpp
+++ b/clang-tools-extra/clangd/Selection.cpp
@@ -684,6 +684,9 @@ public:
     return traverseNode<TypeLoc>(
         &QX, [&] { return TraverseTypeLoc(QX.getUnqualifiedLoc()); });
   }
+  bool TraverseObjCProtocolLoc(ObjCProtocolLoc PL) {
+    return traverseNode(&PL, [&] { return Base::TraverseObjCProtocolLoc(PL); });
+  }
   // Uninteresting parts of the AST that don't have locations within them.
   bool TraverseNestedNameSpecifier(NestedNameSpecifier *) { return true; }
   bool TraverseType(QualType) { return true; }
diff --git a/clang-tools-extra/clangd/unittests/FindTargetTests.cpp b/clang-tools-extra/clangd/unittests/FindTargetTests.cpp
index 4887ee5b5deb..7026f7fced3c 100644
--- a/clang-tools-extra/clangd/unittests/FindTargetTests.cpp
+++ b/clang-tools-extra/clangd/unittests/FindTargetTests.cpp
@@ -946,11 +946,9 @@ TEST_F(TargetDeclTest, ObjC) {
   EXPECT_DECLS("ObjCCategoryImplDecl", "@interface Foo(Ext)");
 
   Code = R"cpp(
-    @protocol Foo
-    @end
-    void test([[id<Foo>]] p);
+    void test(id</*error-ok*/[[InvalidProtocol]]> p);
   )cpp";
-  EXPECT_DECLS("ObjCObjectTypeLoc", "@protocol Foo");
+  EXPECT_DECLS("ParmVarDecl", "id p");
 
   Code = R"cpp(
     @class C;
@@ -966,7 +964,7 @@ TEST_F(TargetDeclTest, ObjC) {
     @end
     void test(C<[[Foo]]> *p);
   )cpp";
-  EXPECT_DECLS("ObjCObjectTypeLoc", "@protocol Foo");
+  EXPECT_DECLS("ObjCProtocolLoc", "@protocol Foo");
 
   Code = R"cpp(
     @class C;
@@ -976,8 +974,17 @@ TEST_F(TargetDeclTest, ObjC) {
     @end
     void test(C<[[Foo]], Bar> *p);
   )cpp";
-  // FIXME: We currently can't disambiguate between multiple protocols.
-  EXPECT_DECLS("ObjCObjectTypeLoc", "@protocol Foo", "@protocol Bar");
+  EXPECT_DECLS("ObjCProtocolLoc", "@protocol Foo");
+
+  Code = R"cpp(
+    @class C;
+    @protocol Foo
+    @end
+    @protocol Bar
+    @end
+    void test(C<Foo, [[Bar]]> *p);
+  )cpp";
+  EXPECT_DECLS("ObjCProtocolLoc", "@protocol Bar");
 
   Code = R"cpp(
     @interface Foo
diff --git a/clang-tools-extra/clangd/unittests/HoverTests.cpp b/clang-tools-extra/clangd/unittests/HoverTests.cpp
index 0b1203ae8179..0f07aa9d3b42 100644
--- a/clang-tools-extra/clangd/unittests/HoverTests.cpp
+++ b/clang-tools-extra/clangd/unittests/HoverTests.cpp
@@ -2522,6 +2522,22 @@ TEST(Hover, All) {
             HI.Definition = "@property(nonatomic, assign, unsafe_unretained, "
                             "readwrite) int prop1;";
           }},
+      {
+          R"cpp(
+          @protocol MYProtocol
+          @end
+          @interface MYObject
+          @end
+
+          @interface MYObject (Ext) <[[MYProt^ocol]]>
+          @end
+          )cpp",
+          [](HoverInfo &HI) {
+            HI.Name = "MYProtocol";
+            HI.Kind = index::SymbolKind::Protocol;
+            HI.NamespaceScope = "";
+            HI.Definition = "@protocol MYProtocol\n@end";
+          }},
       {R"objc(
         @interface Foo
         @end