fix comptime yield targeting

compiler/checker/comptime.odin

@@ -318,6 +318,7 @@ Ct_State :: struct {
 	bindings:   [dynamic]Ct_Binding,
 	error_refinements: [dynamic]Ct_Error_Refinement,
 	defers:     [dynamic]Ct_Defer,
+	yield_targets: [dynamic]symbol.Id,
 	defer_depth: int,
 	value_return_depth: int,
 	steps:      int,
@@ -361,6 +362,7 @@ ct_state_make :: proc(
 	state.bindings.allocator = checker.allocator
 	state.error_refinements.allocator = checker.allocator
 	state.defers.allocator = checker.allocator
+	state.yield_targets.allocator = checker.allocator
 	state.promoted_cells.allocator = checker.allocator
 	state.promoted_globals.allocator = checker.allocator
 	for value in values {
@@ -401,6 +403,7 @@ ct_state_destroy :: proc(state: ^Ct_State) {
 	delete(state.bindings)
 	delete(state.error_refinements)
 	delete(state.defers)
+	delete(state.yield_targets)
 	delete(state.promoted_cells)
 	delete(state.promoted_globals)
 }
@@ -1271,7 +1274,7 @@ ct_eval_expr :: proc(
 		if expr.left != ast.INVALID_EXPR {
 			return ct_eval_expr(state, expr.left, expected, depth+1)
 		}
-		flow, ok := ct_exec_statements(state, expr.body, true, depth+1)
+		flow, ok := ct_exec_value_source(state, expr.body, symbol.INVALID, false, false, depth+1)
 		if ok && flow.kind == .Yield {
 			return flow.value, ct_flow(.Normal), true
 		}
@@ -3934,7 +3937,9 @@ ct_eval_catch_expr :: proc(state: ^Ct_State, expr: ast.Expr, expected: types.Typ
 		return result, result_flow, result_ok
 	}
 	state.value_return_depth += 1
-	handler, handler_ok := ct_exec_statements(state, expr.body, true, depth+1)
+	handler, handler_ok := ct_exec_value_source(
+		state, expr.body, symbol.INVALID, expr.integer != 0, true, depth+1,
+	)
 	state.value_return_depth -= 1
 	ct_pop_bindings(state, scope_start)
 	if !handler_ok {
@@ -4123,6 +4128,397 @@ ct_return_value :: proc(state: ^Ct_State, expr_id: ast.Expr_Id, span: source.Spa
 	return ct_flow(.Return, value), true
 }
 
+ct_has_yield_target :: proc(state: ^Ct_State, label: symbol.Id) -> bool {
+	for index := len(state.yield_targets) - 1; index >= 0; index -= 1 {
+		if state.yield_targets[index] == label {
+			return true
+		}
+	}
+	return false
+}
+
+ct_statements_exit_value_source :: proc(state: ^Ct_State, statements: []ast.Stmt_Id) -> bool {
+	checker := state.checker
+	for statement_id in statements {
+		if statement_id == ast.INVALID_STMT || int(statement_id) >= len(checker.ast_module.statements) {
+			return false
+		}
+		statement := checker.ast_module.statements[statement_id]
+		#partial switch statement.kind {
+		case .Return, .Break, .Continue:
+			return true
+		case .Yield:
+			if symbol.is_valid(statement.label) && ct_has_yield_target(state, statement.label) {
+				return true
+			}
+		case .If:
+			if statement.else_body != nil &&
+			   ct_statements_exit_value_source(state, statement.body) &&
+			   ct_statements_exit_value_source(state, statement.else_body) {
+				return true
+			}
+		case .Block:
+			if ct_statements_exit_value_source(state, statement.body) {
+				return true
+			}
+		case .Match:
+			exits := len(statement.body) > 0
+			for arm_id in statement.body {
+				arm := checker.ast_module.statements[arm_id]
+				exits = exits && arm.kind == .Match_Arm && ct_statements_exit_value_source(state, arm.body)
+			}
+			if exits {
+				return true
+			}
+		}
+	}
+	return false
+}
+
+ct_validate_ordinary_statements :: proc(state: ^Ct_State, statements: []ast.Stmt_Id) -> bool {
+	checker := state.checker
+	for statement_id in statements {
+		if statement_id == ast.INVALID_STMT || int(statement_id) >= len(checker.ast_module.statements) {
+			return false
+		}
+		statement := checker.ast_module.statements[statement_id]
+		#partial switch statement.kind {
+		case .Declaration, .Assignment:
+			if statement.expr == ast.INVALID_EXPR && !ct_validate_value_source_structure(
+				state, statement.body, statement.label, statement.value_control_flow, false,
+			) {
+				return false
+			}
+		case .Return:
+			if statement.value_control_flow && !ct_validate_value_source_structure(
+				state, statement.body, symbol.INVALID, true, false,
+			) {
+				return false
+			}
+		case .Yield:
+			if !symbol.is_valid(statement.label) {
+				return ct_fail(
+					state, .Not_Comptime, statement.span,
+					"'yield' is only valid as the final statement of a value block",
+				)
+			}
+			if !ct_has_yield_target(state, statement.label) {
+				return ct_failf(
+					state, .Not_Comptime, statement.span,
+					"no enclosing value loop or block is labeled '%s'",
+					symbol_text(checker, statement.label),
+				)
+			}
+			if statement.value_control_flow && !ct_validate_value_source_structure(
+				state, statement.body, symbol.INVALID, true, false,
+			) {
+				return false
+			}
+		case .If:
+			if !ct_validate_ordinary_statements(state, statement.body) ||
+			   !ct_validate_ordinary_statements(state, statement.else_body) {
+				return false
+			}
+		case .While, .For:
+			if !ct_validate_ordinary_statements(state, statement.body) {
+				return false
+			}
+			if statement.update != ast.INVALID_STMT {
+				update := [1]ast.Stmt_Id{statement.update}
+				if !ct_validate_ordinary_statements(state, update[:]) {
+					return false
+				}
+			}
+		case .Block:
+			if !ct_validate_ordinary_statements(state, statement.body) {
+				return false
+			}
+		case .Defer:
+			if statement.update != ast.INVALID_STMT {
+				deferred := [1]ast.Stmt_Id{statement.update}
+				if !ct_validate_ordinary_statements(state, deferred[:]) {
+					return false
+				}
+			}
+		case .Match:
+			for arm_id in statement.body {
+				arm := checker.ast_module.statements[arm_id]
+				if arm.kind == .Match_Arm && !ct_validate_ordinary_statements(state, arm.body) {
+					return false
+				}
+			}
+		}
+	}
+	return true
+}
+
+ct_validate_value_branch_structure :: proc(state: ^Ct_State, statements: []ast.Stmt_Id) -> bool {
+	checker := state.checker
+	if len(statements) == 1 && checker.ast_module.statements[statements[0]].kind == .Expression {
+		return true
+	}
+	n := len(statements)
+	if n > 0 {
+		last := checker.ast_module.statements[statements[n-1]]
+		if last.kind == .Yield && !symbol.is_valid(last.label) {
+			if !ct_validate_ordinary_statements(state, statements[:n-1]) {
+				return false
+			}
+			return !last.value_control_flow || ct_validate_value_source_structure(
+				state, last.body, symbol.INVALID, true, false,
+			)
+		}
+	}
+	if ct_statements_exit_value_source(state, statements) {
+		return ct_validate_ordinary_statements(state, statements)
+	}
+	return ct_fail(
+		state, .Not_Comptime, source.Span{},
+		"a value branch must end with 'yield' or exit on every path (return/break/continue)",
+	)
+}
+
+ct_validate_value_if_structure :: proc(state: ^Ct_State, statement: ast.Stmt) -> bool {
+	checker := state.checker
+	if statement.else_body == nil {
+		return ct_fail(
+			state, .Not_Comptime, statement.span,
+			"an 'if' used as a value must have an 'else' so every path yields",
+		)
+	}
+	if !ct_validate_value_branch_structure(state, statement.body) {
+		return false
+	}
+	if len(statement.else_body) == 1 {
+		else_statement := checker.ast_module.statements[statement.else_body[0]]
+		if else_statement.kind == .If {
+			return ct_validate_value_if_structure(state, else_statement)
+		}
+	}
+	return ct_validate_value_branch_structure(state, statement.else_body)
+}
+
+ct_validate_value_source_structure :: proc(
+	state: ^Ct_State,
+	statements: []ast.Stmt_Id,
+	label: symbol.Id,
+	value_control_flow: bool,
+	allow_exit: bool,
+) -> bool {
+	checker := state.checker
+	if symbol.is_valid(label) {
+		target_start := len(state.yield_targets)
+		append(&state.yield_targets, label)
+		valid := ct_validate_ordinary_statements(state, statements) &&
+		         ct_statements_exit_value_source(state, statements)
+		resize(&state.yield_targets, target_start)
+		if valid {
+			return true
+		}
+		if state.error == .None {
+			return ct_fail(state, .Not_Comptime, source.Span{}, "a labeled value block must 'yield' on every path")
+		}
+		return false
+	}
+	if value_control_flow && len(statements) == 1 {
+		statement := checker.ast_module.statements[statements[0]]
+		#partial switch statement.kind {
+		case .If:
+			return ct_validate_value_if_structure(state, statement)
+		case .For, .While:
+			n := len(statement.body)
+			if !symbol.is_valid(statement.label) {
+				return ct_fail(
+					state, .Not_Comptime, statement.span,
+					"a value loop must label its body (e.g. 'blk:') so a 'yield :blk' can exit it",
+				)
+			}
+			if n == 0 {
+				return ct_fail(
+					state, .Not_Comptime, statement.span,
+					"a value loop's body must end with a 'yield' for when the loop completes",
+				)
+			}
+			last := checker.ast_module.statements[statement.body[n-1]]
+			if last.kind != .Yield || symbol.is_valid(last.label) {
+				return ct_fail(
+					state, .Not_Comptime, statement.span,
+					"a value loop's body must end with a 'yield' for when the loop completes",
+				)
+			}
+			target_start := len(state.yield_targets)
+			append(&state.yield_targets, statement.label)
+			valid := ct_validate_ordinary_statements(state, statement.body[:n-1])
+			if valid && statement.update != ast.INVALID_STMT {
+				update := [1]ast.Stmt_Id{statement.update}
+				valid = ct_validate_ordinary_statements(state, update[:])
+			}
+			resize(&state.yield_targets, target_start)
+			if valid && last.value_control_flow {
+				valid = ct_validate_value_source_structure(
+					state, last.body, symbol.INVALID, true, false,
+				)
+			}
+			return valid
+		case .Match:
+			for arm_id in statement.body {
+				arm := checker.ast_module.statements[arm_id]
+				if arm.kind == .Match_Arm && !ct_validate_value_branch_structure(state, arm.body) {
+					return false
+				}
+			}
+			return true
+		}
+	}
+	n := len(statements)
+	if n > 0 {
+		last := checker.ast_module.statements[statements[n-1]]
+		if last.kind == .Yield && !symbol.is_valid(last.label) {
+			if !ct_validate_ordinary_statements(state, statements[:n-1]) {
+				return false
+			}
+			return !last.value_control_flow || ct_validate_value_source_structure(
+				state, last.body, symbol.INVALID, true, false,
+			)
+		}
+	}
+	if allow_exit {
+		return ct_validate_ordinary_statements(state, statements)
+	}
+	return ct_fail(state, .Not_Comptime, source.Span{}, "a value block must end with an explicit 'yield'")
+}
+
+ct_exec_value_branch :: proc(state: ^Ct_State, statements: []ast.Stmt_Id, depth: int) -> (Ct_Flow, bool) {
+	checker := state.checker
+	if len(statements) == 1 {
+		statement := checker.ast_module.statements[statements[0]]
+		if statement.kind == .Expression {
+			value, flow, ok := ct_eval_expr(state, statement.expr, types.INVALID, depth+1)
+			if !ok || flow.kind != .Normal {
+				return flow, ok
+			}
+			return ct_flow(.Yield, value), true
+		}
+	}
+	flow, ok := ct_exec_statements(state, statements, true, depth+1)
+	if ok && flow.kind == .Normal {
+		return flow, ct_fail(
+			state, .Not_Comptime, source.Span{},
+			"a value branch must end with 'yield' or exit on every path (return/break/continue)",
+		)
+	}
+	return flow, ok
+}
+
+ct_exec_value_loop :: proc(state: ^Ct_State, statement: ast.Stmt, depth: int) -> (Ct_Flow, bool) {
+	checker := state.checker
+	n := len(statement.body)
+	if !symbol.is_valid(statement.label) {
+		return ct_flow(.Normal), ct_fail(
+			state, .Not_Comptime, statement.span,
+			"a value loop must label its body (e.g. 'blk:') so a 'yield :blk' can exit it",
+		)
+	}
+	if n == 0 {
+		return ct_flow(.Normal), ct_fail(
+			state, .Not_Comptime, statement.span,
+			"a value loop's body must end with a 'yield' for when the loop completes",
+		)
+	}
+	fall_stmt := checker.ast_module.statements[statement.body[n-1]]
+	if fall_stmt.kind != .Yield || symbol.is_valid(fall_stmt.label) {
+		return ct_flow(.Normal), ct_fail(
+			state, .Not_Comptime, statement.span,
+			"a value loop's body must end with a 'yield' for when the loop completes",
+		)
+	}
+	fallback := ct_flow(.Normal)
+	fallback_ok := true
+	if fall_stmt.value_control_flow {
+		fallback, fallback_ok = ct_exec_value_source(
+			state, fall_stmt.body, symbol.INVALID, true, false, depth+1,
+		)
+	} else {
+		value, flow, ok := ct_eval_expr(state, fall_stmt.expr, types.INVALID, depth+1)
+		fallback_ok = ok
+		fallback = flow
+		if ok && flow.kind == .Normal {
+			fallback = ct_flow(.Yield, value)
+		}
+	}
+	if !fallback_ok || fallback.kind != .Yield {
+		return fallback, fallback_ok
+	}
+	peeled := statement
+	peeled.body = statement.body[:n-1]
+	target_start := len(state.yield_targets)
+	append(&state.yield_targets, statement.label)
+	defer resize(&state.yield_targets, target_start)
+	flow := ct_flow(.Normal)
+	ok := false
+	if statement.kind == .For {
+		flow, ok = ct_exec_for(state, peeled, false, depth+1)
+	} else {
+		flow, ok = ct_exec_while(state, peeled, false, depth+1)
+	}
+	if !ok {
+		return flow, false
+	}
+	if flow.kind == .Yield && flow.label == statement.label {
+		return ct_flow(.Yield, flow.value), true
+	}
+	if flow.kind == .Normal {
+		return fallback, true
+	}
+	return flow, true
+}
+
+ct_exec_value_source :: proc(
+	state: ^Ct_State,
+	statements: []ast.Stmt_Id,
+	label: symbol.Id,
+	value_control_flow: bool,
+	allow_exit: bool,
+	depth: int,
+) -> (Ct_Flow, bool) {
+	checker := state.checker
+	if !ct_validate_value_source_structure(state, statements, label, value_control_flow, allow_exit) {
+		return ct_flow(.Normal), false
+	}
+	if symbol.is_valid(label) {
+		target_start := len(state.yield_targets)
+		append(&state.yield_targets, label)
+		flow, ok := ct_exec_statements(state, statements, false, depth+1)
+		resize(&state.yield_targets, target_start)
+		if !ok {
+			return flow, false
+		}
+		if flow.kind == .Yield && flow.label == label {
+			return ct_flow(.Yield, flow.value), true
+		}
+		if flow.kind == .Normal {
+			return flow, ct_fail(state, .Not_Comptime, source.Span{}, "a labeled value block must 'yield' on every path")
+		}
+		return flow, true
+	}
+	if value_control_flow && len(statements) == 1 {
+		statement := checker.ast_module.statements[statements[0]]
+		#partial switch statement.kind {
+		case .If:
+			return ct_exec_if(state, statement, true, depth+1)
+		case .For, .While:
+			return ct_exec_value_loop(state, statement, depth+1)
+		case .Match:
+			return ct_exec_match(state, statement, true, depth+1)
+		}
+	}
+	flow, ok := ct_exec_statements(state, statements, true, depth+1)
+	if ok && flow.kind == .Normal && !allow_exit {
+		return flow, ct_fail(state, .Not_Comptime, source.Span{}, "a value block must end with an explicit 'yield'")
+	}
+	return flow, ok
+}
+
 ct_exec_statements :: proc(
 	state: ^Ct_State,
 	statements: []ast.Stmt_Id,
@@ -4139,7 +4535,7 @@ ct_exec_statements :: proc(
 		ct_pop_bindings(state, scope_start)
 		resize(&state.defers, defer_start)
 	}
-	for statement_id in statements {
+	for statement_id, statement_index in statements {
 		if statement_id == ast.INVALID_STMT || int(statement_id) >= len(checker.ast_module.statements) {
 			return ct_flow(.Normal), false
 		}
@@ -4152,7 +4548,9 @@ ct_exec_statements :: proc(
 		#partial switch statement.kind {
 		case .Declaration:
 			if statement.expr == ast.INVALID_EXPR {
-				value_flow, value_ok := ct_exec_statements(state, statement.body, true, depth+1)
+				value_flow, value_ok := ct_exec_value_source(
+					state, statement.body, statement.label, statement.value_control_flow, false, depth+1,
+				)
 				ok = value_ok
 				if ok && value_flow.kind == .Yield {
 					value := value_flow.value
@@ -4163,6 +4561,8 @@ ct_exec_statements :: proc(
 					if ok && statement.name != checker.sink_symbol {
 						ct_bind_value(state, statement.name, state.values[value].type, value, !statement.immutable)
 					}
+				} else if ok && value_flow.kind != .Normal {
+					flow = value_flow
 				} else if ok {
 					ok = ct_fail(state, .Not_Comptime, statement.span, "comptime value block must yield")
 				}
@@ -4189,7 +4589,9 @@ ct_exec_statements :: proc(
 			if yield_returns && state.value_return_depth == 0 {
 				ok = ct_fail(state, .Not_Comptime, statement.span, "'return' is not valid in this comptime block")
 			} else if statement.value_control_flow {
-				value_flow, value_ok := ct_exec_statements(state, statement.body, true, depth+1)
+				value_flow, value_ok := ct_exec_value_source(
+					state, statement.body, symbol.INVALID, true, false, depth+1,
+				)
 				ok = value_ok
 				if ok && value_flow.kind == .Yield {
 					value := value_flow.value
@@ -4205,18 +4607,37 @@ ct_exec_statements :: proc(
 							flow = ct_flow(.Return, value)
 						}
 					}
+				} else if ok {
+					flow = value_flow
 				}
 			} else {
 				flow, ok = ct_return_value(state, statement.expr, statement.span, depth+1)
 			}
 		case .Yield:
-			if !yield_returns {
-				ok = ct_fail(state, .Not_Comptime, statement.span, "'yield' is only valid in a comptime value block")
-			} else if statement.value_control_flow {
-				flow, ok = ct_exec_statements(state, statement.body, true, depth+1)
-				if ok && flow.kind == .Yield && types.is_void(state.values[flow.value].type) {
+			if symbol.is_valid(statement.label) {
+				if !ct_has_yield_target(state, statement.label) {
+					ok = ct_failf(
+						state, .Not_Comptime, statement.span,
+						"no enclosing value loop or block is labeled '%s'",
+						symbol_text(checker, statement.label),
+					)
+				}
+			} else if !yield_returns || statement_index != len(statements)-1 {
+				ok = ct_fail(state, .Not_Comptime, statement.span, "'yield' is only valid as the final statement of a value block")
+			}
+			if ok {
+				if statement.value_control_flow {
+					value_flow, value_ok := ct_exec_value_source(
+						state, statement.body, symbol.INVALID, true, false, depth+1,
+					)
+					ok = value_ok
+					flow = value_flow
+					if ok && value_flow.kind == .Yield {
+						flow.label = statement.label
+						if types.is_void(state.values[flow.value].type) {
 							ok = ct_fail(state, .Not_Comptime, statement.span, "'yield' expression must produce a non-void value")
 						}
+					}
 				} else {
 					value, expr_flow, expr_ok := ct_eval_expr(state, statement.expr, types.INVALID, depth+1)
 					ok = expr_ok
@@ -4229,18 +4650,19 @@ ct_exec_statements :: proc(
 						}
 					}
 				}
+			}
 		case .If:
-			flow, ok = ct_exec_if(state, statement, yield_returns, depth+1)
+			flow, ok = ct_exec_if(state, statement, false, depth+1)
 		case .While:
-			flow, ok = ct_exec_while(state, statement, yield_returns, depth+1)
+			flow, ok = ct_exec_while(state, statement, false, depth+1)
 		case .For:
-			flow, ok = ct_exec_for(state, statement, yield_returns, depth+1)
+			flow, ok = ct_exec_for(state, statement, false, depth+1)
 		case .Break:
 			flow = ct_flow(.Break, INVALID_CT_VALUE, statement.label)
 		case .Continue:
 			flow = ct_flow(.Continue, INVALID_CT_VALUE, statement.label)
 		case .Block:
-			flow, ok = ct_exec_statements(state, statement.body, yield_returns, depth+1)
+			flow, ok = ct_exec_statements(state, statement.body, false, depth+1)
 		case .Defer:
 			if statement.error_only && types.kind(state.result, &checker.module.types) != .Fallible {
 				ok = ct_fail(state, .Not_Comptime, statement.span, "'errdefer' requires an enclosing fallible function")
@@ -4254,7 +4676,7 @@ ct_exec_statements :: proc(
 				})
 			}
 		case .Match:
-			flow, ok = ct_exec_match(state, statement, yield_returns, depth+1)
+			flow, ok = ct_exec_match(state, statement, false, depth+1)
 		case .Match_Arm:
 			ok = ct_fail(state, .Not_Comptime, statement.span, "unexpected match arm outside 'match'")
 		case .Invalid:
@@ -4327,8 +4749,13 @@ ct_exec_assignment :: proc(state: ^Ct_State, statement: ast.Stmt, depth: int) ->
 	}
 	if name == checker.sink_symbol {
 		if statement.expr == ast.INVALID_EXPR {
-			flow, ok := ct_exec_statements(state, statement.body, true, depth+1)
-			return ct_flow(.Normal), ok && (flow.kind == .Yield || flow.kind == .Normal)
+			flow, ok := ct_exec_value_source(
+				state, statement.body, statement.label, statement.value_control_flow, false, depth+1,
+			)
+			if ok && flow.kind == .Yield {
+				return ct_flow(.Normal), true
+			}
+			return flow, ok
 		}
 		_, flow, ok := ct_eval_expr(state, statement.expr, types.INVALID, depth+1)
 		return flow, ok
@@ -4359,7 +4786,9 @@ ct_exec_assignment :: proc(state: ^Ct_State, statement: ast.Stmt, depth: int) ->
 	flow := ct_flow(.Normal)
 	ok := true
 	if statement.expr == ast.INVALID_EXPR {
-		flow, ok = ct_exec_statements(state, statement.body, true, depth+1)
+		flow, ok = ct_exec_value_source(
+			state, statement.body, statement.label, statement.value_control_flow, false, depth+1,
+		)
 		if ok && flow.kind == .Yield {
 			value = flow.value
 			flow = ct_flow(.Normal)
@@ -4411,6 +4840,12 @@ ct_exec_assignment :: proc(state: ^Ct_State, statement: ast.Stmt, depth: int) ->
 
 ct_exec_if :: proc(state: ^Ct_State, statement: ast.Stmt, yield_returns: bool, depth: int) -> (Ct_Flow, bool) {
 	checker := state.checker
+	if yield_returns && statement.else_body == nil {
+		return ct_flow(.Normal), ct_fail(
+			state, .Not_Comptime, statement.span,
+			"an 'if' used as a value must have an 'else' so every path yields",
+		)
+	}
 	if len(statement.captures) == 0 {
 		condition, flow, ok := ct_eval_expr(state, statement.expr, types.BOOL, depth+1)
 		if !ok || flow.kind != .Normal {
@@ -4421,7 +4856,13 @@ ct_exec_if :: proc(state: ^Ct_State, statement: ast.Stmt, yield_returns: bool, d
 			return ct_flow(.Normal), ct_fail(state, .Not_Comptime, statement.span, "'if' condition must be a bool")
 		}
 		body := statement.body if value else statement.else_body
-		return ct_exec_statements(state, body, yield_returns, depth+1)
+		if yield_returns {
+			if !value && len(body) == 1 && checker.ast_module.statements[body[0]].kind == .If {
+				return ct_exec_if(state, checker.ast_module.statements[body[0]], true, depth+1)
+			}
+			return ct_exec_value_branch(state, body, depth+1)
+		}
+		return ct_exec_statements(state, body, false, depth+1)
 	}
 	operands: [dynamic]ast.Expr_Id
 	operands.allocator = checker.allocator
@@ -4466,7 +4907,17 @@ ct_exec_if :: proc(state: ^Ct_State, statement: ast.Stmt, yield_returns: bool, d
 		matched = guard_value
 	}
 	body := statement.body if matched else statement.else_body
-	flow, ok := ct_exec_statements(state, body, yield_returns, depth+1)
+	flow := ct_flow(.Normal)
+	ok := false
+	if yield_returns {
+		if !matched && len(body) == 1 && checker.ast_module.statements[body[0]].kind == .If {
+			flow, ok = ct_exec_if(state, checker.ast_module.statements[body[0]], true, depth+1)
+		} else {
+			flow, ok = ct_exec_value_branch(state, body, depth+1)
+		}
+	} else {
+		flow, ok = ct_exec_statements(state, body, false, depth+1)
+	}
 	ct_pop_bindings(state, scope_start)
 	return flow, ok
 }
@@ -4918,16 +5369,13 @@ ct_exec_match :: proc(state: ^Ct_State, statement: ast.Stmt, yield_returns: bool
 		   selection.tag != INVALID_CT_VALUE {
 			ct_bind_value(state, arm.captures[1], state.values[selection.tag].type, selection.tag, false)
 		}
-		if yield_returns && len(arm.body) == 1 && checker.ast_module.statements[arm.body[0]].kind == .Expression {
-			expr_stmt := checker.ast_module.statements[arm.body[0]]
-			value, expr_flow, expr_ok := ct_eval_expr(state, expr_stmt.expr, types.INVALID, depth+1)
-			ct_pop_bindings(state, scope_start)
-			if !expr_ok || expr_flow.kind != .Normal {
-				return expr_flow, expr_ok
+		arm_flow := ct_flow(.Normal)
+		arm_ok := false
+		if yield_returns {
+			arm_flow, arm_ok = ct_exec_value_branch(state, arm.body, depth+1)
+		} else {
+			arm_flow, arm_ok = ct_exec_statements(state, arm.body, false, depth+1)
 		}
-			return ct_flow(.Yield, value), true
-		}
-		arm_flow, arm_ok := ct_exec_statements(state, arm.body, yield_returns, depth+1)
 		ct_pop_bindings(state, scope_start)
 		return arm_flow, arm_ok
 	}
@@ -5049,7 +5497,13 @@ eval_comptime_statements :: proc(
 ) -> (Constant, bool, bool) {
 	state := ct_state_make(checker, pkg, file, types.INVALID, values, diagnose=false)
 	defer ct_state_destroy(&state)
-	flow, ok := ct_exec_statements(&state, statements, yield_returns, depth)
+	flow := ct_flow(.Normal)
+	ok := false
+	if yield_returns {
+		flow, ok = ct_exec_value_source(&state, statements, symbol.INVALID, false, false, depth)
+	} else {
+		flow, ok = ct_exec_statements(&state, statements, false, depth)
+	}
 	if !ok {
 		#partial switch state.error {
 		case .Overflow:
@@ -5120,7 +5574,7 @@ infer_comptime_expr_type :: proc(
 	if expr.left != ast.INVALID_EXPR {
 		value, flow, ok = ct_eval_expr(&state, expr.left, types.INVALID)
 	} else {
-		flow, ok = ct_exec_statements(&state, expr.body, true)
+		flow, ok = ct_exec_value_source(&state, expr.body, symbol.INVALID, false, false, 0)
 		if ok && flow.kind == .Yield {
 			value = flow.value
 			flow = ct_flow(.Normal)
@@ -5156,7 +5610,7 @@ build_comptime_expr :: proc(
 	if expr.left != ast.INVALID_EXPR {
 		value, flow, ok = ct_eval_expr(&state, expr.left, expected)
 	} else {
-		flow, ok = ct_exec_statements(&state, expr.body, true)
+		flow, ok = ct_exec_value_source(&state, expr.body, symbol.INVALID, false, false, 0)
 		if ok && flow.kind == .Yield {
 			value = flow.value
 			flow = ct_flow(.Normal)

compiler_tests.odin

@@ -5064,7 +5064,7 @@ main func() void {
 		found_expired = found_expired || strings.contains(message, "expired storage")
 		found_quota = found_quota || strings.contains(message, "comptime evaluation exceeded the step quota")
 		found_missing = found_missing || strings.contains(message, "did not return a value")
-		found_yield = found_yield || strings.contains(message, "comptime block must yield a value")
+		found_yield = found_yield || strings.contains(message, "a value block must end with an explicit 'yield'")
 	}
 	testing.expect(t, found_runtime)
 	testing.expect(t, runtime_only_count >= 2)
@@ -7428,6 +7428,64 @@ yield_misuse_is_diagnosed :: proc(t: ^testing.T) {
 	testing.expect(t, misplaced_yield)
 }
 
+@(test)
+comptime_yield_targets_match_value_source_semantics :: proc(t: ^testing.T) {
+	directory := "/tmp/brolang-test-comptime-yield-targets"
+	main_path := "/tmp/brolang-test-comptime-yield-targets/main.bro"
+	output := "/tmp/brolang-test-comptime-yield-targets-output"
+	valid_text := `selected :: ${
+	via_if :: if true {
+		yield 1
+	} else {
+		yield 2
+	}
+	via_label :: done: {
+		if true {
+			yield :done 3
+		}
+		yield :done 4
+	}
+	yield via_if + via_label
+}
+
+main func() i32 { return selected - 4 }
+`
+	_ = os2.remove_all(directory)
+	defer _ = os2.remove_all(directory)
+	defer _ = os.remove(output)
+	testing.expect(t, os.make_directory(directory) == nil)
+	testing.expect(t, os.write_entire_file(main_path, transmute([]byte)valid_text))
+	testing.expect_value(t, compiler_core.compile_package(directory, output), 0)
+	state := run_executable(output)
+	testing.expect_value(t, state.exit_code, 0)
+
+	invalid_text := `invalid :: ${
+	if false {
+		yield 1
+	}
+	yield 2
+}
+
+main func() void {}
+`
+	source_file := source.Source{path="invalid_comptime_yield.bro", text=invalid_text}
+	diagnostics := source.init_diagnostics(&source_file)
+	defer source.destroy_diagnostics(&diagnostics)
+	symbols := symbol.init_table()
+	defer symbol.destroy_table(&symbols)
+	stream := lexer.lex(&source_file, &diagnostics, &symbols)
+	defer delete(stream.items)
+	ast_module := parser.parse(&stream, &source_file, &diagnostics)
+	defer ast.destroy_module(&ast_module)
+	hir_module := checker.check(&ast_module, &diagnostics, &symbols)
+	defer hir.destroy_module(&hir_module)
+	found := false
+	for diagnostic in diagnostics.items {
+		found = found || strings.contains(diagnostic.message, "'yield' is only valid as the final statement of a value block")
+	}
+	testing.expect(t, found)
+}
+
 @(test)
 yield_control_flow_is_diagnosed :: proc(t: ^testing.T) {
 	// Value if/loop/block misuse: an `if` value without an `else`; a branch that neither

examples/programs/comptime_v1/main.bro

@@ -170,10 +170,8 @@ ct_errdefer_check func() i32 {
 }
 
 recover func() i32 {
-	return may_fail(true) catch |e| {
-		match e {
-			.bad: yield 5
-		}
+	return may_fail(true) catch |e| match e {
+		.bad: 5
 	}
 }
 

std/static_string_map/static_string_map.bro

@@ -35,9 +35,10 @@ init func($V type, $N usize, $entries [N]Pair(V)) StaticStringMap(V) {
 			values[i] = entry.1
 		}
 
+		result :: done: {
 			if N == 0 {
 				len_indexes [0]mut u32 = undefined
-			yield StaticStringMap(V) {
+				yield :done StaticStringMap(V) {
 					keys = keys[..],
 					values = values[..],
 					len_indexes = len_indexes[..],
@@ -70,7 +71,7 @@ init func($V type, $N usize, $entries [N]Pair(V)) StaticStringMap(V) {
 				len_indexes[length] = u32(entry_index)
 			}
 
-		yield StaticStringMap(V) {
+			yield :done StaticStringMap(V) {
 				keys = keys[..],
 				values = values[..],
 				len_indexes = len_indexes[..],
@@ -78,6 +79,8 @@ init func($V type, $N usize, $entries [N]Pair(V)) StaticStringMap(V) {
 				max_len = max_len,
 			}
 		}
+		yield result
+	}
 }
 
 get func($V type, map @StaticStringMap(V), key []u8) ?V {