feat(compiler): support continue statement in while and for loops (#39)

src/tiny8/compiler.py

@@ -37,7 +37,7 @@ class Lexer:
         ("STRING", r'"([^"\\]|\\.)*"'),
         (
             "KEYWORD",
-            r"\b(int|char|void|if|else|while|for|break|return|peek|poke|in|out|vram_set|vram_get|vram_clear|vram_puts)\b",
+            r"\b(int|char|void|if|else|while|for|break|continue|return|peek|poke|in|out|vram_set|vram_get|vram_clear|vram_puts)\b",
         ),
         ("ID", r"[a-zA-Z_]\w*"),
         ("OP", r"\+\+|--|<<|>>|==|!=|<=|>=|&&|\|\||[+\-*/%&|^]=|[+\-*/%&|^=!<>]"),
@@ -198,6 +198,11 @@ class BreakStmt(ASTNode):
     """Break statement."""
 
 
+@dataclass
+class ContinueStmt(ASTNode):
+    """Continue statement."""
+
+
 @dataclass
 class ReturnStmt(ASTNode):
     """Return statement."""
@@ -516,6 +521,11 @@ def parse_stmt(self) -> ASTNode:
             self.consume(";")
             return BreakStmt()
 
+        if tok.value == "continue":
+            self.consume("continue")
+            self.consume(";")
+            return ContinueStmt()
+
         if tok.value == "for":
             self.consume("for")
             self.consume("(")
@@ -895,6 +905,7 @@ def __init__(self) -> None:
         self.next_ram = 0x60
         self.current_func = ""
         self.loop_end_stack: list[str] = []
+        self.loop_continue_stack: list[str] = []
 
     def new_label(self, prefix: str = "L") -> str:
         self.label_counter += 1
@@ -939,6 +950,7 @@ def generate(self, program: Program, optimize: bool = True) -> str:
         self.symbols = {}
         self.next_ram = 0x60
         self.loop_end_stack = []
+        self.loop_continue_stack = []
 
         for g in program.globals:
             if isinstance(g, ArrayDecl):
@@ -1078,23 +1090,31 @@ def gen_stmt(self, stmt: ASTNode) -> None:
             if stmt.init:
                 self.gen_stmt(stmt.init)
             lbl_start = self.new_label("for_start")
+            lbl_post = self.new_label("for_post")
             lbl_end = self.new_label("for_end")
             self.loop_end_stack.append(lbl_end)
+            self.loop_continue_stack.append(lbl_post)
             self.emit_label(lbl_start)
             self.gen_expr(stmt.condition)
             self.emit("CPI R16, 0")
             self.emit(f"BREQ {lbl_end}")
             self.gen_block(stmt.body)
+            self.emit_label(lbl_post)
             if stmt.post:
                 self.gen_stmt(stmt.post)
             self.emit(f"JMP {lbl_start}")
             self.emit_label(lbl_end)
+            self.loop_continue_stack.pop()
             self.loop_end_stack.pop()
 
         elif isinstance(stmt, BreakStmt):
             if self.loop_end_stack:
                 self.emit(f"JMP {self.loop_end_stack[-1]}")
 
+        elif isinstance(stmt, ContinueStmt):
+            if self.loop_continue_stack:
+                self.emit(f"JMP {self.loop_continue_stack[-1]}")
+
         elif isinstance(stmt, IfStmt):
             else_lbl = self.new_label("else")
             end_lbl = self.new_label("endif")
@@ -1117,6 +1137,7 @@ def gen_stmt(self, stmt: ASTNode) -> None:
             loop_lbl = self.new_label("while_start")
             end_lbl = self.new_label("while_end")
             self.loop_end_stack.append(end_lbl)
+            self.loop_continue_stack.append(loop_lbl)
             self.emit_label(loop_lbl)
             self.gen_expr(stmt.condition)
             self.emit("CPI R16, 0")
@@ -1127,6 +1148,7 @@ def gen_stmt(self, stmt: ASTNode) -> None:
 
             self.emit(f"JMP {loop_lbl}")
             self.emit_label(end_lbl)
+            self.loop_continue_stack.pop()
             self.loop_end_stack.pop()
 
         elif isinstance(stmt, ReturnStmt):

tests/test_compiler.py

@@ -3,7 +3,7 @@
 import pytest
 
 from tiny8 import CPU, assemble
-from tiny8.compiler import compile_c, optimize_assembly
+from tiny8.compiler import ContinueStmt, Lexer, Parser, compile_c, optimize_assembly
 
 
 def test_compiler_variable_and_arithmetic():
@@ -119,6 +119,105 @@ def test_compiler_for_loop_and_break():
     assert cpu.read_ram(0x84) == 6
 
 
+def test_compiler_while_loop_continue():
+    c_code = """
+    int main() {
+        int sum = 0;
+        int i = 0;
+        while (i < 6) {
+            i = i + 1;
+            if (i == 3) {
+                continue;
+            }
+            sum = sum + i;
+        }
+        poke(0x84, sum);
+        return 0;
+    }
+    """
+    asm = compile_c(c_code)
+    asm_res = assemble(asm)
+    cpu = CPU()
+    cpu.load_program(asm_res)
+    cpu.run(max_steps=300, show_progress=False)
+
+    # 1 + 2 + 4 + 5 + 6 = 18
+    assert cpu.read_ram(0x84) == 18
+
+
+def test_compiler_for_loop_continue():
+    c_code = """
+    int main() {
+        int sum = 0;
+        int i = 0;
+        for (i = 0; i < 6; i = i + 1) {
+            if (i == 3) {
+                continue;
+            }
+            sum = sum + i;
+        }
+        poke(0x84, sum);
+        return 0;
+    }
+    """
+    asm = compile_c(c_code)
+    asm_res = assemble(asm)
+    cpu = CPU()
+    cpu.load_program(asm_res)
+    cpu.run(max_steps=300, show_progress=False)
+
+    # 0 + 1 + 2 + 4 + 5 = 12
+    assert cpu.read_ram(0x84) == 12
+
+
+def test_compiler_nested_loops_continue_and_break():
+    c_code = """
+    int main() {
+        int sum = 0;
+        int i = 0;
+        int j = 0;
+        for (i = 0; i < 4; i = i + 1) {
+            if (i == 1) {
+                continue;
+            }
+            j = 0;
+            while (j < 4) {
+                j = j + 1;
+                if (j == 2) {
+                    continue;
+                }
+                if (j == 4) {
+                    break;
+                }
+                sum = sum + 1;
+            }
+        }
+        poke(0x84, sum);
+        return 0;
+    }
+    """
+    asm = compile_c(c_code)
+    asm_res = assemble(asm)
+    cpu = CPU()
+    cpu.load_program(asm_res)
+    cpu.run(max_steps=1000, show_progress=False)
+
+    # i: 0, 2, 3 (i=1 skipped)
+    # for each i: j=1 -> sum++, j=2 -> continue, j=3 -> sum++, j=4 -> break
+    # sum = 3 * 2 = 6
+    assert cpu.read_ram(0x84) == 6
+
+
+def test_compiler_continue_ast_and_lexer():
+    tokens = Lexer("continue;").tokenize()
+    assert len(tokens) == 2
+    assert tokens[0].type == "KEYWORD"
+    assert tokens[0].value == "continue"
+
+    prog = Parser(Lexer("void main() { continue; }").tokenize()).parse()
+    assert isinstance(prog.functions[0].body.statements[0], ContinueStmt)
+
+
 def test_compiler_array():
     c_code = """
     int main() {