aster/aster-core/src/vm/compiler.rs

//! AST → bytecode compiler for the Aster VM.
//!
//! Walks the AST recursively, emitting stack-based bytecode instructions.
//! Handles local variable resolution (slot indices), upvalue capture for
//! closures, jump backpatching for control flow, and constant pool management.

use crate::ast::*;
use crate::ast::expr::{Literal, UnaryOp, BinaryOp, LogicalOp, AssignOp};
use crate::error::RuntimeError;
use crate::interpreter::Value;
use super::opcode::*;

use std::rc::Rc;
use std::cell::RefCell;

// ============================================================================
// FunctionProto — compiled function blueprint
// ============================================================================

#[derive(Debug, Clone)]
pub struct FunctionProto {
    pub name: Option<String>,
    pub arity: u8,
    pub code: Vec<u8>,
    pub constants: Vec<Value>,
    pub upvalue_count: u8,
    pub lines: Vec<(usize, usize)>, // (bytecode_offset, source_line)
}

impl FunctionProto {
    pub fn new(name: Option<String>) -> Self {
        Self {
            name,
            arity: 0,
            code: Vec::new(),
            constants: Vec::new(),
            upvalue_count: 0,
            lines: Vec::new(),
        }
    }

    fn add_constant(&mut self, val: Value) -> u16 {
        // Check for existing identical constant
        for (i, c) in self.constants.iter().enumerate() {
            if values_eq(c, &val) {
                return i as u16;
            }
        }
        let idx = self.constants.len();
        self.constants.push(val);
        idx as u16
    }
}

fn values_eq(a: &Value, b: &Value) -> bool {
    match (a, b) {
        (Value::Number(x), Value::Number(y)) => (x - y).abs() < f64::EPSILON,
        (Value::String(x), Value::String(y)) => x == y,
        (Value::Bool(x), Value::Bool(y)) => x == y,
        (Value::Nil, Value::Nil) => true,
        _ => false,
    }
}

// ============================================================================
// Compiler
// ============================================================================

struct Local {
    name: String,
    depth: u8,       // scope depth where declared; 0 = uninitialized
    is_captured: bool,
    is_const: bool,
}

struct Upvalue {
    index: u8,
    is_local: bool,  // true = captured from enclosing fn's local; false = from upvalue
}

struct LoopContext {
    start_ip: usize,             // bytecode offset of loop condition
    break_patches: Vec<usize>,   // jump instruction offsets that need break dest
    continue_patches: Vec<usize>,// jump instruction offsets that need continue dest
    scope_depth: u8,             // scope depth when loop started
}

pub struct Compiler {
    function: FunctionProto,
    locals: Vec<Local>,
    upvalues: Vec<Upvalue>,
    scope_depth: u8,
    loop_stack: Vec<LoopContext>,
    /// Index into an outer compiler array for upvalue resolution
    enclosing_idx: Option<usize>,
}

impl Compiler {
    pub fn new(name: Option<String>) -> Self {
        Self {
            function: FunctionProto::new(name),
            locals: Vec::new(),
            upvalues: Vec::new(),
            scope_depth: 0,
            loop_stack: Vec::new(),
            enclosing_idx: None,
        }
    }

    /// Convenience: compile a list of statements into a FunctionProto.
    pub fn compile(stmts: &[Stmt]) -> Result<FunctionProto, RuntimeError> {
        let mut compiler = Self::new(None);
        for stmt in stmts {
            compiler.compile_stmt(stmt)?;
        }
        // Implicit return nil at end of function
        compiler.emit_op(OpCode::LoadNil);
        compiler.emit_op(OpCode::Return);
        Ok(compiler.function)
    }

    /// Compile a list of statements (for use from nested compilers).
    fn compile_stmts(&mut self, stmts: &[Stmt]) -> Result<(), RuntimeError> {
        for stmt in stmts {
            self.compile_stmt(stmt)?;
        }
        Ok(())
    }

    // ========================================================================
    // Statement compilation
    // ========================================================================

