we / crates / js / src / compiler.rs

//! AST → register-based bytecode compiler.
//!
//! Walks the AST produced by the parser and emits bytecode instructions.
//! Uses a simple greedy register allocator: each new temporary gets the next
//! available register, and registers are freed when no longer needed.

use crate::ast::*;
use crate::bytecode::*;
use crate::JsError;

/// Compiler state for a single function scope.
struct FunctionCompiler {
    builder: BytecodeBuilder,
    /// Maps local variable names to their register slots.
    locals: Vec<Local>,
    /// Next free register index.
    next_reg: u8,
    /// Stack of loop contexts for break/continue.
    loop_stack: Vec<LoopCtx>,
}

#[derive(Debug, Clone)]
struct Local {
    name: String,
    reg: Reg,
}

struct LoopCtx {
    /// Label, if this is a labeled loop.
    label: Option<String>,
    /// Patch positions for break jumps.
    break_patches: Vec<usize>,
    /// Patch positions for continue jumps (patched after body compilation).
    continue_patches: Vec<usize>,
}

impl FunctionCompiler {
    fn new(name: String, param_count: u8) -> Self {
        Self {
            builder: BytecodeBuilder::new(name, param_count),
            locals: Vec::new(),
            next_reg: 0,
            loop_stack: Vec::new(),
        }
    }

    /// Allocate a register, updating the high-water mark.
    fn alloc_reg(&mut self) -> Reg {
        let r = self.next_reg;
        self.next_reg = self.next_reg.checked_add(1).expect("register overflow");
        if self.next_reg > self.builder.func.register_count {
            self.builder.func.register_count = self.next_reg;
        }
        r
    }

    /// Free the last allocated register (must be called in reverse order).
    fn free_reg(&mut self, r: Reg) {
        debug_assert_eq!(
            r,
            self.next_reg - 1,
            "registers must be freed in reverse order"
        );
        self.next_reg -= 1;
    }

    /// Look up a local variable by name.
    fn find_local(&self, name: &str) -> Option<Reg> {
        self.locals
            .iter()
            .rev()
            .find(|l| l.name == name)
            .map(|l| l.reg)
    }

    /// Define a local variable.
    fn define_local(&mut self, name: &str) -> Reg {
        let reg = self.alloc_reg();
        self.locals.push(Local {
            name: name.to_string(),
            reg,
        });
        reg
    }
}

/// Compile a parsed program into a top-level bytecode function.
pub fn compile(program: &Program) -> Result<Function, JsError> {
    let mut fc = FunctionCompiler::new("<main>".into(), 0);

    // Reserve r0 for the implicit return value.
    let result_reg = fc.alloc_reg();
    fc.builder.emit_reg(Op::LoadUndefined, result_reg);

    compile_stmts(&mut fc, &program.body, result_reg)?;

    fc.builder.emit_reg(Op::Return, result_reg);
    Ok(fc.builder.finish())
}

fn compile_stmts(
    fc: &mut FunctionCompiler,
    stmts: &[Stmt],
    result_reg: Reg,
) -> Result<(), JsError> {
    for stmt in stmts {
        compile_stmt(fc, stmt, result_reg)?;
    }
    Ok(())
}

fn compile_stmt(fc: &mut FunctionCompiler, stmt: &Stmt, result_reg: Reg) -> Result<(), JsError> {
    match &stmt.kind {
        StmtKind::Expr(expr) => {
            // Expression statement: compile and store result in result_reg.
            compile_expr(fc, expr, result_reg)?;
        }

        StmtKind::Block(stmts) => {
            let saved_locals = fc.locals.len();
            let saved_next = fc.next_reg;
            compile_stmts(fc, stmts, result_reg)?;
            // Pop locals from this block.
            fc.locals.truncate(saved_locals);
            fc.next_reg = saved_next;
        }

        StmtKind::VarDecl {
            kind: _,
            declarators,
        } => {
            for decl in declarators {
                compile_var_declarator(fc, decl)?;
            }
        }

        StmtKind::FunctionDecl(func_def) => {
            compile_function_decl(fc, func_def)?;
        }

        StmtKind::If {
            test,
            consequent,
            alternate,
        } => {
            compile_if(fc, test, consequent, alternate.as_deref(), result_reg)?;
        }

        StmtKind::While { test, body } => {
            compile_while(fc, test, body, None, result_reg)?;
        }

        StmtKind::DoWhile { body, test } => {
            compile_do_while(fc, body, test, None, result_reg)?;
        }

        StmtKind::For {
            init,
            test,
            update,
            body,
        } => {
            compile_for(
                fc,
                init.as_ref(),
                test.as_ref(),
                update.as_ref(),
                body,
                None,
                result_reg,
            )?;
        }

        StmtKind::ForIn { left, right, body } => {
            let saved_locals = fc.locals.len();
            let saved_next = fc.next_reg;

            // Evaluate the RHS object.
            let obj_r = fc.alloc_reg();
            compile_expr(fc, right, obj_r)?;

            // ForInInit: collect enumerable keys into an array.
            let keys_r = fc.alloc_reg();
            fc.builder.emit_reg_reg(Op::ForInInit, keys_r, obj_r);
            // obj_r is no longer needed but we don't free it (LIFO constraint).

            // Index register (starts at 0).
            let idx_r = fc.alloc_reg();
            fc.builder.emit_load_int8(idx_r, 0);

            // Key and done registers (reused each iteration).
            let key_r = fc.alloc_reg();
            let done_r = fc.alloc_reg();

            let loop_start = fc.builder.offset();

            // ForInNext: get next key or done flag.
            fc.builder
                .emit_reg4(Op::ForInNext, key_r, done_r, keys_r, idx_r);

            // Exit loop if done.
            let exit_patch = fc.builder.emit_cond_jump(Op::JumpIfTrue, done_r);

            // Bind the loop variable.
            match left {
                ForInOfLeft::VarDecl { kind: _, pattern } => {
                    if let PatternKind::Identifier(name) = &pattern.kind {
                        let var_r = fc.define_local(name);
                        fc.builder.emit_reg_reg(Op::Move, var_r, key_r);
                    }
                }
                ForInOfLeft::Pattern(pattern) => {
                    if let PatternKind::Identifier(name) = &pattern.kind {
                        if let Some(local) = fc.find_local(name) {
                            fc.builder.emit_reg_reg(Op::Move, local, key_r);
                        } else {
                            let ni = fc.builder.add_name(name);
                            fc.builder.emit_store_global(ni, key_r);
                        }
                    }
                }
            }

            // Compile the body.
            fc.loop_stack.push(LoopCtx {
                label: None,
                break_patches: Vec::new(),
                continue_patches: Vec::new(),
            });
            compile_stmt(fc, body, result_reg)?;

            // Increment index: idx = idx + 1.
            // Use a temp register for the constant 1. Since we allocate it
            // after the loop body, we can't free it with LIFO either — the
            // saved_next restoration handles cleanup.
            let one_r = fc.alloc_reg();
            fc.builder.emit_load_int8(one_r, 1);
            fc.builder.emit_reg3(Op::Add, idx_r, idx_r, one_r);

            // Jump back to loop start.
            fc.builder.emit_jump_to(loop_start);

            // Patch exit and continue jumps.
            fc.builder.patch_jump(exit_patch);
            let ctx = fc.loop_stack.pop().unwrap();
            for patch in ctx.break_patches {
                fc.builder.patch_jump(patch);
            }
            for patch in ctx.continue_patches {
                fc.builder.patch_jump_to(patch, loop_start);
            }

            // Restore locals/regs — frees all temporaries at once.
            fc.locals.truncate(saved_locals);
            fc.next_reg = saved_next;
        }

        StmtKind::ForOf {
            left,
            right,
            body,
            is_await: _,
        } => {
            let _ = left;
            let tmp = fc.alloc_reg();
            compile_expr(fc, right, tmp)?;
            fc.free_reg(tmp);
            compile_stmt(fc, body, result_reg)?;
        }

        StmtKind::Return(expr) => {
            let ret_reg = fc.alloc_reg();
            if let Some(e) = expr {
                compile_expr(fc, e, ret_reg)?;
            } else {
                fc.builder.emit_reg(Op::LoadUndefined, ret_reg);
            }
            fc.builder.emit_reg(Op::Return, ret_reg);
            fc.free_reg(ret_reg);
        }

        StmtKind::Throw(expr) => {
            let tmp = fc.alloc_reg();
            compile_expr(fc, expr, tmp)?;
            fc.builder.emit_reg(Op::Throw, tmp);
            fc.free_reg(tmp);
        }

        StmtKind::Break(label) => {
            // Find the matching loop context.
            let idx = find_loop_ctx(&fc.loop_stack, label.as_deref())
                .ok_or_else(|| JsError::SyntaxError("break outside of loop".into()))?;
            let patch = fc.builder.emit_jump(Op::Jump);
            fc.loop_stack[idx].break_patches.push(patch);
        }

        StmtKind::Continue(label) => {
            let idx = find_loop_ctx(&fc.loop_stack, label.as_deref())
                .ok_or_else(|| JsError::SyntaxError("continue outside of loop".into()))?;
            let patch = fc.builder.emit_jump(Op::Jump);
            fc.loop_stack[idx].continue_patches.push(patch);
        }

        StmtKind::Labeled { label, body } => {
            // If body is a loop, propagate the label.
            match &body.kind {
                StmtKind::While { test, body: inner } => {
                    compile_while(fc, test, inner, Some(label.clone()), result_reg)?;
                }
                StmtKind::DoWhile { body: inner, test } => {
                    compile_do_while(fc, inner, test, Some(label.clone()), result_reg)?;
                }
                StmtKind::For {
                    init,
                    test,
                    update,
                    body: inner,
                } => {
                    compile_for(
                        fc,
                        init.as_ref(),
                        test.as_ref(),
                        update.as_ref(),
                        inner,
                        Some(label.clone()),
                        result_reg,
                    )?;
                }
                _ => {
                    compile_stmt(fc, body, result_reg)?;
                }
            }
        }

        StmtKind::Switch {
            discriminant,
            cases,
        } => {
            compile_switch(fc, discriminant, cases, result_reg)?;
        }

        StmtKind::Try {
            block,
            handler,
            finalizer,
        } => {
            // Simplified: compile blocks sequentially.
            // Real try/catch needs exception table support from the VM.
            compile_stmts(fc, block, result_reg)?;
            if let Some(catch) = handler {
                compile_stmts(fc, &catch.body, result_reg)?;
            }
            if let Some(fin) = finalizer {
                compile_stmts(fc, fin, result_reg)?;
            }
        }

        StmtKind::Empty | StmtKind::Debugger => {
            // No-op.
        }

        StmtKind::With { object, body } => {
            // Compile `with` as: evaluate object (discard), then run body.
            // Proper `with` scope requires VM support.
            let tmp = fc.alloc_reg();
            compile_expr(fc, object, tmp)?;
            fc.free_reg(tmp);
            compile_stmt(fc, body, result_reg)?;
        }

        StmtKind::Import { .. } => {
            // Module imports are resolved before execution; no bytecode needed.
        }

        StmtKind::Export(export) => {
            compile_export(fc, export, result_reg)?;
        }

        StmtKind::ClassDecl(class_def) => {
            compile_class_decl(fc, class_def)?;
        }
    }
    Ok(())
}

