Refactor into multipass compiler

9 files changed, 357 insertions, 210 deletions

diff --git a/src/asm_pass1.py b/src/asm_pass1.py
new file mode 100644
index 0000000..6537aea
--- /dev/null
+++ b/src/asm_pass1.py
@@ -0,0 +1,52 @@
+import importlib
+from dataclasses import dataclass
+from assembler_ops import *
+from assembler_eval import *
+
+@dataclass
+class Result1Line:
+  line_num: int
+  parts: list[str]
+
+@dataclass
+class Result1:
+  lines: list[Result1Line]
+  options: Options
+  operations: OpExpansionDict
+
+def pass_1(lines: list[str]) -> Result1:
+  options = Options()
+  result_lines: list[Result1Line] = []
+  operations: OpExpansionDict = default_expansions.copy()
+
+  for (line_num, line) in enumerate(lines):
+    line = line.split(";")[0].strip()
+    if line == "": continue
+    keyword, *args = line.lower().split()
+    match keyword:
+      case "@opt":
+        opt_name, opt_value = args
+        match opt_name:
+          case "stack_pointer": options.stack_pointer = opt_value
+          case "csr_scratch": options.csr_scratch = opt_value
+          case _: raise Exception("Unknown @opt: " + opt_name)
+      case "@use":
+        module_name, ops = args[0].split(":")
+        ops_split = ops.split(",")
+        module = importlib.import_module(module_name)
+        mod_expansions: OpExpansionDict = module.expansions
+        for op in mod_expansions:
+          expansion = mod_expansions[op]
+          if ops == "*" or op in ops_split:
+            operations[op] = expansion
+      case _:
+        result_lines.append(Result1Line(
+          line_num=line_num,
+          parts=[keyword, *args]
+        ))
+
+  return Result1(
+    operations=operations,
+    options=options,
+    lines=result_lines
+  )
diff --git a/src/asm_pass2.py b/src/asm_pass2.py
new file mode 100644
index 0000000..77b8fa6
--- /dev/null
+++ b/src/asm_pass2.py
@@ -0,0 +1,39 @@
+from dataclasses import dataclass
+from assembler_ops import *
+from assembler_eval import *
+from asm_pass1 import *
+
+@dataclass
+class Result2Line:
+  line_num: int
+  parts: list[str]
+
+@dataclass
+class Result2:
+  lines: list[Result2Line]
+  options: Options
+  operations: OpExpansionDict
+
+def pass_2(result1: Result1) -> Result2:
+  options = Options()
+  result_lines: list[Result2Line] = []
+
+  for line in result1.lines:
+    keyword, *args = line.parts
+    if not keyword.startswith("@"):
+      fn = result1.operations[keyword]
+      output: list[list[str]] = fn(*args)
+    else:
+      output = [line.parts]
+
+    for parts in output:
+      result_lines.append(Result2Line(
+        line_num=line.line_num,
+        parts=parts
+      ))
+
+  return Result2(
+    operations=result1.operations,
+    options=options,
+    lines=result_lines
+  )
diff --git a/src/asm_pass3.py b/src/asm_pass3.py
new file mode 100644
index 0000000..5365e3b
--- /dev/null
+++ b/src/asm_pass3.py
@@ -0,0 +1,46 @@
+from dataclasses import dataclass
+from assembler_ops import *
+from assembler_eval import *
+from asm_pass2 import *
+
+@dataclass
+class Result3Line:
+  line_num: int
+  address: int
+  parts: list[str]
+
+@dataclass
+class Result3:
+  lines: list[Result3Line]
+  options: Options
+  operations: OpExpansionDict
+  labels: dict[str, int]
+
+def pass_3(result2: Result2) -> Result3:
+  result_lines: list[Result3Line] = []
+  labels: dict[str, int] = {}
+  address = 0
+
+  for line in result2.lines:
+    keyword, *args = line.parts
+    match keyword:
+      case "@address":
+        address = eval(args[0])
+        continue
+      case "@label":
+        labels[args[0]] = address
+        continue
+      case _:
+        result_lines.append(Result3Line(
+          line_num=line.line_num,
+          parts=line.parts,
+          address=address,
+        ))
+        address += 1
+
+  return Result3(
+    operations=result2.operations,
+    options=result2.options,
+    lines=result_lines,
+    labels=labels,
+  )
diff --git a/src/asm_pass4.py b/src/asm_pass4.py
new file mode 100644
index 0000000..96d9e65
--- /dev/null
+++ b/src/asm_pass4.py
@@ -0,0 +1,64 @@
+from dataclasses import dataclass
+from assembler_ops import *
+from assembler_eval import *
+from asm_pass3 import *
+
+@dataclass
+class Result4Line:
+  line_num: int
+  address: int
+  word: int
+  text: str
+
+@dataclass
+class Result4:
+  lines: list[Result4Line]
+  options: Options
+  operations: OpExpansionDict
+  labels: dict[str, int]
+
+def translate_opt(arg: str, options: Options) -> str:
+  match arg:
+    case "__STACK_POINTER": return options.stack_pointer
+    case "__CSR_SCRATCH": return options.csr_scratch
+    case _: return arg
+
+def pass_4(result3: Result3) -> Result4:
+  result_lines: list[Result4Line] = []
+  address = 0
+
+  for line in result3.lines:
+    keyword, *args = line.parts
+    meta = Meta(
+      address=address,
+    )
+
+    args = list(map(lambda a: translate_opt(a, result3.options), args))
+    text = " ".join([keyword, *args])
+
+    if keyword in operations:
+      fn = operations[keyword]
+      word = fn(meta, result3.labels, *args)
+      result_lines.append(Result4Line(
+        line_num=line.line_num,
+        address=line.address,
+        word=word,
+        text=text,
+      ))
+    else:
+      try:
+        result_lines.append(Result4Line(
+          line_num=line.line_num,
+          address=line.address,
+          word=eval_expr(result3.labels, keyword),
+          text=text,
+        ))
+      except:
+        raise Exception(f"Invalid assembly at {line.line_num + 1}\n\n{line}")
+
+  return Result4(
+    operations=result3.operations,
+    labels=result3.labels,
+    options=result3.options,
+    lines=result_lines,
+  )
diff --git a/src/assembler.py b/src/assembler.py
index a4ddde4..fc7bc45 100755
--- a/src/assembler.py
+++ b/src/assembler.py
@@ -2,10 +2,12 @@
 # Assemble ATK16 assembly to bytecode
 
 import sys
