[clangd] Improve handling of Objective-C protocols in types

Improve support for Objective-C protocols for types/type locs

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

clang-tools-extra/clangd/FindTarget.cpp

@@ -432,10 +432,13 @@ public:
         Outer.add(OIT->getDecl(), Flags);
       }
       void VisitObjCObjectType(const ObjCObjectType *OOT) {
-        // FIXME: ObjCObjectTypeLoc has no children for the protocol list, so
-        // there is no node in id<Foo> that refers to ObjCProtocolDecl Foo.
-        if (OOT->isObjCQualifiedId() && OOT->getNumProtocols() == 1)
-          Outer.add(OOT->getProtocol(0), Flags);
+        // 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());
@@ -813,30 +816,34 @@ refInTypeLoc(TypeLoc L, const HeuristicResolver *Resolver) {
     Visitor(const HeuristicResolver *Resolver) : Resolver(Resolver) {}
 
     const HeuristicResolver *Resolver;
-    llvm::Optional<ReferenceLoc> Ref;
+    llvm::SmallVector<ReferenceLoc> Refs;
 
     void VisitElaboratedTypeLoc(ElaboratedTypeLoc L) {
       // We only know about qualifier, rest if filled by inner locations.
+      size_t InitialSize = Refs.size();
       Visit(L.getNamedTypeLoc().getUnqualifiedLoc());
-      // Fill in the qualifier.
-      if (!Ref)
-        return;
-      assert(!Ref->Qualifier.hasQualifier() && "qualifier already set");
-      Ref->Qualifier = L.getQualifierLoc();
+      size_t NewSize = Refs.size();
+      // Add qualifier for the newly-added refs.
+      for (unsigned I = InitialSize; I < NewSize; ++I) {
+        ReferenceLoc *Ref = &Refs[I];
+        // Fill in the qualifier.
+        assert(!Ref->Qualifier.hasQualifier() && "qualifier already set");
+        Ref->Qualifier = L.getQualifierLoc();
+      }
     }
 
     void VisitTagTypeLoc(TagTypeLoc L) {
-      Ref = ReferenceLoc{NestedNameSpecifierLoc(),
-                         L.getNameLoc(),
-                         /*IsDecl=*/false,
-                         {L.getDecl()}};
+      Refs.push_back(ReferenceLoc{NestedNameSpecifierLoc(),
+                                  L.getNameLoc(),
+                                  /*IsDecl=*/false,
+                                  {L.getDecl()}});
     }
 
     void VisitTemplateTypeParmTypeLoc(TemplateTypeParmTypeLoc L) {
-      Ref = ReferenceLoc{NestedNameSpecifierLoc(),
-                         L.getNameLoc(),
-                         /*IsDecl=*/false,
-                         {L.getDecl()}};
+      Refs.push_back(ReferenceLoc{NestedNameSpecifierLoc(),
+                                  L.getNameLoc(),
+                                  /*IsDecl=*/false,
+                                  {L.getDecl()}});
     }
 
     void VisitTemplateSpecializationTypeLoc(TemplateSpecializationTypeLoc L) {
@@ -848,63 +855,71 @@ refInTypeLoc(TypeLoc L, const HeuristicResolver *Resolver) {
       //    1. valias with mask 'Alias'.
       //    2. 'vector<int>' with mask 'Underlying'.
       //  we want to return only #1 in this case.
-      Ref = ReferenceLoc{
+      Refs.push_back(ReferenceLoc{
           NestedNameSpecifierLoc(), L.getTemplateNameLoc(), /*IsDecl=*/false,
           explicitReferenceTargets(DynTypedNode::create(L.getType()),
-                                   DeclRelation::Alias, Resolver)};
+                                   DeclRelation::Alias, Resolver)});
     }
     void VisitDeducedTemplateSpecializationTypeLoc(
         DeducedTemplateSpecializationTypeLoc L) {
-      Ref = ReferenceLoc{
+      Refs.push_back(ReferenceLoc{
           NestedNameSpecifierLoc(), L.getNameLoc(), /*IsDecl=*/false,
           explicitReferenceTargets(DynTypedNode::create(L.getType()),
-                                   DeclRelation::Alias, Resolver)};
+                                   DeclRelation::Alias, Resolver)});
     }
 
     void VisitInjectedClassNameTypeLoc(InjectedClassNameTypeLoc TL) {
-      Ref = ReferenceLoc{NestedNameSpecifierLoc(),
-                         TL.getNameLoc(),
-                         /*IsDecl=*/false,
-                         {TL.getDecl()}};
+      Refs.push_back(ReferenceLoc{NestedNameSpecifierLoc(),
+                                  TL.getNameLoc(),
+                                  /*IsDecl=*/false,
+                                  {TL.getDecl()}});
     }
 
     void VisitDependentTemplateSpecializationTypeLoc(
         DependentTemplateSpecializationTypeLoc L) {
-      Ref = ReferenceLoc{L.getQualifierLoc(), L.getTemplateNameLoc(),
-                         /*IsDecl=*/false,
-                         explicitReferenceTargets(
-                             DynTypedNode::create(L.getType()), {}, Resolver)};
+      Refs.push_back(
+          ReferenceLoc{L.getQualifierLoc(), L.getTemplateNameLoc(),
+                       /*IsDecl=*/false,
+                       explicitReferenceTargets(
+                           DynTypedNode::create(L.getType()), {}, Resolver)});
     }
 
     void VisitDependentNameTypeLoc(DependentNameTypeLoc L) {
-      Ref = ReferenceLoc{L.getQualifierLoc(), L.getNameLoc(), /*IsDecl=*/false,
-                         explicitReferenceTargets(
-                             DynTypedNode::create(L.getType()), {}, Resolver)};
+      Refs.push_back(
+          ReferenceLoc{L.getQualifierLoc(), L.getNameLoc(),
+                       /*IsDecl=*/false,
+                       explicitReferenceTargets(
+                           DynTypedNode::create(L.getType()), {}, Resolver)});
     }
 
     void VisitTypedefTypeLoc(TypedefTypeLoc L) {
-      Ref = ReferenceLoc{NestedNameSpecifierLoc(),
-                         L.getNameLoc(),
-                         /*IsDecl=*/false,
-                         {L.getTypedefNameDecl()}};
+      Refs.push_back(ReferenceLoc{NestedNameSpecifierLoc(),
+                                  L.getNameLoc(),
+                                  /*IsDecl=*/false,
+                                  {L.getTypedefNameDecl()}});
     }
 
     void VisitObjCInterfaceTypeLoc(ObjCInterfaceTypeLoc L) {
-      Ref = ReferenceLoc{NestedNameSpecifierLoc(),
-                         L.getNameLoc(),
-                         /*IsDecl=*/false,
-                         {L.getIFaceDecl()}};
+      Refs.push_back(ReferenceLoc{NestedNameSpecifierLoc(),
+                                  L.getNameLoc(),
+                                  /*IsDecl=*/false,
+                                  {L.getIFaceDecl()}});
     }
 
-    // FIXME: add references to protocols in ObjCObjectTypeLoc and maybe
-    // ObjCObjectPointerTypeLoc.
+    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};
   V.Visit(L.getUnqualifiedLoc());