// ── Variable declarations ───────────────────────────────────

fn compile_var_declarator(fc: &mut FunctionCompiler, decl: &VarDeclarator) -> Result<(), JsError> {
    match &decl.pattern.kind {
        PatternKind::Identifier(name) => {
            let reg = fc.define_local(name);
            if let Some(init) = &decl.init {
                compile_expr(fc, init, reg)?;
            } else {
                fc.builder.emit_reg(Op::LoadUndefined, reg);
            }
        }
        _ => {
            // Destructuring: evaluate init, then bind patterns.
            let tmp = fc.alloc_reg();
            if let Some(init) = &decl.init {
                compile_expr(fc, init, tmp)?;
            } else {
                fc.builder.emit_reg(Op::LoadUndefined, tmp);
            }
            compile_destructuring_pattern(fc, &decl.pattern, tmp)?;
            fc.free_reg(tmp);
        }
    }
    Ok(())
}

fn compile_destructuring_pattern(
    fc: &mut FunctionCompiler,
    pattern: &Pattern,
    src: Reg,
) -> Result<(), JsError> {
    match &pattern.kind {
        PatternKind::Identifier(name) => {
            let reg = fc.define_local(name);
            fc.builder.emit_reg_reg(Op::Move, reg, src);
        }
        PatternKind::Object {
            properties,
            rest: _,
        } => {
            for prop in properties {
                let key_name = match &prop.key {
                    PropertyKey::Identifier(s) | PropertyKey::String(s) => s.clone(),
                    _ => {
                        return Err(JsError::SyntaxError(
                            "computed destructuring keys not yet supported".into(),
                        ));
                    }
                };
                let val_reg = fc.alloc_reg();
                let name_idx = fc.builder.add_name(&key_name);
                fc.builder.emit_get_prop_name(val_reg, src, name_idx);
                compile_destructuring_pattern(fc, &prop.value, val_reg)?;
                fc.free_reg(val_reg);
            }
        }
        PatternKind::Array { elements, rest: _ } => {
            for (i, elem) in elements.iter().enumerate() {
                if let Some(pat) = elem {
                    let idx_reg = fc.alloc_reg();
                    if i <= 127 {
                        fc.builder.emit_load_int8(idx_reg, i as i8);
                    } else {
                        let ci = fc.builder.add_constant(Constant::Number(i as f64));
                        fc.builder.emit_reg_u16(Op::LoadConst, idx_reg, ci);
                    }
                    let val_reg = fc.alloc_reg();
                    fc.builder.emit_reg3(Op::GetProperty, val_reg, src, idx_reg);
                    compile_destructuring_pattern(fc, pat, val_reg)?;
                    fc.free_reg(val_reg);
                    fc.free_reg(idx_reg);
                }
            }
        }
        PatternKind::Assign { left, right } => {
            // Default value: if src is undefined, use default.
            let val_reg = fc.alloc_reg();
            fc.builder.emit_reg_reg(Op::Move, val_reg, src);
            // Check if undefined, if so use default.
            let check_reg = fc.alloc_reg();
            let undef_reg = fc.alloc_reg();
            fc.builder.emit_reg(Op::LoadUndefined, undef_reg);
            fc.builder
                .emit_reg3(Op::StrictEq, check_reg, val_reg, undef_reg);
            fc.free_reg(undef_reg);
            let patch = fc.builder.emit_cond_jump(Op::JumpIfFalse, check_reg);
            fc.free_reg(check_reg);
            // Is undefined → evaluate default.
            compile_expr(fc, right, val_reg)?;
            fc.builder.patch_jump(patch);
            compile_destructuring_pattern(fc, left, val_reg)?;
            fc.free_reg(val_reg);
        }
    }
    Ok(())
}

// ── Function declarations ───────────────────────────────────

fn compile_function_decl(fc: &mut FunctionCompiler, func_def: &FunctionDef) -> Result<(), JsError> {
    let name = func_def.id.clone().unwrap_or_default();
    let inner = compile_function_body(func_def)?;
    let func_idx = fc.builder.add_function(inner);

    let reg = fc.define_local(&name);
    fc.builder.emit_reg_u16(Op::CreateClosure, reg, func_idx);

    // Also store as global so inner/recursive calls via LoadGlobal can find it.
    if !name.is_empty() {
        let name_idx = fc.builder.add_name(&name);
        fc.builder.emit_store_global(name_idx, reg);
    }
    Ok(())
}

fn compile_function_body(func_def: &FunctionDef) -> Result<Function, JsError> {
    let name = func_def.id.clone().unwrap_or_default();
    let param_count = func_def.params.len().min(255) as u8;
    let mut inner = FunctionCompiler::new(name, param_count);

    // Allocate registers for parameters.
    for p in &func_def.params {
        if let PatternKind::Identifier(name) = &p.kind {
            inner.define_local(name);
        } else {
            // Destructuring param: allocate a register for the whole param,
            // then destructure from it.
            let _ = inner.alloc_reg();
        }
    }

    // Result register for the function body.
    let result_reg = inner.alloc_reg();
    inner.builder.emit_reg(Op::LoadUndefined, result_reg);

    compile_stmts(&mut inner, &func_def.body, result_reg)?;

    // Implicit return undefined.
    inner.builder.emit_reg(Op::Return, result_reg);
    Ok(inner.builder.finish())
}

