`#unroll(N) for`

3 files changed, 48 insertions, 5 deletions

diff --git a/core/odin/ast/ast.odin b/core/odin/ast/ast.odin
index f62feec8c..3b8998b31 100644
--- a/core/odin/ast/ast.odin
+++ b/core/odin/ast/ast.odin
@@ -432,10 +432,13 @@ Range_Stmt :: struct {
 	reverse:   bool,
 }
 
-Inline_Range_Stmt :: struct {
+Inline_Range_Stmt :: Unroll_Range_Stmt
+
+Unroll_Range_Stmt :: struct {
 	using node: Stmt,
 	label:     ^Expr,
-	inline_pos: tokenizer.Pos,
+	unroll_pos: tokenizer.Pos,
+	args:       []^Expr,
 	for_pos:    tokenizer.Pos,
 	val0:       ^Expr,
 	val1:       ^Expr,
diff --git a/core/odin/ast/clone.odin b/core/odin/ast/clone.odin
index 67f7ffa95..b7501e6ca 100644
--- a/core/odin/ast/clone.odin
+++ b/core/odin/ast/clone.odin
@@ -242,8 +242,9 @@ clone_node :: proc(node: ^Node) -> ^Node {
 			r.vals = clone(r.vals)
 			r.expr = clone(r.expr)
 			r.body = clone(r.body)
-		case ^Inline_Range_Stmt:
+		case ^Unroll_Range_Stmt:
 			r.label = clone(r.label)
+			r.args = clone(r.args)
 			r.val0 = clone(r.val0)
 			r.val1 = clone(r.val1)
 			r.expr = clone(r.expr)
diff --git a/core/odin/parser/parser.odin b/core/odin/parser/parser.odin
index 5a7440339..63c7e388f 100644
--- a/core/odin/parser/parser.odin
+++ b/core/odin/parser/parser.odin
@@ -1262,11 +1262,49 @@ parse_foreign_decl :: proc(p: ^Parser) -> ^ast.Decl {
 
 
 parse_unrolled_for_loop :: proc(p: ^Parser, inline_tok: tokenizer.Token) -> ^ast.Stmt {
-	for_tok := expect_token(p, .For)
 	val0, val1: ^ast.Expr
 	in_tok: tokenizer.Token
 	expr: ^ast.Expr
 	body: ^ast.Stmt
+	args: [dynamic]^ast.Expr
+
+	if allow_token(p, .Open_Paren) {
+		p.expr_level += 1
+		if p.curr_tok.kind == .Close_Paren {
+			error(p, p.curr_tok.pos, "#unroll expected at least 1 argument, got 0")
+		} else {
+			args = make([dynamic]^ast.Expr)
+			for p.curr_tok.kind != .Close_Paren &&
+			    p.curr_tok.kind != .EOF {
+			    	arg := parse_value(p)
+
+			    	if p.curr_tok.kind == .Eq {
+			    		eq := expect_token(p, .Eq)
+			    		if arg != nil {
+			    			if _, ok := arg.derived.(^ast.Ident); !ok {
+			    				error(p, arg.pos, "expected an identifier for 'key=value'")
+			    			}
+			    		}
+			    		value := parse_value(p)
+			    		fv := ast.new(ast.Field_Value, arg.pos, value)
+			    		fv.field = arg
+			    		fv.sep   = eq.pos
+			    		fv.value = value
+
+			    		arg = fv
+			    	}
+
+			    	append(&args, arg)
+
+				allow_token(p, .Comma) or_break
+			    }
+		}
+
+		p.expr_level -= 1
+		_ = expect_token_after(p, .Close_Paren, "#unroll")
+	}
+
+	for_tok := expect_token(p, .For)
 
 	bad_stmt := false
 
@@ -1309,7 +1347,8 @@ parse_unrolled_for_loop :: proc(p: ^Parser, inline_tok: tokenizer.Token) -> ^ast
 	}
 
 	range_stmt := ast.new(ast.Inline_Range_Stmt, inline_tok.pos, body)
-	range_stmt.inline_pos = inline_tok.pos
+	range_stmt.unroll_pos = inline_tok.pos
+	range_stmt.args = args[:]
 	range_stmt.for_pos = for_tok.pos
 	range_stmt.val0 = val0
 	range_stmt.val1 = val1