Machine Code and the Assembler

#!/usr/bin/env python3
"""
A two-pass assembler for the machine in the Stored-Program Computer blueprint.
Youblob blueprint: Machine Code and the Assembler

Instruction word = opcode*100 + operand.
  0 HLT   1 LOAD  2 STORE  3 ADD  4 SUB  5 JMP  6 JZ  7 OUT

WHY TWO PASSES: a program can jump FORWARD to a label that has not been seen yet.
Pass 1 walks the source assigning addresses and recording where every label lands.
Pass 2 emits code, by which point every label is known. A one-pass assembler must
either forbid forward jumps or patch them up afterwards -- which is exactly what a
linker does, and why the job exists at all.
"""

OPCODES = {"HLT":0, "LOAD":1, "STORE":2, "ADD":3, "SUB":4, "JMP":5, "JZ":6, "OUT":7}


def parse(source):
    """Strip comments and blanks, split each line into (label, mnemonic, operand)."""
    out = []
    for lineno, raw in enumerate(source.splitlines(), 1):
        line = raw.split("#")[0].strip()
        if not line:
            continue
        label = None
        if ":" in line:
            label, _, line = line.partition(":")
            label = label.strip()
            line = line.strip()
        parts = line.split()
        mnem = parts[0].upper() if parts else None
        operand = parts[1] if len(parts) > 1 else None
        out.append((lineno, label, mnem, operand))
    return out


def assemble(source):
    parsed = parse(source)

    # ---- PASS 1: assign addresses, collect labels -------------------------
    symbols, addr = {}, 0
    for lineno, label, mnem, operand in parsed:
        if label:
            if label in symbols:
                raise SyntaxError(f"line {lineno}: label '{label}' defined twice")
            symbols[label] = addr
        if mnem:                       # DATA reserves a word, like any instruction
            addr += 1

    # ---- PASS 2: emit machine code ----------------------------------------
    image = []
    for lineno, label, mnem, operand in parsed:
        if mnem is None:
            continue
        if mnem == "DATA":
            image.append(int(operand or 0))
            continue
        if mnem not in OPCODES:
            raise SyntaxError(f"line {lineno}: unknown mnemonic '{mnem}'")
        op = OPCODES[mnem]
        if operand is None:
            target = 0
        elif operand.isdigit():
            target = int(operand)
        elif operand in symbols:
            target = symbols[operand]      # THE WHOLE POINT: name -> address
        else:
            raise SyntaxError(f"line {lineno}: undefined label '{operand}'")
        image.append(op * 100 + target)

    return image, symbols


PROGRAM = """
# count down from N, printing each value
start:  LOAD  count
loop:   OUT
        SUB   one
        STORE count
        JZ    done          # forward reference -- impossible in one pass
        JMP   loop
done:   HLT
count:  DATA  5
one:    DATA  1
"""

if __name__ == "__main__":
    image, symbols = assemble(PROGRAM)
    print("symbol table:", symbols)
    print("\naddr  word   meaning")
    names = {v: k for k, v in OPCODES.items()}
    for a, w in enumerate(image):
        print(f"{a:>4}  {w:>4}   {names.get(w // 100, 'DATA'):<6}{w % 100}")
    print("\nInsert a line anywhere above and re-run: every address is recomputed")
    print("for you. That is the entire value of the assembler.")