Support `for init; x in y {}` in `core:odin/parser`bill/range-init

4 files changed, 18 insertions, 1 deletions

diff --git a/core/odin/ast/ast.odin b/core/odin/ast/ast.odin
index ec22db434..2cee6e385 100644
--- a/core/odin/ast/ast.odin
+++ b/core/odin/ast/ast.odin
@@ -432,6 +432,7 @@ For_Stmt :: struct {
 Range_Stmt :: struct {
 	using node: Stmt,
 	label:     ^Expr, // possibly nil
+	init:      ^Stmt,
 	for_pos:   tokenizer.Pos,
 	vals:      []^Expr,
 	in_pos:    tokenizer.Pos,
diff --git a/core/odin/ast/clone.odin b/core/odin/ast/clone.odin
index b7501e6ca..df3e1df0d 100644
--- a/core/odin/ast/clone.odin
+++ b/core/odin/ast/clone.odin
@@ -239,6 +239,7 @@ clone_node :: proc(node: ^Node) -> ^Node {
 			r.body = clone(r.body)
 		case ^Range_Stmt:
 			r.label = clone(r.label)
+			r.init = clone(r.init)
 			r.vals = clone(r.vals)
 			r.expr = clone(r.expr)
 			r.body = clone(r.body)
diff --git a/core/odin/ast/walk.odin b/core/odin/ast/walk.odin
index cba040875..24c90c13b 100644
--- a/core/odin/ast/walk.odin
+++ b/core/odin/ast/walk.odin
@@ -222,6 +222,9 @@ walk :: proc(v: ^Visitor, node: ^Node) {
 		if n.label != nil {
 			walk(v, n.label)
 		}
+		if n.init != nil {
+			walk(v, n.init)
+		}
 		for val in n.vals {
 			if val != nil {
 				walk(v, val)
diff --git a/core/odin/parser/parser.odin b/core/odin/parser/parser.odin
index a18942e6b..643673c69 100644
--- a/core/odin/parser/parser.odin
+++ b/core/odin/parser/parser.odin
@@ -881,7 +881,7 @@ parse_for_stmt :: proc(p: ^Parser) -> ^ast.Stmt {
 	body: ^ast.Stmt
 	is_range := false
 
-	if p.curr_tok.kind != .Open_Brace && p.curr_tok.kind != .Do {
+	general_conds: if p.curr_tok.kind != .Open_Brace && p.curr_tok.kind != .Do {
 		prev_level := p.expr_level
 		defer p.expr_level = prev_level
 		p.expr_level = -1
@@ -929,6 +929,17 @@ parse_for_stmt :: proc(p: ^Parser) -> ^ast.Stmt {
 				error(p, p.curr_tok.pos, "Expected ';', followed by a condition expression and post statement, got %s", tokenizer.tokens[p.curr_tok.kind])
 			} else {
 				if p.curr_tok.kind != .Semicolon {
+					if p.curr_tok.kind == .Ident {
+						next_token := peek_token(p)
+						if next_token.kind == .In || next_token.kind == .Comma {
+							cond = parse_simple_stmt(p, {.In})
+							as := cond.derived_stmt.(^ast.Assign_Stmt)
+							assert(as.op.kind == .In)
+							is_range = true
+							break general_conds
+						}
+					}
+
 					cond = parse_simple_stmt(p, nil)
 				}
 
@@ -967,6 +978,7 @@ parse_for_stmt :: proc(p: ^Parser) -> ^ast.Stmt {
 
 		range_stmt := ast.new(ast.Range_Stmt, tok.pos, body)
 		range_stmt.for_pos = tok.pos
+		range_stmt.init = init
 		range_stmt.vals = vals
 		range_stmt.in_pos = assign_stmt.op.pos
 		range_stmt.expr = rhs