Add optimizer, add supported AST nodes

6 files changed, 369 insertions, 206 deletions

diff --git a/asm/atk16.py b/asm/atk16.py
index f1a2c95..6d9bc10 100644
--- a/asm/atk16.py
+++ b/asm/atk16.py
@@ -1,5 +1,4 @@
 from typing import Literal, Any, TypeVar, NewType, Never
-a = ["A", "B", "C"]
 Char = Literal['\x00', '\x01', '\x02', '\x03', '\x04', '\x05', '\x06', '\x07', '\x08', '\t', '\n', '\x0b', '\x0c', '\r', '\x0e', '\x0f', '\x10', '\x11', '\x12', '\x13', '\x14', '\x15', '\x16', '\x17', '\x18', '\x19', '\x1a', '\x1b', '\x1c', '\x1d', '\x1e', '\x1f', ' ', '!', '"', '#', '$', '%', '&', "'", '(', ')', '*', '+', ',', '-', '.', '/', '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', ':', ';', '<', '=', '>', '?', '@', 'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z', '[', '\\', ']', '^', '_', '`', 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z', '{', '|', '}', '~', '\x7f', '\x80', '\x81', '\x82', '\x83', '\x84', '\x85', '\x86', '\x87', '\x88', '\x89', '\x8a', '\x8b', '\x8c', '\x8d', '\x8e', '\x8f', '\x90', '\x91', '\x92', '\x93', '\x94', '\x95', '\x96', '\x97', '\x98', '\x99', '\x9a', '\x9b', '\x9c', '\x9d', '\x9e', '\x9f', '\xa0', '¡', '¢', '£', '¤', '¥', '¦', '§', '¨', '©', 'ª', '«', '¬', '\xad', '®', '¯', '°', '±', '²', '³', '´', 'µ', '¶', '·', '¸', '¹', 'º', '»', '¼', '½', '¾', '¿', 'À', 'Á', 'Â', 'Ã', 'Ä', 'Å', 'Æ', 'Ç', 'È', 'É', 'Ê', 'Ë', 'Ì', 'Í', 'Î', 'Ï', 'Ð', 'Ñ', 'Ò', 'Ó', 'Ô', 'Õ', 'Ö', '×', 'Ø', 'Ù', 'Ú', 'Û', 'Ü', 'Ý', 'Þ', 'ß', 'à', 'á', 'â', 'ã', 'ä', 'å', 'æ', 'ç', 'è', 'é', 'ê', 'ë', 'ì', 'í', 'î', 'ï', 'ð', 'ñ', 'ò', 'ó', 'ô', 'õ', 'ö', '÷', 'ø', 'ù', 'ú', 'û', 'ü', 'ý', 'þ', 'ÿ']
 ConstInt = int
 
@@ -12,9 +11,6 @@ def load(addr: int) -> int:
 def put_char(c: int):
     return NotImplemented
 
-def set_graphics_mode(mode: Literal[0, 1, 2]):
-    return NotImplemented
-
 T = TypeVar("T")
 def call_inline(expr: T) -> T:
     """Inline the function call `expr` at the callsite.
diff --git a/asm/ext_std.py b/asm/ext_std.py
index 3d87372..fd18953 100644
--- a/asm/ext_std.py
+++ b/asm/ext_std.py
@@ -12,14 +12,29 @@ def expand_addi(left: str, imm: str, target: str) -> ExpandResult:
 def expand_subi(left: str, imm: str, target: str) -> ExpandResult:
   return [["ali", "al_minus", left, imm, target]]
 
+def expand_not(reg: str, target: str) -> ExpandResult:
+  return [["alr", "al_xor", reg, "0xFFFF", target]]
+
+def expand_noti(imm: str, target: str) -> ExpandResult:
+  return [["ali", "al_xor", imm, "0xFFFF", target]]
+
 def expand_and(left: str, right: str, target: str) -> ExpandResult:
   return [["alr", "al_and", left, right, target]]
 
+def expand_andi(left: str, imm: str, target: str) -> ExpandResult:
+  return [["ali", "al_and", left, imm, target]]
+
 def expand_or(left: str, right: str, target: str) -> ExpandResult:
   return [["alr", "al_or", left, right, target]]
 
-def expand_xor(left: str, right: str, target: str) -> ExpandResult:
-  return [["alr", "al_xor", left, right, target]]
+def expand_ori(left: str, imm: str, target: str) -> ExpandResult:
+  return [["ali", "al_or", left, imm, target]]
+
+def expand_xor(left: str, imm: str, target: str) -> ExpandResult:
+  return [["alr", "al_xor", left, imm, target]]
+
+def expand_xori(left: str, right: str, target: str) -> ExpandResult:
+  return [["ali", "al_xor", left, right, target]]
 
 def expand_sll(left: str, right: str, target: str) -> ExpandResult:
   return [["alr", "al_sll", left, right, target]]
@@ -105,9 +120,14 @@ expansions: OpExpansionDict = {
   "sub": expand_sub,
   "addi": expand_addi,
   "subi": expand_subi,
+  "not": expand_not,
+  "noti": expand_noti,
   "and": expand_and,
+  "andi": expand_andi,
   "or": expand_or,
+  "ori": expand_ori,
   "xor": expand_xor,
+  "xori": expand_xori,
   "sll": expand_sll,
   "slr": expand_slr,
   "sar": expand_sar,
diff --git a/asm/test_py_src.atk16 b/asm/test_py_src.atk16
index c1bf1af..c035b72 100644
--- a/asm/test_py_src.atk16
+++ b/asm/test_py_src.atk16
@@ -73,88 +73,44 @@
 
 ;; END BOOTSTRAP
 
-@label TEXT_MODE
-  1
-@label int_0
-  23
-@label int_1
-  63488
-@label int_2
-  72
-@label int_3
-  63489
-@label int_4
-  69
-@label int_5
-  63490
-@label int_6
-  76
-@label int_7
-  63491
-@label int_8
-  76
-@label int_9
-  63492
-@label int_10
-  79
+
 
 @label main
-; builtin call store [<ast.Constant object at 0x1045a3d90>, <ast.Name object at 0x1045a3d60>]
-  ldi int_0 RA
-  ldr RA RA
-  spu RA
-  spo RB
-  spo RA
-  str RB RA
-; builtin call store [<ast.Constant object at 0x1045a3c70>, <ast.Constant object at 0x1045a3c40>]
-  ldi int_1 RA
-  ldr RA RA
-  spu RA
-  ldi int_2 RA
-  ldr RA RA
-  spu RA
-  spo RB
-  spo RA
-  str RB RA
-; builtin call store [<ast.Constant object at 0x1045a3b50>, <ast.Constant object at 0x1045a3b20>]
-  ldi int_3 RA
-  ldr RA RA
-  spu RA
-  ldi int_4 RA
-  ldr RA RA
-  spu RA
-  spo RB
-  spo RA
-  str RB RA
-; builtin call store [<ast.Constant object at 0x1045a3a30>, <ast.Constant object at 0x1045a3a00>]
-  ldi int_5 RA
-  ldr RA RA
-  spu RA
-  ldi int_6 RA
-  ldr RA RA
-  spu RA
-  spo RB
-  spo RA
-  str RB RA
-; builtin call store [<ast.Constant object at 0x1045a3910>, <ast.Constant object at 0x1045a38e0>]
-  ldi int_7 RA
-  ldr RA RA
-  spu RA
-  ldi int_8 RA
-  ldr RA RA
+; <ast.If object at 0x1013c1e40>
+; <ast.Constant object at 0x1013c20b0>
+  ldi 1 RA
   spu RA
-  spo RB
   spo RA
-  str RB RA
-; builtin call store [<ast.Constant object at 0x1045a37f0>, <ast.Constant object at 0x1045a37c0>]
-  ldi int_9 RA
-  ldr RA RA
+  addi RA 0 RA
+  bri zero If_false_branch_0
+; <ast.Constant object at 0x1013c2110>
+  ldi 2 RA
   spu RA
-  ldi int_10 RA
-  ldr RA RA
+  jpi If_end_branch_1
+@label If_false_branch_0
+; <ast.Constant object at 0x1013c2170>
+  ldi 3 RA
   spu RA
-  spo RB
+@label If_end_branch_1
+; <ast.While object at 0x1013c2230>
+@label While_test_2
+; <ast.Constant object at 0x1013c2260>
+  ldi 3 RB
+  spu RB
   spo RA
-  str RB RA
-; builtin call asm [<ast.Constant object at 0x1045a2c20>]
-  ldi 1 RA
\ No newline at end of file
+  addi RA 0 RA
+  bri zero While_else_3
+; <ast.Constant object at 0x1013c22c0>
+  ldi 4 RB
+  spu RB
+; <ast.Break object at 0x1013c22f0>
+  jpi While_end_4
+; <ast.Constant object at 0x1013c2350>
+  ldi 5 RB
+  spu RB
+  jpi While_test_2
+@label While_else_3
+; <ast.Constant object at 0x1013c2470>
+  ldi 6 RB
+  spu RB
+@label While_end_4
\ No newline at end of file
diff --git a/asm/test_py_src.atk16_optimized b/asm/test_py_src.atk16_optimized
index a6a77ab..673800b 100644
--- a/asm/test_py_src.atk16_optimized
+++ b/asm/test_py_src.atk16_optimized
@@ -58,58 +58,28 @@
   ldi gr_disabled_mode RB
   str RB RA
   jpi main
-@label TEXT_MODE
-  1
-@label int_0
-  23
-@label int_1
-  63488
-@label int_2
-  72
-@label int_3
-  63489
-@label int_4
-  69
-@label int_5
-  63490
-@label int_6
-  76
-@label int_7
-  63491
-@label int_8
-  76
-@label int_9
-  63492
-@label int_10
-  79
 @label main
-  ldi int_0 RA
-  ldr RA RB
-  spo RA
-  str RB RA
-  ldi int_1 RA
-  ldr RA RA
-  ldi int_2 RB
-  ldr RB RB
-  str RB RA
-  ldi int_3 RA
-  ldr RA RA
-  ldi int_4 RB
-  ldr RB RB
-  str RB RA
-  ldi int_5 RA
-  ldr RA RA
-  ldi int_6 RB
-  ldr RB RB
-  str RB RA
-  ldi int_7 RA
-  ldr RA RA
-  ldi int_8 RB
-  ldr RB RB
-  str RB RA
-  ldi int_9 RA
-  ldr RA RA
-  ldi int_10 RB
-  ldr RB RB
-  str RB RA
-  ldi 1 RA
\ No newline at end of file
+  ldi 1 RA
+  addi RA 0 RA
+  bri zero If_false_branch_0
+  ldi 2 RA
+  spu RA
+  jpi If_end_branch_1
+@label If_false_branch_0
+  ldi 3 RA
+  spu RA
+@label If_end_branch_1
+@label While_test_2
+  ldi 3 RA
+  addi RA 0 RA
+  bri zero While_else_3
+  ldi 4 RB
+  spu RB
+  jpi While_end_4
+  ldi 5 RB
+  spu RB
+  jpi While_test_2
+@label While_else_3
+  ldi 6 RB
+  spu RB
+@label While_end_4
\ No newline at end of file
diff --git a/asm/test_py_src.py b/asm/test_py_src.py
index a1eca3f..48a5cac 100644
--- a/asm/test_py_src.py
+++ b/asm/test_py_src.py
@@ -1,17 +1,29 @@
 import atk16
 
-TEXT_MODE: atk16.ConstInt          = 1
-GRAPHICS_MODE_ADDR: atk16.ConstInt = 0x17
+if True:
+    2
+else:
+    3
 
-atk16.store(GRAPHICS_MODE_ADDR, TEXT_MODE)
-atk16.store(0xF800, 'H')
-atk16.store(0xF801, 'E')
-atk16.store(0xF802, 'L')
-atk16.store(0xF803, 'L')
-atk16.store(0xF804, 'O')
+while 3:
+    4
+    break
+    5
+else:
+    6
 
-a = atk16.call_inline(
-  atk16.load(0xF000)
-)
+# TEXT_MODE: atk16.ConstInt          = 1
+# GRAPHICS_MODE_ADDR: atk16.ConstInt = 0x17
 
-atk16.asm("ldi 1 RA")
+# atk16.store(GRAPHICS_MODE_ADDR, TEXT_MODE)
+# atk16.store(0xF800, 'H')
+# atk16.store(0xF801, 'E')
+# atk16.store(0xF802, 'L')
+# atk16.store(0xF803, 'L')
+# atk16.store(0xF804, 'O')
+
+# a = atk16.call_inline(
+#   atk16.load(0xF000)
+# )
+
+# atk16.asm("ldi 1 RA")
diff --git a/src/ast_compiler.py b/src/ast_compiler.py
index 2208e43..bcd06bf 100644
--- a/src/ast_compiler.py
+++ b/src/ast_compiler.py
@@ -1,7 +1,6 @@
 #!/usr/bin/env python3
 # Generate .atk16 assembly from a subset of Python
 
-from _ast import AnnAssign, Expr, Module
 import sys
 import ast
 
@@ -51,6 +50,7 @@ class Compiler(ast.NodeVisitor):
     self.const_bindings: dict[str, Label] = {}
     self.call_depth: int = 0
     self.unique_name_counter = 0
+    self.latest_break_target: Label | None = None
 
     self.reserved_regs: OrderedDict[RegChar, None] = OrderedDict()
 
@@ -94,38 +94,22 @@ class Compiler(ast.NodeVisitor):
     ])
 
   def emit_builtin_call(self, name: str, args: list[ast.expr]):
-    self.emit(f"; builtin call {name} {args}")
+    self.emit(f"; Builtin call {name} {args}")
     match name:
       case "asm":
         match args:
-          case [ast.Constant(value)]:
-            if type(value) != str:
-              raise Exception(f"asm: invalid arg type: {type(value)}")
-
+          case [ast.Constant(str(value))]:
             self.emit(value)
-          case _: raise Exception(f"asm: invalid args: {args}")
-      case "set_graphics_mode":
-        match args:
-          case [ast.Constant(value)]:
-            if type(value) != int:
-              raise Exception(f"set_graphics_mode: invalid arg type: {type(value)}")
-
-            reg1 = self.alloc_reg()
-            reg2 = self.alloc_reg()
-            self.emit(f"ldi vt_gr_mode_addr {reg1}")
-            self.emit(f"ldr {reg1} {reg1}")
-            self.emit(f"ldi {value} {reg2}")
-            self.emit(f"str {reg2} {reg1}")
-            self.free_reg(reg1)
-            self.free_reg(reg2)
 
-          case _: raise Exception(f"set_graphics_mode: invalid args: {args}")
+          case other: raise Exception(f"asm: invalid args: {other}")
       case "store":
         if len(args) != 2:
           raise Exception("Invalid number of arguments to store: " + str(len(args)))
 
-        self.eval_expr_and_spu(args[0])
-        self.eval_expr_and_spu(args[1])
+        # Evaluate args before call
+        for arg in args:
+          self.visit(arg)
+
         arg1 = self.alloc_reg()
         arg2 = self.alloc_reg()
         self.emit(f"spo {arg2}")
@@ -147,43 +131,237 @@ class Compiler(ast.NodeVisitor):
     self.emit(f"spu {reg}")
     self.free_reg(reg)
 
-  def eval_expr_and_spu(self, expr: ast.expr):
-    match expr:
-      case ast.Constant(int(value)):
+  def visit_Module(self, node: ast.Module):
+    for stmt in node.body:
+      self.visit(stmt)
+
+  def visit_Expr(self, expr: ast.Expr):
+    # if self.call_depth == 0 and type(expr.value) != ast.Call:
+    #   self.emit("; NOP top-level expression")
+    #   return
+
+    self.visit(expr.value)
+
+  def visit_UnaryOp(self, node: ast.UnaryOp):
+    self.emit(f"; {node}")
+    self.visit(node.operand)
+    match node.op:
+      case ast.Not():
+        reg1 = self.alloc_reg()
+        reg2 = self.alloc_reg()
+        self.emit(f"spo {reg1}")
+        self.emit(f"ldi 1 {reg2}")
+        self.emit(f"andi {reg1} 1 {reg1}")
+        self.emit(f"xor {reg1} {reg2} {reg1}")
+        self.emit(f"spu {reg1}")
+        self.free_reg(reg1)
+        self.free_reg(reg2)
+      case ast.Invert(): # aka bitwise not
+        reg1 = self.alloc_reg()
+        self.emit(f"spo {reg1}")
+        self.emit(f"not {reg1}")
+        self.emit(f"spu {reg1}")
+        self.free_reg(reg1)
+      case ast.UAdd(): # +a
+        pass # nop
+      case ast.USub(): # -a
+        reg1 = self.alloc_reg()
+        self.emit(f"spo {reg1}")
+        self.emit(f"not {reg1}")
+        self.emit(f"addi {reg1} 1 {reg1}")
+        self.emit(f"spu {reg1}")
+        self.free_reg(reg1)
+      case other:
+        raise NotImplementedError(f"Unhandled UnaryOp: {other}")
+
+  def visit_BoolOp(self, node: ast.BoolOp):
+    self.emit(f"; {node}")
+
+    match node.op:
+      case ast.And():
+        self.emit("; Boolean and")
+        reg1 = self.alloc_reg()
+        label_short_circuit = self.get_unique_name("And_short_circuit")
+
+        for arg in node.values:
+          self.emit(f"; And operand {arg}")
+          self.visit(arg)
+
+          self.emit(f"spo {reg1}")
+          self.emit(f"addi {reg1} 0 {reg1}")
+          self.emit(f"bri zero {label_short_circuit}")
+
+        self.emit(f"@label {label_short_circuit}")
+        self.emit(f"spu {reg1}")
+        self.free_reg(reg1)
+      case ast.Or():
+        self.emit("; Boolean or")
+        reg1 = self.alloc_reg()
+        reg2 = self.alloc_reg()
+        label_short_circuit = self.get_unique_name("Or_short_circuit")
+
+        for arg in node.values:
+          self.emit(f"; Or operand {arg}")
+          self.visit(arg)
+
+          self.emit(f"spo {reg1}")
+          self.emit(f"subi {reg1} 1 {reg2}")
+          self.emit(f"bri carry {label_short_circuit}")
+
+        self.emit(f"@label {label_short_circuit}")
+        self.emit(f"spu {reg1}")
+        self.free_reg(reg1)
+        self.free_reg(reg2)
+      case other:
+        raise NotImplementedError(f"Unhandled BoolOp: {other}")
+
+  def visit_BinOp(self, node: ast.BinOp):
+    self.emit(f"; {node}")
+
+    self.emit(f"; BinOp lhs {node}")
+    self.visit(node.left)
+    self.emit(f"; BinOp rhs {node}")
+    self.visit(node.right)
+
+    reg1 = self.alloc_reg()
+    reg2 = self.alloc_reg()
+
+    self.emit(f"spo {reg2}")
+    self.emit(f"spo {reg1}")
+
+    match node.op:
+      case ast.Add():
+        self.emit(f"add {reg1} {reg2} {reg1}")
+      case ast.Sub():
+        self.emit(f"sub {reg1} {reg2} {reg1}")
+      case ast.BitAnd():
+        self.emit(f"and {reg1} {reg2} {reg1}")
+      case ast.BitOr():
+        self.emit(f"or {reg1} {reg2} {reg1}")
+      case ast.BitXor():
+        self.emit(f"xor {reg1} {reg2} {reg1}")
+      case ast.LShift():
+        self.emit(f"sll {reg1} {reg2} {reg1}")
+      case ast.RShift():
+        self.emit(f"slr {reg1} {reg2} {reg1}")
+      case other:
+        raise NotImplementedError(f"Unhandled BinOp: {other}")
+
+    self.emit(f"spu {reg1}")
+
+    self.free_reg(reg1)
+    self.free_reg(reg2)
+
+  def visit_Constant(self, node: ast.Constant):
+    self.emit(f"; {node}")
+    match node.value:
+      case bool(value):
+        int_value = 1 if value else 0
+        self.eval_int_constant_and_spu(int_value)
+      case int(value):
         self.eval_int_constant_and_spu(value)
-      case ast.Constant(str(value)):
+      case str(value):
         if len(value) > 1:
           raise Exception("Invalid string, only single char values allowed: " + value)
 
         c = value[0]
         int_value = ord(c)
         self.eval_int_constant_and_spu(int_value)
+      case other:
+        raise NotImplementedError(f"Unhandled Constant: {other}")
 
+  def visit_If(self, node: ast.If):
+    self.emit(f"; {node}")
 
-  def visit_Module(self, node: Module):
-    for stmt in node.body:
-      self.visit(stmt)
+    self.visit(node.test)
+    reg1 = self.alloc_reg()
+    label_false = self.get_unique_name("If_false_branch")
+    label_end = self.get_unique_name("If_end_branch")
+    self.emit(f"spo {reg1}")
+    self.emit(f"addi {reg1} 0 {reg1}")
+    self.emit(f"bri zero {label_false}")
+    self.free_reg(reg1)
+
+    for true_branch_stmt in node.body:
+      self.visit(true_branch_stmt)
+
+    self.emit(f"jpi {label_end}")
+    self.emit(f"@label {label_false}")
+
+    for false_branch_stmt in node.orelse:
+      self.visit(false_branch_stmt)
+
+    self.emit(f"@label {label_end}")
+
+  def visit_While(self, node: ast.While):
+    self.emit(f"; {node}")
+
+    reg1 = self.alloc_reg()
+    label_test = self.get_unique_name("While_test")
+    label_else = self.get_unique_name("While_else")
+    label_end = self.get_unique_name("While_end")
 
-  def visit_Expr(self, expr: Expr):
-    if self.call_depth == 0 and type(expr.value) != ast.Call:
-      self.emit("; NOP top-level expression")
-      return
+    prev_break_target = self.latest_break_target
+    self.latest_break_target = label_end
 
-    match expr.value:
-      case ast.Call(func, args, keywords):
-        match func:
-          case ast.Attribute(ast.Name(id="atk16"), attr):
-            print("FOUND CAPTURED ATK16 CALL: " + attr)
-            self.emit_builtin_call(attr, args)
-          case ast.Name(name):
-            addr = self.const_bindings[name]
-            self.emit(f"csi {addr}")
-          case _:
-            raise NotImplementedError("Unhandled Call: " + str(func))
-      case _:
-        raise NotImplementedError("Unhandled Expr.value: " + str(expr.value))
+    self.emit(f"@label {label_test}")
+    self.visit(node.test)
+    self.emit(f"spo {reg1}")
+    self.emit(f"addi {reg1} 0 {reg1}")
+    self.emit(f"bri zero {label_else}")
 
-  def visit_AnnAssign(self, node: AnnAssign):
+    for body_stmt in node.body:
+      self.visit(body_stmt)
+
+    self.emit(f"jpi {label_test}")
+
+    self.emit(f"@label {label_else}")
+
+    for else_stmt in node.orelse:
+      self.visit(else_stmt)
+
+    self.emit(f"@label {label_end}")
+
+    self.free_reg(reg1)
+    self.latest_break_target = prev_break_target
+
+  def visit_Break(self, node: ast.Break):
+    self.emit(f"; {node}")
+
+    if self.latest_break_target is None:
+      raise Exception("Invalid break: no break target defined, i.e. no place to break out to")
+
+    self.emit(f"jpi {self.latest_break_target}")
+
+  def visit_Call(self, node: ast.Call):
+    self.emit(f"; {node}")
+    match node.func:
+      case ast.Attribute(ast.Name(id="atk16"), attr):
+        self.emit_builtin_call(attr, node.args)
+      case ast.Name(name):
+        # Evaluate args before call
+        for arg in node.args:
+          self.visit(arg)
+
+        addr = self.const_bindings[name]
+        self.emit(f"csi {addr}")
+      case other:
+        raise NotImplementedError(f"Unhandled Call: {other}")
+
+  def visit_Assign(self, node: ast.Assign) -> Any:
+    raise NotImplementedError(f"TODO assign {node}")
+
+  def generic_visit(self, node: ast.AST) -> Any:
+    raise NotImplementedError(f"type {type(node)}, value: {node}")
+
+  def visit_Import(self, node: ast.Import) -> Any:
+    match node.names:
+      case [ast.alias(name="atk16")]:
+        return
+      case other:
+        raise NotImplementedError(f"Unsupported import {other}")
+
+  def visit_AnnAssign(self, node: ast.AnnAssign):
     match node:
       case ast.AnnAssign(
         target=ast.Name(name),
@@ -209,9 +387,12 @@ class Optimizer:
     asm = [row for row in asm if not len(row) == 0]
     asm = [row.split() for row in asm]
     asm = self.compact_spu_spo_pattern(asm)
-    asm = self.compact_load_mov_pattern(asm)
-    asm = self.compact_load_mov_pattern(asm)
+    asm = self.compact_target_mov_pattern(asm)
+    asm = self.compact_target_mov_pattern(asm)
+    asm = self.compact_mov_source_pattern(asm)
+    asm = self.compact_mov_source_pattern(asm)
     asm = self.compact_spu_load_spo_pattern(asm)
+    # asm = self.convert_alr_to_ali(asm)
     result = "\n".join([format_asm_row(" ".join(row)) for row in asm])
 
     return result
@@ -252,7 +433,8 @@ class Optimizer:
 
     return result
 
-  def compact_load_mov_pattern(self, asm: list[list[str]]) -> list[list[str]]:
+  def compact_target_mov_pattern(self, asm: list[list[str]]) -> list[list[str]]:
+    ops_with_target = ["ldi", "ldr", "add", "sub", "addi", "subi", "and", "or", "xor", "sll", "slr", "sar", "slli", "slri", "sari", "inc", "dec", "mov", "ali", "alr", "lpc"]
     i = 0
     result: list[list[str]] = []
     while i < len(asm):
@@ -263,12 +445,39 @@ class Optimizer:
         result.append(current)
         continue
 
-      if (current[0] == "ldi" or current[0] == "ldr") and next[0] == "mov":
-        load_op, load_from, load_target_reg = current[0], current[1], current[2]
+      if current[0] in ops_with_target and next[0] == "mov":
+
+        op_op, op_operands, op_target_reg = current[0], current[1:len(current) - 1], current[len(current) - 1]
         mov_from_reg, mov_to_reg = next[1], next[2]
 
-        if load_target_reg == mov_from_reg:
-          ret = [load_op, load_from, mov_to_reg]
+        if op_target_reg == mov_from_reg:
+          ret = [op_op, *op_operands, mov_to_reg]
+          result.append(ret)
+          i += 1
+          continue
+
+      result.append(current)
+
+    return result
+
+  def compact_mov_source_pattern(self, asm: list[list[str]]) -> list[list[str]]:
+    ops_with_source = ["str", "ldr", "add", "sub", "addi", "subi", "and", "or", "xor", "sll", "slr", "sar", "slli", "slri", "sari", "mov", "ali", "alr"]
+    i = 0
+    result: list[list[str]] = []
+    while i < len(asm):
+      current = asm[i]
+      next = asm[i + 1] if i + 1 < len(asm) else None
+      i += 1
+      if current[0].startswith("@") or next is None:
+        result.append(current)
+        continue
+
+      if current[0] == "mov" and next[0] in ops_with_source:
+        mov_from_reg, mov_to_reg = current[1], current[2]
+        op_op, op_source_reg, op_operands = next[0], next[1], next[2:]
+
+        if op_source_reg == mov_to_reg:
+          ret = [op_op, op_source_reg, *op_operands]
           result.append(ret)
           i += 1
           continue