Nand2Tetris — Chapter 8

# VM Translator 2
# Input: .vm file or directory
# Output: .asm file
# Scope: program flow, function call / return
# movement:
# Call Sys.init after setting SP=256 in bootstrap
# call saves return-address and LCL/ARG/THIS/THAT and moves to a new frame
# return restores FRAME=LCL, RET=*(FRAME-5) and then jumps to the caller location.

class VMWriter:
    def __init__(self):
        self.call_id = 0
        self.current_function = ""

    def scoped(self, label):
        if self.current_function:
            return f"{self.current_function}${label}"
        return label

    def bootstrap(self):
        return [
            "@256",
            "D=A",
            "@SP",
            "M=D",
        ] + self.write_call("Sys.init", 0)

    def write_label(self, label):
        return [f"({self.scoped(label)})"]

    def write_goto(self, label):
        return [
            f"@{self.scoped(label)}",
            "0;JMP",
        ]

    def write_if(self, label):
        return [
            "@SP",
            "AM=M-1",
            "D=M",
            f"@{self.scoped(label)}",
            "D;JNE",
        ]

    def write_function(self, name, local_count):
        self.current_function = name
        code = [f"({name})"]

        for _ in range(int(local_count)):
            code += [
                "@SP",
                "A=M",
                "M=0",
                "@SP",
                "M=M+1",
            ]

        return code

    def push_register(self, register):
        return [
            f"@{register}",
            "D=M",
            "@SP",
            "A=M",
            "M=D",
            "@SP",
            "M=M+1",
        ]

    def write_call(self, name, arg_count):
        return_label = f"{name}$ret.{self.call_id}"
        self.call_id += 1

        code = [
            f"@{return_label}",
            "D=A",
            "@SP",
            "A=M",
            "M=D",
            "@SP",
            "M=M+1",
        ]

        for register in ["LCL", "ARG", "THIS", "THAT"]:
            code += self.push_register(register)

        code += [
            "@SP",
            "D=M",
            f"@{int(arg_count) + 5}",
            "D=D-A",
            "@ARG",
            "M=D",
            "@SP",
            "D=M",
            "@LCL",
            "M=D",
            f"@{name}",
            "0;JMP",
            f"({return_label})",
        ]

        return code

    def write_return(self):
        return [
            "@LCL",
            "D=M",
            "@R13",
            "M=D",
            "@5",
            "A=D-A",
            "D=M",
            "@R14",
            "M=D",
            "@SP",
            "AM=M-1",
            "D=M",
            "@ARG",
            "A=M",
            "M=D",
            "@ARG",
            "D=M+1",
            "@SP",
            "M=D",
            "@R13",
            "AM=M-1",
            "D=M",
            "@THAT",
            "M=D",
            "@R13",
            "AM=M-1",
            "D=M",
            "@THIS",
            "M=D",
            "@R13",
            "AM=M-1",
            "D=M",
            "@ARG",
            "M=D",
            "@R13",
            "AM=M-1",
            "D=M",
            "@LCL",
            "M=D",
            "@R14",
            "A=M",
            "0;JMP",
        ]