    fn compile_stmt(&mut self, stmt: &Stmt) -> Result<(), RuntimeError> {
        match stmt {
            Stmt::Let { name, initializer, mutable } => {
                self.compile_let(name, initializer, *mutable)?;
            }
            Stmt::ExprStmt(expr) => {
                self.compile_expr(expr)?;
                self.emit_op(OpCode::Pop);
            }
            Stmt::Block(stmts) => {
                self.begin_scope();
                self.compile_stmts(stmts)?;
                self.end_scope();
            }
            Stmt::If { condition, then_branch, else_branch } => {
                self.compile_if(condition, then_branch, else_branch.as_deref())?;
            }
            Stmt::While { condition, body } => {
                self.compile_while(condition, body)?;
            }
            Stmt::For { initializer, condition, step, body } => {
                self.compile_for(initializer.as_deref(), condition.as_ref(), step.as_ref(), body)?;
            }
            Stmt::ForIn { var_name, iterable, body } => {
                self.compile_for_in(var_name, iterable, body)?;
            }
            Stmt::Function { name, params, body } => {
                self.compile_function_decl(name, params, body)?;
            }
            Stmt::Return(expr) => {
                if let Some(e) = expr {
                    self.compile_expr(e)?;
                } else {
                    self.emit_op(OpCode::LoadNil);
                }
                self.emit_op(OpCode::Return);
            }
            Stmt::Break => {
                self.compile_break()?;
            }
            Stmt::Continue => {
                self.compile_continue()?;
            }
        }
        Ok(())
    }

    fn compile_let(&mut self, name: &str, initializer: &Expr, mutable: bool) -> Result<(), RuntimeError> {
        self.compile_expr(initializer)?;
        if self.scope_depth == 0 {
            // Global variable
            let name_idx = self.add_string_constant(name);
            emit_u16(&mut self.function.code, OpCode::DefineGlobal, name_idx);
        } else {
            // Local variable
            let slot = self.locals.len() as u8;
            self.locals.push(Local {
                name: name.to_string(),
                depth: self.scope_depth,
                is_captured: false,
                is_const: !mutable,
            });
            emit_u8(&mut self.function.code, OpCode::StoreLocal, slot);
        }
        Ok(())
    }

    fn compile_if(&mut self, condition: &Expr, then_branch: &Stmt, else_branch: Option<&Stmt>) -> Result<(), RuntimeError> {
        self.compile_expr(condition)?;
        let else_jump = self.emit_jump(OpCode::JumpIfFalse);
        self.compile_stmt(then_branch)?;
        if let Some(else_stmt) = else_branch {
            let end_jump = self.emit_jump(OpCode::Jump);
            self.patch_jump(else_jump);
            self.compile_stmt(else_stmt)?;
            self.patch_jump(end_jump);
        } else {
            self.patch_jump(else_jump);
        }
        Ok(())
    }

    fn compile_while(&mut self, condition: &Expr, body: &Stmt) -> Result<(), RuntimeError> {
        let start_ip = self.function.code.len();
        self.compile_expr(condition)?;
        let exit_jump = self.emit_jump(OpCode::JumpIfFalse);

        self.loop_stack.push(LoopContext {
            start_ip,
            break_patches: Vec::new(),
            continue_patches: Vec::new(),
            scope_depth: self.scope_depth,
        });

        self.compile_stmt(body)?;
        self.emit_loop_jump(start_ip);

        self.patch_jump(exit_jump);

        let loop_ctx = self.loop_stack.pop().unwrap();
        for patch in loop_ctx.break_patches {
            self.patch_jump(patch);
        }
        for patch in loop_ctx.continue_patches {
            self.patch_jump_to(patch, start_ip);
        }
        Ok(())
    }

    fn compile_for(&mut self, initializer: Option<&Stmt>, condition: Option<&Expr>, step: Option<&Expr>, body: &Stmt) -> Result<(), RuntimeError> {
        // For loop: new scope for init
        self.begin_scope();

        if let Some(init) = initializer {
            self.compile_stmt(init)?;
        }

        let start_ip = self.function.code.len();
        if let Some(cond) = condition {
            self.compile_expr(cond)?;
        } else {
            self.emit_op(OpCode::LoadTrue);
        }
        let exit_jump = self.emit_jump(OpCode::JumpIfFalse);

        self.loop_stack.push(LoopContext {
            start_ip,
            break_patches: Vec::new(),
            continue_patches: Vec::new(),
            scope_depth: self.scope_depth,
        });

        self.compile_stmt(body)?;

        // continue lands here (after body, before step)
        let continue_ip = self.function.code.len();
        let loop_ctx = self.loop_stack.pop().unwrap();
        for patch in loop_ctx.continue_patches {
            self.patch_jump_to(patch, continue_ip);
        }

        if let Some(s) = step {
            self.compile_expr(s)?;
            self.emit_op(OpCode::Pop);
        }

        self.emit_loop_jump(start_ip);
        self.patch_jump(exit_jump);

        for patch in loop_ctx.break_patches {
            self.patch_jump(patch);
        }

        self.end_scope();
        Ok(())
    }

