compiler.py

"""
Eva Compiler - Main compilation functions
"""

from eva_value import *
from opcodes import *
from parser import *
from scope import *
from compiler_helpers import *

# Scope Analysis 

def analyze(state, exp, scope):
    """Analyze scopes and variable allocations"""
    if is_symbol_ast(exp):
        # Boolean literals
        if exp in ['true', 'false', 'null']:
            return
        # Variables
        else:
            scope_maybe_promote(scope, exp)
    
    elif is_list_ast(exp):
        if len(exp) == 0:
            return
        
        tag = exp[0]
        
        if not is_symbol_ast(tag):
            for e in exp:
                analyze(state, e, scope)
            return
        
        op = tag
        
        if op == 'begin':
            # New block scope
            new_scope = make_scope(
                SCOPE_BLOCK if scope else SCOPE_GLOBAL,
                scope
            )
            state['scope_info'][id(exp)] = new_scope
            
            for i in range(1, len(exp)):
                analyze(state, exp[i], new_scope)
        
        elif op == 'var':
            # Variable declaration
            scope_add_local(scope, exp[1])
            analyze(state, exp[2], scope)
        
        elif op == 'def':
            # Function declaration
            fn_name = exp[1]
            scope_add_local(scope, fn_name)
            
            new_scope = make_scope(SCOPE_FUNCTION, scope)
            state['scope_info'][id(exp)] = new_scope
            
            # Add function name for recursion
            scope_add_local(new_scope, fn_name)
            
            # Parameters
            params = exp[2]
            for param in params:
                scope_add_local(new_scope, param)
            
            # Body
            analyze(state, exp[3], new_scope)
        
        elif op == 'lambda':
            # Lambda function
            new_scope = make_scope(SCOPE_FUNCTION, scope)
            state['scope_info'][id(exp)] = new_scope
            
            # Parameters
            params = exp[1]
            for param in params:
                scope_add_local(new_scope, param)
            
            # Body
            analyze(state, exp[2], new_scope)
        
        elif op == 'class':
            # Class declaration
            class_name = exp[1]
            new_scope = make_scope(SCOPE_CLASS, scope)
            state['scope_info'][id(exp)] = new_scope
            scope_add_local(scope, class_name)
            
            # Body
            for i in range(3, len(exp)):
                analyze(state, exp[i], scope)
        
        elif op == 'prop':
            analyze(state, exp[1], scope)
        
        else:
            # Default: analyze all sub-expressions
            for i in range(1, len(exp)):
                analyze(state, exp[i], scope)

# Code Generation 

