codeemitter.py

from assembly.ast import *
from typing import List


class CodeEmitter:
    """
    A class for emitting executable x64 assembly code using AT&T syntax.

    This class takes an operating system as input and generates assembly code
    for functions and instructions. It handles specific syntax differences
    between operating systems like Linux and Darwin (macOS).

    Attributes:
        operating_system (str): The target operating system ('Linux' or 'Darwin').
        function_name_prefix (str): Prefix for function names based on the operating system.

    Methods:
        visit_program(program): Translates a program node into assembly code.
        visit_fun(function): Translates a function node into assembly code.
        visit_mov(mov): Translates a 'Mov' instruction into assembly code.
    """

    def __init__(self, operating_system: str):
        """
        Initializes the CodeEmitter with the target operating system.

        Args:
            operating_system (str): The target operating system ('Linux' or 'Darwin').
        """
        self.operating_system: str = operating_system
        self.function_name_prefix: str = '_' if operating_system == "Darwin" else ''
        self.ending_line: str = '\t.section .note.GNU-stack,"",@progbits\n' if self.operating_system == 'Linux' else ''

    def _visit(self, assembly_ast_node):
        return assembly_ast_node.accept(self)
    
    def parse(self, assembly_ast: AssemblyASTNode) -> str:
        """
        Parses the assembly AST and returns code for assembler

        Returns:
            Code as string for assembler
        """
        return self._visit(assembly_ast) or ''

    def visit_program(self, program: Program) -> str:
        """
        Translates a program node into assembly code.

        Args:
            program (Program): The program node containing the function definition.

        Returns:
            str: The assembly code for the program.
        """
        function_code: str = self._visit(program.func_def)
        # Executable stack not needed on Linux
        return function_code + self.ending_line

    def visit_fun(self, function: Function) -> str:
        """
        Translates a function node into assembly code.

        Args:
            function (Function): The function node containing instructions.

        Returns:
            str: The assembly code for the function.
        """
        func_name: str = self.function_name_prefix + function.name
        header: str = f"\t.global {func_name}\n"
        label: str = f"{func_name}:\n"
        instructions: str = ''
        for instruct in function.instructions:
            if isinstance(instruct, Mov):
                instructions += f"\t{self._visit(instruct)}\n"
            elif isinstance(instruct, Ret):
                instructions += "\tret\n"
            else:
                raise RuntimeError(
                    f"Failed to translate node '{
                        instruct}' into a valid assembly instruction"
                )
        return header + label + instructions

    def visit_mov(self, mov: Mov) -> str:
        """
        Translates a 'Mov' instruction into assembly code.

        Args:
            mov (Mov): The 'Mov' instruction node with source and destination operands.

        Returns:
            str: The assembly code for the 'Mov' instruction.

        Raises:
            RuntimeError: If the source or destination operand cannot be translated.
        """
        src: str = ""
        dst: str = ""
        if isinstance(mov.source, Immediate):
            src = f"${mov.source.value}"
        elif isinstance(mov.source, Register):
            src = f"%{mov.source.name}"
        else:
            raise RuntimeError(
                f"Code Emitter: Failed to translate operand '{
                    mov.source}' inside move instruction"
            )

        if isinstance(mov.destination, Register):
            dst = f"%{mov.destination.name}"
        else:
            raise RuntimeError(
                f"Code Emitter: Failed to translate operand '{
                    mov.destination}' inside move instruction"
            )
        return f"movl\t{src}, {dst}"