add defer

compiler/checker/checker.odin

@@ -68,9 +68,18 @@ Build_Ctx :: struct {
 	global_reads: ^[dynamic]hir.Global_Id,
 	calls:        ^[dynamic]hir.Function_Id,
 	problematic:  ^bool,
-	// Number of enclosing loops being built. `break`/`continue` are only valid
-	// when this is > 0; bumped around loop-body builds in `build_block`.
-	loop_depth:   int,
+	// `defer` lowering. Deferred statements are built once at the `defer` site and
+	// their hir stmt ids stored here as a flat stack across scopes (one entry per
+	// deferred statement); they are replayed (appended) at each scope exit in LIFO
+	// order. `loop_defer_starts` records `len(defers)` at each enclosing loop body
+	// entry: `break`/`continue` flush down to that mark (and need `len > loop_floor`
+	// to be valid). `defer_depth`/`loop_floor` guard control flow inside a deferred
+	// statement: `return` is rejected while `defer_depth > 0`, and `break`/`continue`
+	// only see loops opened within the defer (those past `loop_floor`).
+	defers:            ^[dynamic][]hir.Stmt_Id,
+	loop_defer_starts: ^[dynamic]int,
+	defer_depth:       int,
+	loop_floor:        int,
 }
 
 Constant_Kind :: enum {
@@ -737,6 +746,11 @@ mark_block_imports_used :: proc(checker: ^Checker, statements: []ast.Stmt_Id, fi
 		case .For:
 			mark_expr_imports_used(checker, statement.expr, file)
 			mark_block_imports_used(checker, statement.body, file)
+		case .Block:
+			mark_block_imports_used(checker, statement.body, file)
+		case .Defer:
+			deferred := [1]ast.Stmt_Id{statement.update}
+			mark_block_imports_used(checker, deferred[:], file)
 		case .Break, .Continue:
 		case .Invalid:
 		}
@@ -1875,6 +1889,11 @@ infer_statements :: proc(
 			}
 			infer_statements(checker, statement.body, locals, local_types, pkg, file, demanded, result, result_hint)
 			resize(locals, capture_start)
+		case .Block:
+			infer_statements(checker, statement.body, locals, local_types, pkg, file, demanded, result, result_hint)
+		case .Defer:
+			deferred := [1]ast.Stmt_Id{statement.update}
+			infer_statements(checker, deferred[:], locals, local_types, pkg, file, demanded, result, result_hint)
 		}
 	}
 	resize(locals, scope_start)
@@ -3962,6 +3981,17 @@ make_link_name :: proc(checker: ^Checker, id: Spec_Id) -> string {
 	return fmt.aprintf("%s", strings.to_string(builder), allocator = checker.allocator)
 }
 
