[clang] Add `ObjCProtocolLoc` to represent protocol references

[dgoldman]

Summary

Add ObjCProtocolLoc which behaves like TypeLoc but for
ObjCProtocolDecl references.

RecursiveASTVisitor now synthesizes ObjCProtocolLoc during traversal
and the ObjCProtocolLoc can be stored in a DynTypedNode.

In a follow up patch, I'll update clangd to make use of this
to properly support protocol references for hover + goto definition.

[dgoldman]

Minor lint fixes

[sammccall]

Thanks @akyrtzi @egorzhdan I was going to suggest that I think someone at Apple should probably look at changes here first.

The goal here is to ensure that any tokens in the code that reference a decl have some AST node (which can be stored in DynTypedNode).
This is provides a nice simple model which clangd's go-to-definition etc rely on: token -> ast node -> target.
Clang almost follows this rule but references to protocols have no dedicated AST node.

clang/include/clang/AST/RecursiveASTVisitor.h

331

RAV changes should have tests

clang/include/clang/AST/TypeLoc.h

2612

These fields shouldn't be public in AST nodes, provide accessors instead (I think they can be read-only).

2621

will we ever have invalid instances?

[egorzhdan]

Thank you @dgoldman, this approach looks good to me.
I don't have anything to add other than what @sammccall has already commented.

[dgoldman]

Add tests and public accessors

clang/include/clang/AST/TypeLoc.h

2621

From what I can tell, nope. This is explicitly used here though: https://github.com/llvm/llvm-project/blob/ceb5dc55c2e395c085cd9d16c4152c5a1923d7e2/clang/lib/AST/ParentMapContext.cpp#L405

[dgoldman]

PTAL, thanks!

[sammccall]

I'm a bit concerned about the parent map vs ASTMatchFinder's view being out of sync: we'll have a ProtocolLoc node that has a parent but the parent doesn't have the child.

I'm not sure this breaks anything immediately, but I think we should also make these nodes visible to matchers, even if there's no specific matcher yet.

clang/include/clang/AST/RecursiveASTVisitor.h

1350

i => I etc

clang/include/clang/AST/TypeLoc.h

2611

Also add a public getLocation() accessor for ProtocolLoc? (That returns SourceLocation)

2617

Return non-const pointer, matching other AST nodes. (But method should be const)

2620

This comment doesn't really say anything.

Maybe replace with "the source range is just the protocol name"?

2621

From what I can tell, nope. This is explicitly used here though: https://github.com/llvm/llvm-project/blob/ceb5dc55c2e395c085cd9d16c4152c5a1923d7e2/clang/lib/AST/ParentMapContext.cpp#L405

Well it's not used today, you're adding the use in this patch.

It can be solved in some other way, like pulling out an isNull template and then overloading it for ProtocolLoc.
I don't think we should add API surface to AST nodes to satisfy brittle templates

[dgoldman]

Template and variable casing fixes

[dgoldman]

In D119363#3323778, @sammccall wrote:

I'm a bit concerned about the parent map vs ASTMatchFinder's view being out of sync: we'll have a ProtocolLoc node that has a parent but the parent doesn't have the child.

I'm not sure this breaks anything immediately, but I think we should also make these nodes visible to matchers, even if there's no specific matcher yet.

Hmm I can try to do it - what do I need to modify to make them visible to matchers?

clang/include/clang/AST/TypeLoc.h

2611

this was already there? did you want something else?

2617

Why - Is there a reason to make it mutable? TypeLoc's getTypePtr is const (although this is a decl not a type).

2621

Gotcha, I'm relatively new to templates but took a stab at isNull - LMK what you think/if there's a better way.

[sammccall]

In D119363#3323867, @dgoldman wrote:

In D119363#3323778, @sammccall wrote:

I'm a bit concerned about the parent map vs ASTMatchFinder's view being out of sync: we'll have a ProtocolLoc node that has a parent but the parent doesn't have the child.

I'm not sure this breaks anything immediately, but I think we should also make these nodes visible to matchers, even if there's no specific matcher yet.

Hmm I can try to do it - what do I need to modify to make them visible to matchers?

I don't remember specifically, I think ASTMatchFinderImpl has a RecursiveASTVisitor that you need to extend? I'm not sure actually how you would observe these nodes cleanly without matchers (hacks like has(hasParent()) maybe?) So maybe it's best to ignore it in this patch and add basic matchers in a next one

