code-generation

name	code-generation
description	Target code emission and optimization. Use when generating WebAssembly Text format (WAT), translating AST to output code, or implementing compiler backends.

Code Generation

Generate target code from Abstract Syntax Trees, focusing on WebAssembly Text format (WAT).

Quick Start

class CodeGenerator:
    """Generate WebAssembly Text format from AST."""

    def __init__(self):
        self.output = []
        self.indent_level = 0

    def emit(self, code):
        """Emit a line of code with proper indentation."""
        indent = '  ' * self.indent_level
        self.output.append(indent + code)

    def generate(self, ast):
        """Generate WAT from AST."""
        self.emit('(module')
        self.indent_level += 1

        for func in ast['functions']:
            self.generate_function(func)

        self.indent_level -= 1
        self.emit(')')
        return '\n'.join(self.output)

    def generate_function(self, node):
        name = node['name']
        params = node['params']
        return_type = node['return_type']

        # Function signature
        param_str = ' '.join([f'(param ${p["name"]} i32)' for p in params])
        self.emit(f'(func ${name} {param_str} (result i32)')
        self.indent_level += 1

        # Generate body
        for stmt in node['body']:
            self.generate_statement(stmt)

        self.indent_level -= 1
        self.emit(')')

        # Export main function
        if name == 'main':
            self.emit(f'(export "main" (func ${name}))')

WAT Format Basics

# WebAssembly Text Format structure:
# (module
#   (func $name (param $p1 i32) (result i32)
#     (local $var i32)
#     (local.set $var (i32.const 10))
#     (i32.add (local.get $p1) (local.get $var))
#   )
#   (export "name" (func $name))
# )

def generate_module(functions):
    """Generate complete WAT module."""
    lines = ['(module']
    for func in functions:
        lines.extend(generate_function_wat(func))
    lines.append(')')
    return '\n'.join(lines)

def generate_function_wat(func):
    """Generate WAT for a single function."""
    name = func['name']
    params = ' '.join([f'(param ${p["name"]} i32)' for p in func['params']])
    result = '(result i32)' if func['return_type'] else ''

    lines = [f'  (func ${name} {params} {result}']

    # Local variables
    for var in func.get('locals', []):
        lines.append(f'    (local ${var["name"]} i32)')

    # Function body
    for stmt in func['body']:
        lines.extend(['    ' + line for line in generate_statement_wat(stmt)])

    lines.append('  )')

    # Export if main
    if name == 'main':
        lines.append(f'  (export "main" (func ${name}))')

    return lines

Expression Code Generation

def generate_expression(node):
    """Generate WAT expression from AST node."""
    node_type = node['type']

    if node_type == 'Number':
        return f'(i32.const {node["value"]})'

    elif node_type == 'Identifier':
        return f'(local.get ${node["name"]})'

    elif node_type == 'BinaryExpression':
        op = node['operator']
        left = generate_expression(node['left'])
        right = generate_expression(node['right'])

        # Map operators to WAT instructions
        wat_ops = {
            '+': 'i32.add',
            '-': 'i32.sub',
            '*': 'i32.mul',
            '/': 'i32.div_s',
            '<': 'i32.lt_s',
            '>': 'i32.gt_s',
            '==': 'i32.eq',
        }

        wat_op = wat_ops.get(op)
        if not wat_op:
            raise ValueError(f"Unsupported operator: {op}")

        return f'({wat_op} {left} {right})'

    elif node_type == 'CallExpression':
        callee = node['callee']
        args = ' '.join([generate_expression(arg) for arg in node['arguments']])
        return f'(call ${callee} {args})'

    else:
        raise ValueError(f"Unknown expression type: {node_type}")

Statement Code Generation