-from typing import Callable
-import importlib
 from assembler_ops import *
 from assembler_eval import *
+from asm_pass1 import pass_1
+from asm_pass2 import pass_2
+from asm_pass3 import pass_3
+from asm_pass4 import pass_4
 
 if len(sys.argv) != 3:
   print("usage: assembler.py <infile> <outfile> # read from file")
@@ -47,91 +49,27 @@ def parse(line: str) -> list[str]:
   result.append(acc)
   return list(filter(lambda x: len(x) > 0, result))
 
-# 1st pass, gather labels, set options, eval operations
-options = Options()
-address: int = 0
+result1 = pass_1(src_lines)
+result2 = pass_2(result1)
+result3 = pass_3(result2)
+result4 = pass_4(result3)
 
-for (lineNo, line) in enumerate(src_lines):
-  line = line.split(";")[0].strip()
-  if line == "": continue
-  keyword, *args = line.lower().split()
-  match keyword:
-    case "@address":
-      address = eval(args[0])
-      continue
-    case "@label":
-      labels[args[0]] = address
-      # todo bugi: 1st pass ei ota monirivisiä käskyjä huomioon
-      continue
-    case "@opt":
-      opt_name, opt_value = args
-      match opt_name:
-        case "stack_pointer": options.stack_pointer = opt_value
-        case "csr_scratch": options.csr_scratch = opt_value
-        case _: raise Exception("Unknown @opt: " + opt_name)
-    case "@use":
-      module_name, ops = args[0].split(":")
-      ops_split = ops.split(",")
-      module = importlib.import_module(module_name)
-      mod_operations: dict[str, Callable[..., list[int]]] = module.operations
-      for op in mod_operations:
-        callable = mod_operations[op]
-        if ops == "*" or op in ops_split:
-            operations[op] = callable
-    case _:
-      pass
-
-  address += 1
-
-# 2nd pass
-
-address: int = 0
 nop = bytearray([0b1000_0000, 0])
 result = bytearray()
 # initially one nop
 result.extend(nop)
 