clang/include/clang/AST/TypeLoc.h

2611

Nope sorry I'm just going blind it seems

2617

A bunch of stuff operates on non-const nodes only, RecursiveASTVisitor for example.
And for whatever reason this is the convention, rather than a const+non-const variant of everything.

TypeLoc's getTypePtr is const (although this is a decl not a type).

Yeah types are always const, this is related to the fact that they're interned so mutating them wouldn't make sense.

clang/lib/AST/ParentMapContext.cpp

404

t -> Node (case)

404

These should be static I think so they're not exposed

405

No need for template<>, this can just be a normal function that forms an overload set along with the template

[dgoldman]

Template fixes

[dgoldman]

In D119363#3324024, @sammccall wrote:

In D119363#3323867, @dgoldman wrote:

In D119363#3323778, @sammccall wrote:

I'm a bit concerned about the parent map vs ASTMatchFinder's view being out of sync: we'll have a ProtocolLoc node that has a parent but the parent doesn't have the child.

I'm not sure this breaks anything immediately, but I think we should also make these nodes visible to matchers, even if there's no specific matcher yet.

Hmm I can try to do it - what do I need to modify to make them visible to matchers?

I don't remember specifically, I think ASTMatchFinderImpl has a RecursiveASTVisitor that you need to extend? I'm not sure actually how you would observe these nodes cleanly without matchers (hacks like has(hasParent()) maybe?) So maybe it's best to ignore it in this patch and add basic matchers in a next one

It looks like ASTMatchFinder.cpp has a MatchASTVisitor which I think is what you meant, but yeah I think it's probably best to do in a follow up unless you think it'll break something here.

[sammccall]

Great!
This is worth mentioning in the clang release notes. (Along with any matchers that match these new nodes)

clang/include/clang/AST/ASTFwd.h

 class CXXCtorInitializer;
 class OMPClause;
 #define GEN_CLANG_CLAUSE_CLASS
 #define CLAUSE_CLASS(Enum, Str, Class) class Class;
 #include "llvm/Frontend/OpenMP/OMP.inc"
 class Attr;
 #define ATTR(A) class A##Attr;
 #include "clang/Basic/AttrList.inc"
+class ObjCProtocolLoc;
 
 } // end namespace clang
 
 #endif

clang/include/clang/AST/ASTTypeTraits.h

 #include "clang/AST/TypeNodes.inc"
     NKI_OMPClause,
 #define GEN_CLANG_CLAUSE_CLASS
 #define CLAUSE_CLASS(Enum, Str, Class) NKI_##Class,
 #include "llvm/Frontend/OpenMP/OMP.inc"
     NKI_Attr,
 #define ATTR(A) NKI_##A##Attr,
 #include "clang/Basic/AttrList.inc"
+    NKI_ObjCProtocolLoc,
     NKI_NumberOfKinds
   };
 
   /// Use getFromNodeKind<T>() to construct the kind.
   ASTNodeKind(NodeKindId KindId) : KindId(KindId) {}
 
   /// Returns \c true if \c Base is a base kind of (or same as) \c
   ///   Derived.

 #define TYPELOC(CLASS, PARENT) KIND_TO_KIND_ID(CLASS##TypeLoc)
 #include "clang/AST/TypeLocNodes.def"
 KIND_TO_KIND_ID(TypeLoc)
 KIND_TO_KIND_ID(Decl)
 KIND_TO_KIND_ID(Stmt)
 KIND_TO_KIND_ID(Type)
 KIND_TO_KIND_ID(OMPClause)
 KIND_TO_KIND_ID(Attr)
+KIND_TO_KIND_ID(ObjCProtocolLoc)
 KIND_TO_KIND_ID(CXXBaseSpecifier)
 #define DECL(DERIVED, BASE) KIND_TO_KIND_ID(DERIVED##Decl)
 #include "clang/AST/DeclNodes.inc"
 #define STMT(DERIVED, BASE) KIND_TO_KIND_ID(DERIVED)
 #include "clang/AST/StmtNodes.inc"
 #define TYPE(DERIVED, BASE) KIND_TO_KIND_ID(DERIVED##Type)
 #include "clang/AST/TypeNodes.inc"
 #define GEN_CLANG_CLAUSE_CLASS