-  if (!V.Ref)
-    return {};
-  return {*V.Ref};
+  return V.Refs;
 }
 
 class ExplicitReferenceCollector

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

@@ -943,14 +943,32 @@ TEST_F(TargetDeclTest, ObjC) {
   )cpp";
   EXPECT_DECLS("ObjCObjectTypeLoc", "@protocol Foo");
 
+  Code = R"cpp(
+    @class C;
+    @protocol Foo
+    @end
+    void test([[C]]<Foo> *p);
+  )cpp";
+  EXPECT_DECLS("ObjCInterfaceTypeLoc", "@class C;");
+
   Code = R"cpp(
     @class C;
     @protocol Foo
     @end
     void test(C<[[Foo]]> *p);
   )cpp";
-  // FIXME: there's no AST node corresponding to 'Foo', so we're stuck.
-  EXPECT_DECLS("ObjCObjectTypeLoc");
+  EXPECT_DECLS("ObjCObjectTypeLoc", "@protocol Foo");
+
+  Code = R"cpp(
+    @class C;
+    @protocol Foo
+    @end
+    @protocol Bar
+    @end
+    void test(C<[[Foo]], Bar> *p);
+  )cpp";
+  // FIXME: We currently can't disambiguate between multiple protocols.
+  EXPECT_DECLS("ObjCObjectTypeLoc", "@protocol Foo", "@protocol Bar");
 }
 
 class FindExplicitReferencesTest : public ::testing::Test {
@@ -1610,12 +1628,12 @@ TEST_F(FindExplicitReferencesTest, All) {
             @protocol P
             @end
             void foo() {
-              // FIXME: should reference P
-              $0^I<P> *$1^x;
+              $0^I<$1^P> *$2^x;
             }
           )cpp",
         "0: targets = {I}\n"
-        "1: targets = {x}, decl\n"},
+        "1: targets = {P}\n"
+        "2: targets = {x}, decl\n"},
 
        // Designated initializers.
        {R"cpp(

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

@@ -676,8 +676,7 @@ sizeof...($TemplateParameter[[Elements]]);
         @end
         @interface $Class_decl[[Klass]] <$Interface[[Protocol]]>
         @end
-        // FIXME: protocol list in ObjCObjectType should be highlighted.
-        id<Protocol> $Variable_decl[[x]];
+        id<$Interface[[Protocol]]> $Variable_decl[[x]];
       )cpp",
       R"cpp(
         // ObjC: Categories