feat(compiler): support do-while loop statement (#40)

src/tiny8/compiler.py

@@ -5,7 +5,7 @@
 - Compound assignment operators (+=, -=, *=, /=, %=, &=, |=, ^=)
 - Increment and decrement operators (++, --)
 - Arithmetic (+, -, *, /, %), bitwise (&, |, ^, ~), logical (!), and relational (<, <=, >, >=, ==, !=)
-- Control structures: if / else, while, return
+- Control structures: if / else, while, do-while, for, break, continue, return
 - Functions with register-based parameter passing (R16..R19) and return value in R16
 - Hardware builtins: peek(addr), poke(addr, val), in(port), out(port, val)
 """
@@ -37,7 +37,7 @@ class Lexer:
         ("STRING", r'"([^"\\]|\\.)*"'),
         (
             "KEYWORD",
-            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",
+            r"\b(int|char|void|if|else|while|for|do|break|continue|return|peek|poke|in|out|vram_set|vram_get|vram_clear|vram_puts)\b",
         ),
         ("ID", r"[a-zA-Z_]\w*"),
         ("OP", r"\+\+|--|<<|>>|==|!=|<=|>=|&&|\|\||[+\-*/%&|^]=|[+\-*/%&|^=!<>]"),
@@ -183,6 +183,14 @@ class WhileStmt(ASTNode):
     body: Block
 
 
+@dataclass
+class DoWhileStmt(ASTNode):
+    """Do-while loop."""
+
+    body: Block
+    condition: ASTNode
+
+
 @dataclass
 class ForStmt(ASTNode):
     """For loop."""
@@ -582,6 +590,20 @@ def parse_stmt(self) -> ASTNode:
             )
             return WhileStmt(cond, body)
 
+        if tok.value == "do":
+            self.consume("do")
+            body = (
+                self.parse_block()
+                if self.peek() and self.peek().value == "{"
+                else Block([self.parse_stmt()])
+            )
+            self.consume("while")
+            self.consume("(")
+            cond = self.parse_expr()
+            self.consume(")")
+            self.consume(";")
+            return DoWhileStmt(body, cond)
+
         if tok.value == "return":
             self.consume("return")
             expr = None
@@ -1151,6 +1173,25 @@ def gen_stmt(self, stmt: ASTNode) -> None:
             self.loop_continue_stack.pop()
             self.loop_end_stack.pop()
 
+        elif isinstance(stmt, DoWhileStmt):
+            lbl_start = self.new_label("dowhile_start")
+            lbl_cond = self.new_label("dowhile_cond")
+            lbl_end = self.new_label("dowhile_end")
+            self.loop_end_stack.append(lbl_end)
+            self.loop_continue_stack.append(lbl_cond)
+
+            self.emit_label(lbl_start)
+            self.gen_block(stmt.body)
+
+            self.emit_label(lbl_cond)
+            self.gen_expr(stmt.condition)
+            self.emit("CPI R16, 0")
+            self.emit(f"BRNE {lbl_start}")
+
+            self.emit_label(lbl_end)
+            self.loop_continue_stack.pop()
+            self.loop_end_stack.pop()
+
         elif isinstance(stmt, ReturnStmt):
             if stmt.expr:
                 self.gen_expr(stmt.expr)

tests/test_compiler.py

@@ -3,7 +3,15 @@
 import pytest
 
 from tiny8 import CPU, assemble
-from tiny8.compiler import ContinueStmt, Lexer, Parser, compile_c, optimize_assembly
+from tiny8.compiler import (
+    ContinueStmt,
+    DoWhileStmt,
+    Lexer,
+    Literal,
+    Parser,
+    compile_c,
+    optimize_assembly,
+)
 
 
 def test_compiler_variable_and_arithmetic():
@@ -49,6 +57,113 @@ def test_compiler_while_loop():
     assert cpu.read_ram(0x81) == 15
 
 
+def test_compiler_do_while_basic():
+    c_code = """
+    int main() {
+        int sum = 0;
+        int i = 1;
+        do {
+            sum = sum + i;
+            i = i + 1;
+        } while (i <= 5);
+        poke(0x81, 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 + 3 + 4 + 5 = 15
+    assert cpu.read_ram(0x81) == 15
+
+
+def test_compiler_do_while_false_condition_executes_once():
+    c_code = """
+    int main() {
+        int x = 10;
+        do {
+            x = x + 5;
+        } while (0);
+        poke(0x82, x);
+        return 0;
+    }
+    """
+    asm = compile_c(c_code)
+    asm_res = assemble(asm)
+    cpu = CPU()
+    cpu.load_program(asm_res)
+    cpu.run(max_steps=100, show_progress=False)
+
+    # Executed exactly once: 10 + 5 = 15
+    assert cpu.read_ram(0x82) == 15
+
+
+def test_compiler_do_while_break_and_continue():
+    c_code = """
+    int main() {
+        int sum = 0;
+        int i = 0;
+        do {
+            i = i + 1;
+            if (i == 2) {
+                continue;
+            }
+            if (i == 4) {
+                break;
+            }
+            sum = sum + i;
+        } while (i < 5);
+        poke(0x83, 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)
+
+    # i=1: sum = 1
+    # i=2: continue (sum not added)
+    # i=3: sum = 1 + 3 = 4
+    # i=4: break (loop terminates)
+    assert cpu.read_ram(0x83) == 4
+
+
+def test_compiler_do_while_single_statement():
+    c_code = """
+    int main() {
+        int x = 0;
+        do x = x + 1; while (x < 3);
+        poke(0x84, x);
+        return 0;
+    }
+    """
+    asm = compile_c(c_code)
+    asm_res = assemble(asm)
+    cpu = CPU()
+    cpu.load_program(asm_res)
+    cpu.run(max_steps=200, show_progress=False)
+
+    assert cpu.read_ram(0x84) == 3
+
+
+def test_compiler_do_while_ast_and_lexer():
+    tokens = Lexer("do { x = 1; } while (x < 10);").tokenize()
+    assert tokens[0].type == "KEYWORD"
+    assert tokens[0].value == "do"
+
+    prog = Parser(Lexer("void main() { do { x = 1; } while (1); }").tokenize()).parse()
+    stmt = prog.functions[0].body.statements[0]
+    assert isinstance(stmt, DoWhileStmt)
+    assert len(stmt.body.statements) == 1
+    assert isinstance(stmt.condition, Literal)
+    assert stmt.condition.value == 1
+
+
 def test_compiler_if_else():
     c_code = """
     int main() {