// ── Class declarations ──────────────────────────────────────

fn compile_class_decl(fc: &mut FunctionCompiler, class_def: &ClassDef) -> Result<(), JsError> {
    let name = class_def.id.clone().unwrap_or_default();
    let reg = fc.define_local(&name);

    // Find constructor or create empty one.
    let ctor = class_def.body.iter().find(|m| {
        matches!(
            &m.kind,
            ClassMemberKind::Method {
                kind: MethodKind::Constructor,
                ..
            }
        )
    });

    if let Some(member) = ctor {
        if let ClassMemberKind::Method { value, .. } = &member.kind {
            let inner = compile_function_body(value)?;
            let func_idx = fc.builder.add_function(inner);
            fc.builder.emit_reg_u16(Op::CreateClosure, reg, func_idx);
        }
    } else {
        // No constructor: create a minimal function that returns undefined.
        let mut empty = BytecodeBuilder::new(name.clone(), 0);
        let r = 0u8;
        empty.func.register_count = 1;
        empty.emit_reg(Op::LoadUndefined, r);
        empty.emit_reg(Op::Return, r);
        let func_idx = fc.builder.add_function(empty.finish());
        fc.builder.emit_reg_u16(Op::CreateClosure, reg, func_idx);
    }

    // Compile methods: set them as properties on the constructor's prototype.
    // This is simplified — real class compilation needs prototype chain setup.
    for member in &class_def.body {
        match &member.kind {
            ClassMemberKind::Method {
                key,
                value,
                kind,
                is_static: _,
                computed: _,
            } => {
                if matches!(kind, MethodKind::Constructor) {
                    continue;
                }
                let method_name = match key {
                    PropertyKey::Identifier(s) | PropertyKey::String(s) => s.clone(),
                    _ => continue,
                };
                let inner = compile_function_body(value)?;
                let func_idx = fc.builder.add_function(inner);
                let method_reg = fc.alloc_reg();
                fc.builder
                    .emit_reg_u16(Op::CreateClosure, method_reg, func_idx);
                let name_idx = fc.builder.add_name(&method_name);
                fc.builder.emit_set_prop_name(reg, name_idx, method_reg);
                fc.free_reg(method_reg);
            }
            ClassMemberKind::Property { .. } => {
                // Class fields are set in constructor; skip here.
            }
        }
    }

    Ok(())
}

// ── Export ───────────────────────────────────────────────────

fn compile_export(
    fc: &mut FunctionCompiler,
    export: &ExportDecl,

    result_reg: Reg,
) -> Result<(), JsError> {
    match export {
        ExportDecl::Declaration(stmt) => {
            compile_stmt(fc, stmt, result_reg)?;
        }
        ExportDecl::Default(expr) => {
            compile_expr(fc, expr, result_reg)?;
        }
        ExportDecl::Named { .. } | ExportDecl::AllFrom(_) => {
            // Named re-exports are module-level; no bytecode needed.
        }
    }
    Ok(())
}

// ── Control flow ────────────────────────────────────────────

fn compile_if(
    fc: &mut FunctionCompiler,
    test: &Expr,
    consequent: &Stmt,
    alternate: Option<&Stmt>,

    result_reg: Reg,
) -> Result<(), JsError> {
    let cond = fc.alloc_reg();
    compile_expr(fc, test, cond)?;
    let else_patch = fc.builder.emit_cond_jump(Op::JumpIfFalse, cond);
    fc.free_reg(cond);

    compile_stmt(fc, consequent, result_reg)?;

    if let Some(alt) = alternate {
        let end_patch = fc.builder.emit_jump(Op::Jump);
        fc.builder.patch_jump(else_patch);
        compile_stmt(fc, alt, result_reg)?;
        fc.builder.patch_jump(end_patch);
    } else {
        fc.builder.patch_jump(else_patch);
    }
    Ok(())
}

fn compile_while(
    fc: &mut FunctionCompiler,
    test: &Expr,
    body: &Stmt,
    label: Option<String>,

    result_reg: Reg,
) -> Result<(), JsError> {
    let loop_start = fc.builder.offset();

    let cond = fc.alloc_reg();
    compile_expr(fc, test, cond)?;
    let exit_patch = fc.builder.emit_cond_jump(Op::JumpIfFalse, cond);
    fc.free_reg(cond);

    fc.loop_stack.push(LoopCtx {
        label,
        break_patches: Vec::new(),
        continue_patches: Vec::new(),
    });

    compile_stmt(fc, body, result_reg)?;
    fc.builder.emit_jump_to(loop_start);
    fc.builder.patch_jump(exit_patch);

    let ctx = fc.loop_stack.pop().unwrap();
    for patch in ctx.break_patches {
        fc.builder.patch_jump(patch);
    }
    // In a while loop, continue jumps back to the condition check (loop_start).
    for patch in ctx.continue_patches {
        fc.builder.patch_jump_to(patch, loop_start);
    }
    Ok(())
}

fn compile_do_while(
    fc: &mut FunctionCompiler,
    body: &Stmt,
    test: &Expr,
    label: Option<String>,

    result_reg: Reg,
) -> Result<(), JsError> {
    let loop_start = fc.builder.offset();

    fc.loop_stack.push(LoopCtx {
        label,
        break_patches: Vec::new(),
        continue_patches: Vec::new(),
    });

    compile_stmt(fc, body, result_reg)?;

    // continue in do-while should jump here (the condition check).
    let cond_start = fc.builder.offset();

    let cond = fc.alloc_reg();
    compile_expr(fc, test, cond)?;
    fc.builder
        .emit_cond_jump_to(Op::JumpIfTrue, cond, loop_start);
    fc.free_reg(cond);

    let ctx = fc.loop_stack.pop().unwrap();
    for patch in ctx.break_patches {
        fc.builder.patch_jump(patch);
    }
    for patch in ctx.continue_patches {
        fc.builder.patch_jump_to(patch, cond_start);
    }
    Ok(())
}

fn compile_for(
    fc: &mut FunctionCompiler,
    init: Option<&ForInit>,
    test: Option<&Expr>,
    update: Option<&Expr>,
    body: &Stmt,
    label: Option<String>,

    result_reg: Reg,
) -> Result<(), JsError> {
    let saved_locals = fc.locals.len();
    let saved_next = fc.next_reg;

    // Init.
    if let Some(init) = init {
        match init {
            ForInit::VarDecl {
                kind: _,
                declarators,
            } => {
                for decl in declarators {
                    compile_var_declarator(fc, decl)?;
                }
            }
            ForInit::Expr(expr) => {
                let tmp = fc.alloc_reg();
                compile_expr(fc, expr, tmp)?;
                fc.free_reg(tmp);
            }
        }
    }

    let loop_start = fc.builder.offset();

    // Test.
    let exit_patch = if let Some(test) = test {
        let cond = fc.alloc_reg();
        compile_expr(fc, test, cond)?;
        let patch = fc.builder.emit_cond_jump(Op::JumpIfFalse, cond);
        fc.free_reg(cond);
        Some(patch)
    } else {
        None
    };

    fc.loop_stack.push(LoopCtx {
        label,
        break_patches: Vec::new(),
        continue_patches: Vec::new(),
    });

    compile_stmt(fc, body, result_reg)?;

    // continue in a for-loop should jump here (the update expression).
    let continue_target = fc.builder.offset();

    // Update.
    if let Some(update) = update {
        let tmp = fc.alloc_reg();
        compile_expr(fc, update, tmp)?;
        fc.free_reg(tmp);
    }

    fc.builder.emit_jump_to(loop_start);

    if let Some(patch) = exit_patch {
        fc.builder.patch_jump(patch);
    }

    let ctx = fc.loop_stack.pop().unwrap();
    for patch in ctx.break_patches {
        fc.builder.patch_jump(patch);
    }
    for patch in ctx.continue_patches {
        fc.builder.patch_jump_to(patch, continue_target);
    }

    fc.locals.truncate(saved_locals);
    fc.next_reg = saved_next;
    Ok(())
}

