simple_ass_ Make encode method safer

handle nop

1 files changed, 52 insertions, 23 deletions

diff --git a/nandgame/assembler/simple-assembler.py b/nandgame/assembler/simple_assembler.py
index 5aa05ce..7fb37f4 100755
--- a/nandgame/assembler/simple-assembler.py
+++ b/nandgame/assembler/simple_assembler.py
@@ -32,6 +32,9 @@ Jumps:
 """
 
 import sys
+from typing import Union
+
+import parser_types as pt
 
 ZERO = "#0"
 DEST_NONE = "_"
@@ -41,7 +44,7 @@ ENDIANNESS = "little"
 # mapping from mnemonic to (opcode, two_op)
 MNEMONICS = {
     "and": (0b000, True),
-    "or":  (0b001, True),
+    "or": (0b001, True),
     "xor": (0b010, True),
     "not": (0b011, False),
     "add": (0b100, True),
@@ -52,7 +55,7 @@ MNEMONICS = {
 
 # jump mnemonic -> bits 0..2
 JUMP_ENCODE = {
-    "":    0b000,
+    "": 0b000,
     "jgt": 0b001,
     "jeq": 0b010,
     "jge": 0b011,
@@ -74,11 +77,11 @@ def encode_dest(dest: str) -> int:
 
     bits = 0
     if "A" in dest:
-        bits |= (1 << 5)
+        bits |= 1 << 5
     if "D" in dest:
-        bits |= (1 << 4)
+        bits |= 1 << 4
     if "M" in dest:
-        bits |= (1 << 3)
+        bits |= 1 << 3
     return bits
 
 
@@ -121,11 +124,11 @@ def encode_args_two_op(op1: str, op2: str) -> int:
                 if dec_op1 == op1 and dec_op2 == op2:
                     bits = 0
                     if use_mem:
-                        bits |= (1 << 12)
+                        bits |= 1 << 12
                     if zx:
-                        bits |= (1 << 7)
+                        bits |= 1 << 7
                     if sw:
-                        bits |= (1 << 6)
+                        bits |= 1 << 6
                     return bits
 
     raise ValueError(f"Unsupported operand combination for two-op: {op1}, {op2}")
@@ -145,7 +148,7 @@ def encode_args_one_op(op1: str) -> int:
     bits = 0
 
     if op1 == ZERO:
-        bits |= (1 << 7)  # zx
+        bits |= 1 << 7  # zx
         # sw/use_mem don't matter for arg1 when zx=1, but keep them 0
         return bits
 
@@ -154,30 +157,46 @@ def encode_args_one_op(op1: str) -> int:
         return bits
 
     if op1 == "A":
-        bits |= (1 << 6)  # sw=1
+        bits |= 1 << 6  # sw=1
         # use_mem=0
         return bits
 
     if op1 == "M":
-        bits |= (1 << 6)  # sw=1
-        bits |= (1 << 12)  # use_mem=1
+        bits |= 1 << 6  # sw=1
+        bits |= 1 << 12  # use_mem=1
         return bits
 
     raise ValueError(f"Unsupported operand for one-op: {op1}")
 
+#Arg = Union[str, int, pt.Address, pt.Immediate, pt.Register]
+Arg = Union[str, int, None]
 
-def encode_instruction(mnemonic: str, dest: str, op1: str, op2: str, jump: str) -> int:
+def encode_instruction(
+    mnemonic: str, dest: str, op1: Arg, op2: Arg, jump: str
+) -> int:
     """
     Encode a single instruction into a 16-bit integer.
     """
     mnemonic = mnemonic.strip()
 
     if mnemonic == "hlt":
-        return (0xFFFF & ~0x4000)
+        return 0xFFFF & ~0x4000
 
     dest = dest.strip()
-    op1 = op1.strip()
-    op2 = op2.strip()
+    if op1 is None:
+        op1 = ""
+    elif isinstance(op1, int):
+        op1 = f"#{op1}"
+    else:
+        op1 = str(op1).strip()
+
+    if op2 is None:
+        op2 = ""
+    elif isinstance(op2, int):
+        op2 = f"#{op2}"
+    else:
+        op2 = str(op2).strip()
+
     jump = jump.strip()
 
     # A-instruction: mov A, #imm
@@ -192,10 +211,18 @@ def encode_instruction(mnemonic: str, dest: str, op1: str, op2: str, jump: str)
                 raise ValueError(f"Immediate out of range (0..32767): {value}")
             return value & 0x7FFF
         else:
-            raise ValueError(f"Invalid args to mov.")
+            raise ValueError("Invalid args to mov.")
 
     # C-instruction
-    if mnemonic not in MNEMONICS:
+    if mnemonic in MNEMONICS:
+        pass
+    elif mnemonic == "nop":
+        mnemonic = "and"
+        dest = "_"
+        op1 = "#0"
+        op2 = "A"
+        jump = ""
+    else:
         raise ValueError(f"Unknown mnemonic: {mnemonic}")
 
     opcode, two_op = MNEMONICS[mnemonic]
@@ -206,9 +233,9 @@ def encode_instruction(mnemonic: str, dest: str, op1: str, op2: str, jump: str)
     # opcode bits: low 2 bits in 8..9, high bit in 10 (ar_n_log)
     low2 = opcode & 0b11
     high1 = (opcode >> 2) & 0b1
-    ins |= (low2 << 8)
+    ins |= low2 << 8
     if high1:
-        ins |= (1 << 10)
+        ins |= 1 << 10
 
     # dest bits
     ins |= encode_dest(dest)
@@ -235,7 +262,7 @@ def parse_line(line: str):
     Returns None if the line is empty or comment.
     """
     # strip comments starting with ';' or '#'
-    for sep in (";"):
+    for sep in ";":
         idx = line.find(sep)
         if idx != -1:
             line = line[:idx]
@@ -280,7 +307,9 @@ def parse_line(line: str):
 
 
 def assemble_file(in_filename: str, out_filename: str):
-    with open(in_filename, "r") as fin, open(out_filename, "wb") as fout:
+    with open(in_filename, "r", encoding="ascii") as fin, open(
+        out_filename, "wb"
+    ) as fout:
         lineno = 0
         for line in fin:
             lineno += 1
@@ -292,7 +321,7 @@ def assemble_file(in_filename: str, out_filename: str):
                 ins = encode_instruction(mnemonic, dest, op1, op2, jump)
                 fout.write(ins.to_bytes(2, byteorder=ENDIANNESS))
             except Exception as e:
-                raise SystemExit(f"{in_filename}:{lineno}: {e}")
+                raise SystemExit(f"{in_filename}:{lineno}: {e}") from e
 
 
 def main():