2 files changed, 20 insertions, 7 deletions
diff --git a/lib/AST/DeclObjC.cpp b/lib/AST/DeclObjC.cpp
index 2ec9395644..cd4c6fff9f 100644
--- a/lib/AST/DeclObjC.cpp
+++ b/lib/AST/DeclObjC.cpp
@@ -68,10 +68,9 @@ ObjCForwardProtocolDecl::Create(ASTContext &C, SourceLocation L,
 }
 
 ObjCCategoryDecl *ObjCCategoryDecl::Create(ASTContext &C, SourceLocation L,
-                                           unsigned numRefProtocol, 
                                            IdentifierInfo *Id) {
   void *Mem = C.getAllocator().Allocate<ObjCCategoryDecl>();
-  return new (Mem) ObjCCategoryDecl(L, numRefProtocol, Id);
+  return new (Mem) ObjCCategoryDecl(L, Id);
 }
 
 
@@ -164,6 +163,17 @@ void ObjCProtocolDecl::addMethods(ObjCMethodDecl **insMethods,
   AtEndLoc = endLoc;
 }
 
+void ObjCCategoryDecl::setReferencedProtocolList(ObjCProtocolDecl **List,
+                                                 unsigned NumRPs) {
+  assert(NumReferencedProtocols == 0 && "Protocol list already set");
+  if (NumRPs == 0) return;
+  
+  ReferencedProtocols = new ObjCProtocolDecl*[NumRPs];
+  memcpy(ReferencedProtocols, List, NumRPs*sizeof(ObjCProtocolDecl*));
+  NumReferencedProtocols = NumRPs;
+}
+
+
 /// addMethods - Insert instance and methods declarations into
 /// ObjCCategoryDecl's CatInsMethods and CatClsMethods fields.
 ///
diff --git a/lib/Sema/SemaDeclObjC.cpp b/lib/Sema/SemaDeclObjC.cpp
index 7bdfd9e014..e07d3c59e6 100644
--- a/lib/Sema/SemaDeclObjC.cpp
+++ b/lib/Sema/SemaDeclObjC.cpp
@@ -281,8 +281,7 @@ Sema::DeclTy *Sema::ActOnStartCategoryInterface(
   ObjCInterfaceDecl *IDecl = getObjCInterfaceDecl(ClassName);
   
   ObjCCategoryDecl *CDecl = 
-    ObjCCategoryDecl::Create(Context, AtInterfaceLoc, NumProtoRefs, 
-                             CategoryName);
+    ObjCCategoryDecl::Create(Context, AtInterfaceLoc, CategoryName);
   CDecl->setClassInterface(IDecl);
   
   /// Check that class of this category is already completely declared.
@@ -304,7 +303,8 @@ Sema::DeclTy *Sema::ActOnStartCategoryInterface(
   }
 
   if (NumProtoRefs) {
-    /// Check then save referenced protocols
+    llvm::SmallVector<ObjCProtocolDecl*, 32> RefProtocols;
+    /// Check and then save the referenced protocols.
     for (unsigned int i = 0; i != NumProtoRefs; i++) {
       ObjCProtocolDecl* RefPDecl = ObjCProtocols[ProtoRefNames[i]];
       if (!RefPDecl || RefPDecl->isForwardDecl()) {
@@ -312,10 +312,13 @@ Sema::DeclTy *Sema::ActOnStartCategoryInterface(
              ProtoRefNames[i]->getName(),
              CategoryName->getName());
       }
-      CDecl->setCatReferencedProtocols(i, RefPDecl);
+      if (RefPDecl)
+        RefProtocols.push_back(RefPDecl);
     }
-    CDecl->setLocEnd(EndProtoLoc);
+    if (!RefProtocols.empty())
+      CDecl->setReferencedProtocolList(&RefProtocols[0], RefProtocols.size());
   }
+  CDecl->setLocEnd(EndProtoLoc);
   return CDecl;
 }