    fn compile_for_in(&mut self, var_name: &str, iterable: &Expr, body: &Stmt) -> Result<(), RuntimeError> {
        // Evaluate iterable, set up iterator
        self.compile_expr(iterable)?;
        emit_op(&mut self.function.code, OpCode::ForInInit);
        let exit_jump = emit_i16_placeholder(&mut self.function.code, OpCode::ForInNext);

        // New scope for the loop variable
        self.begin_scope();
        let loop_var_slot = self.locals.len() as u8;
        self.locals.push(Local {
            name: var_name.to_string(),
            depth: self.scope_depth,
            is_captured: false,
            is_const: false,
        });
        emit_u8(&mut self.function.code, OpCode::StoreLocal, loop_var_slot);

        let start_ip = self.function.code.len();

        self.loop_stack.push(LoopContext {
            start_ip,
            break_patches: Vec::new(),
            continue_patches: Vec::new(),
            scope_depth: self.scope_depth,
        });

        self.compile_stmt(body)?;

        // Patch continue → loop back
        let continue_ip = self.function.code.len();
        let loop_ctx = self.loop_stack.pop().unwrap();
        for patch in loop_ctx.continue_patches {
            self.patch_jump_to(patch, continue_ip);
        }

        // Jump back to ForInNext
        self.emit_loop_jump(start_ip - 3); // jump back to the ForInNext instruction
        let exit_ip = self.function.code.len();
        self.patch_jump_to(exit_jump, exit_ip);

        for patch in loop_ctx.break_patches {
            self.patch_jump(patch);
        }

        // Pop the loop variable
        self.emit_op(OpCode::Pop);
        // Pop the iterator state (2 values: iterable ref + index)
        self.emit_op(OpCode::Pop);
        self.emit_op(OpCode::Pop);

        self.end_scope();
        Ok(())
    }

    fn compile_function_decl(&mut self, name: &str, params: &[String], body: &[Stmt]) -> Result<(), RuntimeError> {
        let proto = self.compile_nested_function(Some(name.to_string()), params, body)?;
        let const_idx = self.function.add_constant(Value::Function(Rc::new(
            crate::interpreter::Function {
                params: params.to_vec(),
                body: body.to_vec(),
                env: Rc::new(RefCell::new(crate::interpreter::Env::new(None))),
                name: Some(name.to_string()),
            }
        )));
        // For now, we also need to store the proto for the VM to use.
        // We'll store it as a special constant.
        let proto_idx = self.add_function_proto_constant(proto);

        // Emit Closure opcode with upvalues
        self.emit_closure(proto_idx);

        // Bind to name
        if self.scope_depth == 0 {
            let name_idx = self.add_string_constant(name);
            emit_u16(&mut self.function.code, OpCode::DefineGlobal, name_idx);
        } else {
            let slot = self.locals.len() as u8;
            self.locals.push(Local {
                name: name.to_string(),
                depth: self.scope_depth,
                is_captured: false,
                is_const: false,
            });
            emit_u8(&mut self.function.code, OpCode::StoreLocal, slot);
        }
        Ok(())
    }

    fn compile_break(&mut self) -> Result<(), RuntimeError> {
        let jump_loc = self.emit_jump_placeholder();
        let loop_ctx = self.loop_stack.last_mut().ok_or_else(|| RuntimeError::RuntimeError {
            message: "break outside of loop".into(),
            token: None,
        })?;
        loop_ctx.break_patches.push(jump_loc);
        Ok(())
    }

    fn compile_continue(&mut self) -> Result<(), RuntimeError> {
        let jump_loc = self.emit_jump_placeholder();
        let loop_ctx = self.loop_stack.last_mut().ok_or_else(|| RuntimeError::RuntimeError {
            message: "continue outside of loop".into(),
            token: None,
        })?;
        loop_ctx.continue_patches.push(jump_loc);
        Ok(())
    }

    // ========================================================================
    // Expression compilation
    // ========================================================================

