BUILD
YOUR
COMPILER

#!/usr/bin/env python3
"""
s4c.py - A minimal compiler from S expressions to x86-64 assembly
Demonstrates the 4 phases of compilation:
1. Lexical Analysis (tokenize)
2. Syntax Analysis (parse) 
3. Code Generation (gen)
4. Assembly Output (compile_expr_to_asm)
"""
import sys, argparse

# =============================================================================
# TOKEN TYPES - The vocabulary of our language
# =============================================================================
TOK_NUM, TOK_OP, TOK_LP, TOK_RP, TOK_EOF = 'NUM','OP','LP','RP','EOF'

def tokenize(s):
    """
    PHASE 1: LEXICAL ANALYSIS
    Converts source text into tokens (lexemes)
    Example: "2 + 3 * 4" -> [NUM(2), OP(+), NUM(3), OP(*), NUM(4), EOF]
    """
    i, n, toks = 0, len(s), []  # Current position, length, token list
    
    while i < n:
        c = s[i]
        
        # Skip whitespace
        if c.isspace(): 
            i += 1
            continue
            
        # Parse multi-digit numbers
        if c.isdigit():
            j = i  # Remember start position
            while j < n and s[j].isdigit(): 
                j += 1  # Find end of number
            toks.append((TOK_NUM, int(s[i:j])))  # Convert to integer
            i = j
            continue
            
        # Single-character operators
        if c in '+-*/': 
            toks.append((TOK_OP, c))
            i += 1
            continue
            
        # Parentheses for grouping
        if c == '(': 
            toks.append((TOK_LP, c))
            i += 1
            continue
        if c == ')': 
            toks.append((TOK_RP, c))
            i += 1
            continue
            
        # Unknown character - compilation error
        raise SystemExit(f'Unsupported char: {c}')
    
    # Always end with EOF token
    toks.append((TOK_EOF, ''))
    return toks

class Parser:
    """
    PHASE 2: SYNTAX ANALYSIS
    Builds Abstract Syntax Tree (AST) with correct operator precedence
    Uses recursive descent parsing with these grammar rules:
    
    expr -> term (('+' | '-') term)*     # Lowest precedence
    term -> fact (('*' | '/') fact)*     # Higher precedence  
    fact -> number | '(' expr ')' | unary_op fact  # Highest precedence
    """
    
    def __init__(self, toks): 
        self.toks = toks     # Token stream
        self.i = 0           # Current token index
        
    def cur(self): 
        """Get current token"""
        return self.toks[self.i]
        
    def eat(self, expected_type):
        """Consume token of expected type, or error"""
        if self.cur()[0] == expected_type: 
            self.i += 1
        else: 
            raise SystemExit(f'Expected {expected_type}, got {self.cur()}')
            
    def parse(self): 
        """Entry point - parse expression"""
        return self.expr()
        
    def expr(self):
        """Parse addition/subtraction (lowest precedence)"""
        node = self.term()  # Parse left operand
        
        # Handle left-associative operators: a+b+c = ((a+b)+c)
        while self.cur()[0] == TOK_OP and self.cur()[1] in '+-':
            op = self.cur()[1]
            self.eat(TOK_OP)
            # Create binary operation node: ('bin', operator, left, right)
            node = ('bin', op, node, self.term())
        return node
        
    def term(self):
        """Parse multiplication/division (higher precedence)"""
        node = self.fact()
        
        # Same pattern as expr(), but for */ operators
        while self.cur()[0] == TOK_OP and self.cur()[1] in '*/':
            op = self.cur()[1]
            self.eat(TOK_OP)
            node = ('bin', op, node, self.fact())
        return node
        
    def fact(self):
        """Parse factors: numbers, parentheses, unary operators (highest precedence)"""
        tok = self.cur()
        
        # Unary plus/minus: +5, -3
        if tok[0] == TOK_OP and tok[1] in '+-':
            op = tok[1]
            self.eat(TOK_OP)
            return ('un', op, self.fact())  # Unary operation node
            
        # Number literal
        if tok[0] == TOK_NUM:
            self.eat(TOK_NUM)
            return ('num', tok[1])  # Number node
            
        # Parenthesized expression
        if tok[0] == TOK_LP:
            self.eat(TOK_LP)
            node = self.expr()  # Parse inner expression
            self.eat(TOK_RP)    # Expect closing paren
            return node
            
        # Syntax error
        raise SystemExit(f'Unexpected token: {tok}')