-for (lineNo, line) in enumerate(src_lines):
-  line = line.split(";")[0].strip()
-  if line == "": continue
-  keyword, *args = parse(line.lower())
-
-  meta = Meta(
-    address=address,
-    options=options
-  )
-  if keyword in operations:
-    words = operations[keyword](meta, *args)
-  else:
-    match keyword:
-      # Directives
-      case "@address":
-        address = eval_expr(args[0])
-        continue
-      case "@label" | "@opt" | "@use":
-        continue
-
-      # Default case: evaluate as is (e.g. data word)
-      case _:
-        try:
-          words = [eval_expr(keyword)]
-        except:
-          raise Exception(f"Invalid assembly at {infile_path}:{lineNo + 1}\n\n{line}")
-
-  for (label, label_addr) in labels.items():
-    if address == label_addr:
+for line in result4.lines:
+  for (label, label_addr) in result4.labels.items():
+    if line.address == label_addr:
       print(f"{label}:")
 
-  line_printed = False
-  for word in words:
-    if len(result) < 2 * address + 1:
-      result.extend((2 * address + 1 - len(result)) * nop)
+  if len(result) < 2 * line.address + 1:
+    result.extend((2 * line.address + 1 - len(result)) * nop)
 
-    print(f"{address:>08x}  0x{word:>04x}  {line if not line_printed else '...'}")
-    result[2 * address + 0] = ((word >> 8) & 0xff)
-    result[2 * address + 1] = ((word >> 0) & 0xff)
-    address += 1
-    line_printed = True
+  print(f"{line.address:>08x}  0x{line.word:>04x}  {line.text}")
+  result[2 * line.address + 0] = ((line.word >> 8) & 0xff)
+  result[2 * line.address + 1] = ((line.word >> 0) & 0xff)
 
 with open(outfile_path, "wb") as f:
   f.write(result)
diff --git a/src/assembler_eval.py b/src/assembler_eval.py
index 07f22e7..7cfe336 100644
--- a/src/assembler_eval.py
+++ b/src/assembler_eval.py
@@ -24,13 +24,13 @@ constants: dict[str, str] = {
   "f_sign": "3",
 }
 
-labels: dict[str, int] = {}
+Labels = dict[str, int]
 
 def check_size(bits: int, val: int) -> None:
   if val >= 2 ** bits:
     raise Exception(f"Value does not fit in {bits} bits: {val}")
 
-def eval_symbol(c: str) -> str:
+def eval_symbol(labels: Labels, c: str) -> str:
   if c in labels:
     return str(labels[c])
 
@@ -39,8 +39,9 @@ def eval_symbol(c: str) -> str:
 
   return c
 
-def eval_expr(expr: str, bits: int = 16) -> int:
-  expr = eval_symbol(expr)
+def eval_expr(labels: Labels, expr: str, bits: int = 16) -> int:
+  expr = expr.lower()
+  expr = eval_symbol(labels, expr)
   ret = eval(expr, labels.copy()) # eval as Python expr
   check_size(bits, ret)
   return ret
diff --git a/src/assembler_ops.py b/src/assembler_ops.py
index 2313bfd..e4eeac4 100644
--- a/src/assembler_ops.py
+++ b/src/assembler_ops.py
@@ -4,115 +4,114 @@ from dataclasses import dataclass
 
 @dataclass
 class Options:
-  stack_pointer = "RG"
-  csr_scratch = "RF"
+  stack_pointer = "rg"
+  csr_scratch = "rf"
 
 @dataclass
 class Meta:
   address: int
-  options: Options
 
-def make_alr(meta: Meta, alu_op: str, left: str, right: str, target: str) -> list[int]:
+def make_alr(meta: Meta, labels: Labels, alu_op: str, left: str, right: str, target: str) -> int:
   """ALR 0000 TTTL LLRR RSSS"""
-  target_e = eval_expr(target, bits=3)
-  left_e = eval_expr(left, bits=3)
-  right_e = eval_expr(right, bits=3)
-  alu_op_e = eval_expr(alu_op, bits=3)
+  target_e = eval_expr(labels, target, bits=3)
+  left_e = eval_expr(labels, left, bits=3)
+  right_e = eval_expr(labels, right, bits=3)
+  alu_op_e = eval_expr(labels, alu_op, bits=3)
   word = (0b0000 << 12) + \
               (target_e << 9) + \
               (left_e << 6) + \
               (right_e << 3) + \
               alu_op_e