def generate_statement(node):
    """Generate WAT statement from AST node."""
    node_type = node['type']

    if node_type == 'VariableDeclaration':
        # In WAT, locals are declared at function start
        # Assignment becomes local.set
        name = node['name']
        if node.get('initializer'):
            expr = generate_expression(node['initializer'])
            return f'(local.set ${name} {expr})'
        return ''

    elif node_type == 'ReturnStatement':
        expr = generate_expression(node['expression'])
        return expr  # Last expression in function is return value

    elif node_type == 'IfStatement':
        condition = generate_expression(node['condition'])
        then_stmts = '\n'.join([generate_statement(s) for s in node['then']])

        if node.get('else'):
            else_stmts = '\n'.join([generate_statement(s) for s in node['else']])
            return f'(if (result i32) {condition}\n  (then {then_stmts})\n  (else {else_stmts})\n)'
        else:
            return f'(if {condition}\n  (then {then_stmts})\n)'

    else:
        raise ValueError(f"Unknown statement type: {node_type}")

Complete Code Generator

class WATCodeGenerator:
    """Complete WebAssembly Text format code generator."""

    def __init__(self):
        self.output = []
        self.indent = 0
        self.local_vars = set()

    def emit(self, line):
        self.output.append('  ' * self.indent + line)

    def generate_program(self, ast):
        """Generate complete WAT module from program AST."""
        self.emit('(module')
        self.indent += 1

        for func_node in ast['functions']:
            self.generate_function(func_node)

        self.indent -= 1
        self.emit(')')
        return '\n'.join(self.output)

    def generate_function(self, node):
        """Generate function with locals collected from body."""
        name = node['name']
        params = node['params']
        return_type = node['return_type']

        # Collect local variables from function body
        self.local_vars = set()
        self.collect_locals(node['body'])

        # Generate function signature
        params_str = ' '.join([f'(param ${p["name"]} i32)' for p in params])
        result_str = '(result i32)' if return_type == 'int' else ''
        self.emit(f'(func ${name} {params_str} {result_str}')
        self.indent += 1

        # Declare local variables
        for var in self.local_vars:
            self.emit(f'(local ${var} i32)')

        # Generate function body
        for stmt in node['body']:
            self.generate_statement(stmt)

        self.indent -= 1
        self.emit(')')

        # Export main function
        if name == 'main':
            self.emit(f'(export "main" (func ${name}))')

    def collect_locals(self, statements):
        """Collect local variable declarations."""
        for stmt in statements:
            if stmt['type'] == 'VariableDeclaration':
                self.local_vars.add(stmt['name'])

    def generate_statement(self, stmt):
        if stmt['type'] == 'VariableDeclaration':
            if stmt.get('initializer'):
                expr = self.generate_expression(stmt['initializer'])
                self.emit(f'(local.set ${stmt["name"]} {expr})')

        elif stmt['type'] == 'ReturnStatement':
            expr = self.generate_expression(stmt['expression'])
            self.emit(expr)

    def generate_expression(self, expr):
        if expr['type'] == 'Number':
            return f'(i32.const {expr["value"]})'

        elif expr['type'] == 'Identifier':
            return f'(local.get ${expr["name"]})'

        elif expr['type'] == 'BinaryExpression':
            ops = {'+': 'i32.add', '-': 'i32.sub', '*': 'i32.mul', '/': 'i32.div_s'}
            op = ops[expr['operator']]
            left = self.generate_expression(expr['left'])
            right = self.generate_expression(expr['right'])
            return f'({op} {left} {right})'

        elif expr['type'] == 'CallExpression':
            args = ' '.join([self.generate_expression(arg) for arg in expr['arguments']])
            return f'(call ${expr["callee"]} {args})'

# Usage
generator = WATCodeGenerator()
wat_code = generator.generate_program(ast)
with open('output.wat', 'w') as f:
    f.write(wat_code)

Optimization Techniques

# Constant folding
def constant_fold(expr):
    if expr['type'] == 'BinaryExpression':
        left = constant_fold(expr['left'])
        right = constant_fold(expr['right'])

        if left['type'] == 'Number' and right['type'] == 'Number':
            # Evaluate at compile time
            result = eval(f"{left['value']} {expr['operator']} {right['value']}")
            return {'type': 'Number', 'value': result}

        return {'type': 'BinaryExpression', 'operator': expr['operator'],
                'left': left, 'right': right}
    return expr