fn compile_switch(
    fc: &mut FunctionCompiler,
    discriminant: &Expr,
    cases: &[SwitchCase],

    result_reg: Reg,
) -> Result<(), JsError> {
    let disc_reg = fc.alloc_reg();
    compile_expr(fc, discriminant, disc_reg)?;

    // Use a loop context for break statements.
    fc.loop_stack.push(LoopCtx {
        label: None,
        break_patches: Vec::new(),
        continue_patches: Vec::new(),
    });

    // Phase 1: emit comparison jumps for each non-default case.
    // Store (case_index, patch_position) for each case with a test.
    let mut case_jump_patches: Vec<(usize, usize)> = Vec::new();
    let mut default_index: Option<usize> = None;

    for (i, case) in cases.iter().enumerate() {
        if let Some(test) = &case.test {
            let test_reg = fc.alloc_reg();
            compile_expr(fc, test, test_reg)?;
            let cmp_reg = fc.alloc_reg();
            fc.builder
                .emit_reg3(Op::StrictEq, cmp_reg, disc_reg, test_reg);
            let patch = fc.builder.emit_cond_jump(Op::JumpIfTrue, cmp_reg);
            fc.free_reg(cmp_reg);
            fc.free_reg(test_reg);
            case_jump_patches.push((i, patch));
        } else {
            default_index = Some(i);
        }
    }

    // After all comparisons: jump to default body or end.
    let fallthrough_patch = fc.builder.emit_jump(Op::Jump);

    // Phase 2: emit case bodies in order (fall-through semantics).
    let mut body_offsets: Vec<(usize, usize)> = Vec::new();
    for (i, case) in cases.iter().enumerate() {
        body_offsets.push((i, fc.builder.offset()));
        compile_stmts(fc, &case.consequent, result_reg)?;
    }

    let end_offset = fc.builder.offset();

    // Patch case test jumps to their respective body offsets.
    for (case_idx, patch) in &case_jump_patches {
        let body_offset = body_offsets
            .iter()
            .find(|(i, _)| i == case_idx)
            .map(|(_, off)| *off)
            .unwrap();
        fc.builder.patch_jump_to(*patch, body_offset);
    }

    // Patch fallthrough: jump to default body if present, otherwise to end.
    if let Some(def_idx) = default_index {
        let default_offset = body_offsets
            .iter()
            .find(|(i, _)| *i == def_idx)
            .map(|(_, off)| *off)
            .unwrap();
        fc.builder.patch_jump_to(fallthrough_patch, default_offset);
    } else {
        fc.builder.patch_jump_to(fallthrough_patch, end_offset);
    }

    fc.free_reg(disc_reg);

    let ctx = fc.loop_stack.pop().unwrap();
    for patch in ctx.break_patches {
        fc.builder.patch_jump(patch);
    }
    Ok(())
}

fn find_loop_ctx(stack: &[LoopCtx], label: Option<&str>) -> Option<usize> {
    if let Some(label) = label {
        stack
            .iter()
            .rposition(|ctx| ctx.label.as_deref() == Some(label))
    } else {
        if stack.is_empty() {
            None
        } else {
            Some(stack.len() - 1)
        }
    }
}

// ── Expressions ─────────────────────────────────────────────