    fn compile_expr(&mut self, expr: &Expr) -> Result<(), RuntimeError> {
        match expr {
            Expr::Literal(lit) => self.compile_literal(lit),
            Expr::Variable(name) => self.compile_variable(name),
            Expr::Assign { name, op, value } => self.compile_assign(name, op, value),
            Expr::Get { object, name } => self.compile_get(object, name),
            Expr::Set { object, name, op, value } => self.compile_set(object, name, op, value),
            Expr::ObjectLiteral { properties } => self.compile_object(properties),
            Expr::ArrayLiteral { elements } => self.compile_array(elements),
            Expr::IndexGet { array, index } => self.compile_index_get(array, index),
            Expr::IndexSet { array, index, op, value } => self.compile_index_set(array, index, op, value),
            Expr::Unary { op, right } => self.compile_unary(op, right),
            Expr::Binary { left, op, right } => self.compile_binary(left, op, right),
            Expr::Logical { left, op, right } => self.compile_logical(left, op, right),
            Expr::Ternary { condition, then_branch, else_branch } => {
                self.compile_ternary(condition, then_branch, else_branch)
            }
            Expr::Call { callee, arguments } => self.compile_call(callee, arguments),
            Expr::Lambda { params, body } => self.compile_lambda(params, body),
        }
    }

    fn compile_literal(&mut self, lit: &Literal) -> Result<(), RuntimeError> {
        match lit {
            Literal::Number(n) => {
                let idx = self.function.add_constant(Value::Number(*n));
                emit_u16(&mut self.function.code, OpCode::LoadConst, idx);
            }
            Literal::String(s) => {
                let idx = self.function.add_constant(Value::String(s.clone()));
                emit_u16(&mut self.function.code, OpCode::LoadConst, idx);
            }
            Literal::Bool(true) => self.emit_op(OpCode::LoadTrue),
            Literal::Bool(false) => self.emit_op(OpCode::LoadFalse),
            Literal::Nil => self.emit_op(OpCode::LoadNil),
        }
        Ok(())
    }

    fn compile_variable(&mut self, name: &str) -> Result<(), RuntimeError> {
        // Try to resolve as local
        if let Some(slot) = self.resolve_local(name) {
            emit_u8(&mut self.function.code, OpCode::LoadLocal, slot);
            return Ok(());
        }
        // Try to resolve as upvalue
        if let Some(upvalue_idx) = self.resolve_upvalue(name) {
            // Upvalue access: LoadUpvalue is LoadLocal with slot = upvalue-local-marker
            // For simplicity, we use a convention: upvalues are locals with special marking.
            // Actually, we need a dedicated LoadUpvalue opcode. Let's add it...
            // For now, store upvalues at "local slots" offset by 256.
            emit_u16(&mut self.function.code, OpCode::LoadConst, 0xFFFF); // placeholder
            // We'll handle upvalues properly when we have LoadUpvalue
            // TODO: Add LoadUpvalue opcode
            return Err(RuntimeError::RuntimeError {
                message: format!("Upvalue '{}' not yet supported", name),
                token: None,
            });
        }
        // Fall back to global
        let name_idx = self.add_string_constant(name);
        emit_u16(&mut self.function.code, OpCode::LoadGlobal, name_idx);
        Ok(())
    }

    fn compile_assign(&mut self, name: &str, op: &AssignOp, value: &Expr) -> Result<(), RuntimeError> {
        if *op == AssignOp::Equal {
            // Simple assignment
            self.compile_expr(value)?;
            if let Some(slot) = self.resolve_local(name) {
                emit_u8(&mut self.function.code, OpCode::StoreLocal, slot);
            } else {
                let name_idx = self.add_string_constant(name);
                emit_u16(&mut self.function.code, OpCode::StoreGlobal, name_idx);
            }
        } else {
            // Compound assignment
            self.compile_expr(value)?;
            let compound_op = match op {
                AssignOp::PlusEqual => CompoundOp::PlusEqual,
                AssignOp::MinusEqual => CompoundOp::MinusEqual,
                AssignOp::StarEqual => CompoundOp::StarEqual,
                AssignOp::SlashEqual => CompoundOp::SlashEqual,
                AssignOp::PercentEqual => CompoundOp::PercentEqual,
                AssignOp::Equal => unreachable!(),
            };
            if let Some(slot) = self.resolve_local(name) {
                emit_u8(&mut self.function.code, OpCode::CompoundAssignLocal, slot);
                self.function.code.push(compound_op as u8);
            } else {
                let name_idx = self.add_string_constant(name);
                emit_u16(&mut self.function.code, OpCode::CompoundAssignProp, name_idx);
                self.function.code.push(compound_op as u8);
            }
        }
        Ok(())
    }

