feat(compiler): support compound assignment operators

src/tiny8/compiler.py

@@ -2,6 +2,7 @@
 
 Features:
 - Variable declarations and assignments (int/char)
+- Compound assignment operators (+=, -=, *=, /=, %=, &=, |=, ^=)
 - Arithmetic (+, -, *, /, %), bitwise (&, |, ^, ~), logical (!), and relational (<, <=, >, >=, ==, !=)
 - Control structures: if / else, while, return
 - Functions with register-based parameter passing (R16..R19) and return value in R16
@@ -38,7 +39,7 @@ class Lexer:
             r"\b(int|char|void|if|else|while|for|break|return|peek|poke|in|out|vram_set|vram_get|vram_clear|vram_puts)\b",
         ),
         ("ID", r"[a-zA-Z_]\w*"),
-        ("OP", r"<<|>>|==|!=|<=|>=|&&|\|\||[+\-*/%&|^=!<>]"),
+        ("OP", r"<<|>>|==|!=|<=|>=|&&|\|\||[+\-*/%&|^]=|[+\-*/%&|^=!<>]"),
         ("PUNCT", r"[;,\(\)\{\}\[\]]"),
         ("WS", r"\s+"),
     ]
@@ -136,6 +137,25 @@ class Assign(ASTNode):
     expr: ASTNode
 
 
+@dataclass
+class CompoundAssign(ASTNode):
+    """Compound assignment statement: name op= expr."""
+
+    name: str
+    op: str
+    expr: ASTNode
+
+
+@dataclass
+class ArrayCompoundAssign(ASTNode):
+    """Array element compound assignment: name[index] op= expr."""
+
+    name: str
+    index: ASTNode
+    op: str
+    expr: ASTNode
+
+
 @dataclass
 class IfStmt(ASTNode):
     """If-else statement."""
@@ -285,6 +305,9 @@ def _fold_unary(op: str, operand: ASTNode) -> ASTNode:
     return UnaryOp(op, operand)
 
 
+COMPOUND_ASSIGN_OPS = frozenset({"+=", "-=", "*=", "/=", "%=", "&=", "|=", "^="})
+
+
 class Parser:
     """Recursive descent parser."""
 
@@ -395,6 +418,12 @@ def parse_inline_stmt(self) -> ASTNode:
             next_tok = (
                 self.tokens[self.pos + 1] if self.pos + 1 < len(self.tokens) else None
             )
+            if next_tok and next_tok.value in COMPOUND_ASSIGN_OPS:
+                name = self.consume().value
+                op = self.consume().value[:-1]
+                expr = self.parse_expr()
+                self.consume(";")
+                return CompoundAssign(name, op, expr)
             if next_tok and next_tok.value == "=":
                 name = self.consume().value
                 self.consume("=")
@@ -411,6 +440,11 @@ def parse_inline_post(self) -> ASTNode:
             next_tok = (
                 self.tokens[self.pos + 1] if self.pos + 1 < len(self.tokens) else None
             )
+            if next_tok and next_tok.value in COMPOUND_ASSIGN_OPS:
+                name = self.consume().value
+                op = self.consume().value[:-1]
+                expr = self.parse_expr()
+                return CompoundAssign(name, op, expr)
             if next_tok and next_tok.value == "=":
                 name = self.consume().value
                 self.consume("=")
@@ -537,10 +571,22 @@ def parse_stmt(self) -> ASTNode:
                 self.consume("[")
                 idx_expr = self.parse_expr()
                 self.consume("]")
+                op_tok = self.peek()
+                if op_tok and op_tok.value in COMPOUND_ASSIGN_OPS:
+                    self.consume()
+                    expr = self.parse_expr()
+                    self.consume(";")
+                    return ArrayCompoundAssign(name, idx_expr, op_tok.value[:-1], expr)
                 self.consume("=")
                 expr = self.parse_expr()
                 self.consume(";")
                 return ArrayAssign(name, idx_expr, expr)
+            if next_tok and next_tok.value in COMPOUND_ASSIGN_OPS:
+                name = self.consume().value
+                op = self.consume().value[:-1]
+                expr = self.parse_expr()
+                self.consume(";")
+                return CompoundAssign(name, op, expr)
             if next_tok and next_tok.value == "=":
                 name = self.consume().value
                 self.consume("=")