fn compile_expr(fc: &mut FunctionCompiler, expr: &Expr, dst: Reg) -> Result<(), JsError> {
    match &expr.kind {
        ExprKind::Number(n) => {
            // Optimize small integers.
            let int_val = *n as i64;
            if int_val as f64 == *n && (-128..=127).contains(&int_val) {
                fc.builder.emit_load_int8(dst, int_val as i8);
            } else {
                let ci = fc.builder.add_constant(Constant::Number(*n));
                fc.builder.emit_reg_u16(Op::LoadConst, dst, ci);
            }
        }

        ExprKind::String(s) => {
            let ci = fc.builder.add_constant(Constant::String(s.clone()));
            fc.builder.emit_reg_u16(Op::LoadConst, dst, ci);
        }

        ExprKind::Bool(true) => {
            fc.builder.emit_reg(Op::LoadTrue, dst);
        }

        ExprKind::Bool(false) => {
            fc.builder.emit_reg(Op::LoadFalse, dst);
        }

        ExprKind::Null => {
            fc.builder.emit_reg(Op::LoadNull, dst);
        }

        ExprKind::Identifier(name) => {
            if let Some(local_reg) = fc.find_local(name) {
                if local_reg != dst {
                    fc.builder.emit_reg_reg(Op::Move, dst, local_reg);
                }
            } else {
                // Global lookup.
                let ni = fc.builder.add_name(name);
                fc.builder.emit_load_global(dst, ni);
            }
        }

        ExprKind::This => {
            // `this` is loaded as a global named "this" (the VM binds it).
            let ni = fc.builder.add_name("this");
            fc.builder.emit_load_global(dst, ni);
        }

        ExprKind::Binary { op, left, right } => {
            let lhs = fc.alloc_reg();
            compile_expr(fc, left, lhs)?;
            let rhs = fc.alloc_reg();
            compile_expr(fc, right, rhs)?;
            let bytecode_op = binary_op_to_opcode(*op);
            fc.builder.emit_reg3(bytecode_op, dst, lhs, rhs);
            fc.free_reg(rhs);
            fc.free_reg(lhs);
        }

        ExprKind::Unary { op, argument } => {
            if *op == UnaryOp::Delete {
                // Handle delete specially: need object + key form for member expressions.
                match &argument.kind {
                    ExprKind::Member {
                        object,
                        property,
                        computed,
                    } => {
                        let obj_r = fc.alloc_reg();
                        compile_expr(fc, object, obj_r)?;
                        let key_r = fc.alloc_reg();
                        if *computed {
                            compile_expr(fc, property, key_r)?;
                        } else if let ExprKind::Identifier(name) = &property.kind {
                            let ci = fc.builder.add_constant(Constant::String(name.clone()));
                            fc.builder.emit_reg_u16(Op::LoadConst, key_r, ci);
                        } else {
                            compile_expr(fc, property, key_r)?;
                        }
                        fc.builder.emit_reg3(Op::Delete, dst, obj_r, key_r);
                        fc.free_reg(key_r);
                        fc.free_reg(obj_r);
                    }
                    _ => {
                        // `delete x` on a simple identifier: always true in non-strict mode.
                        fc.builder.emit_reg(Op::LoadTrue, dst);
                    }
                }
            } else {
                let src = fc.alloc_reg();
                compile_expr(fc, argument, src)?;
                match op {
                    UnaryOp::Minus => fc.builder.emit_reg_reg(Op::Neg, dst, src),
                    UnaryOp::Plus => {
                        fc.builder.emit_reg_reg(Op::Move, dst, src);
                    }
                    UnaryOp::Not => fc.builder.emit_reg_reg(Op::LogicalNot, dst, src),
                    UnaryOp::BitwiseNot => fc.builder.emit_reg_reg(Op::BitNot, dst, src),
                    UnaryOp::Typeof => fc.builder.emit_reg_reg(Op::TypeOf, dst, src),
                    UnaryOp::Void => fc.builder.emit_reg_reg(Op::Void, dst, src),
                    UnaryOp::Delete => unreachable!(),
                }
                fc.free_reg(src);
            }
        }

        ExprKind::Update {
            op,
            argument,
            prefix,
        } => {
            // Get current value.
            compile_expr(fc, argument, dst)?;

            let one = fc.alloc_reg();
            fc.builder.emit_load_int8(one, 1);

            if *prefix {
                // ++x / --x: modify first, return modified.
                match op {
                    UpdateOp::Increment => fc.builder.emit_reg3(Op::Add, dst, dst, one),
                    UpdateOp::Decrement => fc.builder.emit_reg3(Op::Sub, dst, dst, one),
                }
                // Store back.
                compile_store(fc, argument, dst)?;
            } else {
                // x++ / x--: return original, then modify.
                let tmp = fc.alloc_reg();
                fc.builder.emit_reg_reg(Op::Move, tmp, dst);
                match op {
                    UpdateOp::Increment => fc.builder.emit_reg3(Op::Add, tmp, tmp, one),
                    UpdateOp::Decrement => fc.builder.emit_reg3(Op::Sub, tmp, tmp, one),
                }
                compile_store(fc, argument, tmp)?;
                fc.free_reg(tmp);
            }
            fc.free_reg(one);
        }

        ExprKind::Logical { op, left, right } => {
            compile_expr(fc, left, dst)?;
            match op {
                LogicalOp::And => {
                    // Short-circuit: if falsy, skip right.
                    let skip = fc.builder.emit_cond_jump(Op::JumpIfFalse, dst);
                    compile_expr(fc, right, dst)?;
                    fc.builder.patch_jump(skip);
                }
                LogicalOp::Or => {
                    let skip = fc.builder.emit_cond_jump(Op::JumpIfTrue, dst);
                    compile_expr(fc, right, dst)?;
                    fc.builder.patch_jump(skip);
                }
                LogicalOp::Nullish => {
                    let skip = fc.builder.emit_cond_jump(Op::JumpIfNullish, dst);
                    // If NOT nullish, skip the right side. Wait — JumpIfNullish
                    // should mean "jump if nullish" so we want: evaluate left,
                    // if NOT nullish skip right.
                    // Let's invert: evaluate left, check if nullish → evaluate right.
                    // We need the jump to skip the "evaluate right" if NOT nullish.
                    // Since JumpIfNullish jumps when nullish, we need the inverse.
                    // Instead: use a two-step approach.
                    //
                    // Actually, rethink: for `a ?? b`:
                    //   1. evaluate a → dst
                    //   2. if dst is NOT null/undefined, jump to end
                    //   3. evaluate b → dst
                    //   end:
                    // JumpIfNullish jumps when IS nullish. So we want jump when NOT nullish.
                    // Let's just use a "not nullish" check.
                    // For now: negate and use JumpIfFalse.
                    // Actually simpler: skip right when not nullish.
                    // JumpIfNullish jumps WHEN nullish. We want to jump over right when NOT nullish.
                    // So:
                    //   evaluate a → dst
                    //   JumpIfNullish dst → evaluate_right
                    //   Jump → end
                    //   evaluate_right: evaluate b → dst
                    //   end:
                    // But we already emitted JumpIfNullish. Let's fix this.
                    // The JumpIfNullish we emitted jumps to "after patch", which is where
                    // we'll put the right-side code. We need another jump to skip right.
                    let end_patch = fc.builder.emit_jump(Op::Jump);
                    fc.builder.patch_jump(skip); // nullish → evaluate right
                    compile_expr(fc, right, dst)?;
                    fc.builder.patch_jump(end_patch);
                }
            }
        }

        ExprKind::Assignment { op, left, right } => {
            if *op == AssignOp::Assign {
                compile_expr(fc, right, dst)?;
                compile_store(fc, left, dst)?;
            } else {
                // Compound assignment: load current, operate, store.
                compile_expr(fc, left, dst)?;
                let rhs = fc.alloc_reg();
                compile_expr(fc, right, rhs)?;
                let arith_op = compound_assign_op(*op);
                fc.builder.emit_reg3(arith_op, dst, dst, rhs);
                fc.free_reg(rhs);
                compile_store(fc, left, dst)?;
            }
        }

        ExprKind::Conditional {
            test,
            consequent,
            alternate,
        } => {
            let cond = fc.alloc_reg();
            compile_expr(fc, test, cond)?;
            let else_patch = fc.builder.emit_cond_jump(Op::JumpIfFalse, cond);
            fc.free_reg(cond);
            compile_expr(fc, consequent, dst)?;
            let end_patch = fc.builder.emit_jump(Op::Jump);
            fc.builder.patch_jump(else_patch);
            compile_expr(fc, alternate, dst)?;
            fc.builder.patch_jump(end_patch);
        }

        ExprKind::Call { callee, arguments } => {
            let func_reg = fc.alloc_reg();
            compile_expr(fc, callee, func_reg)?;

            let args_start = fc.next_reg;
            let arg_count = arguments.len().min(255) as u8;
            for arg in arguments {
                let arg_reg = fc.alloc_reg();
                compile_expr(fc, arg, arg_reg)?;
            }

            fc.builder.emit_call(dst, func_reg, args_start, arg_count);

            // Free argument registers (in reverse).
            for _ in 0..arg_count {
                fc.next_reg -= 1;
            }
            fc.free_reg(func_reg);
        }

        ExprKind::New { callee, arguments } => {
            // For now, compile like a regular call. The VM will differentiate
            // based on the `New` vs `Call` distinction (TODO: add NewCall opcode).
            let func_reg = fc.alloc_reg();
            compile_expr(fc, callee, func_reg)?;

            let args_start = fc.next_reg;
            let arg_count = arguments.len().min(255) as u8;
            for arg in arguments {
                let arg_reg = fc.alloc_reg();
                compile_expr(fc, arg, arg_reg)?;
            }

            fc.builder.emit_call(dst, func_reg, args_start, arg_count);

            for _ in 0..arg_count {
                fc.next_reg -= 1;
            }
            fc.free_reg(func_reg);
        }

        ExprKind::Member {
            object,
            property,
            computed,
        } => {
            let obj_reg = fc.alloc_reg();
            compile_expr(fc, object, obj_reg)?;

            if !computed {
                // Static member: obj.prop → GetPropertyByName.
                if let ExprKind::Identifier(name) = &property.kind {
                    let ni = fc.builder.add_name(name);
                    fc.builder.emit_get_prop_name(dst, obj_reg, ni);
                } else {
                    let key_reg = fc.alloc_reg();
                    compile_expr(fc, property, key_reg)?;
                    fc.builder.emit_reg3(Op::GetProperty, dst, obj_reg, key_reg);
                    fc.free_reg(key_reg);
                }
            } else {
                // Computed member: obj[expr].
                let key_reg = fc.alloc_reg();
                compile_expr(fc, property, key_reg)?;
                fc.builder.emit_reg3(Op::GetProperty, dst, obj_reg, key_reg);
                fc.free_reg(key_reg);
            }
            fc.free_reg(obj_reg);
        }

        ExprKind::Array(elements) => {
            fc.builder.emit_reg(Op::CreateArray, dst);
            for (i, elem) in elements.iter().enumerate() {
                if let Some(el) = elem {
                    let val_reg = fc.alloc_reg();
                    match el {
                        ArrayElement::Expr(e) => compile_expr(fc, e, val_reg)?,
                        ArrayElement::Spread(e) => {
                            // Spread in array: simplified, just compile the expression.
                            compile_expr(fc, e, val_reg)?;
                        }
                    }
                    let idx_reg = fc.alloc_reg();
                    if i <= 127 {
                        fc.builder.emit_load_int8(idx_reg, i as i8);
                    } else {
                        let ci = fc.builder.add_constant(Constant::Number(i as f64));
                        fc.builder.emit_reg_u16(Op::LoadConst, idx_reg, ci);
                    }
                    fc.builder.emit_reg3(Op::SetProperty, dst, idx_reg, val_reg);
                    fc.free_reg(idx_reg);
                    fc.free_reg(val_reg);
                }
            }
        }

        ExprKind::Object(properties) => {
            fc.builder.emit_reg(Op::CreateObject, dst);
            for prop in properties {
                let val_reg = fc.alloc_reg();
                if let Some(value) = &prop.value {
                    compile_expr(fc, value, val_reg)?;
                } else {
                    // Shorthand: `{ x }` means `{ x: x }`.
                    if let PropertyKey::Identifier(name) = &prop.key {
                        if let Some(local) = fc.find_local(name) {
                            fc.builder.emit_reg_reg(Op::Move, val_reg, local);
                        } else {
                            let ni = fc.builder.add_name(name);
                            fc.builder.emit_load_global(val_reg, ni);
                        }
                    } else {
                        fc.builder.emit_reg(Op::LoadUndefined, val_reg);
                    }
                }

                match &prop.key {
                    PropertyKey::Identifier(name) | PropertyKey::String(name) => {
                        let ni = fc.builder.add_name(name);
                        fc.builder.emit_set_prop_name(dst, ni, val_reg);
                    }
                    PropertyKey::Number(n) => {
                        let key_reg = fc.alloc_reg();
                        let ci = fc.builder.add_constant(Constant::Number(*n));
                        fc.builder.emit_reg_u16(Op::LoadConst, key_reg, ci);
                        fc.builder.emit_reg3(Op::SetProperty, dst, key_reg, val_reg);
                        fc.free_reg(key_reg);
                    }
                    PropertyKey::Computed(expr) => {
                        let key_reg = fc.alloc_reg();
                        compile_expr(fc, expr, key_reg)?;
                        fc.builder.emit_reg3(Op::SetProperty, dst, key_reg, val_reg);
                        fc.free_reg(key_reg);
                    }
                }
                fc.free_reg(val_reg);
            }
        }

        ExprKind::Function(func_def) => {
            let inner = compile_function_body(func_def)?;
            let func_idx = fc.builder.add_function(inner);
            fc.builder.emit_reg_u16(Op::CreateClosure, dst, func_idx);
        }

        ExprKind::Arrow {
            params,
            body,
            is_async: _,
        } => {
            let param_count = params.len().min(255) as u8;
            let mut inner = FunctionCompiler::new("<arrow>".into(), param_count);
            for p in params {
                if let PatternKind::Identifier(name) = &p.kind {
                    inner.define_local(name);
                } else {
                    let _ = inner.alloc_reg();
                }
            }
            let result = inner.alloc_reg();
            match body {
                ArrowBody::Expr(e) => {
                    compile_expr(&mut inner, e, result)?;
                }
                ArrowBody::Block(stmts) => {
                    inner.builder.emit_reg(Op::LoadUndefined, result);
                    compile_stmts(&mut inner, stmts, result)?;
                }
            }
            inner.builder.emit_reg(Op::Return, result);
            let inner_func = inner.builder.finish();
            let func_idx = fc.builder.add_function(inner_func);
            fc.builder.emit_reg_u16(Op::CreateClosure, dst, func_idx);
        }

        ExprKind::Class(class_def) => {
            // Class expression: compile like class decl but into dst.
            let name = class_def.id.clone().unwrap_or_default();
            // Find constructor.
            let ctor = class_def.body.iter().find(|m| {
                matches!(
                    &m.kind,
                    ClassMemberKind::Method {
                        kind: MethodKind::Constructor,
                        ..
                    }
                )
            });
            if let Some(member) = ctor {
                if let ClassMemberKind::Method { value, .. } = &member.kind {
                    let inner = compile_function_body(value)?;
                    let func_idx = fc.builder.add_function(inner);
                    fc.builder.emit_reg_u16(Op::CreateClosure, dst, func_idx);
                }
            } else {
                let mut empty = BytecodeBuilder::new(name, 0);
                let r = 0u8;
                empty.func.register_count = 1;
                empty.emit_reg(Op::LoadUndefined, r);
                empty.emit_reg(Op::Return, r);
                let func_idx = fc.builder.add_function(empty.finish());
                fc.builder.emit_reg_u16(Op::CreateClosure, dst, func_idx);
            }

            // Compile methods as properties on the constructor.
            for member in &class_def.body {
                match &member.kind {
                    ClassMemberKind::Method {
                        key,
                        value,
                        kind,
                        is_static: _,
                        computed: _,
                    } => {
                        if matches!(kind, MethodKind::Constructor) {
                            continue;
                        }
                        let method_name = match key {
                            PropertyKey::Identifier(s) | PropertyKey::String(s) => s.clone(),
                            _ => continue,
                        };
                        let inner = compile_function_body(value)?;
                        let func_idx = fc.builder.add_function(inner);
                        let method_reg = fc.alloc_reg();
                        fc.builder
                            .emit_reg_u16(Op::CreateClosure, method_reg, func_idx);
                        let name_idx = fc.builder.add_name(&method_name);
                        fc.builder.emit_set_prop_name(dst, name_idx, method_reg);
                        fc.free_reg(method_reg);
                    }
                    ClassMemberKind::Property { .. } => {}
                }
            }
        }

        ExprKind::Sequence(exprs) => {
            for e in exprs {
                compile_expr(fc, e, dst)?;
            }
        }

        ExprKind::Spread(inner) => {
            compile_expr(fc, inner, dst)?;
        }

        ExprKind::TemplateLiteral {
            quasis,
            expressions,
        } => {
            // Compile template literal as string concatenation.
            if quasis.len() == 1 && expressions.is_empty() {
                let ci = fc.builder.add_constant(Constant::String(quasis[0].clone()));
                fc.builder.emit_reg_u16(Op::LoadConst, dst, ci);
            } else {
                // Start with first quasi.
                let ci = fc.builder.add_constant(Constant::String(quasis[0].clone()));
                fc.builder.emit_reg_u16(Op::LoadConst, dst, ci);
                for (i, expr) in expressions.iter().enumerate() {
                    let tmp = fc.alloc_reg();
                    compile_expr(fc, expr, tmp)?;
                    fc.builder.emit_reg3(Op::Add, dst, dst, tmp);
                    fc.free_reg(tmp);
                    if i + 1 < quasis.len() {
                        let qi = fc
                            .builder
                            .add_constant(Constant::String(quasis[i + 1].clone()));
                        let tmp2 = fc.alloc_reg();
                        fc.builder.emit_reg_u16(Op::LoadConst, tmp2, qi);
                        fc.builder.emit_reg3(Op::Add, dst, dst, tmp2);
                        fc.free_reg(tmp2);
                    }
                }
            }
        }

        ExprKind::TaggedTemplate { tag, quasi } => {
            // Simplified: call tag with the template as argument.
            let func_reg = fc.alloc_reg();
            compile_expr(fc, tag, func_reg)?;
            let arg_reg = fc.alloc_reg();
            compile_expr(fc, quasi, arg_reg)?;
            fc.builder.emit_call(dst, func_reg, arg_reg, 1);
            fc.free_reg(arg_reg);
            fc.free_reg(func_reg);
        }

        ExprKind::Yield {
            argument,
            delegate: _,
        } => {
            // Yield is a VM-level operation; for now compile the argument.
            if let Some(arg) = argument {
                compile_expr(fc, arg, dst)?;
            } else {
                fc.builder.emit_reg(Op::LoadUndefined, dst);
            }
        }

        ExprKind::Await(inner) => {
            // Await is a VM-level operation; compile the argument.
            compile_expr(fc, inner, dst)?;
        }

        ExprKind::RegExp { .. } => {
            // RegExp literals are created at runtime by the VM.
            fc.builder.emit_reg(Op::LoadUndefined, dst);
        }

        ExprKind::OptionalChain { base } => {
            compile_expr(fc, base, dst)?;
        }
    }
    Ok(())
}

