errdefer and try/defer fix

compiler/checker/checker.odin

@@ -72,6 +72,12 @@ Yield_Target :: struct {
 	defer_floor:     int,
 }
 
+Defer_Entry :: struct {
+	body:       []hir.Stmt_Id,
+	error_only: bool,
+	capture:    hir.Local_Id,
+}
+
 Build_Ctx :: struct {
 	checker:      ^Checker,
 	pkg:          ast.Package_Id,
@@ -96,7 +102,7 @@ Build_Ctx :: struct {
 	// 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,
+	defers:            ^[dynamic]Defer_Entry,
 	loop_defer_starts: ^[dynamic]int,
 	// Parallel to `loop_defer_starts`: the label of each enclosing break target (INVALID
 	// when unlabeled), so a `break :L` / `continue :L` can target an outer one. A labeled
@@ -3686,8 +3692,19 @@ infer_statements :: proc(
 		case .Block:
 			infer_statements(checker, statement.body, locals, local_types, pkg, file, demanded, result, result_hint)
 		case .Defer:
+			capture_start := len(locals^)
+			if statement.error_only && len(statement.captures) > 0 &&
+			   statement.captures[0] != checker.sink_symbol {
+				append(locals, Infer_Local{
+					name=statement.captures[0],
+					type=types.fallible_error(result^, &checker.module.types),
+					declared=types.fallible_error(result^, &checker.module.types),
+					statement=ast.INVALID_STMT,
+				})
+			}
 			deferred := [1]ast.Stmt_Id{statement.update}
 			infer_statements(checker, deferred[:], locals, local_types, pkg, file, demanded, result, result_hint)
+			resize(locals, capture_start)
 		case .Match:
 			// The build pass desugars `match` to an if/else chain, but inference runs first
 			// and must still visit the subject and arm bodies so calls there get specialized
@@ -5508,6 +5525,10 @@ build_compound_expr :: proc(
 			target=hir.INVALID_REF, diagnostic=source.INVALID_DIAGNOSTIC,
 		})
 	case .Try:
+		if checker.current_build_ctx != nil && checker.current_build_ctx.defer_depth > 0 {
+			id := source.add(checker.diagnostics, expr.span, "cannot 'try' inside a 'defer'")
+			return invalid_hir_expr(checker, expr.span, id)
+		}
 		left_expected := expected if checker.ast_module.exprs[expr.left].kind == .Call else types.INVALID
 		channel := build_nested_expr(checker, expr.left, locals, global_reads, calls, left_expected, pkg, file)
 		channel_type := checker.module.exprs[channel].type
@@ -5532,11 +5553,18 @@ build_compound_expr :: proc(
 			id := source.add(checker.diagnostics, expr.span, "'try' error channel cannot be widened to the enclosing error channel")
 			return invalid_hir_expr(checker, expr.span, id, success)
 		}
+		cleanup: []hir.Stmt_Id
+		captures: []hir.Expr_Id
+		if checker.current_build_ctx != nil {
+			cleanup, captures = try_cleanup(checker.current_build_ctx)
+		}
 		return add_hir_expr(checker, hir.Expr{
 			kind=.Try,
 			span=expr.span,
 			type=success,
 			left=channel,
+			body=cleanup,
+			args=captures,
 			target=hir.INVALID_REF,
 			right=hir.INVALID_EXPR,
 			diagnostic=source.INVALID_DIAGNOSTIC,
@@ -6701,17 +6729,51 @@ 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) {
+// Replay eligible cleanup in `[lo, len(defers))` in LIFO order. Error exits run
+// both defer forms; other exits skip errdefer. A direct error return supplies its
+// preserved payload so captured errors can be initialized before each cleanup.
+flush_defers :: proc(
+	ctx: ^Build_Ctx,
+	body: ^[dynamic]hir.Stmt_Id,
+	lo: int,
+	error_exit := false,
+	error_value := hir.INVALID_EXPR,
+) {
 	for i := len(ctx.defers^) - 1; i >= lo; i -= 1 {
-		for stmt_id in ctx.defers^[i] {
+		entry := ctx.defers^[i]
+		if entry.error_only && !error_exit {
+			continue
+		}
+		if error_exit && entry.capture != hir.INVALID_LOCAL && error_value != hir.INVALID_EXPR {
+			append(body, hir.stmt_id(len(ctx.checker.module.statements)))
+			append(&ctx.checker.module.statements, hir.Stmt{
+				kind=.Declaration, local=entry.capture, expr=error_value,
+				diagnostic=source.INVALID_DIAGNOSTIC,
+			})
+		}
+		for stmt_id in entry.body {
 			append(body, stmt_id)
 		}
 	}
 }
 
+// Copy the active error-exit cleanup onto a Try expression. Capture locals are
+// initialized by lowering once the propagated error has been extracted/widened.
+try_cleanup :: proc(ctx: ^Build_Ctx) -> ([]hir.Stmt_Id, []hir.Expr_Id) {
+	body: [dynamic]hir.Stmt_Id
+	body.allocator = ctx.checker.allocator
+	captures: [dynamic]hir.Expr_Id
+	captures.allocator = ctx.checker.allocator
+	for i := len(ctx.defers^) - 1; i >= 0; i -= 1 {
+		entry := ctx.defers^[i]
+		if entry.error_only && entry.capture != hir.INVALID_LOCAL {
+			append(&captures, hir.Expr_Id(entry.capture))
+		}
+		append(&body, ..entry.body)
+	}
+	return body[:], captures[:]
+}
+
 build_block :: proc(
 	ctx: ^Build_Ctx,
 	statements: []ast.Stmt_Id,
@@ -7144,6 +7206,8 @@ build_block :: proc(
 				continue
 			}
 			value := hir.INVALID_EXPR
+			error_exit := false
+			error_value := hir.INVALID_EXPR
 			if statement.value_control_flow {
 				if types.kind(ctx.result, store) == .Fallible {
 					success := types.fallible_success(ctx.result, store)
@@ -7157,13 +7221,12 @@ build_block :: proc(
 			} else if types.kind(ctx.result, store) == .Fallible {
 				success := types.fallible_success(ctx.result, store)
 				error_type := types.fallible_error(ctx.result, store)
-				error_path := false
 				expr_ast := checker.ast_module.exprs[statement.expr]
 				if expr_ast.kind == .Enum_Literal {
 					success_has := types.sum_has_name(store, success, u32(expr_ast.name))
 					error_has := types.sum_has_name(store, error_type, u32(expr_ast.name))
 					if error_has && !success_has {
-						error_path = true
+						error_exit = true
 					} else if error_has && success_has {
 						id := source.add(checker.diagnostics, expr_ast.span, "ambiguous fallible return member")
 						value = invalid_hir_expr(checker, expr_ast.span, id, ctx.result)
@@ -7172,7 +7235,7 @@ build_block :: proc(
 					target_pkg, available := expr_package(checker, expr_ast, ctx.pkg, ctx.file, true)
 					named := types.find_named(store, u32(target_pkg), u32(expr_ast.name), file=u32(expr_lookup_file(expr_ast, ctx.file))) if available else types.INVALID
 					named = types.resolve_alias(named, store)
-					error_path = can_implicitly_convert_type(checker, named, error_type)
+					error_exit = can_implicitly_convert_type(checker, named, error_type)
 				} else {
 					probe := build_expr(
 						checker, statement.expr, ctx.locals^[:], ctx.global_reads, ctx.calls,
@@ -7181,22 +7244,38 @@ build_block :: proc(
 					probe_type := checker.module.exprs[probe].type
 					if can_implicitly_convert_type(checker, probe_type, error_type) &&
 					   !can_implicitly_convert_type(checker, probe_type, success) {
-						error_path = true
+						error_exit = true
 						value = probe
 					} else if can_implicitly_convert_type(checker, probe_type, success) {
 						value = probe
 					}
 				}
 				if value == hir.INVALID_EXPR {
-					expected := error_type if error_path else success
+					expected := error_type if error_exit else success
 					value = build_expr(
 						checker, statement.expr, ctx.locals^[:], ctx.global_reads, ctx.calls,
 						expected, ctx.pkg, ctx.file,
 					)
 				}
-				expected := error_type if error_path else success
+				expected := error_type if error_exit else success
 				value = coerce_expr(checker, value, expected, statement.span)
-				value = fallible_aggregate(checker, statement.span, ctx.result, value, error_path)
+				if error_exit && len(ctx.defers^) > 0 && checker.module.exprs[value].kind != .Invalid {
+					tmp := append_tracked_local(
+						ctx.hir_locals, ctx.local_spans, ctx.local_used, ctx.local_warnable,
+						hir.Local{name=checker.sink_symbol, type=error_type, mutable=false}, source.Span{},
+					)
+					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,
+					})
+					error_value = hir.expr_id(len(checker.module.exprs))
+					append(&checker.module.exprs, hir.Expr{
+						kind=.Local, span=statement.span, type=error_type, target=hir.local_ref(tmp),
+					})
+					value = error_value
+				}
+				value = fallible_aggregate(checker, statement.span, ctx.result, value, error_exit)
 			} else {
 				value = build_expr(
 					checker, statement.expr, ctx.locals^[:], ctx.global_reads, ctx.calls,
@@ -7210,7 +7289,7 @@ build_block :: proc(
 			// 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 {
+				if !error_exit && checker.module.exprs[value].kind != .Invalid {
 					tmp := append_tracked_local(
 						ctx.hir_locals,
 						ctx.local_spans,
@@ -7229,7 +7308,7 @@ build_block :: proc(
 						kind = .Local, span = statement.span, type = ctx.result, target = hir.local_ref(tmp),
 					})
 				}
-				flush_defers(ctx, &body, 0)
+				flush_defers(ctx, &body, 0, error_exit, error_value)
 			}
 			append(&body, hir.stmt_id(len(checker.module.statements)))
 			append(&checker.module.statements, hir.Stmt{
@@ -7753,16 +7832,37 @@ build_block :: proc(
 				ctx.problematic^ = true
 				continue
 			}
+			if statement.error_only && types.kind(ctx.result, store) != .Fallible {
+				id := source.add(checker.diagnostics, statement.span, "'errdefer' requires an enclosing fallible function")
+				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.
+			capture_start := len(ctx.locals^)
+			capture := hir.INVALID_LOCAL
+			if statement.error_only && len(statement.captures) > 0 &&
+			   statement.captures[0] != checker.sink_symbol {
+				capture = append_build_local(
+					ctx, statement.captures[0], types.fallible_error(ctx.result, store), false, statement.span,
+				)
+			}
 			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)
+			resize(ctx.locals, capture_start)
+			append(ctx.defers, Defer_Entry{
+				body=entry, error_only=statement.error_only, capture=capture,
+			})
 		case .Match:
 			build_match(ctx, &body, statement)
 		case .Match_Arm:
@@ -7795,7 +7895,7 @@ build_block :: proc(
 		// 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)
+			delete(ctx.defers^[i].body, checker.allocator)
 		}
 		resize(ctx.defers, defer_start)
 		resize(ctx.locals, scope_start)
@@ -7887,7 +7987,7 @@ build_value_block :: proc(
 	}
 	// Close the scope (build_block left it open for us).
 	for i := defer_start; i < len(ctx.defers^); i += 1 {
-		delete(ctx.defers^[i], checker.allocator)
+		delete(ctx.defers^[i].body, checker.allocator)
 	}
 	resize(ctx.defers, defer_start)
 	resize(ctx.locals, scope_start)
@@ -8924,7 +9024,7 @@ block_element_type :: proc(ctx: ^Build_Ctx, block_stmts: []ast.Stmt_Id) -> types
 	result := checker.module.exprs[probe].type if checker.module.exprs[probe].kind != .Invalid else types.INVALID
 	// Discard the throwaway leading build's scope (its hir stmts/locals are dead but stable).
 	for i := defer_start; i < len(ctx.defers^); i += 1 {
-		delete(ctx.defers^[i], checker.allocator)
+		delete(ctx.defers^[i].body, checker.allocator)
 	}
 	resize(ctx.defers, defer_start)
 	resize(ctx.locals, scope_start)
@@ -9310,7 +9410,7 @@ build_function :: proc(checker: ^Checker, id: Spec_Id) {
 			},
 		)
 	}
-	defers: [dynamic][]hir.Stmt_Id
+	defers: [dynamic]Defer_Entry
 	defers.allocator = checker.allocator
 	loop_defer_starts: [dynamic]int
 	loop_defer_starts.allocator = checker.allocator
@@ -9392,7 +9492,7 @@ build_function :: proc(checker: ^Checker, id: Spec_Id) {
 		},
 	)
 	for entry in defers {
-		delete(entry, checker.allocator)
+		delete(entry.body, checker.allocator)
 	}
 	delete(defers)
 	delete(loop_defer_starts)