# =============================================================================
# PHASE 3: CODE GENERATION
# =============================================================================
reg_cycle = ['rax', 'rbx', 'rcx', 'rdx']  # Available x86-64 registers

def gen(node, code, it=[0]):
    """
    CODE GENERATION - Tree walking code generator
    Recursively traverses AST and emits x86-64 assembly instructions
    Uses register allocation with cycling through available registers
    """
    
    def nextreg(): 
        """Get next available register (cycles through rax->rbx->rcx->rdx->rax...)"""
        r = reg_cycle[it[0] % len(reg_cycle)]
        it[0] += 1
        return r
    
    # Number literal: load immediate value
    if node[0] == 'num':
        r = nextreg()
        code.append(f'mov {r}, {node[1]}')  # mov rax, 42
        return r
        
    # Unary operation: +expr or -expr  
    if node[0] == 'un':
        r = gen(node[2], code)  # Generate code for operand
        if node[1] == '+': 
            return r  # Unary + is no-op
        if node[1] == '-': 
            code.append(f'neg {r}')  # Two's complement negation
            return r
            
    # Binary operation: left op right
    if node[0] == 'bin':
        # Generate code for both operands (left-to-right evaluation)
        l = gen(node[2], code)  # Left operand in register l
        r = gen(node[3], code)  # Right operand in register r
        
        # Emit operation instruction
        if node[1] == '+': 
            code.append(f'add {l}, {r}')    # l += r
        if node[1] == '-': 
            code.append(f'sub {l}, {r}')    # l -= r  
        if node[1] == '*': 
            code.append(f'imul {l}, {r}')   # l *= r (signed multiply)
        if node[1] == '/': 
            # Division is complex in x86-64: requires rax/rdx register pair
            code += [
                f'mov rax, {l}',    # Move dividend to rax
                'cqo',              # Sign-extend rax into rdx:rax  
                f'idiv {r}',        # Signed divide rdx:rax by r
                f'mov {l}, rax'     # Move quotient back to result register
            ]
        return l  # Result is in left register
        
    # Unknown AST node type
    raise SystemExit('Unknown node')

# =============================================================================
# PHASE 4: ASSEMBLY OUTPUT  
# =============================================================================

# Program wrapper - creates complete Linux executable
HEAD = 'global _start\nsection .text\n_start:'  # ELF entry point
TAIL = '\n    mov rdi, rax\n    mov rax, 60\n    syscall\n'  # Linux exit syscall

def compile_expr_to_asm(expr):
    """
    MAIN COMPILATION PIPELINE
    Takes S expression string, returns complete x86-64 assembly program
    """
    # Phase 1: Tokenize source code
    toks = tokenize(expr)
    
    # Phase 2: Parse tokens into AST  
    ast = Parser(toks).parse()
    
    # Phase 3: Generate assembly code
    body = []           # Assembly instruction list
    it = [0]           # Register counter (mutable for nested calls)
    result_reg = gen(ast, body, it)
    
    # Ensure final result is in rax (required for exit code)
    if result_reg != 'rax': 
        body.append(f'mov rax, {result_reg}')
    
    # Phase 4: Wrap in complete program
    # Indent body instructions and combine with header/footer
    asm = [HEAD] + ['    ' + line for line in body] + [TAIL]
    return '\n'.join(asm)

def main():
    """Command-line interface - reads S file, writes assembly file"""
    ap = argparse.ArgumentParser(description='S4C: S to x86-64 Compiler')
    ap.add_argument('src', help='S source file (MVP: a single expression)')
    ap.add_argument('-o', '--out', default='out.asm', help='Output assembly file')
    args = ap.parse_args()
    
    # Read source expression from file
    expr = open(args.src).read().strip()
    
    # Compile to assembly
    asm = compile_expr_to_asm(expr)
    
    # Write assembly output
    open(args.out, 'w').write(asm)
    print(f'[ok] wrote {args.out}')

if __name__ == '__main__':
    main()