/// Compile a store operation (assignment target).
fn compile_store(fc: &mut FunctionCompiler, target: &Expr, src: Reg) -> Result<(), JsError> {
    match &target.kind {
        ExprKind::Identifier(name) => {
            if let Some(local) = fc.find_local(name) {
                if local != src {
                    fc.builder.emit_reg_reg(Op::Move, local, src);
                }
            } else {
                let ni = fc.builder.add_name(name);
                fc.builder.emit_store_global(ni, src);
            }
        }
        ExprKind::Member {
            object,
            property,
            computed,
        } => {
            let obj_reg = fc.alloc_reg();
            compile_expr(fc, object, obj_reg)?;
            if !computed {
                if let ExprKind::Identifier(name) = &property.kind {
                    let ni = fc.builder.add_name(name);
                    fc.builder.emit_set_prop_name(obj_reg, ni, src);
                } else {
                    let key_reg = fc.alloc_reg();
                    compile_expr(fc, property, key_reg)?;
                    fc.builder.emit_reg3(Op::SetProperty, obj_reg, key_reg, src);
                    fc.free_reg(key_reg);
                }
            } else {
                let key_reg = fc.alloc_reg();
                compile_expr(fc, property, key_reg)?;
                fc.builder.emit_reg3(Op::SetProperty, obj_reg, key_reg, src);
                fc.free_reg(key_reg);
            }
            fc.free_reg(obj_reg);
        }
        _ => {
            // Other assignment targets (destructuring) not handled here.
        }
    }
    Ok(())
}

