codegen.py

# codegen.py
# Target Code Generator for CPL
# TAC Instructions → x86-style Assembly (AT&T syntax, simplified)
#
# WHY x86 assembly?
# The problem statement says "cross-compile to C / Assembly-like code".
# We generate a clean, readable x86-like assembly that demonstrates the
# concept — it's not meant to run directly on hardware, but it follows
# real x86 conventions closely enough to be credible.
#
# REGISTER CONVENTION we use:
#   %eax  — accumulator (return values, intermediate results)
#   %ebx  — second operand in binary ops
#   %esp  — stack pointer (for function calls)
#   Variables live in memory: -4(%ebp), -8(%ebp), etc. (stack frame)
#
# APPROACH:
# We maintain a variable→stack-offset map.
# Each new variable gets the next available slot on the stack.
# Temporaries and variables are treated the same way.
#
# WHAT GETS GENERATED:
#   .section .data     ← string literals
#   .section .text     ← code
#   Each function is a labeled block with a standard prologue/epilogue.

from ir import TACInstruction


class CodeGenerator:
    """
    Converts a list of TACInstruction objects into x86-like assembly text.

    Usage:
        cg  = CodeGenerator()
        asm = cg.generate(tac_list)
        print(asm)
    """

    def __init__(self):
        self.output      = []       # lines of assembly
        self.var_offset  = {}       # variable name → stack offset (e.g. -4, -8)
        self.stack_size  = 0        # current stack usage in bytes
        self.strings     = {}       # string literal → label name (for .data section)
        self.str_count   = 0

    # ── helpers ───────────────────────────────────────────

    def emit(self, line: str):
        self.output.append(line)

    def get_var_offset(self, name: str) -> int:
        """
        Allocate a stack slot for 'name' if it doesn't have one yet.
        WHY stack? Local variables live on the stack in real x86 code.
        Each int/float takes 4 bytes, so we go -4, -8, -12, ...
        """
        if name not in self.var_offset:
            self.stack_size     += 4
            self.var_offset[name] = -self.stack_size
        return self.var_offset[name]

    def mem(self, name: str) -> str:
        """Return the memory reference for a variable: e.g. -8(%ebp)"""
        offset = self.get_var_offset(name)
        return f"{offset}(%ebp)"

    def load_to_eax(self, name):
        """Load a variable or immediate value into %eax."""
        if isinstance(name, (int, float, bool)):
            val = int(name) if isinstance(name, bool) else name
            self.emit(f"    movl    ${val}, %eax")
        else:
            self.emit(f"    movl    {self.mem(name)}, %eax")

    def store_from_eax(self, name: str):
        """Store %eax into a variable's stack slot."""
        self.emit(f"    movl    %eax, {self.mem(name)}")

    # ── main entry ────────────────────────────────────────

    def generate(self, instructions: list) -> str:
        """
        Two-pass generation:
        Pass 1: collect string literals → .data section
        Pass 2: emit code → .text section
        """
        # Pass 1: find strings
        for instr in instructions:
            if instr.op == 'ASSIGN' and isinstance(instr.arg1, str) and \
               instr.arg1.startswith('"') or \
               (instr.op == 'ASSIGN' and isinstance(instr.arg1, str) and
                not instr.arg1.startswith('t') and not instr.arg1.startswith('L')
                and not instr.arg1[0].isalpha()):
                pass    # string detection handled inline below

        # Header
        self.emit("# CPL Compiled Output — x86 AT&T Assembly")
        self.emit("# Generated by cplc (Bootstrap-Enabled Cross Compiler)")
        self.emit("")
        self.emit(".section .data")
        self.emit('    fmt_int:    .string "%d\\n"')    # printf format for ints
        self.emit('    fmt_float:  .string "%f\\n"')    # printf format for floats
        self.emit('    fmt_str:    .string "%s\\n"')    # printf format for strings
        self.emit('    fmt_input:  .string "%d"')       # scanf format
        self.emit("")
        self.emit(".section .text")
        self.emit(".global main")
        self.emit("")

        # Check if there are function declarations
        has_main = any(i.op == 'FUNC_BEGIN' and i.result == 'main'
                       for i in instructions)

        # If no explicit main function, wrap everything in main
        if not has_main:
            self._emit_func_prologue('main')
            for instr in instructions:
                self._emit_instruction(instr)
            self._emit_func_epilogue('main')
        else:
            for instr in instructions:
                self._emit_instruction(instr)

        return "\n".join(self.output)

    # ── function prologue / epilogue ──────────────────────

    def _emit_func_prologue(self, name: str):
        """
        Standard x86 function entry:
          push %ebp         ← save caller's base pointer
          mov %esp, %ebp   ← set our own base pointer
          sub $N, %esp     ← reserve space for local variables

        WHY? This is the "stack frame" setup — every real C function does this.
        %ebp becomes our reference point for accessing local variables.
        """
        self.emit(f"{name}:")
        self.emit(f"    pushl   %ebp")
        self.emit(f"    movl    %esp, %ebp")
        # Reserve 64 bytes for locals (we'll use as much as we need)
        self.emit(f"    subl    $64, %esp")
        self.emit("")

    def _emit_func_epilogue(self, name: str):
        """
        Standard x86 function exit:
          mov %ebp, %esp   ← restore stack pointer
          pop %ebp         ← restore caller's base pointer
          ret              ← return to caller

        For main, we also call exit(0) properly via  movl $0, %eax + ret.
        """
        self.emit("")
        if name == 'main':
            self.emit(f"    movl    $0, %eax")      # return 0 from main
        self.emit(f"    movl    %ebp, %esp")
        self.emit(f"    popl    %ebp")
        self.emit(f"    ret")
        self.emit("")

    # ── instruction emitter ───────────────────────────────

    def _emit_instruction(self, instr: TACInstruction):
        op = instr.op

        if op == 'FUNC_BEGIN':
            self.var_offset = {}       # new function = fresh stack frame
            self.stack_size = 0
            self._emit_func_prologue(instr.result)

        elif op == 'FUNC_END':
            self._emit_func_epilogue(instr.result)

        elif op == 'LABEL':
            # Labels become assembly labels (no indentation)
            self.emit(f"{instr.result}:")

        elif op == 'ASSIGN':
            # result = arg1
            self.load_to_eax(instr.arg1)
            self.store_from_eax(instr.result)

        elif op == 'BINOP':
            self._emit_binop(instr)

        elif op == 'UNARY':
            self._emit_unary(instr)

        elif op == 'JUMP':
            self.emit(f"    jmp     {instr.result}")

        elif op == 'JUMPIF':
            # if arg1 goto result
            self.load_to_eax(instr.arg1)
            self.emit(f"    cmpl    $0, %eax")
            self.emit(f"    jne     {instr.result}")

        elif op == 'JUMPIFNOT':
            # ifnot arg1 goto result
            self.load_to_eax(instr.arg1)
            self.emit(f"    cmpl    $0, %eax")
            self.emit(f"    je      {instr.result}")

        elif op == 'PRINT':
            self._emit_print(instr.arg1)

        elif op == 'INPUT':
            self._emit_input(instr.arg1)

        elif op == 'RETURN':
            if instr.arg1:
                self.load_to_eax(instr.arg1)
            self.emit(f"    movl    %ebp, %esp")
            self.emit(f"    popl    %ebp")
            self.emit(f"    ret")

        elif op == 'PARAM':
            # Push argument onto stack before function call
            self.load_to_eax(instr.arg1)
            self.emit(f"    pushl   %eax")

        elif op == 'CALL':
            self.emit(f"    call    {instr.arg1}")
            # Clean up pushed arguments (arg2 = number of args, each 4 bytes)
            if instr.arg2 and instr.arg2 > 0:
                self.emit(f"    addl    ${instr.arg2 * 4}, %esp")
            self.store_from_eax(instr.result)

        elif op == 'ARRAY_GET':
            # result = array[index]
            # Load index → compute address → load value
            self.load_to_eax(instr.arg2)
            self.emit(f"    imull   $4, %eax")         # index * 4 (each element = 4 bytes)
            arr_off = self.get_var_offset(instr.arg1)
            self.emit(f"    leal    {arr_off}(%ebp), %ebx")
            self.emit(f"    movl    (%ebx,%eax,1), %eax")
            self.store_from_eax(instr.result)

        elif op == 'ARRAY_SET':
            # array[index] = value
            self.load_to_eax(instr.arg1)
            self.emit(f"    imull   $4, %eax")
            arr_off = self.get_var_offset(instr.result)
            self.emit(f"    leal    {arr_off}(%ebp), %ebx")
            self.emit(f"    pushl   %eax")
            self.load_to_eax(instr.arg2)
            self.emit(f"    popl    %ecx")
            self.emit(f"    movl    %eax, (%ebx,%ecx,1)")

    def _emit_binop(self, instr: TACInstruction):
        """
        Emit x86 for: result = arg1 OP arg2

        Pattern:
          load arg1 → %eax
          load arg2 → %ebx
          apply op  → result in %eax
          store %eax → result
        """
        self.load_to_eax(instr.arg1)
        self.emit(f"    movl    %eax, %ebx")    # save arg1 in %ebx
        self.load_to_eax(instr.arg2)            # arg2 → %eax

        op = instr.op_sym

        if op == '+':
            self.emit(f"    addl    %ebx, %eax")
        elif op == '-':
            # a - b: we need ebx - eax (arg1 - arg2)
            self.emit(f"    subl    %eax, %ebx")
            self.emit(f"    movl    %ebx, %eax")
        elif op == '*':
            self.emit(f"    imull   %ebx, %eax")
        elif op == '/':
            # idivl requires dividend in %eax:%edx
            self.emit(f"    movl    %ebx, %ecx")    # save arg1
            self.emit(f"    movl    %eax, %ebx")    # arg2 → divisor
            self.emit(f"    movl    %ecx, %eax")    # arg1 → dividend
            self.emit(f"    cdq")                    # sign-extend eax into edx:eax
            self.emit(f"    idivl   %ebx")           # eax = eax / ebx
        elif op == '%':
            self.emit(f"    movl    %ebx, %ecx")
            self.emit(f"    movl    %eax, %ebx")
            self.emit(f"    movl    %ecx, %eax")
            self.emit(f"    cdq")
            self.emit(f"    idivl   %ebx")
            self.emit(f"    movl    %edx, %eax")    # remainder is in %edx
        elif op in ('==', '!=', '<', '>', '<=', '>='):
            # Comparison: use cmp + setXX to get 0/1 in %eax
            self.emit(f"    cmpl    %eax, %ebx")    # compare arg1 vs arg2
            cmp_map = {
                '==': 'sete', '!=': 'setne',
                '<':  'setl', '>':  'setg',
                '<=': 'setle', '>=': 'setge',
            }
            self.emit(f"    {cmp_map[op]}  %al")
            self.emit(f"    movzbl  %al, %eax")     # zero-extend byte to 32-bit
        elif op == '&&':
            # Logical AND: both nonzero → 1
            self.emit(f"    cmpl    $0, %ebx")
            self.emit(f"    setne   %cl")
            self.emit(f"    cmpl    $0, %eax")
            self.emit(f"    setne   %al")
            self.emit(f"    andb    %cl, %al")
            self.emit(f"    movzbl  %al, %eax")
        elif op == '||':
            self.emit(f"    orl     %ebx, %eax")
            self.emit(f"    setne   %al")
            self.emit(f"    movzbl  %al, %eax")

        self.store_from_eax(instr.result)

    def _emit_unary(self, instr: TACInstruction):
        self.load_to_eax(instr.arg1)
        if instr.op_sym == '-':
            self.emit(f"    negl    %eax")
        elif instr.op_sym == '!':
            self.emit(f"    cmpl    $0, %eax")
            self.emit(f"    sete    %al")
            self.emit(f"    movzbl  %al, %eax")
        self.store_from_eax(instr.result)

    def _emit_print(self, var: str):
        """
        Call printf with the variable's value.
        We use the int format by default.
        WHY printf? It's the standard C library call — realistic and simple.
        """
        self.load_to_eax(var)
        self.emit(f"    pushl   %eax")
        self.emit(f"    pushl   $fmt_int")
        self.emit(f"    call    printf")
        self.emit(f"    addl    $8, %esp")

    def _emit_input(self, var: str):
        """
        Call scanf to read an int from stdin into the variable.
        """
        offset = self.get_var_offset(var)
        self.emit(f"    leal    {offset}(%ebp), %eax")
        self.emit(f"    pushl   %eax")
        self.emit(f"    pushl   $fmt_input")
        self.emit(f"    call    scanf")
        self.emit(f"    addl    $8, %esp")