    fn compile_get(&mut self, object: &Expr, name: &str) -> Result<(), RuntimeError> {
        self.compile_expr(object)?;
        let name_idx = self.add_string_constant(name);
        emit_u16(&mut self.function.code, OpCode::GetProperty, name_idx);
        Ok(())
    }

    fn compile_set(&mut self, object: &Expr, name: &str, op: &AssignOp, value: &Expr) -> Result<(), RuntimeError> {
        if *op == AssignOp::Equal {
            self.compile_expr(object)?;
            self.compile_expr(value)?;
            let name_idx = self.add_string_constant(name);
            emit_u16(&mut self.function.code, OpCode::SetProperty, name_idx);
        } else {
            // Compound property set: object.name op= value
            // Strategy: load object, dup, get property as current value,
            // load rhs, apply op, set property
            self.compile_expr(object)?;
            self.emit_op(OpCode::Dup);
            let name_idx = self.add_string_constant(name);
            emit_u16(&mut self.function.code, OpCode::GetProperty, name_idx); // current value
            self.compile_expr(value)?; // rhs
            let compound_op = assign_op_to_compound(op);
            self.function.code.push(compound_op as u8);
            // Now stack: object, current_val, rhs → set property
            emit_u16(&mut self.function.code, OpCode::CompoundAssignProp, name_idx);
            self.function.code.push(compound_op as u8);
        }
        Ok(())
    }

    fn compile_object(&mut self, properties: &[(String, Expr)]) -> Result<(), RuntimeError> {
        emit_op(&mut self.function.code, OpCode::NewObject);
        for (key, value_expr) in properties {
            self.emit_op(OpCode::Dup);
            self.compile_expr(value_expr)?;
            let name_idx = self.add_string_constant(key);
            emit_u16(&mut self.function.code, OpCode::SetProperty, name_idx);
            // SetProperty leaves the value on stack; we need the object, so Pop the value
            self.emit_op(OpCode::Pop);
        }
        Ok(())
    }

    fn compile_array(&mut self, elements: &[Expr]) -> Result<(), RuntimeError> {
        for e in elements {
            self.compile_expr(e)?;
        }
        emit_u16(&mut self.function.code, OpCode::NewArray, elements.len() as u16);
        Ok(())
    }

    fn compile_index_get(&mut self, array: &Expr, index: &Expr) -> Result<(), RuntimeError> {
        self.compile_expr(array)?;
        self.compile_expr(index)?;
        emit_op(&mut self.function.code, OpCode::GetIndex);
        Ok(())
    }

    fn compile_index_set(&mut self, array: &Expr, index: &Expr, op: &AssignOp, value: &Expr) -> Result<(), RuntimeError> {
        self.compile_expr(array)?;
        self.compile_expr(index)?;
        self.compile_expr(value)?;
        if *op == AssignOp::Equal {
            emit_op(&mut self.function.code, OpCode::SetIndex);
        } else {
            let compound_op = assign_op_to_compound(op);
            emit_u8(&mut self.function.code, OpCode::CompoundAssignIndex, compound_op as u8);
        }
        Ok(())
    }

    fn compile_unary(&mut self, op: &UnaryOp, right: &Expr) -> Result<(), RuntimeError> {
        self.compile_expr(right)?;
        match op {
            UnaryOp::Negate => self.emit_op(OpCode::Negate),
            UnaryOp::Not => self.emit_op(OpCode::Not),
        }
        Ok(())
    }

    fn compile_binary(&mut self, left: &Expr, op: &BinaryOp, right: &Expr) -> Result<(), RuntimeError> {
        self.compile_expr(left)?;
        self.compile_expr(right)?;
        let opcode = match op {
            BinaryOp::Add => OpCode::Add,
            BinaryOp::Sub => OpCode::Sub,
            BinaryOp::Mul => OpCode::Mul,
            BinaryOp::Div => OpCode::Div,
            BinaryOp::Mod => OpCode::Mod,
            BinaryOp::Greater => OpCode::Greater,
            BinaryOp::GreaterEqual => OpCode::GreaterEqual,
            BinaryOp::Less => OpCode::Less,
            BinaryOp::LessEqual => OpCode::LessEqual,
            BinaryOp::Equal => OpCode::Equal,
            BinaryOp::NotEqual => OpCode::NotEqual,
        };
        self.emit_op(opcode);
        Ok(())
    }

