Code Generation & Register Allocation

Code generation is the final phase — it translates the optimized IR into actual machine code for the target architecture. This is where all the abstract work becomes real instructions the CPU executes.

Three sub-problems:

1. Instruction Selection:
Map IR operations to target machine instructions. Not always 1:1:
- IR "t = a + b" → x86 "add rax, rbx" (but which registers?)
- IR "t = a * 4" → x86 "lea rax, [rax*4]" (load effective address — faster than imul)
- IR "t = a[i]" → x86 "mov rax, [rbx + rcx*8]" (memory addressing modes)
Modern compilers use tree pattern matching (BURG — Bottom-Up Rewriting Grammar) to find optimal instruction sequences.

2. Register Allocation:
Map the unlimited IR temporaries to the finite CPU registers.
x86-64 has 16 general-purpose registers. LLVM IR has infinite virtual registers.
Graph Coloring approach:
- Build interference graph: nodes = variables, edge = "live at same time" (can't share register)
- Graph k-coloring where k = number of registers
- If a variable can't get a register → SPILL to memory (store to stack, load when needed)
- Spilling is expensive (extra load/store instructions)

Liveness analysis: variable v is LIVE at point p if there exists a path from p to a use of v not going through a definition of v. Computed by dataflow analysis (backward analysis):
- LIVEOUT(b) = ∪ LIVEIN(successor blocks)
- LIVEIN(b) = USE(b) ∪ (LIVEOUT(b) - DEF(b))

3. Instruction Scheduling:
Reorder instructions to avoid pipeline stalls (without changing semantics).
After a memory load, the result isn't available for 3-4 cycles on modern CPUs (load-use hazard). Scheduler tries to put independent instructions between the load and the use.

Calling Convention implementation:
- Save caller-saved registers before call, restore after
- Set up stack frame: push rbp, mov rbp rsp, sub rsp N (for local variables)
- Put arguments in right registers/stack positions
- Clean up stack on return

Object code and linking:
- Assembler: converts assembly text to object file (.o) — binary machine code with relocation entries
- Linker: combines .o files, resolves external references (printf, malloc), produces executable
- Dynamic linker (ld.so): at load time, resolves shared library symbols

Just-In-Time (JIT) compilation:
- Start by interpreting (fast startup)
- Profile which functions are "hot" (called frequently)
- Compile only hot functions to native code at runtime
- Can use runtime information unavailable to ahead-of-time compilers (actual types, branch frequencies)
- Used by: V8 (JavaScript), JVM (Java), PyPy (Python), LuaJIT

1# Code generation: TAC → x86-64 assembly
2
3from dataclasses import dataclass
4
5@dataclass
6class AsmInstr:
7    op:      str
8    dst:     str = ""
9    src:     str = ""
10    comment: str = ""
11    def __str__(self):
12        parts = f"    {self.op}"
13        if self.dst: parts += f" {self.dst}"
14        if self.src: parts += f", {self.src}"
15        if self.comment: parts += f"    ; {self.comment}"
16        return parts
17
18@dataclass
19class AsmLabel:
20    name: str
21    def __str__(self): return f"{self.name}:"
22
23# ── Simple code generator ──────────────────────────────
24class CodeGen:
25    # x86-64 general purpose registers
26    REGISTERS = ["rax", "rbx", "rcx", "rdx", "rsi", "rdi", "r8", "r9", "r10", "r11"]
27
28    def __init__(self):
29        self.assembly    = []
30        self.reg_map     = {}    # variable → register
31        self.free_regs   = list(self.REGISTERS)
32        self.stack_vars  = {}    # spilled variables → stack offset
33        self.stack_size  = 0
34
35    def emit(self, instr):
36        self.assembly.append(instr)
37
38    def get_reg(self, var) -> str:
39        if var in self.reg_map:
40            return self.reg_map[var]
41        if self.free_regs:
42            reg = self.free_regs.pop(0)
43            self.reg_map[var] = reg
44            return reg
45        # Spill: move oldest variable to stack
46        spill_var = next(iter(self.reg_map))
47        spill_reg = self.reg_map.pop(spill_var)
48        self.stack_size += 8
49        offset = self.stack_size
50        self.stack_vars[spill_var] = offset
51        self.emit(AsmInstr("mov", f"[rbp-{offset}]", spill_reg, f"spill {spill_var}"))
52        self.reg_map[var] = spill_reg
53        return spill_reg
54
55    def get_operand(self, var) -> str:
56        """Get register or memory location for a variable"""
57        try:
58            float(var)   # is it a constant?
59            return var
60        except ValueError:
61            pass
62        if var in self.reg_map:
63            return self.reg_map[var]
64        if var in self.stack_vars:
65            return f"[rbp-{self.stack_vars[var]}]"
66        return f"[{var}]"   # global variable
67
68    def gen_prologue(self, func_name):
69        self.emit(AsmLabel(f"_{func_name}"))
70        self.emit(AsmInstr("push", "rbp",          comment="save frame pointer"))
71        self.emit(AsmInstr("mov",  "rbp", "rsp",   comment="set up frame"))
72
73    def gen_epilogue(self):
74        self.emit(AsmInstr("pop",  "rbp",          comment="restore frame pointer"))
75        self.emit(AsmInstr("ret",                  comment="return"))
76
77    def gen_instr(self, instr):
78        if isinstance(instr, TACBinOp):
79            dst = self.get_reg(instr.result)
80            src1 = self.get_operand(instr.left)
81            src2 = self.get_operand(instr.right)
82
83            op_map = {"+": "add", "-": "sub", "*": "imul",
84                      "&": "and", "|": "or", "^": "xor",
85                      "<<": "shl", ">>": "shr"}
86
87            # Load first operand into dst register
88            if src1 != dst:
89                self.emit(AsmInstr("mov", dst, src1, f"{instr.result} = {instr.left}"))
90
91            # Apply operation
92            x86op = op_map.get(instr.op, instr.op)
93            self.emit(AsmInstr(x86op, dst, src2, f"{instr.left} {instr.op} {instr.right}"))
94
95        elif isinstance(instr, TACCopy):
96            dst = self.get_reg(instr.result)
97            src = self.get_operand(instr.value)
98            self.emit(AsmInstr("mov", dst, src, f"{instr.result} = {instr.value}"))
99
100        elif isinstance(instr, TACLabel):
101            self.emit(AsmLabel(instr.name))
102
103        elif isinstance(instr, TACJump):
104            self.emit(AsmInstr("jmp", instr.label))
105
106        elif isinstance(instr, TACCondJump):
107            cond = self.get_operand(instr.condition)
108            self.emit(AsmInstr("cmp",  cond, "0",         "test condition"))
109            self.emit(AsmInstr("jne",  instr.true_label,  "jump if true"))
110            self.emit(AsmInstr("jmp",  instr.false_label, "jump if false"))
111
112        elif isinstance(instr, TACReturn):
113            val = self.get_operand(instr.value)
114            if val != "rax":
115                self.emit(AsmInstr("mov", "rax", val, "return value in rax"))
116
117    def generate(self, tac_instrs, func_name="main"):
118        self.gen_prologue(func_name)
119        for instr in tac_instrs:
120            self.gen_instr(instr)
121        self.gen_epilogue()
122        return self.assembly
123
124# ── Generate assembly for: result = 2 + 3 * x ────────
125tac = [
126    TACBinOp("t1", "3",  "*", "x"),
127    TACBinOp("t2", "2",  "+", "t1"),
128    TACCopy("result", "t2"),
129    TACReturn("result"),
130]
131
132gen  = CodeGen()
133asm  = gen.generate(tac, "compute")
134print("Generated x86-64 Assembly:")
135print()
136for instr in asm:
137    print(instr)
138
139print("
140Register assignments:", gen.reg_map)
141
142# ── Real compilation pipeline via subprocess ──────────
143import subprocess, tempfile, os
144c_code = "int add(int a, int b) { return a + b; }"
145with tempfile.NamedTemporaryFile(suffix=".c", mode="w", delete=False) as f:
146    f.write(c_code)
147    fname = f.name
148
149result = subprocess.run(
150    ["gcc", "-O2", "-S", "-o", "/dev/stdout", fname],
151    capture_output=True, text=True
152)
153if result.returncode == 0:
154    print(f"
155GCC output for '{c_code}':")
156    for line in result.stdout.split("
157"):
158        if line.strip() and not line.startswith(".") and not line.startswith("#"):
159            print(f"  {line}")
160else:
161    print("
162(gcc not available — but this is what it would output)")
163    print("  add:")
164    print("    lea eax, [rdi+rsi]   ; result = a + b (single instruction!)")
165    print("    ret")
166os.unlink(fname)