+// Replay the deferred statements in frames `[lo, len(defers))` into `body`,
+// innermost-most-recent first (LIFO across frames); each frame's own statements
+// keep their forward order. Used at every scope-exit path in `build_block`.
+flush_defers :: proc(ctx: ^Build_Ctx, body: ^[dynamic]hir.Stmt_Id, lo: int) {
+	for i := len(ctx.defers^) - 1; i >= lo; i -= 1 {
+		for stmt_id in ctx.defers^[i] {
+			append(body, stmt_id)
+		}
+	}
+}
+
 build_block :: proc(
 	ctx: ^Build_Ctx,
 	statements: []ast.Stmt_Id,
@@ -3971,6 +4001,7 @@ build_block :: proc(
 	body: [dynamic]hir.Stmt_Id
 	body.allocator = checker.allocator
 	scope_start := len(ctx.locals^)
+	defer_start := len(ctx.defers^)
 	duplicate_start := scope_start if duplicate_scope_start < 0 else duplicate_scope_start
 	for statement_id in statements {
 		statement := checker.ast_module.statements[statement_id]
@@ -4224,6 +4255,16 @@ build_block :: proc(
 			})
 			ctx.problematic^ = ctx.problematic^ || checker.module.exprs[value].kind == .Invalid
 		case .Return:
+			if ctx.defer_depth > 0 {
+				id := source.add(checker.diagnostics, statement.span, "cannot 'return' inside a 'defer'")
+				append(&body, hir.stmt_id(len(checker.module.statements)))
+				append(&checker.module.statements, hir.Stmt{
+					kind = .Trap, span = statement.span, expr = hir.INVALID_EXPR,
+					local = hir.INVALID_LOCAL, diagnostic = id,
+				})
+				ctx.problematic^ = true
+				continue
+			}
 			if statement.expr == ast.INVALID_EXPR {
 				if !types.is_void(ctx.result) {
 					id := source.add(checker.diagnostics, statement.span, "'return _' is only valid in a void function")
@@ -4234,6 +4275,7 @@ build_block :: proc(
 					})
 					ctx.problematic^ = true
 				} else {
+					flush_defers(ctx, &body, 0)
 					append(&body, hir.stmt_id(len(checker.module.statements)))
 					append(&checker.module.statements, hir.Stmt{
 						kind = .Return, span = statement.span, expr = hir.INVALID_EXPR,
@@ -4257,12 +4299,32 @@ build_block :: proc(
 				ctx.result, ctx.pkg, ctx.file,
 			)
 			value = coerce_expr(checker, value, ctx.result, statement.span)
+			ctx.problematic^ = ctx.problematic^ || checker.module.exprs[value].kind == .Invalid
+			// Run deferred statements before returning, but capture the return value
+			// first (spill it to a temp) so a defer that mutates the returned local
+			// can't change what is returned — Zig evaluates the return value, then
+			// runs defers.
+			if len(ctx.defers^) > 0 {
+				if checker.module.exprs[value].kind != .Invalid {
+					tmp := hir.local_id(len(ctx.hir_locals^))
+					append(ctx.hir_locals, hir.Local{name = checker.sink_symbol, type = ctx.result, mutable = false})
+					append(&body, hir.stmt_id(len(checker.module.statements)))
+					append(&checker.module.statements, hir.Stmt{
+						kind = .Declaration, span = statement.span, local = tmp, expr = value,
+						diagnostic = source.INVALID_DIAGNOSTIC,
+					})
+					value = hir.expr_id(len(checker.module.exprs))
+					append(&checker.module.exprs, hir.Expr{
+						kind = .Local, span = statement.span, type = ctx.result, target = hir.local_ref(tmp),
+					})
+				}
+				flush_defers(ctx, &body, 0)
+			}
 			append(&body, hir.stmt_id(len(checker.module.statements)))
 			append(&checker.module.statements, hir.Stmt{
 				kind = .Return, span = statement.span, expr = value,
 				local = hir.INVALID_LOCAL, diagnostic = source.INVALID_DIAGNOSTIC,
 			})
-			ctx.problematic^ = ctx.problematic^ || checker.module.exprs[value].kind == .Invalid
 		case .Expression:
 			value := build_expr(checker, statement.expr, ctx.locals^[:], ctx.global_reads, ctx.calls, types.INVALID, ctx.pkg, ctx.file)
 			if !types.is_void(checker.module.exprs[value].type) {
@@ -4443,9 +4505,9 @@ build_block :: proc(
 				condition = invalid_hir_expr(checker, statement.span, id, types.BOOL)
 				ctx.problematic^ = true
 			}
-			ctx.loop_depth += 1
+			append(ctx.loop_defer_starts, len(ctx.defers^))
 			loop_body := build_block(ctx, statement.body)
-			ctx.loop_depth -= 1
+			pop(ctx.loop_defer_starts)
 			update := hir.INVALID_STMT
 			if statement.update != ast.INVALID_STMT {
 				update_ast := [1]ast.Stmt_Id{statement.update}
@@ -4537,9 +4599,9 @@ build_block :: proc(
 					append(ctx.locals, Build_Local{name=statement.index_name, type=types.USIZE, mutable=false, id=index_local})
 				}
 			}
-			ctx.loop_depth += 1
+			append(ctx.loop_defer_starts, len(ctx.defers^))
 			loop_body := build_block(ctx, statement.body, capture_start)
-			ctx.loop_depth -= 1
+			pop(ctx.loop_defer_starts)
 			resize(ctx.locals, capture_start)
 
 			append(&body, hir.stmt_id(len(checker.module.statements)))
@@ -4571,7 +4633,9 @@ build_block :: proc(
 				ctx.problematic^ = true
 			}
 		case .Break, .Continue:
-			if ctx.loop_depth == 0 {
+			// Inside a `defer`, `loop_floor` hides the enclosing loops so only loops
+			// opened within the defer count.
+			if len(ctx.loop_defer_starts^) <= ctx.loop_floor {
 				keyword := "break" if statement.kind == .Break else "continue"
 				id := source.addf(checker.diagnostics, statement.span, "'%s' outside of a loop", keyword)
 				append(&body, hir.stmt_id(len(checker.module.statements)))
@@ -4582,12 +4646,50 @@ build_block :: proc(
 				ctx.problematic^ = true
 				continue
 			}
+			// Exit the loop body and any blocks between here and it: run their
+			// deferred statements down to and including the innermost loop body.
+			flush_defers(ctx, &body, ctx.loop_defer_starts^[len(ctx.loop_defer_starts^) - 1])
 			append(&body, hir.stmt_id(len(checker.module.statements)))
 			append(&checker.module.statements, hir.Stmt{
 				kind = .Break if statement.kind == .Break else .Continue,
 				span = statement.span, expr = hir.INVALID_EXPR,
 				local = hir.INVALID_LOCAL, diagnostic = source.INVALID_DIAGNOSTIC,
 			})
+		case .Block:
+			// A bare `{ ... }` scope: build it (its own locals/defers are scoped by
+			// the recursive call) and splice its statements in.
+			block := build_block(ctx, statement.body)
+			for stmt in block {
+				append(&body, stmt)
+			}
+			delete(block, checker.allocator)
+		case .Defer:
+			deferred := checker.ast_module.statements[statement.update]
+			if deferred.kind == .Return || deferred.kind == .Break ||
+			   deferred.kind == .Continue || deferred.kind == .Defer {
+				keyword := "return"
+				if deferred.kind == .Break    { keyword = "break" }
+				if deferred.kind == .Continue { keyword = "continue" }
+				if deferred.kind == .Defer    { keyword = "defer" }
+				id := source.addf(checker.diagnostics, statement.span, "cannot defer a '%s' statement", keyword)
+				append(&body, hir.stmt_id(len(checker.module.statements)))
+				append(&checker.module.statements, hir.Stmt{
+					kind = .Trap, span = statement.span, expr = hir.INVALID_EXPR,
+					local = hir.INVALID_LOCAL, diagnostic = id,
+				})
+				ctx.problematic^ = true
+				continue
+			}
+			// Build the deferred statement once, guarded so a `return` inside it is
+			// rejected and `break`/`continue` only target loops opened within the
+			// defer; its hir is replayed at each scope exit, not emitted here.
+			saved_floor := ctx.loop_floor
+			ctx.defer_depth += 1
+			ctx.loop_floor = len(ctx.loop_defer_starts^)
+			entry := build_block(ctx, []ast.Stmt_Id{statement.update})
+			ctx.loop_floor = saved_floor
+			ctx.defer_depth -= 1
+			append(ctx.defers, entry)
 		case .Invalid:
 			append(&body, hir.stmt_id(len(checker.module.statements)))
 			append(&checker.module.statements, hir.Stmt{
@@ -4597,6 +4699,18 @@ build_block :: proc(
 			ctx.problematic^ = true
 		}
 	}
+	// Normal fall-through exit: run this block's own deferred statements, unless
+	// every path already exited early (return/break/continue) — that would only
+	// emit unreachable duplicates.
+	if !all_paths_exit(&checker.module, body[:]) {
+		flush_defers(ctx, &body, defer_start)
+	}
+	// Free this block's deferred-statement entry slices (their stmt ids were already
+	// replayed at every path that can leave this block) and pop the frame.
+	for i := defer_start; i < len(ctx.defers^); i += 1 {
+		delete(ctx.defers^[i], checker.allocator)
+	}
+	resize(ctx.defers, defer_start)
 	resize(ctx.locals, scope_start)
 	return body[:]
 }
@@ -4652,6 +4766,34 @@ loop_body_breaks :: proc(module: ^hir.Module, stmts: []hir.Stmt_Id) -> bool {
 	return false
 }
 
+// Like `all_paths_return`, but also treats `break`/`continue` as terminating the
+// block. Used only to decide whether `build_block` may skip the fall-through defer
+// flush (a block that always exits early would otherwise emit unreachable copies).
+all_paths_exit :: proc(module: ^hir.Module, stmts: []hir.Stmt_Id) -> bool {
+	for id in stmts {
+		statement := module.statements[id]
+		#partial switch statement.kind {
+		case .Return, .Trap, .Break, .Continue:
+			return true
+		case .If:
+			if statement.else_body != nil &&
+			   all_paths_exit(module, statement.then_body) &&
+			   all_paths_exit(module, statement.else_body) {
+				return true
+			}
+		case .While:
+			if statement.expr != hir.INVALID_EXPR && int(statement.expr) < len(module.exprs) {
+				condition := module.exprs[statement.expr]
+				if condition.kind == .Bool && condition.integer != 0 &&
+				   !loop_body_breaks(module, statement.then_body) {
+					return true
+				}
+			}
+		}
+	}
+	return false
+}
+
 build_function :: proc(checker: ^Checker, id: Spec_Id) {
 	spec := checker.specs[id]
 	function := checker.ast_module.functions[spec.template]
@@ -4751,6 +4893,10 @@ build_function :: proc(checker: ^Checker, id: Spec_Id) {
 			},
 		)
 	}
+	defers: [dynamic][]hir.Stmt_Id
+	defers.allocator = checker.allocator
+	loop_defer_starts: [dynamic]int
+	loop_defer_starts.allocator = checker.allocator
 	ctx := Build_Ctx{
 		checker = checker,
 		pkg = function.pkg,
@@ -4762,6 +4908,8 @@ build_function :: proc(checker: ^Checker, id: Spec_Id) {
 		global_reads = &global_reads,
 		calls = &calls,
 		problematic = &problematic,
+		defers = &defers,
+		loop_defer_starts = &loop_defer_starts,
 	}
 	block := build_block(&ctx, function.body)
 	returns := all_paths_return(&checker.module, block)
@@ -4806,6 +4954,11 @@ build_function :: proc(checker: ^Checker, id: Spec_Id) {
 			diagnostic = source.INVALID_DIAGNOSTIC,
 		},
 	)
+	for entry in defers {
+		delete(entry, checker.allocator)
+	}
+	delete(defers)
+	delete(loop_defer_starts)
 	delete(locals)
 }