added: index variables in `for` loops

diff --git a/compiler/include/astnodes.h b/compiler/include/astnodes.h
index a8d5fc554cba5c36ea6c6597c95a080d2ce7eea8..f920df4cac1b584aa7b7961067d1af51ea67f44d 100644 (file)
--- a/compiler/include/astnodes.h
+++ b/compiler/include/astnodes.h
@@ -835,6 +835,7 @@ struct AstFor           {
 
     // NOTE: Local defining the iteration variable
     AstLocal* var;
+    AstLocal* index_var;
 
     // NOTE: This can be any expression, but it is checked that
     // it is of a type that we know how to iterate over.

diff --git a/compiler/src/checker.c b/compiler/src/checker.c
index a9bece890894c7566f0a0443ad50e1107ad8b548..1760594a1f7af066f3e613dfc5c472b832b6af69 100644 (file)
--- a/compiler/src/checker.c
+++ b/compiler/src/checker.c
@@ -346,6 +346,15 @@ CheckStatus check_for(AstFor* fornode) {
         fornode->var->type = given_type;
     }
 
+    if (fornode->index_var) {
+        fornode->index_var->flags |= Ast_Flag_Cannot_Take_Addr;
+        CHECK(expression, (AstTyped **) &fornode->index_var);
+
+        if (!type_is_integer(fornode->index_var->type)) {
+            ERROR_(fornode->index_var->token->pos, "Index for a for loop must be an integer type, but it is a '%s'.", type_get_name(fornode->index_var->type));
+        }
+    }
+
     if (fornode->by_pointer)
         fornode->var->flags |= Ast_Flag_Cannot_Take_Addr;
 

diff --git a/compiler/src/parser.c b/compiler/src/parser.c
index 8036464834419f44a1de5ceb466175fd4e2a399a..a570a200c55945758d9aa138b8d0f6c7299fae17 100644 (file)
--- a/compiler/src/parser.c
+++ b/compiler/src/parser.c
@@ -1348,16 +1348,38 @@ static AstFor* parse_for_stmt(OnyxParser* parser) {
         for_node->by_pointer = 1;
     }
 
+    //
+    // For loops can take on a lot of shapes.
+    //     for value in iter
+    //     for value: i64 in iter
+    //     for value, index in iter
+    //     for value: i64, index in iter
+    //     for value: i64, index: i32 in iter
+    //
     if (next_tokens_are(parser, 2, Token_Type_Symbol, Token_Type_Keyword_In)
         || next_tokens_are(parser, 2, Token_Type_Symbol, ':')
+        || next_tokens_are(parser, 2, Token_Type_Symbol, ',')
     ) {
-        OnyxToken* local_sym = expect_token(parser, Token_Type_Symbol);
-        AstLocal* var_node = make_local(parser->allocator, local_sym, NULL);
+        for_node->var = make_local(
+            parser->allocator,
+            expect_token(parser, Token_Type_Symbol),
+            NULL
+        );
 
-        for_node->var = var_node;
-        if (peek_token(0)->type == ':') {
-            expect_token(parser, ':');
-            var_node->type_node = parse_type(parser);
+        if (consume_token_if_next(parser, ':')) {
+            for_node->var->type_node = parse_type(parser);
+        }
+
+        if (consume_token_if_next(parser, ',')) {
+            for_node->index_var = make_local(
+                parser->allocator,
+                expect_token(parser, Token_Type_Symbol),
+                (AstType *) &basic_type_u32
+            );
+
+            if (consume_token_if_next(parser, ':')) {
+                for_node->index_var->type_node = parse_type(parser);
+            }
         }
 
         expect_token(parser, Token_Type_Keyword_In);

diff --git a/compiler/src/symres.c b/compiler/src/symres.c
index c7844b5873b7fd1c4722b412eb8c5ba84963ddb8..eb50b9b7e66213510fea4fc1e68fef4d3d52fe48 100644 (file)
--- a/compiler/src/symres.c
+++ b/compiler/src/symres.c
@@ -825,6 +825,12 @@ static SymresStatus symres_for(AstFor* fornode) {
     scope_enter(fornode->scope);
     SYMRES(expression, &fornode->iter);
     SYMRES(local, &fornode->var);
+
+    // Right now, the index variable is optional
+    if (fornode->index_var) {
+        SYMRES(local, &fornode->index_var);
+    }
+
     SYMRES(block, fornode->stmt);
     scope_leave();
 

diff --git a/compiler/src/wasm_emit.c b/compiler/src/wasm_emit.c
index f00a4e13b386420c235fdc224e094a74a5810f30..addef11400a58a003c3608ebcbd2c2bf019570ed 100644 (file)
--- a/compiler/src/wasm_emit.c
+++ b/compiler/src/wasm_emit.c
@@ -13,6 +13,7 @@
 
 
 
+#include "types.h"
 #define BH_DEBUG
 #include "wasm_emit.h"
 #include "utils.h"
@@ -1313,7 +1314,54 @@ EMIT_FUNC(while, AstIfWhile* while_node) {
     *pcode = code;
 }
 
-EMIT_FUNC(for_range, AstFor* for_node, u64 iter_local) {
+EMIT_FUNC(for__prologue, AstFor* for_node, u64 iter_local, i64 index_local) {
+    bh_arr(WasmInstruction) code = *pcode;
+
+    if (for_node->has_first) {
+        for_node->first_local = local_raw_allocate(mod->local_alloc, WASM_TYPE_INT32);
+        WIL(for_node->token, WI_I32_CONST, 1);
+        WIL(for_node->token, WI_LOCAL_SET, for_node->first_local);
+    }
+
+    if (index_local != -1) {
+        if (type_is_small_integer(for_node->index_var->type)) {
+            WIL(for_node->token, WI_I32_CONST, 0);
+        } else {
+            WIL(for_node->token, WI_I64_CONST, 0);
+        }
+
+        WIL(for_node->token, WI_LOCAL_SET, index_local);
+    }
+
+    *pcode = code;
+}
+
+EMIT_FUNC(for__epilogue, AstFor* for_node, u64 iter_local, i64 index_local) {
+    bh_arr(WasmInstruction) code = *pcode;
+    
+    if (for_node->has_first) {
+        WIL(NULL, WI_I32_CONST, 0);
+        WIL(NULL, WI_LOCAL_SET, for_node->first_local);
+    }
+
+    if (index_local != -1) {
+        WIL(for_node->token, WI_LOCAL_GET, index_local);
+
+        if (type_is_small_integer(for_node->index_var->type)) {
+            WIL(for_node->token, WI_I32_CONST, 1);
+            WI(for_node->token, WI_I32_ADD);
+        } else {
+            WIL(for_node->token, WI_I64_CONST, 1);
+            WI(for_node->token, WI_I64_ADD);
+        }
+
+        WIL(for_node->token, WI_LOCAL_SET, index_local);
+    }
+
+    *pcode = code;
+}
+
+EMIT_FUNC(for_range, AstFor* for_node, u64 iter_local, i64 index_local) {
     bh_arr(WasmInstruction) code = *pcode;
 
     // NOTE: There are some aspects of the code below that rely on the
@@ -1325,11 +1373,9 @@ EMIT_FUNC(for_range, AstFor* for_node, u64 iter_local) {
     AstStructLiteral *range = (AstStructLiteral *) for_node->iter;
     u64 offset = 0;
 
-    StructMember low_mem, high_mem, step_mem;
-    type_lookup_member(builtin_range_type_type, "low", &low_mem);
+    StructMember high_mem, step_mem;
     type_lookup_member(builtin_range_type_type, "high", &high_mem);
     type_lookup_member(builtin_range_type_type, "step", &step_mem);
-    u64 low_local  = local_raw_allocate(mod->local_alloc, onyx_type_to_wasm_type(low_mem.type));
     u64 high_local = local_raw_allocate(mod->local_alloc, onyx_type_to_wasm_type(high_mem.type));
     u64 step_local = local_raw_allocate(mod->local_alloc, onyx_type_to_wasm_type(step_mem.type));
 
@@ -1337,14 +1383,9 @@ EMIT_FUNC(for_range, AstFor* for_node, u64 iter_local) {
 
     WIL(for_node->token, WI_LOCAL_SET, step_local);
     WIL(for_node->token, WI_LOCAL_SET, high_local);
-    WIL(for_node->token, WI_LOCAL_TEE, low_local);
     WIL(for_node->token, WI_LOCAL_SET, iter_local);
 
-    if (for_node->has_first) {
-        for_node->first_local = local_raw_allocate(mod->local_alloc, WASM_TYPE_INT32);
-        WIL(for_node->token, WI_I32_CONST, 1);
-        WIL(for_node->token, WI_LOCAL_SET, for_node->first_local);
-    }
+    emit_for__prologue(mod, &code, for_node, iter_local, index_local);
 
     emit_enter_structured_block(mod, &code, SBT_Breakable_Block, for_node->token);
     emit_enter_structured_block(mod, &code, SBT_Basic_Loop, for_node->token);
@@ -1393,10 +1434,7 @@ EMIT_FUNC(for_range, AstFor* for_node, u64 iter_local) {
     WI(for_node->token, WI_I32_ADD);
     WIL(for_node->token, WI_LOCAL_SET, iter_local);
 
-    if (for_node->has_first) {
-        WIL(for_node->token, WI_I32_CONST, 0);
-        WIL(for_node->token, WI_LOCAL_SET, for_node->first_local);
-    }
+    emit_for__epilogue(mod, &code, for_node, iter_local, index_local);
 
     if (bh_arr_last(code).type != WI_JUMP)
         WID(for_node->token, WI_JUMP, 0x00);
@@ -1405,14 +1443,13 @@ EMIT_FUNC(for_range, AstFor* for_node, u64 iter_local) {
     emit_leave_structured_block(mod, &code);
 
     if (for_node->has_first) local_raw_free(mod->local_alloc, WASM_TYPE_INT32);
-    local_raw_free(mod->local_alloc, onyx_type_to_wasm_type(low_mem.type));
     local_raw_free(mod->local_alloc, onyx_type_to_wasm_type(high_mem.type));
     local_raw_free(mod->local_alloc, onyx_type_to_wasm_type(step_mem.type));
 
     *pcode = code;
 }
 
-EMIT_FUNC(for_slice, AstFor* for_node, u64 iter_local) {
+EMIT_FUNC(for_slice, AstFor* for_node, u64 iter_local, i64 index_local) {
     bh_arr(WasmInstruction) code = *pcode;
 
     u64 end_ptr_local, ptr_local;
@@ -1442,11 +1479,7 @@ EMIT_FUNC(for_slice, AstFor* for_node, u64 iter_local) {
     WI(for_node->token, WI_PTR_ADD);
     WIL(for_node->token, WI_LOCAL_SET, end_ptr_local);
 
-    if (for_node->has_first) {
-        for_node->first_local = local_raw_allocate(mod->local_alloc, WASM_TYPE_INT32);
-        WIL(for_node->token, WI_I32_CONST, 1);
-        WIL(for_node->token, WI_LOCAL_SET, for_node->first_local);
-    }
+    emit_for__prologue(mod, &code, for_node, iter_local, index_local);
 
     emit_enter_structured_block(mod, &code, SBT_Breakable_Block, for_node->token);
     emit_enter_structured_block(mod, &code, SBT_Basic_Loop, for_node->token);
@@ -1476,10 +1509,7 @@ EMIT_FUNC(for_slice, AstFor* for_node, u64 iter_local) {
     WI(for_node->token, WI_PTR_ADD);
     WIL(for_node->token, WI_LOCAL_SET, ptr_local);
 
-    if (for_node->has_first) {
-        WIL(NULL, WI_I32_CONST, 0);
-        WIL(NULL, WI_LOCAL_SET, for_node->first_local);
-    }
+    emit_for__epilogue(mod, &code, for_node, iter_local, index_local);
 
     if (bh_arr_last(code).type != WI_JUMP)
         WID(for_node->token, WI_JUMP, 0x00);
@@ -1494,7 +1524,7 @@ EMIT_FUNC(for_slice, AstFor* for_node, u64 iter_local) {
     *pcode = code;
 }
 
-EMIT_FUNC(for_iterator, AstFor* for_node, u64 iter_local) {
+EMIT_FUNC(for_iterator, AstFor* for_node, u64 iter_local, i64 index_local) {
     bh_arr(WasmInstruction) code = *pcode;
 
     // Allocate temporaries for iterator contents
@@ -1532,11 +1562,7 @@ EMIT_FUNC(for_iterator, AstFor* for_node, u64 iter_local) {
         bh_arr_push(mod->for_remove_info, remove_info);
     }
 
-    if (for_node->has_first) {
-        for_node->first_local = local_raw_allocate(mod->local_alloc, WASM_TYPE_INT32);
-        WIL(for_node->token, WI_I32_CONST, 1);
-        WIL(for_node->token, WI_LOCAL_SET, for_node->first_local);
-    }
+    emit_for__prologue(mod, &code, for_node, iter_local, index_local);
 
     AstLocal* var = for_node->var;
     b32 it_is_local = (b32) ((iter_local & LOCAL_IS_WASM) != 0);
@@ -1614,10 +1640,7 @@ EMIT_FUNC(for_iterator, AstFor* for_node, u64 iter_local) {
 
     emit_block(mod, &code, for_node->stmt, 0);
 
-    if (for_node->has_first) {
-        WIL(NULL, WI_I32_CONST, 0);
-        WIL(NULL, WI_LOCAL_SET, for_node->first_local);
-    }
+    emit_for__epilogue(mod, &code, for_node, iter_local, index_local);
 
     WID(for_node->token, WI_JUMP, 0x00);
 
@@ -1645,22 +1668,33 @@ EMIT_FUNC(for, AstFor* for_node) {
     u64 iter_local = local_allocate(mod->local_alloc, (AstTyped *) var);
     bh_imap_put(&mod->local_map, (u64) var, iter_local);
 
+    i64 index_local = -1;
+    if (for_node->index_var) {
+        index_local = local_allocate(mod->local_alloc, (AstTyped *) for_node->index_var);
+        bh_imap_put(&mod->local_map, (u64) for_node->index_var, index_local);
+    }
+
     debug_enter_symbol_frame(mod);
     debug_introduce_symbol(mod, var->token,
         local_is_wasm_local((AstTyped *) var) ? DSL_REGISTER : DSL_STACK,
         iter_local, var->type);
 
+    if (for_node->index_var) {
+        // index variables must be register allocated.
+        debug_introduce_symbol(mod, for_node->index_var->token, DSL_REGISTER, index_local, for_node->index_var->type);
+    }
+
     emit_expression(mod, &code, for_node->iter);
 
     switch (for_node->loop_type) {
-        case For_Loop_Range:    emit_for_range(mod, &code, for_node, iter_local); break;
+        case For_Loop_Range:    emit_for_range(mod, &code, for_node, iter_local, index_local); break;
 
         // NOTE: For static arrays, simply outputing the size
         // of the array right after the pointer to the start
         // of the array essentially makes it a slice.
         case For_Loop_Array:
             WIL(NULL, WI_I32_CONST, for_node->iter->type->Array.count);
-            emit_for_slice(mod, &code, for_node, iter_local);
+            emit_for_slice(mod, &code, for_node, iter_local, index_local);
             break;
 
         // NOTE: A dynamic array is just a slice with a capacity and allocator on the end.
@@ -1671,8 +1705,8 @@ EMIT_FUNC(for, AstFor* for_node) {
             emit_load_slice(mod, &code);
             // fallthrough
 
-        case For_Loop_Slice:    emit_for_slice(mod, &code, for_node, iter_local); break;
-        case For_Loop_Iterator: emit_for_iterator(mod, &code, for_node, iter_local); break;
+        case For_Loop_Slice:    emit_for_slice(mod, &code, for_node, iter_local, index_local); break;
+        case For_Loop_Iterator: emit_for_iterator(mod, &code, for_node, iter_local, index_local); break;
         default: onyx_report_error(for_node->token->pos, Error_Critical, "Invalid for loop type. You should probably not be seeing this...");
     }