-  return [word]
+  return word
 
-def make_ali(meta: Meta, alu_op: str, left: str, imm: str, target: str) -> list[int]:
+def make_ali(meta: Meta, labels: Labels, alu_op: str, left: str, imm: str, target: str) -> int:
   """ALI 0001 TTTL LLII ISSS"""
-  target_e = eval_expr(target, bits=3)
-  left_e = eval_expr(left, bits=3)
-  imm_e = eval_expr(imm, bits=3)
-  alu_op_e = eval_expr(alu_op, bits=3)
+  target_e = eval_expr(labels, target, bits=3)
+  left_e = eval_expr(labels, left, bits=3)
+  imm_e = eval_expr(labels, imm, bits=3)
+  alu_op_e = eval_expr(labels, alu_op, bits=3)
   word = (0b0001 << 12) + \
               (target_e << 9) + \
               (left_e << 6) + \
               (imm_e << 3) + \
               alu_op_e
-  return [word]
+  return word
 
-def make_ldr(meta: Meta, to_reg: str, addr_reg: str) -> list[int]:
+def make_ldr(meta: Meta, labels: Labels, to_reg: str, addr_reg: str) -> int:
   """LDR 0010 TTTR RRXX XXXX"""
-  to_reg_e = eval_expr(to_reg, bits=3)
-  addr_reg_e = eval_expr(addr_reg, bits=3)
+  to_reg_e = eval_expr(labels, to_reg, bits=3)
+  addr_reg_e = eval_expr(labels, addr_reg, bits=3)
   word = (0b0010 << 12) + \
               (to_reg_e << 9) + \
               (addr_reg_e << 6)
-  return [word]
+  return word
 
-def make_str(meta: Meta, from_reg: str, addr_reg: str) -> list[int]:
+def make_str(meta: Meta, labels: Labels, from_reg: str, addr_reg: str) -> int:
   """STR 0011 TTTR RRXX XXXX"""
-  from_reg_e = eval_expr(from_reg, bits=3)
-  addr_reg_e = eval_expr(addr_reg, bits=3)
+  from_reg_e = eval_expr(labels, from_reg, bits=3)
+  addr_reg_e = eval_expr(labels, addr_reg, bits=3)
   word = (0b0011 << 12) + \
               (from_reg_e << 9) + \
               (addr_reg_e << 6)
-  return [word]
+  return word
 
-def make_ldi(meta: Meta, to_reg: str, imm: str) -> list[int]:
+def make_ldi(meta: Meta, labels: Labels, to_reg: str, imm: str) -> int:
   """LDI 0100 TTTI IIII IIII"""
-  to_reg_e = eval_expr(to_reg, bits=3)
-  imm_e = eval_expr(imm, bits=9)
+  to_reg_e = eval_expr(labels, to_reg, bits=3)
+  imm_e = eval_expr(labels, imm, bits=9)
   word = (0b0100 << 12) + \
               (to_reg_e << 9) + \
               imm_e
-  return [word]
+  return word
 
-def make_jpr(meta: Meta, addr_reg: str) -> list[int]:
+def make_jpr(meta: Meta, labels: Labels, addr_reg: str) -> int:
   """JPR 0101 XXXR RRXX XXXX"""
-  addr_reg_e = eval_expr(addr_reg, bits=3)
+  addr_reg_e = eval_expr(labels, addr_reg, bits=3)
   word = (0b0101 << 12) + \
               (addr_reg_e << 6)
-  return [word]
+  return word
 
-def make_jpi(meta: Meta, imm: str) -> list[int]:
+def make_jpi(meta: Meta, labels: Labels, imm: str) -> int:
   """JPI 0110 XXXI IIII IIII"""
-  imm_e = eval_expr(imm, bits=9)
+  imm_e = eval_expr(labels, imm, bits=9)
   imm_e = imm_e - meta.address - 1
   word = (0b0110 << 12) + \
               imm_e
-  return [word]
+  return word
 