    fn compile_logical(&mut self, left: &Expr, op: &LogicalOp, right: &Expr) -> Result<(), RuntimeError> {
        self.compile_expr(left)?;
        match op {
            LogicalOp::And => {
                // Short-circuit: if left is falsy, skip right and return left
                let end_jump = self.emit_jump(OpCode::PopJumpIfFalse);
                self.emit_op(OpCode::Pop); // discard left (truthy)
                self.compile_expr(right)?;
                self.patch_jump(end_jump);
            }
            LogicalOp::Or => {
                // Short-circuit: if left is truthy, skip right and return left
                self.emit_op(OpCode::Dup);
                let end_jump = self.emit_jump(OpCode::JumpIfTrue);
                self.emit_op(OpCode::Pop); // discard left (falsy)
                self.compile_expr(right)?;
                self.patch_jump(end_jump);
            }
        }
        Ok(())
    }

    fn compile_ternary(&mut self, condition: &Expr, then_branch: &Expr, else_branch: &Expr) -> Result<(), RuntimeError> {
        self.compile_expr(condition)?;
        let else_jump = self.emit_jump(OpCode::JumpIfFalse);
        self.compile_expr(then_branch)?;
        let end_jump = self.emit_jump(OpCode::Jump);
        self.patch_jump(else_jump);
        self.compile_expr(else_branch)?;
        self.patch_jump(end_jump);
        Ok(())
    }

    fn compile_call(&mut self, callee: &Expr, arguments: &[Expr]) -> Result<(), RuntimeError> {
        self.compile_expr(callee)?;
        for arg in arguments {
            self.compile_expr(arg)?;
        }
        emit_u8(&mut self.function.code, OpCode::Call, arguments.len() as u8);
        Ok(())
    }

    fn compile_lambda(&mut self, params: &[String], body: &[Stmt]) -> Result<(), RuntimeError> {
        let proto = self.compile_nested_function(None, params, body)?;
        let proto_idx = self.add_function_proto_constant(proto);
        self.emit_closure(proto_idx);
        Ok(())
    }

    /// Compile a nested function (lambda or function declaration) and return its proto.
    fn compile_nested_function(&mut self, name: Option<String>, params: &[String], body: &[Stmt]) -> Result<FunctionProto, RuntimeError> {
        let mut child = Compiler::new(name);
        child.enclosing_idx = Some(0); // placeholder — we handle upvalues differently

        // Add params as locals
        for param in params {
            let slot = child.locals.len() as u8;
            child.locals.push(Local {
                name: param.clone(),
                depth: 1, // params are at scope depth 1 (function body)
                is_captured: false,
                is_const: false,
            });
            // Params are already on stack from Call; StoreLocal just peeks
            emit_u8(&mut child.function.code, OpCode::StoreLocal, slot);
        }
        child.function.arity = params.len() as u8;

        // Compile the body
        child.scope_depth = 1;
        for stmt in body {
            child.compile_stmt(stmt)?;
        }
        // Implicit return nil
        child.emit_op(OpCode::LoadNil);
        child.emit_op(OpCode::Return);

        // Resolve upvalues: for each variable reference in the child that wasn't
        // resolved locally, check if it exists in the parent's locals/upvalues.
        // (We handle this lazily in compile_variable for now — if not local, try upvalue.)

        child.function.upvalue_count = child.upvalues.len() as u8;
        Ok(child.function)
    }

    // ========================================================================
    // Scope management
    // ========================================================================

    fn begin_scope(&mut self) {
        self.scope_depth += 1;
    }

    fn end_scope(&mut self) {
        self.scope_depth -= 1;
        // Pop locals that are going out of scope
        let mut pop_count = 0u8;
        while let Some(local) = self.locals.last() {
            if local.depth > self.scope_depth {
                if local.is_captured {
                    // Close upvalue instead of pop
                    // (For now, just pop — upvalue closing handled in VM)
                }
                self.locals.pop();
                pop_count += 1;
            } else {
                break;
            }
        }
        for _ in 0..pop_count {
            self.emit_op(OpCode::Pop);
        }
    }