def gen(state, exp):
    """Generate bytecode for expression"""
    co = state['co']
    
    # Number
    if is_number_ast(exp):
        emit(co, OP_CONST)
        emit(co, alloc_numeric_const(co, exp))
    
    # String
    elif is_string_ast(exp):
        emit(co, OP_CONST)
        emit(co, alloc_string_const(co, exp))
    
    # Symbol
    elif is_symbol_ast(exp):
        # Boolean
        if exp in ['true', 'false']:
            emit(co, OP_CONST)
            emit(co, alloc_boolean_const(co, exp == 'true'))
        # Variable
        else:
            var_name = exp
            current_scope = state['scope_stack'][-1]
            op_code_getter = scope_get_name_getter(current_scope, var_name)
            
            emit(co, op_code_getter)
            
            if op_code_getter == OP_GET_LOCAL:
                emit(co, co_get_local_index(co, var_name))
            elif op_code_getter == OP_GET_CELL:
                emit(co, co_get_cell_index(co, var_name))
            else:  # Global
                if not global_exists(state['globals'], var_name):
                    raise Exception(f"Reference error: {var_name} doesn't exist")
                emit(co, global_get_index(state['globals'], var_name))
    
    # List
    elif is_list_ast(exp):
        if len(exp) == 0:
            return
        
        tag = exp[0]
        
        # Special forms
        if is_symbol_ast(tag):
            op = tag
            
            # Binary operations
            if op in ['+', '-', '*', '/']:
                gen(state, exp[1])
                gen(state, exp[2])
                if op == '+':
                    emit(co, OP_ADD)
                elif op == '-':
                    emit(co, OP_SUB)
                elif op == '*':
                    emit(co, OP_MUL)
                elif op == '/':
                    emit(co, OP_DIV)
            
            # Comparison
            elif op in COMPARE_OPS:
                gen(state, exp[1])
                gen(state, exp[2])
                emit(co, OP_COMPARE)
                emit(co, COMPARE_OPS[op])
            
            # If statement
            elif op == 'if':
                gen(state, exp[1])  # condition
                emit(co, OP_JMP_IF_FALSE)
                emit(co, 0)
                emit(co, 0)
                
                else_jmp_addr = get_offset(co) - 2
                
                gen(state, exp[2])  # consequent
                emit(co, OP_JMP)
                emit(co, 0)
                emit(co, 0)
                
                end_addr = get_offset(co) - 2
                
                # Patch else branch
                else_branch_addr = get_offset(co)
                patch_jump_address(co, else_jmp_addr, else_branch_addr)
                
                # Alternate
                if len(exp) == 4:
                    gen(state, exp[3])
                
                # Patch end
                end_branch_addr = get_offset(co)
                patch_jump_address(co, end_addr, end_branch_addr)
            
            # While loop
            elif op == 'while':
                loop_start = get_offset(co)
                
                gen(state, exp[1])  # test
                emit(co, OP_JMP_IF_FALSE)
                emit(co, 0)
                emit(co, 0)
                
                loop_end_jmp = get_offset(co) - 2
                
                gen(state, exp[2])  # body
                
                emit(co, OP_JMP)
                emit(co, 0)
                emit(co, 0)
                
                patch_jump_address(co, get_offset(co) - 2, loop_start)
                
                # Patch end
                loop_end = get_offset(co) + 1
                patch_jump_address(co, loop_end_jmp, loop_end)
            
            # Variable declaration
            elif op == 'var':
                var_name = exp[1]
                current_scope = state['scope_stack'][-1]
                op_code_setter = scope_get_name_setter(current_scope, var_name)
                
                # Special treatment for lambda
                if is_lambda(exp[2]):
                    compile_function(state, exp[2], var_name, exp[2][1], exp[2][2])
                else:
                    gen(state, exp[2])
                
                if op_code_setter == OP_SET_GLOBAL:
                    # Global variable
                    global_define(state['globals'], var_name)
                    emit(co, OP_SET_GLOBAL)
                    emit(co, global_get_index(state['globals'], var_name))
                elif op_code_setter == OP_SET_CELL:
                    # Cell variable
                    co['cell_names'].append(var_name)
                    emit(co, OP_SET_CELL)
                    emit(co, len(co['cell_names']) - 1)
                    emit(co, OP_POP)
                else:
                    # Local variable
                    co_add_local(co, var_name)
            
            # Set variable
            elif op == 'set':
                if is_prop(exp[1]):
                    # Property set
                    gen(state, exp[2])  # value
                    gen(state, exp[1][1])  # instance
                    emit(co, OP_SET_PROP)
                    emit(co, alloc_string_const(co, exp[1][2]))
                else:
                    var_name = exp[1]
                    current_scope = state['scope_stack'][-1]
                    op_code_setter = scope_get_name_setter(current_scope, var_name)
                    
                    gen(state, exp[2])
                    
                    if op_code_setter == OP_SET_LOCAL:
                        emit(co, OP_SET_LOCAL)
                        emit(co, co_get_local_index(co, var_name))
                    elif op_code_setter == OP_SET_CELL:
                        emit(co, OP_SET_CELL)
                        emit(co, co_get_cell_index(co, var_name))
                    else:
                        # Global
                        global_index = global_get_index(state['globals'], var_name)
                        if global_index == -1:
                            raise Exception(f"Reference error: {var_name} is not defined")
                        emit(co, OP_SET_GLOBAL)
                        emit(co, global_index)
            
            # Begin block
            elif op == 'begin':
                scope = state['scope_info'][id(exp)]
                state['scope_stack'].append(scope)
                block_enter(co)
                
                for i in range(1, len(exp)):
                    is_last = (i == len(exp) - 1)
                    is_decl = is_declaration(exp[i])
                    
                    gen(state, exp[i])
                    
                    if not is_last and not is_decl:
                        emit(co, OP_POP)
                
                block_exit(co)
                state['scope_stack'].pop()
            
            # Function definition
            elif op == 'def':
                fn_name = exp[1]
                
                compile_function(state, exp, fn_name, exp[2], exp[3])
                
                if state['class_object'] is None:
                    if is_global_scope(co):
                        global_define(state['globals'], fn_name)
                        emit(co, OP_SET_GLOBAL)
                        emit(co, global_get_index(state['globals'], fn_name))
                    else:
                        co_add_local(co, fn_name)
            
            # Lambda
            elif op == 'lambda':
                compile_function(state, exp, 'lambda', exp[1], exp[2])
            
            # Class definition
            elif op == 'class':
                name = exp[1]
                super_class = None if exp[2] == 'null' else get_class_by_name(state, exp[2])
                
                cls = make_class(name, super_class)
                state['class_objects'].append(cls)
                
                co_add_constant(co, cls)
                
                # Register as global
                global_define(state['globals'], name)
                global_set(state['globals'], global_get_index(state['globals'], name), cls)
                
                # Compile class body
                if len(exp) > 3:
                    prev_class = state['class_object']
                    state['class_object'] = cls
                    
                    scope = state['scope_info'][id(exp)]
                    state['scope_stack'].append(scope)
                    
                    for i in range(3, len(exp)):
                        gen(state, exp[i])
                    
                    state['scope_stack'].pop()
                    state['class_object'] = prev_class
                
                # Update constructor to return 'self'
                constr_fn = cls['properties']['constructor']
                constr_fn['code']['code'].insert(-3, OP_POP)
                constr_fn['code']['code'].insert(-3, OP_GET_LOCAL)
                constr_fn['code']['code'].insert(-3, 1)
            
            # New instance
            elif op == 'new':
                class_name = exp[1]
                cls = get_class_by_name(state, class_name)
                
                if cls is None:
                    raise Exception(f"Unknown class: {class_name}")
                
                # Get class
                emit(co, OP_GET_GLOBAL)
                emit(co, global_get_index(state['globals'], class_name))
                
                # New instance
                emit(co, OP_NEW)
                
                # Arguments
                for i in range(2, len(exp)):
                    gen(state, exp[i])
                
                # Call constructor
                emit(co, OP_CALL)
                emit(co, cls['properties']['constructor']['code']['arity'])
            
            # Property access
            elif op == 'prop':
                gen(state, exp[1])  # instance
                emit(co, OP_GET_PROP)
                emit(co, alloc_string_const(co, exp[2]))
            
            # Super
            elif op == 'super':
                class_name = exp[1]
                cls = get_class_by_name(state, class_name)
                
                if cls is None:
                    raise Exception(f"Unknown class: {class_name}")
                
                if cls['super_class'] is None:
                    raise Exception(f"Class {cls['name']} doesn't have super class")
                
                emit(co, OP_GET_GLOBAL)
                emit(co, global_get_index(state['globals'], cls['super_class']['name']))
            
            # Function call
            else:
                function_call(state, exp)
        
        # Lambda call
        else:
            function_call(state, exp)

