4 files changed, 79 insertions, 7 deletions

diff --git a/lib/CodeGen/CGObjC.cpp b/lib/CodeGen/CGObjC.cpp
index 26f678cc4e..5376d0f3c7 100644
--- a/lib/CodeGen/CGObjC.cpp
+++ b/lib/CodeGen/CGObjC.cpp
@@ -281,10 +281,10 @@ void CodeGenFunction::EmitObjCForCollectionStmt(const ObjCForCollectionStmt &S)
   llvm::AllocaInst *StatePtr = CreateTempAlloca(ConvertType(StateTy), 
                                                 "state.ptr");
   StatePtr->setAlignment(getContext().getTypeAlign(StateTy) >> 3);  
-  EmitMemSetToZero(StatePtr,StateTy);
+  EmitMemSetToZero(StatePtr, StateTy);
   
   // Number of elements in the items array.
-  static const unsigned NumItems = 2;
+  static const unsigned NumItems = 16;
   
   // Get selector
   llvm::SmallVector<IdentifierInfo*, 3> II;
@@ -323,23 +323,64 @@ void CodeGenFunction::EmitObjCForCollectionStmt(const ObjCForCollectionStmt &S)
   Builder.CreateStore(CountRV.getScalarVal(), LimitPtr);
   
   llvm::BasicBlock *NoElements = llvm::BasicBlock::Create("noelements");
-  llvm::BasicBlock *LoopStart = llvm::BasicBlock::Create("loopstart");
+  llvm::BasicBlock *SetStartMutations = 
+    llvm::BasicBlock::Create("setstartmutations");
   
   llvm::Value *Limit = Builder.CreateLoad(LimitPtr);
   llvm::Value *Zero = llvm::Constant::getNullValue(UnsignedLongLTy);
 
   llvm::Value *IsZero = Builder.CreateICmpEQ(Limit, Zero, "iszero");
-  Builder.CreateCondBr(IsZero, NoElements, LoopStart);
+  Builder.CreateCondBr(IsZero, NoElements, SetStartMutations);
 
+  EmitBlock(SetStartMutations);
+  
+  llvm::Value *StartMutationsPtr = 
+    CreateTempAlloca(UnsignedLongLTy);
+  
+  llvm::Value *StateMutationsPtrPtr = 
+    Builder.CreateStructGEP(StatePtr, 2, "mutationsptr.ptr");
+  llvm::Value *StateMutationsPtr = Builder.CreateLoad(StateMutationsPtrPtr, 
+                                                      "mutationsptr");
+  
+  llvm::Value *StateMutations = Builder.CreateLoad(StateMutationsPtr, 
+                                                   "mutations");
+  
+  Builder.CreateStore(StateMutations, StartMutationsPtr);
+  
+  llvm::BasicBlock *LoopStart = llvm::BasicBlock::Create("loopstart");
   EmitBlock(LoopStart);
 
-  llvm::BasicBlock *LoopBody = llvm::BasicBlock::Create("loopbody");
-
   llvm::Value *CounterPtr = CreateTempAlloca(UnsignedLongLTy, "counter.ptr");
   Builder.CreateStore(Zero, CounterPtr);
   
+  llvm::BasicBlock *LoopBody = llvm::BasicBlock::Create("loopbody"); 
   EmitBlock(LoopBody);
 
+  StateMutationsPtr = Builder.CreateLoad(StateMutationsPtrPtr, "mutationsptr");
+  StateMutations = Builder.CreateLoad(StateMutationsPtr, "statemutations");
+
+  llvm::Value *StartMutations = Builder.CreateLoad(StartMutationsPtr, 
+                                                   "mutations");
+  llvm::Value *MutationsEqual = Builder.CreateICmpEQ(StateMutations, 
+                                                     StartMutations,
+                                                     "tobool");
+  
+  
+  llvm::BasicBlock *WasMutated = llvm::BasicBlock::Create("wasmutated");
+  llvm::BasicBlock *WasNotMutated = llvm::BasicBlock::Create("wasnotmutated");
+  
+  Builder.CreateCondBr(MutationsEqual, WasNotMutated, WasMutated);
+  
+  EmitBlock(WasMutated);
+  llvm::Value *V =
+    Builder.CreateBitCast(Collection, 
+                          ConvertType(getContext().getObjCIdType()),
+                          "tmp");
+  Builder.CreateCall(CGM.getObjCRuntime().EnumerationMutationFunction(),
+                     V);
+  
+  EmitBlock(WasNotMutated);
+  
   llvm::Value *StateItemsPtr = 
     Builder.CreateStructGEP(StatePtr, 1, "stateitems.ptr");
 