    // ========================================================================
    // Variable resolution
    // ========================================================================

    /// Find a local variable by name, returning its slot index.
    fn resolve_local(&self, name: &str) -> Option<u8> {
        for (i, local) in self.locals.iter().enumerate().rev() {
            if local.name == name && local.depth > 0 {
                return Some(i as u8);
            }
        }
        None
    }

    /// Try to resolve a variable as an upvalue from enclosing functions.
    fn resolve_upvalue(&mut self, name: &str) -> Option<u8> {
        // For now, we don't have enclosing compiler access.
        // Upvalues will be fully implemented in a follow-up.
        None
    }

    // ========================================================================
    // Bytecode emission helpers
    // ========================================================================

    fn emit_op(&mut self, op: OpCode) {
        emit_op(&mut self.function.code, op);
    }

    fn emit_jump(&mut self, op: OpCode) -> usize {
        let loc = self.function.code.len();
        emit_i16(&mut self.function.code, op, 0x7FFF); // placeholder
        loc
    }

    fn emit_jump_placeholder(&mut self) -> usize {
        let loc = self.function.code.len();
        emit_i16(&mut self.function.code, OpCode::Jump, 0x7FFF);
        loc
    }

    fn emit_loop_jump(&mut self, target: usize) {
        let offset = target as isize - self.function.code.len() as isize;
        emit_i16(&mut self.function.code, OpCode::Jump, offset as i16);
    }

    fn emit_closure(&mut self, proto_idx: u16) {
        let code = &mut self.function.code;
        code.push(OpCode::Closure as u8);
        code.push((proto_idx & 0xFF) as u8);
        code.push(((proto_idx >> 8) & 0xFF) as u8);
        // No upvalues yet
        code.push(0u8); // upvalue count = 0 for now
    }

    fn patch_jump(&mut self, jump_loc: usize) {
        let offset = (self.function.code.len() - jump_loc) as i16;
        let code = &mut self.function.code;
        code[jump_loc + 1] = (offset & 0xFF) as u8;
        code[jump_loc + 2] = ((offset >> 8) & 0xFF) as u8;
    }

    fn patch_jump_to(&mut self, jump_loc: usize, target: usize) {
        let offset = target as isize - jump_loc as isize;
        let code = &mut self.function.code;
        code[jump_loc + 1] = (offset & 0xFF) as u8;
        code[jump_loc + 2] = ((offset >> 8) & 0xFF) as u8;
    }

    // ========================================================================
    // Constant pool helpers
    // ========================================================================

    fn add_string_constant(&mut self, s: &str) -> u16 {
        self.function.add_constant(Value::String(s.to_string()))
    }

    fn add_function_proto_constant(&mut self, proto: FunctionProto) -> u16 {
        // Store compiled proto as a special marker.
        // We use Value::Nil as placeholder since FunctionProto isn't a Value.
        // The VM will look up the proto from a separate table.
        // For now, store the proto's index in a side table.
        // Actually, let's store it as a string tag that the VM can recognize.
        // We'll use a dedicated proto storage: just append to a Vec.
        // But FunctionProto isn't a Value... let's store it inline.
        // HACK: store proto in constants with a special Object wrapping
        let idx = self.function.constants.len() as u16;
        // Use a Value::Object with a special marker
        // The VM will need to handle this
        let mut map = std::collections::HashMap::new();
        map.insert("__proto__".to_string(), Value::String(format!("proto_{}", idx)));
        self.function.constants.push(Value::Object(Rc::new(RefCell::new(map))));
        idx
    }
}

// ============================================================================
// Helpers
// ============================================================================

fn assign_op_to_compound(op: &AssignOp) -> CompoundOp {
    match op {
        AssignOp::PlusEqual => CompoundOp::PlusEqual,
        AssignOp::MinusEqual => CompoundOp::MinusEqual,
        AssignOp::StarEqual => CompoundOp::StarEqual,
        AssignOp::SlashEqual => CompoundOp::SlashEqual,
        AssignOp::PercentEqual => CompoundOp::PercentEqual,
        AssignOp::Equal => unreachable!(),
    }
}

fn emit_i16_placeholder(code: &mut Vec<u8>, op: OpCode) -> usize {
    let loc = code.len();
    emit_i16(code, op, 0x7FFF);
    loc
}