@@ -960,6 +1006,32 @@ def gen_stmt(self, stmt: ASTNode) -> None:
                 self.emit(f"LDI R31, {addr >> 8}")
             self.emit("ST R30, R16")
 
+        elif isinstance(stmt, CompoundAssign):
+            self.gen_stmt(
+                Assign(
+                    stmt.name,
+                    BinaryOp(op=stmt.op, left=Var(stmt.name), right=stmt.expr),
+                )
+            )
+
+        elif isinstance(stmt, ArrayCompoundAssign):
+            # Evaluate the index once into a hidden slot so side effects in the
+            # index expression are not repeated by the read-modify-write.
+            tmp = self.new_label("__cidx")
+            self.gen_stmt(Assign(tmp, stmt.index))
+            idx = Var(tmp)
+            self.gen_stmt(
+                ArrayAssign(
+                    stmt.name,
+                    idx,
+                    BinaryOp(
+                        op=stmt.op,
+                        left=ArrayAccess(stmt.name, idx),
+                        right=stmt.expr,
+                    ),
+                )
+            )
+
         elif isinstance(stmt, ForStmt):
             if stmt.init:
                 self.gen_stmt(stmt.init)

tests/test_compiler.py

@@ -421,3 +421,106 @@ def test_peephole_remove_jump_to_next_label(self):
         ]
         opt = optimize_assembly(lines)
         assert opt == ["loop:", "    NOP"]
+
+
+def _run_compound_program(c_code, max_steps=5000):
+    cpu = CPU()
+    cpu.load_program(assemble(compile_c(c_code)))
+    cpu.run(max_steps=max_steps, show_progress=False)
+    return cpu
+
+
+def test_compiler_compound_assign_scalar_operators():
+    cpu = _run_compound_program(
+        """
+        int main() {
+            int a = 10;
+            a += 5;
+            poke(0x80, a);
+            a -= 3;
+            poke(0x81, a);
+            a *= 2;
+            poke(0x82, a);
+            a /= 5;
+            poke(0x83, a);
+            a %= 3;
+            poke(0x84, a);
+            a |= 6;
+            poke(0x85, a);
+            a &= 5;
+            poke(0x86, a);
+            a ^= 3;
+            poke(0x87, a);
+            return 0;
+        }
+        """
+    )
+    assert [cpu.read_ram(0x80 + i) for i in range(8)] == [15, 12, 24, 4, 1, 7, 5, 6]
+
+
+def test_compiler_compound_assign_uses_whole_right_hand_side():
+    cpu = _run_compound_program(
+        """
+        int main() {
+            int a = 3;
+            int b = 4;
+            a += b * 2;
+            poke(0x80, a);
+            return 0;
+        }
+        """
+    )
+    assert cpu.read_ram(0x80) == 11
+
+
+def test_compiler_compound_assign_in_for_loop_step():
+    cpu = _run_compound_program(
+        """
+        int main() {
+            int sum = 0;
+            int i = 0;
+            for (i = 0; i < 10; i += 2) {
+                sum += i;
+            }
+            poke(0x80, sum);
+            return 0;
+        }
+        """
+    )
+    assert cpu.read_ram(0x80) == 20
+
+
+def test_compiler_compound_assign_array_elements():
+    cpu = _run_compound_program(
+        """
+        int main() {
+            int arr[3] = {1, 2, 3};
+            arr[1] += 10;
+            arr[2] *= 4;
+            poke(0x80, arr[0]);
+            poke(0x81, arr[1]);
+            poke(0x82, arr[2]);
+            return 0;
+        }
+        """
+    )
+    assert [cpu.read_ram(0x80 + i) for i in range(3)] == [1, 12, 12]
+
+
+def test_compiler_compound_assign_array_index_evaluated_once():
+    # The index expression writes to RAM as a side effect, so evaluating it twice
+    # would read one element and write another.
+    cpu = _run_compound_program(
+        """
+        int main() {
+            int arr[2] = {5, 7};
+            arr[poke(0x90, peek(0x90) + 1) - 1] += 10;
+            poke(0x80, arr[0]);
+            poke(0x81, arr[1]);
+            return 0;
+        }
+        """
+    )
+    assert cpu.read_ram(0x90) == 1
+    assert cpu.read_ram(0x80) == 15
+    assert cpu.read_ram(0x81) == 7