From 328d893cb58d4b96f12207eb49da01273deda6ce Mon Sep 17 00:00:00 2001
From: gingerBill <bill@gingerbill.org>
Date: Fri, 10 Jan 2025 12:14:43 +0000
Subject: `#unroll(N) for`

---
 src/parser.cpp | 58 +++++++++++++++++++++++++++++++++++++++++++++++-----------
 1 file changed, 47 insertions(+), 11 deletions(-)

(limited to 'src/parser.cpp')

diff --git a/src/parser.cpp b/src/parser.cpp
index 03c5a5962..94f8fd42c 100644
--- a/src/parser.cpp
+++ b/src/parser.cpp
@@ -348,10 +348,11 @@ gb_internal Ast *clone_ast(Ast *node, AstFile *f) {
 		n->RangeStmt.body  = clone_ast(n->RangeStmt.body, f);
 		break;
 	case Ast_UnrollRangeStmt:
-		n->UnrollRangeStmt.val0  = clone_ast(n->UnrollRangeStmt.val0, f);
-		n->UnrollRangeStmt.val1  = clone_ast(n->UnrollRangeStmt.val1, f);
-		n->UnrollRangeStmt.expr  = clone_ast(n->UnrollRangeStmt.expr, f);
-		n->UnrollRangeStmt.body  = clone_ast(n->UnrollRangeStmt.body, f);
+		n->UnrollRangeStmt.args = clone_ast_array(n->UnrollRangeStmt.args, f);
+		n->UnrollRangeStmt.val0 = clone_ast(n->UnrollRangeStmt.val0, f);
+		n->UnrollRangeStmt.val1 = clone_ast(n->UnrollRangeStmt.val1, f);
+		n->UnrollRangeStmt.expr = clone_ast(n->UnrollRangeStmt.expr, f);
+		n->UnrollRangeStmt.body = clone_ast(n->UnrollRangeStmt.body, f);
 		break;
 	case Ast_CaseClause:
 		n->CaseClause.list  = clone_ast_array(n->CaseClause.list, f);
@@ -1037,15 +1038,16 @@ gb_internal Ast *ast_range_stmt(AstFile *f, Token token, Slice<Ast *> vals, Toke
 	return result;
 }
 
-gb_internal Ast *ast_unroll_range_stmt(AstFile *f, Token unroll_token, Token for_token, Ast *val0, Ast *val1, Token in_token, Ast *expr, Ast *body) {
+gb_internal Ast *ast_unroll_range_stmt(AstFile *f, Token unroll_token, Slice<Ast *> args, Token for_token, Ast *val0, Ast *val1, Token in_token, Ast *expr, Ast *body) {
 	Ast *result = alloc_ast_node(f, Ast_UnrollRangeStmt);
 	result->UnrollRangeStmt.unroll_token = unroll_token;
+	result->UnrollRangeStmt.args      = args;
 	result->UnrollRangeStmt.for_token = for_token;
-	result->UnrollRangeStmt.val0 = val0;
-	result->UnrollRangeStmt.val1 = val1;
-	result->UnrollRangeStmt.in_token = in_token;
-	result->UnrollRangeStmt.expr  = expr;
-	result->UnrollRangeStmt.body  = body;
+	result->UnrollRangeStmt.val0      = val0;
+	result->UnrollRangeStmt.val1      = val1;
+	result->UnrollRangeStmt.in_token  = in_token;
+	result->UnrollRangeStmt.expr      = expr;
+	result->UnrollRangeStmt.body      = body;
 	return result;
 }
 
@@ -5137,6 +5139,40 @@ gb_internal Ast *parse_attribute(AstFile *f, Token token, TokenKind open_kind, T
 
 
 gb_internal Ast *parse_unrolled_for_loop(AstFile *f, Token unroll_token) {
+	Array<Ast *> args = {};
+
+	if (allow_token(f, Token_OpenParen)) {
+		f->expr_level++;
+		if (f->curr_token.kind == Token_CloseParen) {
+			syntax_error(f->curr_token, "#unroll expected at least 1 argument, got 0");
+		} else {
+			args = array_make<Ast *>(ast_allocator(f));
+			while (f->curr_token.kind != Token_CloseParen &&
+			       f->curr_token.kind != Token_EOF) {
+				Ast *arg = nullptr;
+				arg = parse_value(f);
+
+				if (f->curr_token.kind == Token_Eq) {
+					Token eq = expect_token(f, Token_Eq);
+					if (arg != nullptr && arg->kind != Ast_Ident) {
+						syntax_error(arg, "Expected an identifier for 'key=value'");
+					}
+					Ast *value = parse_value(f);
+					arg = ast_field_value(f, arg, value, eq);
+				}
+
+				array_add(&args, arg);
+
+				if (!allow_field_separator(f)) {
+					break;
+				}
+			}
+		}
+		f->expr_level--;
+		Token close = expect_closing(f, Token_CloseParen, str_lit("#unroll"));
+		gb_unused(close);
+	}
+
 	Token for_token = expect_token(f, Token_for);
 	Ast *val0 = nullptr;
 	Ast *val1 = nullptr;
@@ -5180,7 +5216,7 @@ gb_internal Ast *parse_unrolled_for_loop(AstFile *f, Token unroll_token) {
 	if (bad_stmt) {
 		return ast_bad_stmt(f, unroll_token, f->curr_token);
 	}
-	return ast_unroll_range_stmt(f, unroll_token, for_token, val0, val1, in_token, expr, body);
+	return ast_unroll_range_stmt(f, unroll_token, slice_from_array(args), for_token, val0, val1, in_token, expr, body);
 }
 
 gb_internal Ast *parse_stmt(AstFile *f) {
-- 
cgit v1.2.3