assembler.py

from collections import defaultdict
from collections.abc import Iterable

from instructions import ImpliedInst, Instruction
from label import Label
from modes import Addr, Val


def _assemble_mode(mode: Addr | Val) -> Iterable[int]:
    match mode:
        case Val():
            yield mode & 0xFF
        case Addr():
            yield mode & 0xFF
            yield mode >> 8


type AnyInst = Instruction | ImpliedInst


def assemble(elem: list[AnyInst | Label], start: int = 0x0000) -> bytearray:
    result = bytearray()

    label_addresses: dict[Label, int] = {}
    label_holes: dict[Label, set[int]] = defaultdict(set)

    address = start

    for elm in elem:
        start_len = len(result)
        match elm:
            case Label() as lbl:
                label_addresses[lbl] = address
            case int() as inst:
                result.append(inst)
            case Instruction(opcode=opcode, mode=Label() as lbl):
                result.append(opcode)
                label_holes[lbl].add(address + 1)
                result.extend((0x00, 0x00))
            case Instruction(opcode=opcode, mode=Val() | Addr() as mode):
                result.append(opcode)
                result.extend(_assemble_mode(mode))
        address += len(result) - start_len

    for lbl, holes in label_holes.items():
        assert lbl in label_addresses, f"Label {lbl} not found"

        print(f"{lbl.name} = {Addr(label_addresses[lbl])}")

        label_addr = tuple(_assemble_mode(mode=Addr(label_addresses[lbl])))

        for addr in holes:
            result[addr - start : addr - start + 2] = label_addr

    return result