fn binary_op_to_opcode(op: BinaryOp) -> Op {
    match op {
        BinaryOp::Add => Op::Add,
        BinaryOp::Sub => Op::Sub,
        BinaryOp::Mul => Op::Mul,
        BinaryOp::Div => Op::Div,
        BinaryOp::Rem => Op::Rem,
        BinaryOp::Exp => Op::Exp,
        BinaryOp::Eq => Op::Eq,
        BinaryOp::Ne => Op::NotEq,
        BinaryOp::StrictEq => Op::StrictEq,
        BinaryOp::StrictNe => Op::StrictNotEq,
        BinaryOp::Lt => Op::LessThan,
        BinaryOp::Le => Op::LessEq,
        BinaryOp::Gt => Op::GreaterThan,
        BinaryOp::Ge => Op::GreaterEq,
        BinaryOp::Shl => Op::ShiftLeft,
        BinaryOp::Shr => Op::ShiftRight,
        BinaryOp::Ushr => Op::UShiftRight,
        BinaryOp::BitAnd => Op::BitAnd,
        BinaryOp::BitOr => Op::BitOr,
        BinaryOp::BitXor => Op::BitXor,
        BinaryOp::In => Op::In,
        BinaryOp::Instanceof => Op::InstanceOf,
    }
}

