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

[dgoldman]

Summary

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.

[dgoldman]

Minor lint fixes

[sammccall]

LG, thanks!
You might want to check out how hover behaves and add a test there.

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

948–949

This new test case looks interesting, but can we keep the old one too?

[dgoldman]

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

948–949

It now fails, I thought this was expected since the selection is now split up into multiple nodes - should it continue to work as it did previously?

llvm-project/clang-tools-extra/clangd/unittests/FindTargetTests.cpp:953: Failure
Value of: assertNodeAndPrintDecls("ObjCObjectTypeLoc")
Expected: has 1 element and that element is equal to @protocol Foo Rel=
  Actual: {}

    @protocol Foo
    @end
    void test([[id<Foo>]] p);

[dgoldman]

add hover test

[sammccall]

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

948–949

Oops, I wasn't reading carefully enough, you're right!

[dgoldman]

Rebase on top of latest ProtocolLoc changes

clang-tools-extra/clangd/FindTarget.cpp

struct Visitor : public TypeVisitor<Visitor> {
         }
       }
       void VisitTemplateTypeParmType(const TemplateTypeParmType *TTPT) {
         Outer.add(TTPT->getDecl(), Flags);
       }
       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());
   }
 
   void add(const NestedNameSpecifier *NNS, RelSet Flags) {
     if (!NNS)
       return;
     debug(*NNS, Flags);
allTargetDecls(const DynTypedNode &N, const HeuristicResolver *Resolver) {
   else if (const QualType *QT = N.get<QualType>())
     Finder.add(*QT, Flags);
   else if (const CXXCtorInitializer *CCI = N.get<CXXCtorInitializer>())
     Finder.add(CCI, Flags);
   else if (const TemplateArgumentLoc *TAL = N.get<TemplateArgumentLoc>())
     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();
 }
 
 llvm::SmallVector<const NamedDecl *, 1>
 targetDecl(const DynTypedNode &N, DeclRelationSet Mask,
            const HeuristicResolver *Resolver) {
   llvm::SmallVector<const NamedDecl *, 1> Result;
   for (const auto &Entry : allTargetDecls(N, Resolver)) {
struct Visitor : ConstDeclVisitor<Visitor> {
     void VisitObjCMethodDecl(const ObjCMethodDecl *OMD) {
       // The name may have several tokens, we can only report the first.
       Refs.push_back(ReferenceLoc{NestedNameSpecifierLoc(),
                                   OMD->getSelectorStartLoc(),
                                   /*IsDecl=*/true,
                                   {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(),
                                   /*IsDecl=*/false,
                                   {OCD->getClassInterface()}});
       Refs.push_back(ReferenceLoc{NestedNameSpecifierLoc(),
                                   OCD->getCategoryNameLoc(),
                                   /*IsDecl=*/true,
                                   {OCD}});
     }
 
     void VisitObjCCategoryImplDecl(const ObjCCategoryImplDecl *OCID) {
       Refs.push_back(ReferenceLoc{NestedNameSpecifierLoc(),
                                   OCID->getLocation(),
                                   /*IsDecl=*/false,
                                   {OCID->getClassInterface()}});
       Refs.push_back(ReferenceLoc{NestedNameSpecifierLoc(),
                                   OCID->getCategoryNameLoc(),
                                   /*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};
   V.Visit(D);
   return V.Refs;
 }
 
 llvm::SmallVector<ReferenceLoc> refInStmt(const Stmt *S,
struct Visitor : TypeLocVisitor<Visitor> {
     }
 
     void VisitObjCInterfaceTypeLoc(ObjCInterfaceTypeLoc L) {
       Refs.push_back(ReferenceLoc{NestedNameSpecifierLoc(),
                                   L.getNameLoc(),
                                   /*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};
   V.Visit(L.getUnqualifiedLoc());
   return V.Refs;
 }
 
 class ExplicitReferenceCollector
if (!L.getNestedNameSpecifier())
       return true;
     visitNode(DynTypedNode::create(L));
     // Inner type is missing information about its qualifier, skip it.
     if (auto TL = L.getTypeLoc())
       TypeLocsToSkip.insert(TL.getBeginLoc());
     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);
   }
 
 private:
   /// Obtain information about a reference directly defined in \p N. Does not
   /// recurse into child nodes, e.g. do not expect references for constructor
if (const CXXCtorInitializer *CCI = N.get<CXXCtorInitializer>()) {
       // the typeLoc.
       if (CCI->isAnyMemberInitializer()) {
         return {ReferenceLoc{NestedNameSpecifierLoc(),
                              CCI->getMemberLocation(),
                              /*IsDecl=*/false,
                              {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 {};
   }
 
   void visitNode(DynTypedNode N) {
     for (auto &R : explicitReference(N))
       reportReference(std::move(R), N);
   }

clang-tools-extra/clangd/Selection.cpp

public:
   // TraverseTypeLoc is not called for the inner UnqualTypeLoc.
   // This means we'd never see 'int' in 'const int'! Work around that here.
   // (The reason for the behavior is to avoid traversing the nested Type twice,
   // but we ignore TraverseType anyway).
   bool TraverseQualifiedTypeLoc(QualifiedTypeLoc QX) {
     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; }
 
   // The DeclStmt for the loop variable claims to cover the whole range
   // inside the parens, this causes the range-init expression to not be hit.
   // Traverse the loop VarDecl instead, which has the right source range.
   bool TraverseCXXForRangeStmt(CXXForRangeStmt *S) {

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

Code = R"cpp(
     @end
     @interface Foo (Ext)
     @end
     @implementation [[Foo]] (Ext)
     @end
   )cpp";
   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;
     @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";
-  EXPECT_DECLS("ObjCObjectTypeLoc", "@protocol Foo");
+  EXPECT_DECLS("ObjCProtocolLoc", "@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");
+  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
     + (id)sharedInstance;
     @end
     @implementation Foo
     + (id)sharedInstance { return 0; }
     @end

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

           [](HoverInfo &HI) {
             HI.Name = "prop1";
             HI.Kind = index::SymbolKind::InstanceProperty;
             HI.LocalScope = "MYProtocol::";
             HI.NamespaceScope = "";
             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
 
         @implementation Foo(Private)
         + (int)somePrivateMethod {
           int [[res^ult]] = 2;
           return result;