1 files changed, 15 insertions, 16 deletions
diff --git a/lib/Serialization/ASTWriter.cpp b/lib/Serialization/ASTWriter.cpp
index deaec02271..a98c4cc82b 100644
--- a/lib/Serialization/ASTWriter.cpp
+++ b/lib/Serialization/ASTWriter.cpp
@@ -2995,7 +2995,6 @@ void ASTWriter::WriteASTCore(Sema &SemaRef, MemorizeStatCalls *StatCalls,
   // Resolve any declaration pointers within the declaration updates block and
   // chained Objective-C categories block to declaration IDs.
   ResolveDeclUpdatesBlocks();
-  ResolveChainedObjCCategories();
   
   // Form the record of special types.
   RecordData SpecialTypes;
@@ -3183,7 +3182,7 @@ void ASTWriter::ResolveDeclUpdatesBlocks() {
     const Decl *D = I->first;
     UpdateRecord &URec = I->second;
     
-    if (DeclsToRewrite.count(D))
+    if (isRewritten(D))
       continue; // The decl will be written completely
 
     unsigned Idx = 0, N = URec.size();
@@ -3216,7 +3215,7 @@ void ASTWriter::WriteDeclUpdatesBlocks() {
     const Decl *D = I->first;
     UpdateRecord &URec = I->second;
 
-    if (DeclsToRewrite.count(D))
+    if (isRewritten(D))
       continue; // The decl will be written completely,no need to store updates.
 
     uint64_t Offset = Stream.GetCurrentBitNo();
@@ -3243,17 +3242,6 @@ void ASTWriter::WriteDeclReplacementsBlock() {
   Stream.EmitRecord(DECL_REPLACEMENTS, Record);
 }
 
-void ASTWriter::ResolveChainedObjCCategories() {
-  for (SmallVector<ChainedObjCCategoriesData, 16>::iterator
-       I = LocalChainedObjCCategories.begin(),
-       E = LocalChainedObjCCategories.end(); I != E; ++I) {
-    ChainedObjCCategoriesData &Data = *I;
-    Data.InterfaceID = GetDeclRef(Data.Interface);
-    Data.TailCategoryID = GetDeclRef(Data.TailCategory);
-  }
-
-}
-
 void ASTWriter::WriteChainedObjCCategories() {
   if (LocalChainedObjCCategories.empty())
     return;
@@ -3263,13 +3251,16 @@ void ASTWriter::WriteChainedObjCCategories() {
          I = LocalChainedObjCCategories.begin(),
          E = LocalChainedObjCCategories.end(); I != E; ++I) {
     ChainedObjCCategoriesData &Data = *I;
+    if (isRewritten(Data.Interface))
+      continue;
+
     serialization::DeclID
         HeadCatID = getDeclID(Data.Interface->getCategoryList());
     assert(HeadCatID != 0 && "Category not written ?");
 
-    Record.push_back(Data.InterfaceID);
+    Record.push_back(GetDeclRef(Data.Interface));
     Record.push_back(HeadCatID);
-    Record.push_back(Data.TailCategoryID);
+    Record.push_back(GetDeclRef(Data.TailCategory));
   }
   Stream.EmitRecord(OBJC_CHAINED_CATEGORIES, Record);
 }
@@ -4136,3 +4127,11 @@ void ASTWriter::AddedObjCCategoryToInterface(const ObjCCategoryDecl *CatD,
   ChainedObjCCategoriesData Data =  { IFD, CatD, 0, 0 };
   LocalChainedObjCCategories.push_back(Data);
 }
+
+void ASTWriter::CompletedObjCForwardRef(const ObjCContainerDecl *D) {
+  assert(!WritingAST && "Already writing the AST!");
+  if (!D->isFromASTFile())
+    return; // Declaration not imported from PCH.
+
+  RewriteDecl(D);
+}