@@ -384,7 +425,7 @@ void CodeGenFunction::EmitObjCForCollectionStmt(const ObjCForCollectionStmt &S)
   llvm::BasicBlock *FetchMore = llvm::BasicBlock::Create("fetchmore");
   
   llvm::Value *IsLess = Builder.CreateICmpULT(Counter, Limit, "isless");
-  Builder.CreateCondBr(IsLess, LoopBody, FetchMore);
+  Builder.CreateCondBr(IsLess, LoopStart, FetchMore);
 
   // Fetch more elements.
   EmitBlock(FetchMore);
diff --git a/lib/CodeGen/CGObjCGNU.cpp b/lib/CodeGen/CGObjCGNU.cpp
index 4565fbfeb0..4d69aed816 100644
--- a/lib/CodeGen/CGObjCGNU.cpp
+++ b/lib/CodeGen/CGObjCGNU.cpp
@@ -120,6 +120,7 @@ public:
                                            const ObjCProtocolDecl *PD);
   virtual void GenerateProtocol(const ObjCProtocolDecl *PD);
   virtual llvm::Function *ModuleInitFunction();
+  virtual llvm::Function *EnumerationMutationFunction();
 };
 } // end anonymous namespace
 
@@ -958,6 +959,13 @@ llvm::Function *CGObjCGNU::GenerateMethod(const ObjCMethodDecl *OMD) {
   return Method;
 }
 
+llvm::Function *CGObjCGNU::EnumerationMutationFunction()
+{
+  assert(0 && "No enumeration mutation function in the GNU runtime!");
+  
+  return 0;
+}
+
 CodeGen::CGObjCRuntime *CodeGen::CreateGNUObjCRuntime(CodeGen::CodeGenModule &CGM){
   return new CGObjCGNU(CGM);
 }
diff --git a/lib/CodeGen/CGObjCMac.cpp b/lib/CodeGen/CGObjCMac.cpp
index ef4196e11f..798c9d3c4b 100644
--- a/lib/CodeGen/CGObjCMac.cpp
+++ b/lib/CodeGen/CGObjCMac.cpp
@@ -130,6 +130,7 @@ public:
   /// MethodListPtrTy - LLVM type for struct objc_method_list *.
   const llvm::Type *MethodListPtrTy;
 
+  llvm::Function *EnumerationMutationFn;
 public:
   ObjCTypesHelper(CodeGen::CodeGenModule &cgm);
   ~ObjCTypesHelper();
@@ -370,6 +371,7 @@ public:
   virtual void GenerateProtocol(const ObjCProtocolDecl *PD);
 
   virtual llvm::Function *ModuleInitFunction();
+  virtual llvm::Function *EnumerationMutationFunction();
 };
 } // end anonymous namespace
 
@@ -1379,6 +1381,11 @@ llvm::Function *CGObjCMac::ModuleInitFunction() {
   return NULL;
 }
 
+llvm::Function *CGObjCMac::EnumerationMutationFunction()
+{
+  return ObjCTypes.EnumerationMutationFn;
+}
+
 /* *** Private Interface *** */
 
 /// EmitImageInfo - Emit the image info marker used to encode some module
@@ -1934,6 +1941,18 @@ ObjCTypesHelper::ObjCTypesHelper(CodeGen::CodeGenModule &cgm)
                            llvm::Function::ExternalLinkage,
                            "objc_msgSendSuper_stret",
                            &CGM.getModule());
+  
+  // Enumeration mutation.
+  
+  Params.clear();
+  Params.push_back(ObjectPtrTy);
+  EnumerationMutationFn = 
+    llvm::Function::Create(llvm::FunctionType::get(llvm::Type::VoidTy,
+                                                   Params,
+                                                   false),
+                           llvm::Function::ExternalLinkage,
+                           "objc_enumerationMutation",
+                           &CGM.getModule());
 }
 
 ObjCTypesHelper::~ObjCTypesHelper() {
diff --git a/lib/CodeGen/CGObjCRuntime.h b/lib/CodeGen/CGObjCRuntime.h
index c0044e9a94..4bc293afac 100644
--- a/lib/CodeGen/CGObjCRuntime.h
+++ b/lib/CodeGen/CGObjCRuntime.h
@@ -122,6 +122,10 @@ public:
   virtual llvm::Value *GetClass(BuilderType &Builder, 
                                 const ObjCInterfaceDecl *OID) = 0;
 
+  /// EnumerationMutationFunction - Return the function that's called by the
+  /// compiler when a mutation is detected during foreach iteration.
+  virtual llvm::Function *EnumerationMutationFunction() = 0;
+    
   /// If instance variable addresses are determined at runtime then this should
   /// return true, otherwise instance variables will be accessed directly from
   /// the structure.  If this returns true then @defs is invalid for this