fn compound_assign_op(op: AssignOp) -> Op {
    match op {
        AssignOp::AddAssign => Op::Add,
        AssignOp::SubAssign => Op::Sub,
        AssignOp::MulAssign => Op::Mul,
        AssignOp::DivAssign => Op::Div,
        AssignOp::RemAssign => Op::Rem,
        AssignOp::ExpAssign => Op::Exp,
        AssignOp::ShlAssign => Op::ShiftLeft,
        AssignOp::ShrAssign => Op::ShiftRight,
        AssignOp::UshrAssign => Op::UShiftRight,
        AssignOp::BitAndAssign => Op::BitAnd,
        AssignOp::BitOrAssign => Op::BitOr,
        AssignOp::BitXorAssign => Op::BitXor,
        AssignOp::AndAssign => Op::BitAnd, // logical AND assignment uses short-circuit; simplified here
        AssignOp::OrAssign => Op::BitOr,   // likewise
        AssignOp::NullishAssign => Op::Move, // simplified
        AssignOp::Assign => unreachable!(),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::parser::Parser;

    /// Helper: parse and compile source, return the top-level function.
    fn compile_src(src: &str) -> Function {
        let program = Parser::parse(src).expect("parse failed");
        compile(&program).expect("compile failed")
    }

    #[test]
    fn test_compile_number_literal() {
        let f = compile_src("42;");
        let dis = f.disassemble();
        assert!(dis.contains("LoadInt8 r0, 42"), "got:\n{dis}");
        assert!(dis.contains("Return r0"));
    }

    #[test]
    fn test_compile_large_number() {
        let f = compile_src("3.14;");
        let dis = f.disassemble();
        assert!(dis.contains("LoadConst r0, #0"), "got:\n{dis}");
        assert!(
            f.constants.contains(&Constant::Number(3.14)),
            "constants: {:?}",
            f.constants
        );
    }

    #[test]
    fn test_compile_string() {
        let f = compile_src("\"hello\";");
        let dis = f.disassemble();
        assert!(dis.contains("LoadConst r0, #0"));
        assert!(f.constants.contains(&Constant::String("hello".into())));
    }

    #[test]
    fn test_compile_bool_null() {
        let f = compile_src("true; false; null;");
        let dis = f.disassemble();
        assert!(dis.contains("LoadTrue r0"));
        assert!(dis.contains("LoadFalse r0"));
        assert!(dis.contains("LoadNull r0"));
    }

    #[test]
    fn test_compile_binary_arithmetic() {
        let f = compile_src("1 + 2;");
        let dis = f.disassemble();
        assert!(dis.contains("Add r0, r1, r2"), "got:\n{dis}");
    }

    #[test]
    fn test_compile_nested_arithmetic() {
        let f = compile_src("(1 + 2) * 3;");
        let dis = f.disassemble();
        assert!(dis.contains("Add"), "got:\n{dis}");
        assert!(dis.contains("Mul"), "got:\n{dis}");
    }

    #[test]
    fn test_compile_var_decl() {
        let f = compile_src("var x = 10; x;");
        let dis = f.disassemble();
        // x should get a register, then be loaded from that register.
        assert!(dis.contains("LoadInt8"), "got:\n{dis}");
        assert!(
            dis.contains("Move") || dis.contains("LoadInt8"),
            "got:\n{dis}"
        );
    }

    #[test]
    fn test_compile_let_const() {
        let f = compile_src("let a = 1; const b = 2; a + b;");
        let dis = f.disassemble();
        assert!(dis.contains("Add"), "got:\n{dis}");
    }

    #[test]
    fn test_compile_if_else() {
        let f = compile_src("if (true) { 1; } else { 2; }");
        let dis = f.disassemble();
        assert!(dis.contains("JumpIfFalse"), "got:\n{dis}");
        assert!(dis.contains("Jump"), "got:\n{dis}");
    }

    #[test]
    fn test_compile_while() {
        let f = compile_src("var i = 0; while (i < 10) { i = i + 1; }");
        let dis = f.disassemble();
        assert!(dis.contains("LessThan"), "got:\n{dis}");
        assert!(dis.contains("JumpIfFalse"), "got:\n{dis}");
        assert!(
            dis.contains("Jump"),
            "backward jump should be present: {dis}"
        );
    }

    #[test]
    fn test_compile_do_while() {
        let f = compile_src("var i = 0; do { i = i + 1; } while (i < 5);");
        let dis = f.disassemble();
        assert!(dis.contains("JumpIfTrue"), "got:\n{dis}");
    }

    #[test]
    fn test_compile_for_loop() {
        let f = compile_src("for (var i = 0; i < 10; i = i + 1) { i; }");
        let dis = f.disassemble();
        assert!(dis.contains("LessThan"), "got:\n{dis}");
        assert!(dis.contains("JumpIfFalse"), "got:\n{dis}");
    }

    #[test]
    fn test_compile_function_decl() {
        let f = compile_src("function add(a, b) { return a + b; }");
        let dis = f.disassemble();
        assert!(dis.contains("CreateClosure"), "got:\n{dis}");
        assert!(!f.functions.is_empty(), "should have nested function");
        let inner = &f.functions[0];
        assert_eq!(inner.name, "add");
        assert_eq!(inner.param_count, 2);
        let inner_dis = inner.disassemble();
        assert!(inner_dis.contains("Add"), "inner:\n{inner_dis}");
        assert!(inner_dis.contains("Return"), "inner:\n{inner_dis}");
    }

    #[test]
    fn test_compile_function_call() {
        let f = compile_src("function f() { return 42; } f();");
        let dis = f.disassemble();
        assert!(dis.contains("Call"), "got:\n{dis}");
    }

    #[test]
    fn test_compile_arrow_function() {
        let f = compile_src("var add = (a, b) => a + b;");
        let dis = f.disassemble();
        assert!(dis.contains("CreateClosure"), "got:\n{dis}");
        let inner = &f.functions[0];
        assert_eq!(inner.param_count, 2);
    }

    #[test]
    fn test_compile_assignment() {
        let f = compile_src("var x = 1; x = x + 2;");
        let dis = f.disassemble();
        assert!(dis.contains("Add"), "got:\n{dis}");
        assert!(
            dis.contains("Move"),
            "assignment should produce Move:\n{dis}"
        );
    }

    #[test]
    fn test_compile_compound_assignment() {
        let f = compile_src("var x = 10; x += 5;");
        let dis = f.disassemble();
        assert!(dis.contains("Add"), "got:\n{dis}");
    }

    #[test]
    fn test_compile_member_access() {
        let f = compile_src("var obj = {}; obj.x;");
        let dis = f.disassemble();
        assert!(dis.contains("CreateObject"), "got:\n{dis}");
        assert!(dis.contains("GetPropertyByName"), "got:\n{dis}");
    }

    #[test]
    fn test_compile_computed_member() {
        let f = compile_src("var arr = []; arr[0];");
        let dis = f.disassemble();
        assert!(dis.contains("GetProperty"), "got:\n{dis}");
    }

    #[test]
    fn test_compile_object_literal() {
        let f = compile_src("var obj = { a: 1, b: 2 };");
        let dis = f.disassemble();
        assert!(dis.contains("CreateObject"), "got:\n{dis}");
        assert!(dis.contains("SetPropertyByName"), "got:\n{dis}");
    }

    #[test]
    fn test_compile_array_literal() {
        let f = compile_src("[1, 2, 3];");
        let dis = f.disassemble();
        assert!(dis.contains("CreateArray"), "got:\n{dis}");
        assert!(dis.contains("SetProperty"), "got:\n{dis}");
    }

    #[test]
    fn test_compile_conditional() {
        let f = compile_src("true ? 1 : 2;");
        let dis = f.disassemble();
        assert!(dis.contains("JumpIfFalse"), "got:\n{dis}");
    }

    #[test]
    fn test_compile_logical_and() {
        let f = compile_src("true && false;");
        let dis = f.disassemble();
        assert!(dis.contains("JumpIfFalse"), "short-circuit:\n{dis}");
    }

    #[test]
    fn test_compile_logical_or() {
        let f = compile_src("false || true;");
        let dis = f.disassemble();
        assert!(dis.contains("JumpIfTrue"), "short-circuit:\n{dis}");
    }

    #[test]
    fn test_compile_typeof() {
        let f = compile_src("typeof 42;");
        let dis = f.disassemble();
        assert!(dis.contains("TypeOf"), "got:\n{dis}");
    }

    #[test]
    fn test_compile_unary_minus() {
        let f = compile_src("-42;");
        let dis = f.disassemble();
        assert!(dis.contains("Neg"), "got:\n{dis}");
    }

    #[test]
    fn test_compile_not() {
        let f = compile_src("!true;");
        let dis = f.disassemble();
        assert!(dis.contains("LogicalNot"), "got:\n{dis}");
    }

    #[test]
    fn test_compile_return() {
        let f = compile_src("function f() { return 42; }");
        let inner = &f.functions[0];
        let dis = inner.disassemble();
        assert!(dis.contains("Return"), "got:\n{dis}");
    }

    #[test]
    fn test_compile_empty_return() {
        let f = compile_src("function f() { return; }");
        let inner = &f.functions[0];
        let dis = inner.disassemble();
        assert!(dis.contains("LoadUndefined"), "got:\n{dis}");
        assert!(dis.contains("Return"), "got:\n{dis}");
    }

    #[test]
    fn test_compile_throw() {
        let f = compile_src("function f() { throw 42; }");
        let inner = &f.functions[0];
        let dis = inner.disassemble();
        assert!(dis.contains("Throw"), "got:\n{dis}");
    }

    #[test]
    fn test_compile_this() {
        let f = compile_src("this;");
        let dis = f.disassemble();
        assert!(dis.contains("LoadGlobal"), "got:\n{dis}");
        assert!(f.names.contains(&"this".to_string()));
    }

    #[test]
    fn test_compile_global_var() {
        let f = compile_src("console;");
        let dis = f.disassemble();
        assert!(dis.contains("LoadGlobal"), "got:\n{dis}");
        assert!(f.names.contains(&"console".to_string()));
    }

    #[test]
    fn test_compile_template_literal() {
        let f = compile_src("`hello`;");
        assert!(
            f.constants.contains(&Constant::String("hello".into())),
            "constants: {:?}",
            f.constants
        );
    }

    #[test]
    fn test_compile_switch() {
        let f = compile_src("switch (1) { case 1: 42; break; case 2: 99; break; }");
        let dis = f.disassemble();
        assert!(dis.contains("StrictEq"), "got:\n{dis}");
    }

    #[test]
    fn test_compile_class() {
        let f = compile_src("class Foo { constructor() {} greet() { return 1; } }");
        let dis = f.disassemble();
        assert!(dis.contains("CreateClosure"), "got:\n{dis}");
    }

    #[test]
    fn test_compile_update_prefix() {
        let f = compile_src("var x = 0; ++x;");
        let dis = f.disassemble();
        assert!(dis.contains("Add"), "got:\n{dis}");
    }

    #[test]
    fn test_compile_comparison() {
        let f = compile_src("1 === 2;");
        let dis = f.disassemble();
        assert!(dis.contains("StrictEq"), "got:\n{dis}");
    }

    #[test]
    fn test_compile_bitwise() {
        let f = compile_src("1 & 2;");
        let dis = f.disassemble();
        assert!(dis.contains("BitAnd"), "got:\n{dis}");
    }

    #[test]
    fn test_compile_void() {
        let f = compile_src("void 0;");
        let dis = f.disassemble();
        assert!(dis.contains("Void"), "got:\n{dis}");
    }

    #[test]
    fn test_disassembler_output_format() {
        let f = compile_src("var x = 42; x + 1;");
        let dis = f.disassemble();
        // Should contain function header.
        assert!(dis.contains("function <main>"));
        // Should contain code section.
        assert!(dis.contains("code:"));
        // Should have hex offsets.
        assert!(dis.contains("0000"));
    }

    #[test]
    fn test_register_allocation_is_minimal() {
        // `var a = 1; var b = 2; a + b;` should use few registers.
        let f = compile_src("var a = 1; var b = 2; a + b;");
        // r0 = result, r1 = a, r2 = b, r3/r4 = temps for addition
        assert!(
            f.register_count <= 6,
            "too many registers: {}",
            f.register_count
        );
    }

    #[test]
    fn test_nested_function_closure() {
        let f = compile_src("function outer() { function inner() { return 1; } return inner; }");
        assert_eq!(f.functions.len(), 1);
        let outer = &f.functions[0];
        assert_eq!(outer.name, "outer");
        assert_eq!(outer.functions.len(), 1);
        let inner = &outer.functions[0];
        assert_eq!(inner.name, "inner");
    }

    #[test]
    fn test_for_with_no_parts() {
        // `for (;;) { break; }` — infinite loop with immediate break.
        let f = compile_src("for (;;) { break; }");
        let dis = f.disassemble();
        assert!(dis.contains("Jump"), "got:\n{dis}");
    }

    #[test]
    fn test_for_continue_targets_update() {
        // `continue` in a for-loop must jump to the update expression, not back
        // to the condition check. Verify the continue jump goes to the Add (i + 1)
        // rather than to the LessThan condition.
        let f = compile_src("for (var i = 0; i < 10; i = i + 1) { continue; }");
        let dis = f.disassemble();
        // The for-loop should contain: LessThan (test), JumpIfFalse (exit),
        // Jump (continue), Add (update), Jump (back to test).
        assert!(dis.contains("LessThan"), "missing test: {dis}");
        assert!(dis.contains("Add"), "missing update: {dis}");
        // There should be at least 2 Jump instructions (continue + back-edge).
        let jump_count = dis.matches("Jump ").count();
        assert!(
            jump_count >= 2,
            "expected >= 2 jumps for continue + back-edge, got {jump_count}: {dis}"
        );
    }

    #[test]
    fn test_do_while_continue_targets_condition() {
        // `continue` in do-while must jump to the condition, not the body start.
        let f = compile_src("var i = 0; do { i = i + 1; continue; } while (i < 5);");
        let dis = f.disassemble();
        assert!(dis.contains("LessThan"), "missing condition: {dis}");
        assert!(dis.contains("JumpIfTrue"), "missing back-edge: {dis}");
    }

    #[test]
    fn test_switch_default_case() {
        // Default case must not corrupt bytecode.
        let f = compile_src("switch (1) { case 1: 10; break; default: 20; break; }");
        let dis = f.disassemble();
        assert!(dis.contains("StrictEq"), "missing case test: {dis}");
        // The first instruction should NOT be corrupted.
        assert!(
            dis.contains("LoadUndefined r0"),
            "first instruction corrupted: {dis}"
        );
    }

    #[test]
    fn test_switch_only_default() {
        // Switch with only a default case.
        let f = compile_src("switch (42) { default: 99; }");
        let dis = f.disassemble();
        // Should compile without panicking and contain the default body.
        assert!(dis.contains("LoadInt8"), "got:\n{dis}");
    }

    #[test]
    fn test_class_empty_constructor_has_return() {
        // A class without an explicit constructor should produce a function with Return.
        let f = compile_src("class Foo {}");
        assert!(!f.functions.is_empty(), "should have constructor function");
        let ctor = &f.functions[0];
        let dis = ctor.disassemble();
        assert!(
            dis.contains("Return"),
            "empty constructor must have Return: {dis}"
        );
    }

    #[test]
    fn test_class_expression_compiles_methods() {
        // Class expression should compile methods, not just the constructor.
        let f = compile_src("var C = class { greet() { return 1; } };");
        let dis = f.disassemble();
        assert!(
            dis.contains("SetPropertyByName"),
            "method should be set as property: {dis}"
        );
    }
}