expand yield to if-statements and loops

compiler/checker/checker.odin

@@ -57,6 +57,18 @@ Build_Local :: struct {
 // nested control-flow blocks (if/else) can be built recursively. `locals` is a
 // scope stack: each block records its entry length and truncates back to it on
 // exit, while `hir_locals` keeps every allocated slot for the function.
+// A labeled value-loop currently being built. A `yield :label x` inside the loop
+// body assigns `x` to the loop's result `slot` (typed `slot_type`) and `break`s.
+// Pushed by `build_value_loop` while its body is built; innermost is last.
+Yield_Target :: struct {
+	label:           symbol.Id,
+	slot:            hir.Local_Id,
+	slot_type:       types.Type,
+	// True when the loop also yields `none` (a `{T, none}` set → `?T`); set from a
+	// pure-AST scan, used to pick the slot's element type on the first concrete yield.
+	result_optional: bool,
+}
+
 Build_Ctx :: struct {
 	checker:      ^Checker,
 	pkg:          ast.Package_Id,
@@ -68,6 +80,8 @@ Build_Ctx :: struct {
 	global_reads: ^[dynamic]hir.Global_Id,
 	calls:        ^[dynamic]hir.Function_Id,
 	problematic:  ^bool,
+	// Stack of labeled value-loops being built (innermost last); see Yield_Target.
+	yield_targets: ^[dynamic]Yield_Target,
 	// `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
@@ -4041,7 +4055,7 @@ build_block :: proc(
 				if typed {
 					expected = type_from_syntax(statement.type)
 				}
-				value, value_type := build_value_block(ctx, &body, statement.body, expected, statement.span)
+				value, value_type := build_value_source(ctx, &body, statement.body, expected, statement.span)
 				if _, found := find_build_local(ctx.locals^[duplicate_start:], statement.name); found {
 					id := source.addf(
 						checker.diagnostics, statement.span,
@@ -4252,7 +4266,7 @@ build_block :: proc(
 				} else if statement.expr == ast.INVALID_EXPR {
 					// `target = { ... yield v }`: build the value block against the
 					// target's type (build_value_block coerces internally).
-					value, _ = build_value_block(ctx, &body, statement.body, target_type, statement.span)
+					value, _ = build_value_source(ctx, &body, statement.body, target_type, statement.span)
 				} else {
 					value = build_expr(
 						checker, statement.expr, ctx.locals^[:], ctx.global_reads, ctx.calls,
@@ -4272,7 +4286,7 @@ build_block :: proc(
 			if statement.name == checker.sink_symbol {
 				value: hir.Expr_Id
 				if statement.expr == ast.INVALID_EXPR {
-					value, _ = build_value_block(ctx, &body, statement.body, types.INVALID, statement.span)
+					value, _ = build_value_source(ctx, &body, statement.body, types.INVALID, statement.span)
 				} else {
 					value = build_expr(checker, statement.expr, ctx.locals^[:], ctx.global_reads, ctx.calls, types.INVALID, ctx.pkg, ctx.file)
 				}
@@ -4316,7 +4330,7 @@ build_block :: proc(
 			}
 			value: hir.Expr_Id
 			if statement.expr == ast.INVALID_EXPR {
-				value, _ = build_value_block(ctx, &body, statement.body, local.type, statement.span)
+				value, _ = build_value_source(ctx, &body, statement.body, local.type, statement.span)
 			} else {
 				value = build_expr(
 					checker, statement.expr, ctx.locals^[:], ctx.global_reads, ctx.calls,
@@ -4740,13 +4754,71 @@ build_block :: proc(
 			}
 			delete(block, checker.allocator)
 		case .Yield:
-			// A legitimate yield is peeled off by build_value_block as the value
-			// block's final statement; reaching it here means it is misplaced
-			// (nested in an if/loop/inner block, or in a non-value block).
-			// ponytail: yield from if/loops/labeled blocks is a later milestone.
+			// A labeled `yield :blk x` exits the value-loop labeled `blk`: assign the
+			// result slot, then `break` (which flushes defers down to the loop body and
+			// branches to its exit). HIR holds no `.Yield` — it becomes Assignment + Break.
+			if symbol.is_valid(statement.label) {
+				target_index := -1
+				for i := len(ctx.yield_targets^) - 1; i >= 0; i -= 1 {
+					if ctx.yield_targets^[i].label == statement.label {
+						target_index = i
+						break
+					}
+				}
+				if target_index < 0 {
+					id := source.addf(checker.diagnostics, statement.span,
+						"no enclosing value loop is labeled '%s'", symbol_text(checker, statement.label))
+					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
+				}
+				// ponytail: a `yield :blk` lowers to milestone-18 `break`, which targets the
+				// innermost loop only, so the label must name the innermost value loop.
+				if target_index != len(ctx.yield_targets^) - 1 {
+					id := source.addf(checker.diagnostics, statement.span,
+						"'yield :%s' must target the innermost loop", symbol_text(checker, statement.label))
+					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
+				}
+				target := &ctx.yield_targets^[target_index]
+				yielded := build_expr(checker, statement.expr, ctx.locals^[:], ctx.global_reads, ctx.calls, target.slot_type, ctx.pkg, ctx.file)
+				yielded = resolve_loop_slot(ctx, target, yielded, checker.module.exprs[yielded].type if yielded != hir.INVALID_EXPR else types.INVALID, statement.span)
+				if target.slot == hir.INVALID_LOCAL || yielded == hir.INVALID_EXPR {
+					id := source.add(checker.diagnostics, statement.span,
+						"could not determine the value loop's yield type; annotate the binding or yield a concrete value first")
+					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
+				}
+				// slot = value (the slot is un-nameable, so no defer can mutate it; no spill).
+				emit_slot_assign(checker, &body, target.slot, yielded, statement.span)
+				// Exit the loop: flush defers down to the loop body, then break.
+				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, span = statement.span, expr = hir.INVALID_EXPR,
+					local = hir.INVALID_LOCAL, diagnostic = source.INVALID_DIAGNOSTIC,
+				})
+				continue
+			}
+			// An unlabeled yield reaching here is misplaced: a legitimate trailing yield
+			// is peeled by the value builders (value block / if branch / loop fall-through).
 			id := source.add(
 				checker.diagnostics, statement.span,
-				"'yield' is only valid as the final statement of a value block",
+				"'yield' is only valid as the final statement of a value block, or as 'yield :label' inside a labeled value loop",
 			)
 			append(&body, hir.stmt_id(len(checker.module.statements)))
 			append(&checker.module.statements, hir.Stmt{
@@ -4887,6 +4959,309 @@ build_value_block :: proc(
 	return value, value_type
 }
 
+// build_value_source feeds a declaration/assignment RHS into the right value builder:
+// a `{ ... }` block, an `if` whose branches yield, or a `for`/`while` whose iterations
+// yield. All three return the produced value and its type for the enclosing binding.
+build_value_source :: proc(
+	ctx: ^Build_Ctx,
+	body: ^[dynamic]hir.Stmt_Id,
+	body_stmts: []ast.Stmt_Id,
+	expected: types.Type,
+	span: source.Span,
+) -> (value: hir.Expr_Id, value_type: types.Type) {
+	checker := ctx.checker
+	if len(body_stmts) == 1 {
+		#partial switch checker.ast_module.statements[body_stmts[0]].kind {
+		case .If:
+			return build_value_if(ctx, body, body_stmts[0], expected, span)
+		case .For, .While:
+			return build_value_loop(ctx, body, body_stmts[0], expected, span)
+		}
+	}
+	return build_value_block(ctx, body, body_stmts, expected, span)
+}
+
+// new_value_slot allocates a fresh, un-nameable mutable local to hold a value-if/loop
+// result. Branches/iterations assign it; the construct's value is a read of it.
+new_value_slot :: proc(ctx: ^Build_Ctx, slot_type: types.Type) -> hir.Local_Id {
+	slot := hir.local_id(len(ctx.hir_locals^))
+	append(ctx.hir_locals, hir.Local{name = ctx.checker.sink_symbol, type = slot_type, mutable = true})
+	return slot
+}
+
+// slot_read builds a `.Local` read of a result slot.
+slot_read :: proc(checker: ^Checker, slot: hir.Local_Id, slot_type: types.Type, span: source.Span) -> hir.Expr_Id {
+	id := hir.expr_id(len(checker.module.exprs))
+	append(&checker.module.exprs, hir.Expr{
+		kind = .Local, span = span, type = slot_type, target = hir.local_ref(slot),
+	})
+	return id
+}
+
+// emit_slot_assign appends a bare-local `slot = value` assignment to `out`.
+emit_slot_assign :: proc(checker: ^Checker, out: ^[dynamic]hir.Stmt_Id, slot: hir.Local_Id, value: hir.Expr_Id, span: source.Span) {
+	append(out, hir.stmt_id(len(checker.module.statements)))
+	append(&checker.module.statements, hir.Stmt{
+		kind = .Assignment, span = span, expr = value, local = slot,
+		target = hir.INVALID_EXPR, diagnostic = source.INVALID_DIAGNOSTIC,
+	})
+}
+
+// build_value_if turns `if c { … yield A } else { … yield B }` into a result slot
+// each branch assigns, read after the if. Every path must yield: a mandatory `else`,
+// each branch ends in `yield`, and all branches share a type (the first establishes it
+// when untyped; later branches coerce). HIR holds an ordinary `.If` + a `.Local` read.
+build_value_if :: proc(
+	ctx: ^Build_Ctx,
+	body: ^[dynamic]hir.Stmt_Id,
+	if_id: ast.Stmt_Id,
+	expected: types.Type,
+	span: source.Span,
+) -> (value: hir.Expr_Id, value_type: types.Type) {
+	checker := ctx.checker
+	slot := hir.INVALID_LOCAL
+	slot_type := types.INVALID
+	if is_runtime_type(checker, expected) {
+		slot_type = expected
+		slot = new_value_slot(ctx, slot_type)
+	}
+	subtree: [dynamic]hir.Stmt_Id
+	subtree.allocator = checker.allocator
+	ok := emit_value_if(ctx, &subtree, if_id, &slot, &slot_type, span)
+	if !ok || slot == hir.INVALID_LOCAL {
+		for s in subtree {
+			append(body, s)
+		}
+		delete(subtree)
+		ctx.problematic^ = true
+		return invalid_hir_expr(checker, span, source.INVALID_DIAGNOSTIC), types.INVALID
+	}
+	// The slot's poison declaration precedes the if; every path assigns it.
+	append(body, hir.stmt_id(len(checker.module.statements)))
+	append(&checker.module.statements, hir.Stmt{
+		kind = .Declaration, span = span, local = slot, expr = hir.INVALID_EXPR,
+		diagnostic = source.INVALID_DIAGNOSTIC,
+	})
+	for s in subtree {
+		append(body, s)
+	}
+	delete(subtree)
+	value = slot_read(checker, slot, slot_type, span)
+	return value, slot_type
+}
+
+// emit_value_if builds one `if`/`else if`/`else` level of a value-if, appending the
+// assembled `.If` to `out`. `slot`/`slot_type` thread through so the first branch can
+// fix an untyped slot and `else if` chains share it.
+emit_value_if :: proc(
+	ctx: ^Build_Ctx,
+	out: ^[dynamic]hir.Stmt_Id,
+	if_id: ast.Stmt_Id,
+	slot: ^hir.Local_Id,
+	slot_type: ^types.Type,
+	span: source.Span,
+) -> bool {
+	checker := ctx.checker
+	if_stmt := checker.ast_module.statements[if_id]
+	// ponytail: an unwrap `if` (captures) as a value source is a later milestone.
+	if len(if_stmt.captures) > 0 {
+		source.add(checker.diagnostics, if_stmt.span, "an unwrap 'if' cannot yet be used as a value")
+		ctx.problematic^ = true
+		return false
+	}
+	condition := build_expr(checker, if_stmt.expr, ctx.locals^[:], ctx.global_reads, ctx.calls, types.BOOL, ctx.pkg, ctx.file)
+	if checker.module.exprs[condition].kind != .Invalid && !types.is_bool(checker.module.exprs[condition].type) {
+		id := source.add(checker.diagnostics, if_stmt.span, "'if' condition must be a bool")
+		condition = invalid_hir_expr(checker, if_stmt.span, id, types.BOOL)
+		ctx.problematic^ = true
+	}
+	then_body: [dynamic]hir.Stmt_Id
+	then_body.allocator = checker.allocator
+	if !emit_value_branch(ctx, &then_body, if_stmt.body, slot, slot_type, span) {
+		delete(then_body)
+		return false
+	}
+	if if_stmt.else_body == nil {
+		source.add(checker.diagnostics, if_stmt.span, "an 'if' used as a value must have an 'else' so every path yields")
+		delete(then_body)
+		ctx.problematic^ = true
+		return false
+	}
+	else_body: [dynamic]hir.Stmt_Id
+	else_body.allocator = checker.allocator
+	branch_ok := true
+	if len(if_stmt.else_body) == 1 && checker.ast_module.statements[if_stmt.else_body[0]].kind == .If {
+		branch_ok = emit_value_if(ctx, &else_body, if_stmt.else_body[0], slot, slot_type, span)
+	} else {
+		branch_ok = emit_value_branch(ctx, &else_body, if_stmt.else_body, slot, slot_type, span)
+	}
+	if !branch_ok {
+		delete(then_body)
+		delete(else_body)
+		return false
+	}
+	append(out, hir.stmt_id(len(checker.module.statements)))
+	append(&checker.module.statements, hir.Stmt{
+		kind = .If, span = if_stmt.span, expr = condition, guard = hir.INVALID_EXPR,
+		then_body = then_body[:], else_body = else_body[:],
+		local = hir.INVALID_LOCAL, target = hir.INVALID_EXPR, diagnostic = source.INVALID_DIAGNOSTIC,
+	})
+	return true
+}
+
+// emit_value_branch builds one branch of a value-if as a value block (leading stmts +
+// trailing yield) and appends `slot = <value>`. The first branch of an untyped value-if
+// fixes the slot type; later branches coerce to it (a mismatch is the "same type" error).
+emit_value_branch :: proc(
+	ctx: ^Build_Ctx,
+	out: ^[dynamic]hir.Stmt_Id,
+	branch_stmts: []ast.Stmt_Id,
+	slot: ^hir.Local_Id,
+	slot_type: ^types.Type,
+	span: source.Span,
+) -> bool {
+	checker := ctx.checker
+	value, vtype := build_value_block(ctx, out, branch_stmts, slot_type^, span)
+	if checker.module.exprs[value].kind == .Invalid {
+		return false
+	}
+	if slot^ == hir.INVALID_LOCAL {
+		slot_type^ = vtype
+		slot^ = new_value_slot(ctx, slot_type^)
+	} else {
+		value = coerce_expr(checker, value, slot_type^, span)
+	}
+	emit_slot_assign(checker, out, slot^, value, span)
+	return true
+}
+
+// loop_yields_none reports whether any `yield` that targets this loop (a labeled
+// `yield :blk` inside `if`/block branches, or the trailing fall-through) yields the
+// literal `none` — making the loop's result optional. Pure AST walk; does not descend
+// into nested loops or value sources, whose yields belong to them.
+loop_yields_none :: proc(checker: ^Checker, stmts: []ast.Stmt_Id) -> bool {
+	for id in stmts {
+		s := checker.ast_module.statements[id]
+		#partial switch s.kind {
+		case .Yield:
+			if s.expr != ast.INVALID_EXPR && checker.ast_module.exprs[s.expr].kind == .None {
+				return true
+			}
+		case .If:
+			if loop_yields_none(checker, s.body) || loop_yields_none(checker, s.else_body) {
+				return true
+			}
+		case .Block:
+			if loop_yields_none(checker, s.body) {
+				return true
+			}
+		}
+	}
+	return false
+}
+
+// resolve_loop_slot fixes a value-loop's result slot from its first concrete yield (an
+// optional element type when the loop also yields `none`) and coerces `value` into it.
+// Returns INVALID when the type can't be fixed yet (a `none`/invalid first yield).
+resolve_loop_slot :: proc(ctx: ^Build_Ctx, target: ^Yield_Target, value: hir.Expr_Id, vtype: types.Type, span: source.Span) -> hir.Expr_Id {
+	checker := ctx.checker
+	if target.slot == hir.INVALID_LOCAL {
+		if !is_runtime_type(checker, vtype) {
+			return hir.INVALID_EXPR
+		}
+		target.slot_type = types.optional(&checker.module.types, vtype) if target.result_optional else vtype
+		target.slot = new_value_slot(ctx, target.slot_type)
+	}
+	return coerce_expr(checker, value, target.slot_type, span)
+}
+
+// build_value_loop turns a labeled `for/while ... blk: { … }` whose body ends in a
+// fall-through `yield` (and may early-exit via `yield :blk x`) into a result slot:
+// the fall-through value initializes the slot before the loop, each `yield :blk x`
+// desugars (in build_block) to `slot = x; break`, and the construct's value is a read
+// of the slot after the loop. Reuses the ordinary `.For`/`.While` build via a peeled
+// copy; no new HIR. The yielded type is the annotation when typed, else the first
+// concrete yield's type (optional when any yield is `none`).
+build_value_loop :: proc(
+	ctx: ^Build_Ctx,
+	body: ^[dynamic]hir.Stmt_Id,
+	loop_id: ast.Stmt_Id,
+	expected: types.Type,
+	span: source.Span,
+) -> (value: hir.Expr_Id, value_type: types.Type) {
+	checker := ctx.checker
+	loop_stmt := checker.ast_module.statements[loop_id]
+	n := len(loop_stmt.body)
+	last_is_fallthrough := n > 0 &&
+		checker.ast_module.statements[loop_stmt.body[n - 1]].kind == .Yield &&
+		!symbol.is_valid(checker.ast_module.statements[loop_stmt.body[n - 1]].label)
+	if !symbol.is_valid(loop_stmt.label) {
+		id := source.add(checker.diagnostics, span,
+			"a value loop must label its body (e.g. 'blk:') so a 'yield :blk' can exit it")
+		ctx.problematic^ = true
+		return invalid_hir_expr(checker, span, id), types.INVALID
+	}
+	if !last_is_fallthrough {
+		id := source.add(checker.diagnostics, span,
+			"a value loop's body must end with a 'yield' for when the loop completes")
+		ctx.problematic^ = true
+		return invalid_hir_expr(checker, span, id), types.INVALID
+	}
+	fall_stmt := checker.ast_module.statements[loop_stmt.body[n - 1]]
+
+	result_optional := loop_yields_none(checker, loop_stmt.body)
+	slot := hir.INVALID_LOCAL
+	slot_type := types.INVALID
+	if is_runtime_type(checker, expected) {
+		slot_type = expected
+		slot = new_value_slot(ctx, slot_type)
+		result_optional = types.is_optional(slot_type, &checker.module.types)
+	}
+	append(ctx.yield_targets, Yield_Target{
+		label = loop_stmt.label, slot = slot, slot_type = slot_type, result_optional = result_optional,
+	})
+
+	// Build the loop with the fall-through peeled off, reusing the normal For/While arm.
+	// The peeled body is a fresh copy so destroy_module won't double-free the original.
+	peeled := loop_stmt
+	peeled_body := make([]ast.Stmt_Id, n - 1, checker.ast_module.allocator)
+	copy(peeled_body, loop_stmt.body[:n - 1])
+	peeled.body = peeled_body
+	peeled_id := ast.stmt_id(len(checker.ast_module.statements))
+	append(&checker.ast_module.statements, peeled)
+	loop_block := build_block(ctx, []ast.Stmt_Id{peeled_id})
+
+	target := pop(ctx.yield_targets)
+
+	// The fall-through value initializes the slot before the loop (loop captures are
+	// out of scope here), so the loop completing leaves it as the result.
+	fall_value := build_expr(checker, fall_stmt.expr, ctx.locals^[:], ctx.global_reads, ctx.calls, target.slot_type, ctx.pkg, ctx.file)
+	fall_value = resolve_loop_slot(ctx, &target, fall_value, checker.module.exprs[fall_value].type if fall_value != hir.INVALID_EXPR else types.INVALID, fall_stmt.span)
+	if target.slot == hir.INVALID_LOCAL || fall_value == hir.INVALID_EXPR {
+		for s in loop_block {
+			append(body, s)
+		}
+		delete(loop_block, checker.allocator)
+		id := source.add(checker.diagnostics, span,
+			"could not determine the value loop's yield type; annotate the binding")
+		ctx.problematic^ = true
+		return invalid_hir_expr(checker, span, id), types.INVALID
+	}
+
+	append(body, hir.stmt_id(len(checker.module.statements)))
+	append(&checker.module.statements, hir.Stmt{
+		kind = .Declaration, span = span, local = target.slot, expr = fall_value,
+		diagnostic = source.INVALID_DIAGNOSTIC,
+	})
+	for s in loop_block {
+		append(body, s)
+	}
+	delete(loop_block, checker.allocator)
+
+	value = slot_read(checker, target.slot, target.slot_type, span)
+	return value, target.slot_type
+}
+
 // Reports whether every control-flow path through `stmts` terminates (returns or traps),
 // so the end of the block is unreachable. A `.Return` or `.Trap` terminates outright; an
 // `.If` terminates only when it has an `else` and both arms terminate. A literal
@@ -5069,6 +5444,8 @@ build_function :: proc(checker: ^Checker, id: Spec_Id) {
 	defers.allocator = checker.allocator
 	loop_defer_starts: [dynamic]int
 	loop_defer_starts.allocator = checker.allocator
+	yield_targets: [dynamic]Yield_Target
+	yield_targets.allocator = checker.allocator
 	ctx := Build_Ctx{
 		checker = checker,
 		pkg = function.pkg,
@@ -5082,6 +5459,7 @@ build_function :: proc(checker: ^Checker, id: Spec_Id) {
 		problematic = &problematic,
 		defers = &defers,
 		loop_defer_starts = &loop_defer_starts,
+		yield_targets = &yield_targets,
 	}
 	block := build_block(&ctx, function.body)
 	returns := all_paths_return(&checker.module, block)
@@ -5131,6 +5509,7 @@ build_function :: proc(checker: ^Checker, id: Spec_Id) {
 	}
 	delete(defers)
 	delete(loop_defer_starts)
+	delete(yield_targets)
 	delete(locals)
 }