private:
   /// Note that we can store \c Decls, \c Stmts, \c Types,
   /// \c NestedNameSpecifiers and \c CXXCtorInitializer by pointer as they are
   /// guaranteed to be unique pointers pointing to dedicated storage in the AST.
   /// \c QualTypes, \c NestedNameSpecifierLocs, \c TypeLocs,
   /// \c TemplateArguments and \c TemplateArgumentLocs on the other hand do not
   /// have storage or unique pointers and thus need to be stored by value.
   llvm::AlignedCharArrayUnion<const void *, TemplateArgument,
                               TemplateArgumentLoc, NestedNameSpecifierLoc,
-                              QualType, TypeLoc>
+                              QualType, TypeLoc, ObjCProtocolLoc>
       Storage;
 };
 
 template <typename T>
 struct DynTypedNode::BaseConverter<
     T, std::enable_if_t<std::is_base_of<Decl, T>::value>>
     : public DynCastPtrConverter<T, Decl> {};
 

 struct DynTypedNode::BaseConverter<
     T, std::enable_if_t<std::is_base_of<TypeLoc, T>::value>>
     : public DynCastValueConverter<T, TypeLoc> {};
 
 template <>
 struct DynTypedNode::BaseConverter<CXXBaseSpecifier, void>
     : public PtrConverter<CXXBaseSpecifier> {};
 
