diff --git a/src/tiny8/compiler.py b/src/tiny8/compiler.py index 72c305c..8e45ec1 100644 --- a/src/tiny8/compiler.py +++ b/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): diff --git a/tests/test_compiler.py b/tests/test_compiler.py index fdef0ed..441410e 100644 --- a/tests/test_compiler.py +++ b/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() {