Fix #622 on both backends

2 files changed, 8 insertions, 4 deletions

diff --git a/src/ir.cpp b/src/ir.cpp
index 65e565bc5..e2c9f90ed 100644
--- a/src/ir.cpp
+++ b/src/ir.cpp
@@ -3307,10 +3307,9 @@ void ir_emit_defer_stmts(irProcedure *proc, irDeferExitKind kind, irBlock *block
 	while (i --> 0) {
 		irDefer d = proc->defer_stmts[i];
 
-		// TODO(bill, 2020-03-05): Why was this added?
-		// if (proc->context_stack.count >= d.context_stack_count) {
-		// 	proc->context_stack.count = d.context_stack_count;
-		// }
+		isize prev_context_stack_count = proc->context_stack.count;
+		defer (proc->context_stack.count = prev_context_stack_count);
+		proc->context_stack.count = d.context_stack_count;
 
 		if (kind == irDeferExit_Default) {
 			if (proc->scope_index == d.scope_index &&
diff --git a/src/llvm_backend.cpp b/src/llvm_backend.cpp
index 816d81f33..06209a35e 100644
--- a/src/llvm_backend.cpp
+++ b/src/llvm_backend.cpp
@@ -6432,6 +6432,11 @@ void lb_build_defer_stmt(lbProcedure *p, lbDefer d) {
 		return;
 	}
 
+	isize prev_context_stack_count = p->context_stack.count;
+	defer (p->context_stack.count = prev_context_stack_count);
+	p->context_stack.count = d.context_stack_count;
+
+
 	lbBlock *b = lb_create_block(p, "defer");
 	if (last_instr == nullptr || !LLVMIsATerminatorInst(last_instr)) {
 		lb_emit_jump(p, b);