def function_call(state, exp):
    """Generate code for function call"""
    co = state['co']
    
    gen(state, exp[0])  # function
    
    for i in range(1, len(exp)):
        gen(state, exp[i])  # arguments
    
    emit(co, OP_CALL)
    emit(co, len(exp) - 1)

def compile_function(state, exp, fn_name, params_exp, body):
    """Compile a function"""
    scope_info = state['scope_info'][id(exp)]
    state['scope_stack'].append(scope_info)
    
    params = params_exp
    arity = len(params)
    
    # Save previous code object
    prev_co = state['co']
    
    # Create new code object
    full_name = fn_name
    if state['class_object'] is not None:
        full_name = state['class_object']['name'] + '.' + fn_name
    
    co = make_code(full_name, arity)
    state['co'] = co
    state['code_objects'].append(co)
    
    # Set up free and cell variables
    co['free_count'] = len(scope_info['free'])
    co['cell_names'] = list(scope_info['free']) + list(scope_info['cells'])
    
    # Store code object as constant
    co_add_constant(prev_co, co)
    
    # Add function name as local (for recursion)
    co_add_local(co, fn_name)
    
    # Add parameters as locals
    for param in params:
        co_add_local(co, param)
        
        # If param is captured by cell, emit code
        cell_index = co_get_cell_index(co, param)
        if cell_index != -1:
            emit(co, OP_SET_CELL)
            emit(co, cell_index)
    
    # Compile body
    prev_class = state['class_object']
    state['class_object'] = None
    gen(state, body)
    state['class_object'] = prev_class
    
    # Scope exit if not a block
    if not is_block(body):
        emit(co, OP_SCOPE_EXIT)
        emit(co, arity + 1)
    
    # Return
    emit(co, OP_RETURN)
    
    # Handle class methods
    if state['class_object'] is not None:
        fn = make_function(co)
        state['co'] = prev_co
        state['class_object']['properties'][fn_name] = fn
    
    # Simple functions (no free variables)
    elif len(scope_info['free']) == 0:
        fn = make_function(co)
        state['co'] = prev_co
        
        co_add_constant(prev_co, fn)
        emit(prev_co, OP_CONST)
        emit(prev_co, len(prev_co['constants']) - 1)
    
    # Closures (with free variables)
    else:
        state['co'] = prev_co
        
        # Load free variables
        for free_var in scope_info['free']:
            emit(prev_co, OP_LOAD_CELL)
            emit(prev_co, co_get_cell_index(prev_co, free_var))
        
        emit(prev_co, OP_CONST)
        emit(prev_co, len(prev_co['constants']) - 1)
        
        emit(prev_co, OP_MAKE_FUNCTION)
        emit(prev_co, len(scope_info['free']))
    
    state['scope_stack'].pop()

# Main Compile Function

def compile_program(state, program):
    """Compile Eva program"""
    # Create main code object
    co = make_code('main')
    state['co'] = co
    state['code_objects'].append(co)
    
    # Create main function
    main_fn = make_function(co)
    state['main'] = main_fn
    
    # Parse program
    ast = parse('(begin ' + program + ')')
    
    # Scope analysis
    analyze(state, ast, None)
    
    # Generate code
    gen(state, ast)
    
    # Halt
    emit(co, OP_HALT)
    
    return main_fn