microscript: fix loops

1 files changed, 12 insertions, 3 deletions

diff --git a/uscript/parser.c b/uscript/parser.c
index f2c601a..d064c6c 100644
--- a/uscript/parser.c
+++ b/uscript/parser.c
@@ -779,8 +779,12 @@ void loop_stat(struct parser *p)
 
         int exit_jump = -1;
 
+        bool is_iter_loop = false;
+
         if (!consume(p, ':')) {
                 if (consume(p, TOKEN_LET)) {
+                        is_iter_loop = true;
+
                         expect(p, TOKEN_IDENT, "expected identifier");
                         struct token ident = p->prev;
 
@@ -792,7 +796,7 @@ void loop_stat(struct parser *p)
 
                         begin_scope(p);
                         loop.start = da_len(p->fp->proto->bytecode) - 1;
-                        
+
                         parser_add_byte(p, BC_GET_LOCAL);
                         parser_add_byte(p, (u8)iter);
                         parser_add_byte(p, BC_CALL);
@@ -808,8 +812,6 @@ void loop_stat(struct parser *p)
                 }
 
                 expect(p, ':', "expected ':' to begin 'loop' block");
-        } else {
-                begin_scope(p);
         }
 
         if (consume(p, '<')) {
@@ -819,11 +821,17 @@ void loop_stat(struct parser *p)
                 expect(p, '>', "expected '>' after loop label");
         }
 
+        if (!is_iter_loop)
+                begin_scope(p);
+
         while (p->cur.kind != TOKEN_END && p->cur.kind != TOKEN_EOF)
                 stat(p);
         end_scope(p);
         add_loop(p, loop.start);
 
+        if (is_iter_loop)
+                parser_add_byte(p, BC_POP);
+
         if (exit_jump != -1) {
                 end_jump(p, exit_jump);
                 parser_add_byte(p, BC_POP);
@@ -903,6 +911,7 @@ void next_stat(struct parser *p)
 
         struct loop *loop = loop_label(p);
         pop_scope(p, loop->scope);
+
         add_loop(p, loop->start);
 }