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