3 files changed, 34 insertions, 15 deletions

diff --git a/lib/CodeGen/CGObjC.cpp b/lib/CodeGen/CGObjC.cpp
index a116316641..efd9299d1a 100644
--- a/lib/CodeGen/CGObjC.cpp
+++ b/lib/CodeGen/CGObjC.cpp
@@ -361,6 +361,11 @@ void CodeGenFunction::EmitObjCForCollectionStmt(const ObjCForCollectionStmt &S)
     return;
   }
 
+  // We want to ensure that any vlas between here and when we
+  // push the break and continue context below can be destroyed
+  // when we break 
+  llvm::Value *saveBreakStackDepth = StackDepth;
+
   if (const DeclStmt *SD = dyn_cast<DeclStmt>(S.getElement())) {
     EmitStmt(SD);
     assert(HaveInsertPoint() && "DeclStmt destroyed insert point!");
@@ -446,10 +451,6 @@ void CodeGenFunction::EmitObjCForCollectionStmt(const ObjCForCollectionStmt &S)
   llvm::BasicBlock *LoopStart = createBasicBlock("loopstart");
   EmitBlock(LoopStart);
 
-  // We want to ensure there are no vlas between here and when we
-  // push the break and continue context below.
-  llvm::Value *saveStackDepth = StackDepth;
-
   llvm::Value *CounterPtr = CreateTempAlloca(UnsignedLongLTy, "counter.ptr");
   Builder.CreateStore(Zero, CounterPtr);
   
@@ -519,9 +520,8 @@ void CodeGenFunction::EmitObjCForCollectionStmt(const ObjCForCollectionStmt &S)
   llvm::BasicBlock *LoopEnd = createBasicBlock("loopend");
   llvm::BasicBlock *AfterBody = createBasicBlock("afterbody");
   
-  assert (StackDepth == saveStackDepth && "vla unhandled in for");
-
-  BreakContinuePush(LoopEnd, AfterBody);
+  // Ensure any vlas created between there and here, are undone
+  BreakContinuePush(LoopEnd, AfterBody, saveBreakStackDepth, StackDepth);
 
   EmitStmt(S.getBody());
   
diff --git a/lib/CodeGen/CGStmt.cpp b/lib/CodeGen/CGStmt.cpp
index c65b8114f3..ac97c6329d 100644
--- a/lib/CodeGen/CGStmt.cpp
+++ b/lib/CodeGen/CGStmt.cpp
@@ -466,6 +466,8 @@ void CodeGenFunction::EmitForStmt(const ForStmt &S) {
 
   EmitBlock(CondBlock);
 
+  llvm::Value *saveStackDepth = StackDepth;
+
   // Evaluate the condition if present.  If not, treat it as a
   // non-zero-constant according to 6.8.5.3p2, aka, true.
   if (S.getCond()) {
@@ -491,8 +493,10 @@ void CodeGenFunction::EmitForStmt(const ForStmt &S) {
     ContinueBlock = CondBlock;  
   
   // Store the blocks to use for break and continue.
-  BreakContinuePush(AfterFor, ContinueBlock);
-  
+  // Ensure any vlas created between there and here, are undone
+  BreakContinuePush(AfterFor, ContinueBlock,
+                    saveStackDepth, saveStackDepth);
+
   // If the condition is true, execute the body of the for stmt.
   EmitStmt(S.getBody());
 
@@ -708,6 +712,9 @@ void CodeGenFunction::EmitSwitchStmt(const SwitchStmt &S) {
   llvm::SwitchInst *SavedSwitchInsn = SwitchInsn;
   llvm::BasicBlock *SavedCRBlock = CaseRangeBlock;
 
+  // Ensure any vlas created inside are destroyed on break.
+  llvm::Value *saveBreakStackDepth = StackDepth;
+
   // Create basic block to hold stuff that comes after switch
   // statement. We also need to create a default block now so that
   // explicit case ranges tests can have a place to jump to on
@@ -723,9 +730,14 @@ void CodeGenFunction::EmitSwitchStmt(const SwitchStmt &S) {
   // All break statements jump to NextBlock. If BreakContinueStack is non empty
   // then reuse last ContinueBlock.
   llvm::BasicBlock *ContinueBlock = NULL;
-  if (!BreakContinueStack.empty())
+  llvm::Value *saveContinueStackDepth = NULL;
+  if (!BreakContinueStack.empty()) {
     ContinueBlock = BreakContinueStack.back().ContinueBlock;
-  BreakContinuePush(NextBlock, ContinueBlock);
+    saveContinueStackDepth = BreakContinueStack.back().SaveContinueStackDepth;
+  }
+  // Ensure any vlas created between there and here, are undone
+  BreakContinuePush(NextBlock, ContinueBlock,
+                    saveBreakStackDepth, saveContinueStackDepth);
 
   // Emit switch body.
   EmitStmt(S.getBody());
diff --git a/lib/CodeGen/CodeGenFunction.h b/lib/CodeGen/CodeGenFunction.h
index 4ff3badd99..068627139d 100644
--- a/lib/CodeGen/CodeGenFunction.h
+++ b/lib/CodeGen/CodeGenFunction.h
@@ -156,6 +156,12 @@ private:
   /// condition has possibly started a vla.
   void BreakContinuePush(llvm::BasicBlock *bb, llvm::BasicBlock *cb) {
     BreakContinueStack.push_back(BreakContinue(bb, cb, StackDepth,
+                                               StackDepth,
+                                               ObjCEHStack.size()));
+  }
+  void BreakContinuePush(llvm::BasicBlock *bb, llvm::BasicBlock *cb,
+                         llvm::Value *bsd, llvm::Value *csd) {
+    BreakContinueStack.push_back(BreakContinue(bb, cb, bsd, csd,
                                                ObjCEHStack.size()));
   }
 
@@ -169,13 +175,14 @@ private:
   // of the eh stack.
   struct BreakContinue {
     BreakContinue(llvm::BasicBlock *bb, llvm::BasicBlock *cb,
-                  llvm::Value *sd, size_t ehss)
-      : BreakBlock(bb), ContinueBlock(cb), SaveStackDepth(sd),
-        EHStackSize(ehss) {}
+                  llvm::Value *bsd, llvm::Value *csd, size_t ehss)
+      : BreakBlock(bb), ContinueBlock(cb),  SaveBreakStackDepth(bsd),
+        SaveContinueStackDepth(csd), EHStackSize(ehss) {}
       
     llvm::BasicBlock *BreakBlock;
     llvm::BasicBlock *ContinueBlock;
-    llvm::Value *SaveStackDepth;
+    llvm::Value *SaveBreakStackDepth;
+    llvm::Value *SaveContinueStackDepth;
     size_t EHStackSize;
   }; 
   llvm::SmallVector<BreakContinue, 8> BreakContinueStack;