expand yield to if-statements and loops

compiler/parser/parser.odin

@@ -49,6 +49,23 @@ previous :: proc(parser: ^Parser) -> token.Token {
 	return parser.tokens.items[max(parser.cursor-1, 0)]
 }
 
+peek :: proc(parser: ^Parser) -> token.Token {
+	return parser.tokens.items[min(parser.cursor+1, len(parser.tokens.items)-1)]
+}
+
+// A loop body may be labeled `blk: { ... }` so a nested `yield :blk x` can exit
+// it past an enclosing `if`. Consumes and returns the label when the next tokens
+// are `Identifier Colon`; otherwise leaves the cursor untouched.
+parse_optional_loop_label :: proc(parser: ^Parser) -> symbol.Id {
+	if current(parser).kind == .Identifier && peek(parser).kind == .Colon {
+		name := advance(parser) // the label name
+		advance(parser)         // consume ':'
+		skip_newlines(parser)
+		return name.symbol
+	}
+	return symbol.INVALID
+}
+
 advance :: proc(parser: ^Parser) -> token.Token {
 	result := current(parser)
 	if result.kind != .Eof {
@@ -992,11 +1009,23 @@ parse_return :: proc(parser: ^Parser) -> ast.Stmt_Id {
 parse_yield :: proc(parser: ^Parser) -> ast.Stmt_Id {
 	start := advance(parser) // consume 'yield'
 	skip_newlines(parser)
+	// `yield :blk x` targets the loop labeled `blk`; a bare `yield x` targets
+	// the directly-enclosing value block / if branch. No expression starts with
+	// ':', so a leading colon is unambiguously a label.
+	label := symbol.INVALID
+	if _, ok := allow(parser, .Colon); ok {
+		if name, name_ok := allow(parser, .Identifier); name_ok {
+			label = name.symbol
+		} else {
+			source.add(parser.diagnostics, current(parser).span, "expected a loop label after ':'")
+		}
+	}
 	expr := parse_expression(parser)
 	id := ast.stmt_id(len(parser.module.statements))
 	append(&parser.module.statements, ast.Stmt{
 		kind=.Yield,
 		span=span_from(start.span, parser.module.exprs[expr].span),
+		label=label,
 		expr=expr,
 		diagnostic=source.INVALID_DIAGNOSTIC,
 	})
@@ -1078,6 +1107,19 @@ starts_declared_type :: proc(parser: ^Parser) -> bool {
 	return cursor < len(parser.tokens.items) && is_type_token(parser.tokens.items[cursor].kind)
 }
 
+// A declaration/assignment RHS may be a value-producing control-flow construct:
+// an `if`/`for`/`while` whose branches/iterations `yield`. Returns the parsed
+// statement (to be carried as a one-element block-init `body`) and true when the
+// current token opens one.
+parse_value_control_flow :: proc(parser: ^Parser) -> (ast.Stmt_Id, bool) {
+	#partial switch current(parser).kind {
+	case .Keyword_If:    return parse_if(parser), true
+	case .Keyword_For:   return parse_for(parser), true
+	case .Keyword_While: return parse_while(parser), true
+	}
+	return ast.INVALID_STMT, false
+}
+
 parse_statement :: proc(parser: ^Parser) -> ast.Stmt_Id {
 	if current(parser).kind == .Keyword_Return {
 		return parse_return(parser)
@@ -1146,6 +1188,25 @@ parse_statement :: proc(parser: ^Parser) -> ast.Stmt_Id {
 				})
 				return id
 			}
+			// A value-producing `if`/`for`/`while`: carried as a one-element block-init
+			// body, the same signal a value block uses (`expr` invalid).
+			if cf, is_cf := parse_value_control_flow(parser); is_cf {
+				body := make([]ast.Stmt_Id, 1, parser.module.allocator)
+				body[0] = cf
+				id := ast.stmt_id(len(parser.module.statements))
+				append(&parser.module.statements, ast.Stmt{
+					kind=kind,
+					span=span_from(name.span, previous(parser).span),
+					name=name.symbol,
+					type=type_syntax,
+					immutable=immutable,
+					target=ast.INVALID_EXPR,
+					expr=ast.INVALID_EXPR,
+					body=body,
+					diagnostic=source.INVALID_DIAGNOSTIC,
+				})
+				return id
+			}
 			expr := parse_expression(parser)
 			id := ast.stmt_id(len(parser.module.statements))
 			append(&parser.module.statements, ast.Stmt{
@@ -1181,6 +1242,21 @@ parse_statement :: proc(parser: ^Parser) -> ast.Stmt_Id {
 			})
 			return id
 		}
+		// A value-producing `if`/`for`/`while` assigned to a complex target.
+		if cf, is_cf := parse_value_control_flow(parser); is_cf {
+			body := make([]ast.Stmt_Id, 1, parser.module.allocator)
+			body[0] = cf
+			id := ast.stmt_id(len(parser.module.statements))
+			append(&parser.module.statements, ast.Stmt{
+				kind=.Assignment,
+				span=span_from(parser.module.exprs[expr].span, previous(parser).span),
+				target=expr,
+				expr=ast.INVALID_EXPR,
+				body=body,
+				diagnostic=source.INVALID_DIAGNOSTIC,
+			})
+			return id
+		}
 		value := parse_expression(parser)
 		id := ast.stmt_id(len(parser.module.statements))
 		append(&parser.module.statements, ast.Stmt{
@@ -1494,6 +1570,7 @@ parse_for :: proc(parser: ^Parser) -> ast.Stmt_Id {
 		}
 	}
 	skip_newlines(parser)
+	label := parse_optional_loop_label(parser)
 	body := parse_block(parser)
 
 	id := ast.stmt_id(len(parser.module.statements))
@@ -1502,6 +1579,7 @@ parse_for :: proc(parser: ^Parser) -> ast.Stmt_Id {
 		span=span_from(start.span, previous(parser).span),
 		name=item_name,
 		index_name=index_name,
+		label=label,
 		pointer_capture=pointer_capture,
 		expr=iterable,
 		body=body,
@@ -1526,6 +1604,7 @@ parse_while :: proc(parser: ^Parser) -> ast.Stmt_Id {
 		update = parse_while_update(parser)
 		skip_newlines(parser)
 	}
+	label := parse_optional_loop_label(parser)
 	body := parse_block(parser)
 
 	id := ast.stmt_id(len(parser.module.statements))
@@ -1534,6 +1613,7 @@ parse_while :: proc(parser: ^Parser) -> ast.Stmt_Id {
 		span=span_from(start.span, previous(parser).span),
 		expr=condition,
 		body=body,
+		label=label,
 		update=update,
 		diagnostic=source.INVALID_DIAGNOSTIC,
 	})