nand2tetris: HACK assembler

assembler.py

#!/usr/bin/env python3

import sys


def check_args():
    if len(sys.argv) != 2:
        print("Usage: python3 assembler.py <file.asm>")
        sys.exit(1)

    return sys.argv[1]

def load_file(input_file):
    with open(input_file, "r") as f:
        return f.readlines()

def remove_comments(lines):
    return [line.split("//")[0] for line in lines]

def remove_whitespace(lines):
    return [line.strip() for line in lines if line.strip()]

def save_binary(output):
    for line in output:
        print(line)

def assemble_a_instruction(line, labels):
    instruction = line[1:]
    if instruction in labels:
        return f"0{labels[instruction]:015b}"

    return f"0{int(instruction):015b}"

def jump_to_binary(jump):
    jumps = {
        "": 0,
        "JGT": 1,
        "JEQ": 2,
        "JGE": 3,
        "JLT": 4,
        "JNE": 5,
        "JLE": 6,
        "JMP": 7,
    }

    return jumps[jump]

def dest_to_binary(dest):
    dests = {
        "": 0,
        "M": 1,
        "D": 2,
        "MD": 3,
        "A": 4,
        "AM": 5,
        "AD": 6,
        "AMD": 7,
    }

    return dests[dest]

def comp_to_binary(comp):
    comps_left = {
        "0": "0101010",
        "1": "0111111",
        "-1": "0111010",
        "D": "0001100",
        "A": "0110000",
        "!D": "0001101",
        "!A": "1110001",
        "-D": "0001111",
        "-A": "0110011",
        "D+1": "0011111",
        "A+1": "0110111",
        "D-1": "0001110",
        "A-1": "0110010",
        "D+A": "0000010",
        "D-A": "0010011",
        "A-D": "0000111",
        "D&A": "0000000",
        "D|A": "0010101",
    }

    comps_right = {
        "M": "1110000",
        "!M": "1110001",
        "-M": "1110011",
        "M+1": "1110111",
        "M-1": "1110010",
        "D+M": "1000010",
        "D-M": "1010011",
        "M-D": "1000111",
        "D&M": "1000000",
        "D|M": "1010101",
    }

    if comp in comps_left:
        return f"{comps_left[comp]}"
    return f"{comps_right[comp]}"

def assemble_c_instruction(line):
    dest, comp, jump = "", "", ""

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

    bcomp = comp_to_binary(comp)
    bdest = dest_to_binary(dest)
    bjump = jump_to_binary(jump)

    bdest= f"{bdest:03b}"
    bjump = f"{bjump:03b}"

    return f"111{bcomp}{bdest}{bjump}"

def assemble_l_instruction(line, labels):
    label_name = line.strip("()")
    if label_name not in labels:
        labels[label_name] = len(labels)

    return None, labels

def assemble_line(line, labels):
    if line.startswith('@'):
        return assemble_a_instruction(line, labels)
    return assemble_c_instruction(line)

def preprocess_labels(lines, labels):
    label_count = 0
    for i, line in enumerate(lines):
        if not line.startswith('('):
            continue
        label_name = line.strip("()")
        labels[label_name] = i - label_count
        label_count += 1

    lines = [line for line in lines if not line.startswith('(')]

    return lines, labels

def preprocess_variables(lines, labels):
    for line in lines:
        if not line.startswith('@'):
            continue

        var_name = line.strip('@')
        if var_name.isdigit():
             continue

        
        if var_name not in labels:
             labels[var_name] = len(labels)

    return lines, labels 

def assemble(input_file):
    labels = {
        "R0": 0,
        "R1": 1,
        "R2": 2,
        "R3": 3,
        "R4": 4,
        "R5": 5,
        "R6": 6,
        "R7": 7,
        "R8": 8,
        "R9": 9,
        "R10": 10,
        "R11": 11,
        "SCREEN": 16384,
        "KBD": 24576,
    }

    lines = load_file(input_file)
    lines = remove_comments(lines)
    lines = remove_whitespace(lines)
    lines, labels = preprocess_labels(lines, labels)
    lines, labels = preprocess_variables(lines, labels)

    # print(labels)

    output = []
    for line in lines:
        output.append(assemble_line(line, labels))

    save_binary(output)


if __name__ == "__main__":
    input_file = check_args()
    assemble(input_file)