3 files changed, 36 insertions, 1 deletions
diff --git a/lib/CodeGen/CGStmt.cpp b/lib/CodeGen/CGStmt.cpp
index fc7a7c7f8f..19977ec380 100644
--- a/lib/CodeGen/CGStmt.cpp
+++ b/lib/CodeGen/CGStmt.cpp
@@ -188,6 +188,12 @@ void CodeGenFunction::EmitBlock(llvm::BasicBlock *BB, bool IsFinished) {
     return;
   }
 
+  // If necessary, associate the block with the cleanup stack size.
+  if (!CleanupEntries.empty()) {
+    BlockScopes[BB] = CleanupEntries.size() - 1;
+    CleanupEntries.back().Blocks.push_back(BB);
+  }
+  
   CurFn->getBasicBlockList().push_back(BB);
   Builder.SetInsertPoint(BB);
 }
diff --git a/lib/CodeGen/CodeGenFunction.cpp b/lib/CodeGen/CodeGenFunction.cpp
index 838de70218..0e0140c9ad 100644
--- a/lib/CodeGen/CodeGenFunction.cpp
+++ b/lib/CodeGen/CodeGenFunction.cpp
@@ -119,7 +119,11 @@ void CodeGenFunction::FinishFunction(SourceLocation EndLoc) {
 
   assert(BreakContinueStack.empty() &&
          "mismatched push/pop in break/continue stack!");
-
+  assert(BlockScopes.empty() &&
+         "did not remove all blocks from block scope map!");
+  assert(CleanupEntries.empty() &&
+         "mismatched push/pop in cleanup stack!");
+  
   // Emit function epilog (to return). 
   EmitReturnBlock();
 
@@ -537,8 +541,22 @@ void CodeGenFunction::EmitCleanupBlock()
   
   llvm::BasicBlock *CleanupBlock = CE.CleanupBlock;
   
+  std::vector<llvm::BasicBlock *> Blocks;
+  std::swap(Blocks, CE.Blocks);
+
+  std::vector<llvm::BranchInst *> BranchFixups;
+  std::swap(BranchFixups, CE.BranchFixups);
+  
   CleanupEntries.pop_back();
   
   EmitBlock(CleanupBlock);
+  
+  // Remove all blocks from the block scope map.
+  for (size_t i = 0, e = Blocks.size(); i != e; ++i) {
+    assert(BlockScopes.count(Blocks[i]) &&
+           "Did not find block in scope map!");
+    
+    BlockScopes.erase(Blocks[i]);
+  }
 }
 
diff --git a/lib/CodeGen/CodeGenFunction.h b/lib/CodeGen/CodeGenFunction.h
index 8a0220724e..3dbc2c99e7 100644
--- a/lib/CodeGen/CodeGenFunction.h
+++ b/lib/CodeGen/CodeGenFunction.h
@@ -253,11 +253,22 @@ private:
 
     explicit CleanupEntry(llvm::BasicBlock *cb)
       : CleanupBlock(cb) {}
+    
+    ~CleanupEntry() {
+      assert(Blocks.empty() && "Did not empty blocks!");
+      assert(BranchFixups.empty() && "Did not empty branch fixups!");
+    }
+    
   };
   
   /// CleanupEntries - Stack of cleanup entries.
   llvm::SmallVector<CleanupEntry, 8> CleanupEntries;
 
+  typedef llvm::DenseMap<llvm::BasicBlock*, size_t> BlockScopeMap;
+
+  /// BlockScopes - Map of which "cleanup scope" scope basic blocks have.
+  BlockScopeMap BlockScopes;
+  
 public:
   CodeGenFunction(CodeGenModule &cgm);