Nand2Tetris — 第6章

# Hack Assembler
# 入力：.asmファイル
# 出力：.hackファイル
# アクション：
# ステップ1：ラベルをシンボルテーブルに登録する
# ステップ2：A-instructionとC-instructionを16ビットの機械語に変換

import sys


PREDEFINED = {
    "SP": 0,
    "LCL": 1,
    "ARG": 2,
    "THIS": 3,
    "THAT": 4,
    "SCREEN": 16384,
    "KBD": 24576,
}

for i in range(16):
    PREDEFINED[f"R{i}"] = i


DEST = {
    "": "000",
    "M": "001",
    "D": "010",
    "MD": "011",
    "A": "100",
    "AM": "101",
    "AD": "110",
    "AMD": "111",
}

COMP = {
    "0": "0101010",
    "1": "0111111",
    "-1": "0111010",
    "D": "0001100",
    "A": "0110000",
    "M": "1110000",
    "!D": "0001101",
    "!A": "0110001",
    "!M": "1110001",
    "-D": "0001111",
    "-A": "0110011",
    "-M": "1110011",
    "D+1": "0011111",
    "A+1": "0110111",
    "M+1": "1110111",
    "D-1": "0001110",
    "A-1": "0110010",
    "M-1": "1110010",
    "D+A": "0000010",
    "D+M": "1000010",
    "D-A": "0010011",
    "D-M": "1010011",
    "A-D": "0000111",
    "M-D": "1000111",
    "D&A": "0000000",
    "D&M": "1000000",
    "D|A": "0010101",
    "D|M": "1010101",
}

JUMP = {
    "": "000",
    "JGT": "001",
    "JEQ": "010",
    "JGE": "011",
    "JLT": "100",
    "JNE": "101",
    "JLE": "110",
    "JMP": "111",
}


def clean_line(line):
    # コメントとスペースを削除
    line = line.split("//")[0]
    return line.strip()


def load_lines(path):
    lines = []

    with open(path, "r", encoding="utf-8") as file:
        for line in file:
            line = clean_line(line)
            if line:
                lines.append(line)

    return lines


def build_symbol_table(lines):
    symbols = dict(PREDEFINED)
    rom_address = 0

    for line in lines:
        if line.startswith("(") and line.endswith(")"):
            label = line[1:-1]
            symbols[label] = rom_address
        else:
            rom_address += 1

    return symbols


def parse_c_instruction(line):
    # C-instruction 形式: dest=comp;jump
    dest = ""
    jump = ""

    if "=" in line:
        dest, line = line.split("=")

    if ";" in line:
        comp, jump = line.split(";")
    else:
        comp = line

    return dest, comp, jump


def assemble(lines):
    symbols = build_symbol_table(lines)
    next_variable_address = 16
    machine_code = []

    for line in lines:
        if line.startswith("("):
            continue

        if line.startswith("@"):
            symbol = line[1:]

            if symbol.isdigit():
                address = int(symbol)
            else:
                if symbol not in symbols:
                    symbols[symbol] = next_variable_address
                    next_variable_address += 1
                address = symbols[symbol]

            machine_code.append(f"{address:016b}")
            continue

        dest, comp, jump = parse_c_instruction(line)
        machine_code.append("111" + COMP[comp] + DEST[dest] + JUMP[jump])

    return machine_code


def output_path(input_path):
    if input_path.endswith(".asm"):
        return input_path[:-4] + ".hack"
    return input_path + ".hack"


def main():
    if len(sys.argv) != 2:
        print("Usage: python assembler.py Xxx.asm")
        return

    input_path = sys.argv[1]
    lines = load_lines(input_path)
    machine_code = assemble(lines)

    with open(output_path(input_path), "w", encoding="utf-8") as file:
        file.write("\n".join(machine_code))
        file.write("\n")


if __name__ == "__main__":
    main()