-def make_brr(meta: Meta, flag_s: str, addr_reg: str) -> list[int]:
+def make_brr(meta: Meta, labels: Labels, flag_s: str, addr_reg: str) -> int:
   """brr  0111 XFFR RRXX XXXX"""
-  flag_s_e = eval_expr(flag_s, bits=2)
-  addr_reg_e = eval_expr(addr_reg, bits=3)
+  flag_s_e = eval_expr(labels, flag_s, bits=2)
+  addr_reg_e = eval_expr(labels, addr_reg, bits=3)
   word = (0b0111 << 12) + \
               (flag_s_e << 9) + \
               (addr_reg_e << 6)
-  return [word]
+  return word
 
-def make_bri(meta: Meta, flag_s: str, imm: str) -> list[int]:
+def make_bri(meta: Meta, labels: Labels, flag_s: str, imm: str) -> int:
   """BRI 1000 XFFI IIII IIII"""
-  flag_s_e = eval_expr(flag_s, bits=2)
-  imm_e = eval_expr(imm, bits=9)
+  flag_s_e = eval_expr(labels, flag_s, bits=2)
+  imm_e = eval_expr(labels, imm, bits=9)
   imm_e = imm_e - meta.address - 1
   word = (0b1000 << 12) + \
               (flag_s_e << 9) + \
               imm_e
-  return [word]
+  return word
 
-def make_lpc(meta: Meta, target: str) -> list[int]:
+def make_lpc(meta: Meta, labels: Labels, target: str) -> int:
   """LPC 1001 TTTX XXXX XXXX"""
-  target_e = eval_expr(target, bits=3)
+  target_e = eval_expr(labels, target, bits=3)
   word = (0b1001 << 12) + \
               (target_e << 9)
-  return [word]
+  return word
 
-def make_hlt(meta: Meta, ) -> list[int]:
+def make_hlt(meta: Meta, labels: Labels) -> int:
   """HLT 1111 XXXX XXXX XXXX"""
   word = (0b1111 << 12)
-  return [word]
+  return word
 