+template <>
+struct DynTypedNode::BaseConverter<ObjCProtocolLoc, void>
+    : public ValueConverter<ObjCProtocolLoc> {};
+
 // The only operation we allow on unsupported types is \c get.
 // This allows to conveniently use \c DynTypedNode when having an arbitrary
 // AST node that is not supported, but prevents misuse - a user cannot create
 // a DynTypedNode from arbitrary types.
 template <typename T, typename EnablerT> struct DynTypedNode::BaseConverter {
   static const T *get(ASTNodeKind NodeKind, const char Storage[]) {
     return NULL;
   }

clang/include/clang/AST/RecursiveASTVisitor.h

public:
   bool TraverseSynOrSemInitListExpr(InitListExpr *S,
                                     DataRecursionQueue *Queue = nullptr);
 
   /// Recursively visit a reference to a concept with potential arguments.
   ///
   /// \returns false if the visitation was terminated early, true otherwise.
   bool TraverseConceptReference(const ConceptReference &C);
 
+  /// Recursively visit an Objective-C protocol reference with location
+  /// information.
+  ///
+  /// \returns false if the visitation was terminated early, true otherwise.
+  bool TraverseObjCProtocolLoc(ObjCProtocolLoc ProtocolLoc);
+
   // ---- Methods on Attrs ----
 
   // Visit an attribute.
   bool VisitAttr(Attr *A) { return true; }
 
 // Declare Traverse* and empty Visit* for all Attr classes.
 #define ATTR_VISITOR_DECLS_ONLY
 #include "clang/AST/AttrVisitor.inc"
DEF_TRAVERSE_TYPELOC(DependentTemplateSpecializationType, {
   for (unsigned I = 0, E = TL.getNumArgs(); I != E; ++I) {
     TRY_TO(TraverseTemplateArgumentLoc(TL.getArgLoc(I)));
   }
 })
 
 DEF_TRAVERSE_TYPELOC(PackExpansionType,
                      { TRY_TO(TraverseTypeLoc(TL.getPatternLoc())); })
 
-DEF_TRAVERSE_TYPELOC(ObjCTypeParamType, {})
+DEF_TRAVERSE_TYPELOC(ObjCTypeParamType, {
+  for (unsigned I = 0, N = TL.getNumProtocols(); I != N; ++I) {
+    ObjCProtocolLoc ProtocolLoc(TL.getProtocol(I), TL.getProtocolLoc(I));
+    TRY_TO(TraverseObjCProtocolLoc(ProtocolLoc));
+  }
+})
 
 DEF_TRAVERSE_TYPELOC(ObjCInterfaceType, {})
 
 DEF_TRAVERSE_TYPELOC(ObjCObjectType, {
   // We have to watch out here because an ObjCInterfaceType's base
   // type is itself.
   if (TL.getTypePtr()->getBaseType().getTypePtr() != TL.getTypePtr())
     TRY_TO(TraverseTypeLoc(TL.getBaseLoc()));
   for (unsigned i = 0, n = TL.getNumTypeArgs(); i != n; ++i)
     TRY_TO(TraverseTypeLoc(TL.getTypeArgTInfo(i)->getTypeLoc()));
+  for (unsigned I = 0, N = TL.getNumProtocols(); I != N; ++I) {
+    ObjCProtocolLoc ProtocolLoc(TL.getProtocol(I), TL.getProtocolLoc(I));
+    TRY_TO(TraverseObjCProtocolLoc(ProtocolLoc));
+  }
 })
 
 DEF_TRAVERSE_TYPELOC(ObjCObjectPointerType,
                      { TRY_TO(TraverseTypeLoc(TL.getPointeeLoc())); })
 
 DEF_TRAVERSE_TYPELOC(AtomicType, { TRY_TO(TraverseTypeLoc(TL.getValueLoc())); })
 
 DEF_TRAVERSE_TYPELOC(PipeType, { TRY_TO(TraverseTypeLoc(TL.getValueLoc())); })
DEF_TRAVERSE_DECL(
     {// Code in an unnamed namespace shows up automatically in
      // decls_begin()/decls_end().  Thus we don't need to recurse on
      // D->getAnonymousNamespace().
     })
 
 DEF_TRAVERSE_DECL(ObjCCompatibleAliasDecl, {// FIXME: implement
                                            })
 
-DEF_TRAVERSE_DECL(ObjCCategoryDecl, {// FIXME: implement
+DEF_TRAVERSE_DECL(ObjCCategoryDecl, {
   if (ObjCTypeParamList *typeParamList = D->getTypeParamList()) {
     for (auto typeParam : *typeParamList) {
       TRY_TO(TraverseObjCTypeParamDecl(typeParam));
     }
   }
+  for (auto It : llvm::zip(D->protocols(), D->protocol_locs())) {
+    ObjCProtocolLoc ProtocolLoc(std::get<0>(It), std::get<1>(It));
+    TRY_TO(TraverseObjCProtocolLoc(ProtocolLoc));
+  }
 })
 
 DEF_TRAVERSE_DECL(ObjCCategoryImplDecl, {// FIXME: implement
                                         })
 
 DEF_TRAVERSE_DECL(ObjCImplementationDecl, {// FIXME: implement
                                           })
 
-DEF_TRAVERSE_DECL(ObjCInterfaceDecl, {// FIXME: implement
+DEF_TRAVERSE_DECL(ObjCInterfaceDecl, {
   if (ObjCTypeParamList *typeParamList = D->getTypeParamListAsWritten()) {
     for (auto typeParam : *typeParamList) {
       TRY_TO(TraverseObjCTypeParamDecl(typeParam));
     }
   }
 
   if (TypeSourceInfo *superTInfo = D->getSuperClassTInfo()) {
     TRY_TO(TraverseTypeLoc(superTInfo->getTypeLoc()));
   }
+  if (D->isThisDeclarationADefinition()) {
+    for (auto It : llvm::zip(D->protocols(), D->protocol_locs())) {
+      ObjCProtocolLoc ProtocolLoc(std::get<0>(It), std::get<1>(It));
+      TRY_TO(TraverseObjCProtocolLoc(ProtocolLoc));
+    }
+  }
 })
 
-DEF_TRAVERSE_DECL(ObjCProtocolDecl, {// FIXME: implement
-                                    })
+DEF_TRAVERSE_DECL(ObjCProtocolDecl, {
+  if (D->isThisDeclarationADefinition()) {
+    for (auto It : llvm::zip(D->protocols(), D->protocol_locs())) {
+      ObjCProtocolLoc ProtocolLoc(std::get<0>(It), std::get<1>(It));
+      TRY_TO(TraverseObjCProtocolLoc(ProtocolLoc));
+    }
+  }
 
 DEF_TRAVERSE_DECL(ObjCMethodDecl, {
   if (D->getReturnTypeSourceInfo()) {
     TRY_TO(TraverseTypeLoc(D->getReturnTypeSourceInfo()->getTypeLoc()));
   }
   for (ParmVarDecl *Parameter : D->parameters()) {
     TRY_TO(TraverseDecl(Parameter));
   }
bool RecursiveASTVisitor<Derived>::TraverseConceptReference(
   TRY_TO(TraverseDeclarationNameInfo(C.getConceptNameInfo()));
   if (C.hasExplicitTemplateArgs())
     TRY_TO(TraverseTemplateArgumentLocsHelper(
         C.getTemplateArgsAsWritten()->getTemplateArgs(),
         C.getTemplateArgsAsWritten()->NumTemplateArgs));
   return true;
 }
 
+template <typename Derived>
+bool RecursiveASTVisitor<Derived>::TraverseObjCProtocolLoc(
+    ObjCProtocolLoc ProtocolLoc) {
+  return true;
+}
+
 // If shouldVisitImplicitCode() returns false, this method traverses only the
 // syntactic form of InitListExpr.
 // If shouldVisitImplicitCode() return true, this method is called once for
 // each pair of syntactic and semantic InitListExpr, and it traverses the
 // subtrees defined by the two forms. This may cause some of the children to be
 // visited twice, if they appear both in the syntactic and the semantic form.
 //
 // There is no guarantee about which form \p S takes when this method is called.

clang/include/clang/AST/TypeLoc.h

 }
 class BitIntTypeLoc final
     : public InheritingConcreteTypeLoc<TypeSpecTypeLoc, BitIntTypeLoc,
                                        BitIntType> {};
 class DependentBitIntTypeLoc final
     : public InheritingConcreteTypeLoc<TypeSpecTypeLoc, DependentBitIntTypeLoc,
                                        DependentBitIntType> {};
 
+class ObjCProtocolLoc {
+  ObjCProtocolDecl *Protocol = nullptr;
+  SourceLocation Loc = SourceLocation();
+
+public:
+  ObjCProtocolLoc(ObjCProtocolDecl *protocol, SourceLocation loc)
+      : Protocol(protocol), Loc(loc) {}
+  ObjCProtocolDecl *getProtocol() const { return Protocol; }
+  SourceLocation getLocation() const { return Loc; }
+
+  /// The source range is just the protocol name.
+  SourceRange getSourceRange() const LLVM_READONLY {
+    return SourceRange(Loc, Loc);
+  }
+};
+
 } // namespace clang
 
 #endif // LLVM_CLANG_AST_TYPELOC_H

clang/lib/AST/ASTTypeTraits.cpp

 //  a type safe way.
 //
 //===----------------------------------------------------------------------===//
 
 #include "clang/AST/ASTTypeTraits.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/Attr.h"
 #include "clang/AST/DeclCXX.h"
+#include "clang/AST/DeclObjC.h"
 #include "clang/AST/NestedNameSpecifier.h"
 #include "clang/AST/OpenMPClause.h"
 #include "clang/AST/TypeLoc.h"
 
 using namespace clang;
 
 const ASTNodeKind::KindInfo ASTNodeKind::AllKindInfo[] = {
     {NKI_None, "<None>"},

 #include "clang/AST/TypeNodes.inc"
     {NKI_None, "OMPClause"},
 #define GEN_CLANG_CLAUSE_CLASS
 #define CLAUSE_CLASS(Enum, Str, Class) {NKI_OMPClause, #Class},
 #include "llvm/Frontend/OpenMP/OMP.inc"
     {NKI_None, "Attr"},
 #define ATTR(A) {NKI_Attr, #A "Attr"},
 #include "clang/Basic/AttrList.inc"
+    {NKI_None, "ObjCProtocolLoc"},
 };
 
 bool ASTNodeKind::isBaseOf(ASTNodeKind Other, unsigned *Distance) const {
   return isBaseOf(KindId, Other.KindId, Distance);
 }
 
 bool ASTNodeKind::isBaseOf(NodeKindId Base, NodeKindId Derived,
                            unsigned *Distance) {
void DynTypedNode::print(llvm::raw_ostream &OS,
   else if (const Decl *D = get<Decl>())
     D->print(OS, PP);
   else if (const Stmt *S = get<Stmt>())
     S->printPretty(OS, nullptr, PP);
   else if (const Type *T = get<Type>())
     QualType(T, 0).print(OS, PP);
   else if (const Attr *A = get<Attr>())
     A->printPretty(OS, PP);
+  else if (const ObjCProtocolLoc *P = get<ObjCProtocolLoc>())
+    P->getProtocol()->print(OS, PP);
   else
     OS << "Unable to print values of type " << NodeKind.asStringRef() << "\n";
 }
 
 void DynTypedNode::dump(llvm::raw_ostream &OS,
                         const ASTContext &Context) const {
   if (const Decl *D = get<Decl>())
     D->dump(OS);
SourceRange DynTypedNode::getSourceRange() const {
   if (const TemplateArgumentLoc *TAL = get<TemplateArgumentLoc>())
     return TAL->getSourceRange();
   if (const auto *C = get<OMPClause>())
     return SourceRange(C->getBeginLoc(), C->getEndLoc());
   if (const auto *CBS = get<CXXBaseSpecifier>())
     return CBS->getSourceRange();
   if (const auto *A = get<Attr>())
     return A->getRange();
+  if (const ObjCProtocolLoc *P = get<ObjCProtocolLoc>())
+    return P->getSourceRange();
   return SourceRange();
 }

clang/lib/AST/ParentMapContext.cpp

 }
 template <> DynTypedNode createDynTypedNode(const TypeLoc &Node) {
   return DynTypedNode::create(Node);
 }
 template <>
 DynTypedNode createDynTypedNode(const NestedNameSpecifierLoc &Node) {
   return DynTypedNode::create(Node);
 }
+template <> DynTypedNode createDynTypedNode(const ObjCProtocolLoc &Node) {
+  return DynTypedNode::create(Node);
+}
 /// @}
 
 /// A \c RecursiveASTVisitor that builds a map from nodes to their
 /// parents as defined by the \c RecursiveASTVisitor.
 ///
 /// Note that the relationship described here is purely in terms of AST
 /// traversal - there are other relationships (for example declaration context)
 /// in the AST that are better modeled by special matchers.
if (NodeOrVector.isNull()) {
       // types.
       bool Found = ParentStack.back().getMemoizationData() &&
                    llvm::is_contained(*Vector, ParentStack.back());
       if (!Found)
         Vector->push_back(ParentStack.back());
     }
   }
 
+  template <typename T> static bool isNull(T Node) { return !Node; }
+  static bool isNull(ObjCProtocolLoc Node) { return false; }
+
   template <typename T, typename MapNodeTy, typename BaseTraverseFn,
             typename MapTy>
   bool TraverseNode(T Node, MapNodeTy MapNode, BaseTraverseFn BaseTraverse,
                     MapTy *Parents) {
-    if (!Node)
+    if (isNull(Node))
       return true;
     addParent(MapNode, Parents);
     ParentStack.push_back(createDynTypedNode(Node));
     bool Result = BaseTraverse();
     ParentStack.pop_back();
     return Result;
   }
 
return TraverseNode(
         [&] { return VisitorBase::TraverseNestedNameSpecifierLoc(NNSLocNode); },
         &Map.OtherParents);
   }
   bool TraverseAttr(Attr *AttrNode) {
     return TraverseNode(
         AttrNode, AttrNode, [&] { return VisitorBase::TraverseAttr(AttrNode); },
         &Map.PointerParents);
   }
+  bool TraverseObjCProtocolLoc(ObjCProtocolLoc ProtocolLocNode) {
+    return TraverseNode(
+        ProtocolLocNode, DynTypedNode::create(ProtocolLocNode),
+        [&] { return VisitorBase::TraverseObjCProtocolLoc(ProtocolLocNode); },
+        &Map.OtherParents);
+  }
 
   // Using generic TraverseNode for Stmt would prevent data-recursion.
   bool dataTraverseStmtPre(Stmt *StmtNode) {
     addParent(StmtNode, &Map.PointerParents);
     ParentStack.push_back(DynTypedNode::create(*StmtNode));
     return true;
   }
   bool dataTraverseStmtPost(Stmt *StmtNode) {

clang/unittests/AST/RecursiveASTVisitorTest.cpp

enum class VisitEvent {
   StartTraverseFunction,
   EndTraverseFunction,
   StartTraverseAttr,
   EndTraverseAttr,
   StartTraverseEnum,
   EndTraverseEnum,
   StartTraverseTypedefType,
   EndTraverseTypedefType,
+  StartTraverseObjCInterface,
+  EndTraverseObjCInterface,
+  StartTraverseObjCProtocol,
+  EndTraverseObjCProtocol,
+  StartTraverseObjCProtocolLoc,
+  EndTraverseObjCProtocolLoc,
 };
 
 class CollectInterestingEvents
     : public RecursiveASTVisitor<CollectInterestingEvents> {
 public:
   bool TraverseFunctionDecl(FunctionDecl *D) {
     Events.push_back(VisitEvent::StartTraverseFunction);
     bool Ret = RecursiveASTVisitor::TraverseFunctionDecl(D);
public:
   bool TraverseTypedefTypeLoc(TypedefTypeLoc TL) {
     Events.push_back(VisitEvent::StartTraverseTypedefType);
     bool Ret = RecursiveASTVisitor::TraverseTypedefTypeLoc(TL);
     Events.push_back(VisitEvent::EndTraverseTypedefType);
 
     return Ret;
   }
 
+  bool TraverseObjCInterfaceDecl(ObjCInterfaceDecl *ID) {
+    Events.push_back(VisitEvent::StartTraverseObjCInterface);
+    bool Ret = RecursiveASTVisitor::TraverseObjCInterfaceDecl(ID);
+    Events.push_back(VisitEvent::EndTraverseObjCInterface);
+
+    return Ret;
+  }
+
+  bool TraverseObjCProtocolDecl(ObjCProtocolDecl *PD) {
+    Events.push_back(VisitEvent::StartTraverseObjCProtocol);
+    bool Ret = RecursiveASTVisitor::TraverseObjCProtocolDecl(PD);
+    Events.push_back(VisitEvent::EndTraverseObjCProtocol);
+
+    return Ret;
+  }
+
+  bool TraverseObjCProtocolLoc(ObjCProtocolLoc ProtocolLoc) {
+    Events.push_back(VisitEvent::StartTraverseObjCProtocolLoc);
+    bool Ret = RecursiveASTVisitor::TraverseObjCProtocolLoc(ProtocolLoc);
+    Events.push_back(VisitEvent::EndTraverseObjCProtocolLoc);
+
+    return Ret;
+  }
+
   std::vector<VisitEvent> takeEvents() && { return std::move(Events); }
 
 private:
   std::vector<VisitEvent> Events;
 };
 
-std::vector<VisitEvent> collectEvents(llvm::StringRef Code) {
+std::vector<VisitEvent> collectEvents(llvm::StringRef Code,
+                                      const Twine &FileName = "input.cc") {
   CollectInterestingEvents Visitor;
   clang::tooling::runToolOnCode(
       std::make_unique<ProcessASTAction>(
           [&](clang::ASTContext &Ctx) { Visitor.TraverseAST(Ctx); }),
-      Code);
+      Code, FileName);
   return std::move(Visitor).takeEvents();
 }
 } // namespace
 
 TEST(RecursiveASTVisitorTest, AttributesInsideDecls) {
   /// Check attributes are traversed inside TraverseFunctionDecl.
   llvm::StringRef Code = R"cpp(
 __attribute__((annotate("something"))) int foo() { return 10; }
TEST(RecursiveASTVisitorTest, EnumDeclWithBase) {
   )cpp";
 
   EXPECT_THAT(collectEvents(Code),
               ElementsAre(VisitEvent::StartTraverseEnum,
                           VisitEvent::StartTraverseTypedefType,
                           VisitEvent::EndTraverseTypedefType,
                           VisitEvent::EndTraverseEnum));
 }
+
+TEST(RecursiveASTVisitorTest, InterfaceDeclWithProtocols) {
+  // Check interface and its protocols are visited.
+  llvm::StringRef Code = R"cpp(
+  @protocol Foo
+  @end
+  @protocol Bar
+  @end
+
+  @interface SomeObject <Foo, Bar>
+  @end
+  )cpp";
+
+  EXPECT_THAT(collectEvents(Code, "input.m"),
+              ElementsAre(VisitEvent::StartTraverseObjCProtocol,
+                          VisitEvent::EndTraverseObjCProtocol,
+                          VisitEvent::StartTraverseObjCProtocol,
+                          VisitEvent::EndTraverseObjCProtocol,
+                          VisitEvent::StartTraverseObjCInterface,
+                          VisitEvent::StartTraverseObjCProtocolLoc,
+                          VisitEvent::EndTraverseObjCProtocolLoc,
+                          VisitEvent::StartTraverseObjCProtocolLoc,
+                          VisitEvent::EndTraverseObjCProtocolLoc,
+                          VisitEvent::EndTraverseObjCInterface));
+}