-operations: dict[str, Callable[..., list[int]]] = {
-  "alr": make_alr,
+OpWordDict = dict[str, Callable[..., int]]
+operations: OpWordDict = {
   "alr": make_alr,
   "ali": make_ali,
   "ldr": make_ldr,
@@ -125,3 +124,25 @@ operations: dict[str, Callable[..., list[int]]] = {
   "lpc": make_lpc,
   "hlt": make_hlt,
 }
+
+ExpandResult = list[list[str]]
+ExpandFn = Callable[..., ExpandResult]
+OpExpansionDict = dict[str, ExpandFn]
+
+def expand_id(*parts: str) -> ExpandResult:
+  return [list(parts)]
+
+default_expansions: OpExpansionDict = {
+  "alr": lambda *args: expand_id("alr", *args),
+  "ali": lambda *args: expand_id("ali", *args),
+  "ldr": lambda *args: expand_id("ldr", *args),
+  "str": lambda *args: expand_id("str", *args),
+  "ldi": lambda *args: expand_id("ldi", *args),
+  "jpr": lambda *args: expand_id("jpr", *args),
+  "jpi": lambda *args: expand_id("jpi", *args),
+  "brr": lambda *args: expand_id("brr", *args),
+  "bri": lambda *args: expand_id("bri", *args),
+  "lpc": lambda *args: expand_id("lpc", *args),
+  "hlt": lambda *args: expand_id("hlt", *args),
+}
+
diff --git a/src/ext_std.py b/src/ext_std.py
index e78f5cc..a6b175a 100644
--- a/src/ext_std.py
+++ b/src/ext_std.py
@@ -1,102 +1,88 @@
 from assembler_ops import *
 
-def make_add(meta: Meta, left: str, right: str, target: str) -> list[int]:
-  return make_alr(meta, "al_plus", left, right, target)
+def expand_add(left: str, right: str, target: str) -> ExpandResult:
+  return [["alr", "al_plus", left, right, target]]
 
-def make_sub(meta: Meta, left: str, right: str, target: str) -> list[int]:
-  return make_alr(meta, "al_minus", left, right, target)
+def expand_sub(left: str, right: str, target: str) -> ExpandResult:
+  return [["alr", "al_minus", left, right, target]]
 
-def make_addi(meta: Meta, left: str, imm: str, target: str) -> list[int]:
-  return make_ali(meta, "al_plus", left, imm, target)
+def expand_addi(left: str, imm: str, target: str) -> ExpandResult:
+  return [["ali", "al_plus", left, imm, target]]
 
-def make_subi(meta: Meta, left: str, imm: str, target: str) -> list[int]:
-  return make_ali(meta, "al_minus", left, imm, target)
+def expand_subi(left: str, imm: str, target: str) -> ExpandResult:
+  return [["ali", "al_minus", left, imm, target]]
 
-def make_and(meta: Meta, left: str, right: str, target: str) -> list[int]:
-  return make_alr(meta, "al_and", left, right, target)
+def expand_and(left: str, right: str, target: str) -> ExpandResult:
+  return [["alr", "al_and", left, right, target]]
 
-def make_or(meta: Meta, left: str, right: str, target: str) -> list[int]:
-  return make_alr(meta, "al_or", left, right, target)
+def expand_or(left: str, right: str, target: str) -> ExpandResult:
+  return [["alr", "al_or", left, right, target]]
 
-def make_xor(meta: Meta, left: str, right: str, target: str) -> list[int]:
-  return make_alr(meta, "al_xor", left, right, target)
+def expand_xor(left: str, right: str, target: str) -> ExpandResult:
+  return [["alr", "al_xor", left, right, target]]
 
-def make_sll(meta: Meta, left: str, right: str, target: str) -> list[int]:
-  return make_alr(meta, "al_sll", left, right, target)
+def expand_sll(left: str, right: str, target: str) -> ExpandResult:
+  return [["alr", "al_sll", left, right, target]]
 
-def make_slr(meta: Meta, left: str, right: str, target: str) -> list[int]:
-  return make_alr(meta, "al_slr", left, right, target)
+def expand_slr(left: str, right: str, target: str) -> ExpandResult:
+  return [["alr", "al_slr", left, right, target]]
 
-def make_sar(meta: Meta, left: str, right: str, target: str) -> list[int]:
-  return make_alr(meta, "al_sar", left, right, target)
+def expand_sar(left: str, right: str, target: str) -> ExpandResult:
+  return [["alr", "al_sar", left, right, target]]
 
-def make_slli(meta: Meta, left: str, imm: str, target: str) -> list[int]:
-  return make_ali(meta, "al_sll", left, imm, target)
+def expand_slli(left: str, imm: str, target: str) -> ExpandResult:
+  return [["ali", "al_sll", left, imm, target]]
 
-def make_slri(meta: Meta, left: str, imm: str, target: str) -> list[int]:
-  return make_ali(meta, "al_slr", left, imm, target)
+def expand_slri(left: str, imm: str, target: str) -> ExpandResult:
+  return [["ali", "al_slr", left, imm, target]]
 
-def make_sari(meta: Meta, left: str, imm: str, target: str) -> list[int]:
-  return make_ali(meta, "al_sar", left, imm, target)
+def expand_sari(left: str, imm: str, target: str) -> ExpandResult:
+  return [["ali", "al_sar", left, imm, target]]
 
-def make_inc(meta: Meta, reg: str) -> list[int]:
-  return make_ali(meta, "al_plus", reg, "1", reg)
+def expand_inc(reg: str) -> ExpandResult:
+  return [["ali", "al_plus", reg, "1", reg]]
 
-def make_dec(meta: Meta, reg: str) -> list[int]:
-  return make_ali(meta, "al_minus", reg, "1", reg)
+def expand_dec(reg: str) -> ExpandResult:
+  return [["ali", "al_minus", reg, "1", reg]]
 
-def make_mov(meta: Meta, from_reg: str, to_reg: str) -> list[int]:
- return make_ali(meta, "al_plus", from_reg, "0", to_reg)
+def expand_mov(from_reg: str, to_reg: str) -> ExpandResult:
+  return [["ali", "al_plus", from_reg, "0", to_reg]]
 
-def make_spu(meta: Meta, reg: str) -> list[int]:
-  sp_reg = meta.options.stack_pointer
-  str_word = make_str(meta, reg, sp_reg)
-  inc_word = make_inc(meta, sp_reg)
-  return str_word + inc_word
+def expand_spu(reg: str) -> ExpandResult:
+  return [["str", reg, "__STACK_POINTER"]] + expand_inc("__STACK_POINTER")
 
-def make_spo(meta: Meta, reg: str) -> list[int]:
-  sp_reg = meta.options.stack_pointer
-  dec_word = make_dec(meta, sp_reg)
-  ldr_word = make_ldr(meta, reg, sp_reg)
-  return dec_word + ldr_word
+def expand_spo(reg: str):
+  return expand_dec("__STACK_POINTER") + [["ldr", reg, "__STACK_POINTER"]]
 
-def make_csr(meta: Meta, addr_reg: str) -> list[int]:
-  lpc_word = make_lpc(meta, meta.options.csr_scratch)
-  spu_word = make_spu(meta, meta.options.csr_scratch)
-  jpr_word = make_jpr(meta, addr_reg)
-  return lpc_word + spu_word + jpr_word
+def expand_csr(addr_reg: str):
+  return [["lpc", "__STACK_POINTER"]] + expand_spu("__STACK_POINTER") + [["jpr", addr_reg]]
 
-def make_csi(meta: Meta, addr_imm: str) -> list[int]:
-  lpc_word = make_lpc(meta, meta.options.csr_scratch)
-  spu_word = make_spu(meta, meta.options.csr_scratch)
-  jpi_word = make_jpi(meta, addr_imm)
-  return lpc_word + spu_word + jpi_word
+def expand_csi(addr_imm: str):
+  return [["lpc", "__STACK_POINTER"]] + expand_spu("__STACK_POINTER") + [["jpr", addr_imm]]
 
-def make_rsr(meta: Meta) -> list[int]:
-  spo_word = make_spo(meta, meta.options.csr_scratch)
-  jpr_word = make_jpr(meta, meta.options.csr_scratch)
-  return spo_word + jpr_word
+def expand_rsr():
+  return expand_spo("__STACK_POINTER") + [["jpr", "__CSR_SCRATCH"]]
 
-operations = {
-  "add": make_add,
-  "sub": make_sub,
-  "addi": make_addi,
-  "subi": make_subi,
-  "and": make_and,
-  "or": make_or,
-  "xor": make_xor,
-  "sll": make_sll,
-  "slr": make_slr,
-  "sar": make_sar,
-  "slli": make_slli,
-  "slri": make_slri,
-  "sari": make_sari,
-  "inc": make_inc,
-  "dec": make_dec,
-  "mov": make_mov,
-  "spu": make_spu,
-  "spo": make_spo,
-  "csr": make_csr,
-  "csi": make_csi,
-  "rsr": make_rsr,
+expansions: OpExpansionDict = {
+  "add": expand_add,
+  "sub": expand_sub,
+  "addi": expand_addi,
+  "subi": expand_subi,
+  "and": expand_and,
+  "or": expand_or,
+  "xor": expand_xor,
+  "sll": expand_sll,
+  "slr": expand_slr,
+  "sar": expand_sar,
+  "slli": expand_slli,
+  "slri": expand_slri,
+  "sari": expand_sari,
+  "inc": expand_inc,
+  "dec": expand_dec,
+  "mov": expand_mov,
+  "spu": expand_spu,
+  "spo": expand_spo,
+  "csr": expand_csr,
+  "csi": expand_csi,
+  "rsr": expand_rsr,
 }
diff --git a/src/test.atk16 b/src/test.atk16
index 1507430..8e48892 100644
--- a/src/test.atk16
+++ b/src/test.atk16
@@ -13,7 +13,7 @@
   ldi RA 10         ; RA := 10
   ldi RB 20         ; RB := 20
   add RA RB RC      ; RC := RA + RB
-  ldi RD ram_offset ; store address of ram_offset in RD
+  ldi RD @ram_offset ; store address of ram_offset in RD
   ldr RD RD         ; dereference ram_offset address
 @label debug
   str RD RC         ; store RC in RAM
@@ -22,7 +22,7 @@
   mov RC RA         ; RA := result of sum
   ldi RB 30         ; RB := 30
   sub RA RB RC      ; RC := RA - RB
-  bri f_zero success  ; if result is zero, jump to success
+  bri f_zero @success  ; if result is zero, jump to success
 
   ldi RA 2          ; RA := 2 to signal failure
   hlt