aegis_vm_macro 0.1.0

//! AST to Bytecode Compiler
//!
//! Converts Rust expressions to VM bytecode at compile time.

use syn::{ItemFn, Expr, Stmt, BinOp, UnOp, Pat, FnArg, Lit};
use std::collections::BTreeMap; // Deterministic ordering!
use crate::opcodes::{stack, arithmetic, control, native, exec};
use crate::crypto::OpcodeTable;
use crate::mba::MbaTransformer;

/// Compilation error
#[derive(Debug)]
pub struct CompileError(pub String);

impl std::fmt::Display for CompileError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.0)
    }
}

/// Variable location - either in input buffer or in a register
#[derive(Debug, Clone)]
enum VarLocation {
    /// Input buffer offset (for function arguments)
    InputOffset(usize),
    /// Register index (for local variables)
    Register(u8),
}

/// Loop context for break/continue support
#[derive(Debug, Clone)]
struct LoopContext {
    /// Label for continue (jump to condition/increment)
    continue_label: String,
    /// Label for break (jump past loop end)
    break_label: String,
}

/// Compiler state
pub struct Compiler {
    /// Generated bytecode
    bytecode: Vec<u8>,
    /// Function argument name -> input buffer offset
    arg_offsets: BTreeMap<String, usize>,
    /// Local variable name -> register index
    local_registers: BTreeMap<String, u8>,
    /// Next available register for locals (R0-R5, R6-R7 reserved for MBA)
    next_local_reg: u8,
    /// Current input buffer offset for next argument
    next_arg_offset: usize,
    /// Label name -> bytecode offset for jumps (BTreeMap for determinism)
    labels: BTreeMap<String, usize>,
    /// Pending jump fixups (bytecode offset, label name)
    fixups: Vec<(usize, String)>,
    /// Unique label counter
    label_counter: usize,
    /// Stack of active loops for break/continue
    loop_stack: Vec<LoopContext>,
    /// Opcode encoding table (base -> shuffled)
    opcode_table: OpcodeTable,
    /// MBA transformer for arithmetic obfuscation
    mba: MbaTransformer,
    /// Enable MBA transformations (for Paranoid level)
    mba_enabled: bool,
}

impl Compiler {
    /// Create new compiler
    #[allow(dead_code)]
    pub fn new() -> Self {
        Self::with_mba(false)
    }

    /// Create compiler with MBA transformations enabled
    pub fn with_mba(mba_enabled: bool) -> Self {
        // Use build seed for deterministic MBA
        let mba_seed = crate::crypto::get_opcode_table().get_seed();

        Self {
            bytecode: Vec::new(),
            arg_offsets: BTreeMap::new(),
            local_registers: BTreeMap::new(),
            next_local_reg: 0,
            next_arg_offset: 0,
            labels: BTreeMap::new(),
            fixups: Vec::new(),
            label_counter: 0,
            loop_stack: Vec::new(),
            opcode_table: crate::crypto::get_opcode_table(),
            mba: MbaTransformer::new(mba_seed),
            mba_enabled,
        }
    }

    /// Emit an opcode (automatically encoded via shuffle table)
    fn emit_op(&mut self, base_opcode: u8) {
        // Junk Code Insertion (Obfuscation)
        // Use bytecode length as a seed for deterministic pseudo-randomness
        // This ensures reproducible builds while scattering junk code
        let entropy = (self.bytecode.len() as u64)
            .wrapping_mul(0x5deece66d)
            .wrapping_add(0xb);
        
        // 10% chance to insert a simple NOP
        // Removing NOP_N logic as it requires inserting padding bytes which complicates things
        if (entropy % 100) < 10 {
            let nop_shuffled = self.opcode_table.encode(crate::opcodes::special::NOP);
            self.bytecode.push(nop_shuffled);
        } 

        // Emit the actual instruction
        let shuffled = self.opcode_table.encode(base_opcode);
        self.bytecode.push(shuffled);
    }

    /// Emit a single byte
    fn emit(&mut self, byte: u8) {
        self.bytecode.push(byte);
    }

    /// Emit a u16 (little-endian)
    fn emit_u16(&mut self, value: u16) {
        self.bytecode.extend_from_slice(&value.to_le_bytes());
    }

    /// Emit a u64 (little-endian)
    fn emit_u64(&mut self, value: u64) {
        self.bytecode.extend_from_slice(&value.to_le_bytes());
    }

    // =========================================================================
    // MBA-Aware Arithmetic Emission
    // =========================================================================

    /// Emit ADD instruction, potentially with MBA transformation
    fn emit_add(&mut self) {
        if self.mba_enabled {
            // Use MBA transformer to generate obfuscated ADD
            let table = self.opcode_table.clone();
            self.mba.emit_add(&mut self.bytecode, |op| table.encode(op));
        } else {
            self.emit_op(arithmetic::ADD);
        }
    }

    /// Emit SUB instruction, potentially with MBA transformation
    fn emit_sub(&mut self) {
        if self.mba_enabled {
            let table = self.opcode_table.clone();
            self.mba.emit_sub(&mut self.bytecode, |op| table.encode(op));
        } else {
            self.emit_op(arithmetic::SUB);
        }
    }

    /// Emit XOR instruction, potentially with MBA transformation
    fn emit_xor(&mut self) {
        if self.mba_enabled {
            let table = self.opcode_table.clone();
            self.mba.emit_xor(&mut self.bytecode, |op| table.encode(op));
        } else {
            self.emit_op(arithmetic::XOR);
        }
    }

    /// Current bytecode position
    fn pos(&self) -> usize {
        self.bytecode.len()
    }

    /// Generate unique label
    fn unique_label(&mut self, prefix: &str) -> String {
        let label = format!("{}_{}", prefix, self.label_counter);
        self.label_counter += 1;
        label
    }

    /// Mark current position as label
    fn mark_label(&mut self, name: &str) {
        self.labels.insert(name.to_string(), self.pos());
    }

    /// Emit jump with fixup
    fn emit_jump(&mut self, opcode: u8, label: &str) {
        self.emit_op(opcode);
        let fixup_pos = self.pos();
        self.emit_u16(0); // Placeholder
        self.fixups.push((fixup_pos, label.to_string()));
    }

    /// Apply all jump fixups
    fn apply_fixups(&mut self) -> Result<(), CompileError> {
        for (fixup_pos, label) in &self.fixups {
            let target = self.labels.get(label)
                .ok_or_else(|| CompileError(format!("Unknown label: {}", label)))?;

            // Calculate relative offset from position after the jump instruction
            let from = fixup_pos + 2; // After the 2-byte offset
            let offset = (*target as isize) - (from as isize);

            if offset < i16::MIN as isize || offset > i16::MAX as isize {
                return Err(CompileError(format!("Jump offset out of range: {}", offset)));
            }

            let offset_bytes = (offset as i16).to_le_bytes();
            self.bytecode[*fixup_pos] = offset_bytes[0];
            self.bytecode[*fixup_pos + 1] = offset_bytes[1];
        }
        Ok(())
    }

    /// Register a function argument
    fn register_arg(&mut self, name: &str) {
        self.arg_offsets.insert(name.to_string(), self.next_arg_offset);
        self.next_arg_offset += 8; // Each arg is 8 bytes (u64)
    }

    /// Get variable location (argument or local)
    fn get_var_location(&self, name: &str) -> Option<VarLocation> {
        // Check arguments first
        if let Some(&offset) = self.arg_offsets.get(name) {
            return Some(VarLocation::InputOffset(offset));
        }
        // Then check locals
        if let Some(&reg) = self.local_registers.get(name) {
            return Some(VarLocation::Register(reg));
        }
        None
    }

    /// Compile an expression (result pushed to stack)
    fn compile_expr(&mut self, expr: &Expr) -> Result<(), CompileError> {
        match expr {
            // Integer literal
            Expr::Lit(lit) => {
                if let Lit::Int(int_lit) = &lit.lit {
                    let value: u64 = int_lit.base10_parse()
                        .map_err(|e| CompileError(format!("Invalid integer: {}", e)))?;

                    if value <= 255 {
                        self.emit_op(stack::PUSH_IMM8);
                        self.emit(value as u8);
                    } else {
                        self.emit_op(stack::PUSH_IMM);
                        self.emit_u64(value);
                    }
                } else if let Lit::Bool(bool_lit) = &lit.lit {
                    self.emit_op(stack::PUSH_IMM8);
                    self.emit(if bool_lit.value { 1 } else { 0 });
                } else {
                    return Err(CompileError("Unsupported literal type".to_string()));
                }
            }

            // Variable reference
            Expr::Path(path) => {
                if path.path.segments.len() == 1 {
                    let name = path.path.segments[0].ident.to_string();

                    match self.get_var_location(&name) {
                        Some(VarLocation::InputOffset(offset)) => {
                            // Read from input buffer (function argument)
                            self.emit_op(native::NATIVE_READ);
                            self.emit_u16(offset as u16);
                        }
                        Some(VarLocation::Register(reg)) => {
                            // Push register value to stack (local variable)
                            self.emit_op(stack::PUSH_REG);
                            self.emit(reg);
                        }
                        None => {
                            return Err(CompileError(format!("Unknown variable: {}", name)));
                        }
                    }
                } else {
                    return Err(CompileError("Complex paths not supported".to_string()));
                }
            }

            // Binary operation
            Expr::Binary(binary) => {
                // Check for compound assignment operators first (+=, -=, etc.)
                // These need special handling since they modify a variable
                match binary.op {
                    BinOp::AddAssign(_) | BinOp::SubAssign(_) | BinOp::MulAssign(_) |
                    BinOp::DivAssign(_) | BinOp::RemAssign(_) | BinOp::BitXorAssign(_) |
                    BinOp::BitAndAssign(_) | BinOp::BitOrAssign(_) | BinOp::ShlAssign(_) |
                    BinOp::ShrAssign(_) => {
                        self.compile_assign_op(&binary.left, &binary.op, &binary.right)?;
                        // Assignment expression produces unit (0)
                        self.emit_op(stack::PUSH_IMM8);
                        self.emit(0);
                        return Ok(());
                    }
                    _ => {}
                }

                // Compile left and right operands first (for all binary operations)
                // Note: && and || in Rust only work with bool, and our comparisons
                // return 0 or 1, so bitwise AND/OR work correctly for boolean logic
                self.compile_expr(&binary.left)?;
                self.compile_expr(&binary.right)?;

                // Emit operation based on operator type
                // Note: ADD, SUB, XOR use MBA-aware emit methods
                match binary.op {
                    BinOp::Add(_) => self.emit_add(),
                    BinOp::Sub(_) => self.emit_sub(),
                    BinOp::Mul(_) => self.emit_op(arithmetic::MUL),
                    BinOp::BitXor(_) => self.emit_xor(),
                    BinOp::BitAnd(_) => self.emit_op(arithmetic::AND),
                    BinOp::BitOr(_) => self.emit_op(arithmetic::OR),
                    BinOp::Shl(_) => self.emit_op(arithmetic::SHL),
                    BinOp::Shr(_) => self.emit_op(arithmetic::SHR),
                    BinOp::Div(_) => self.emit_op(arithmetic::DIV),
                    BinOp::Rem(_) => self.emit_op(arithmetic::MOD),
                    // Logical AND/OR: Since Rust's && and || only work with bool,
                    // and our comparisons return 0 or 1, bitwise ops work correctly
                    BinOp::And(_) => self.emit_op(arithmetic::AND),
                    BinOp::Or(_) => self.emit_op(arithmetic::OR),
                    BinOp::Eq(_) => {
                        // a == b -> (a XOR b) == 0 -> push 1 if zero, else 0
                        self.emit_xor();
                        let else_label = self.unique_label("eq_else");
                        let end_label = self.unique_label("eq_end");
                        self.emit_jump(control::JNZ, &else_label);
                        // True case: push 1
                        self.emit_op(stack::PUSH_IMM8);
                        self.emit(1);
                        self.emit_jump(control::JMP, &end_label);
                        // False case: push 0
                        self.mark_label(&else_label);
                        self.emit_op(stack::PUSH_IMM8);
                        self.emit(0);
                        self.mark_label(&end_label);
                    }
                    BinOp::Ne(_) => {
                        // a != b -> (a XOR b) != 0
                        self.emit_xor();
                        let else_label = self.unique_label("ne_else");
                        let end_label = self.unique_label("ne_end");
                        self.emit_jump(control::JZ, &else_label);
                        // True case: push 1
                        self.emit_op(stack::PUSH_IMM8);
                        self.emit(1);
                        self.emit_jump(control::JMP, &end_label);
                        // False case: push 0
                        self.mark_label(&else_label);
                        self.emit_op(stack::PUSH_IMM8);
                        self.emit(0);
                        self.mark_label(&end_label);
                    }
                    BinOp::Lt(_) => {
                        // a < b: CMP sets flags (and pushes values back), then JLT
                        self.emit_op(control::CMP);
                        self.emit_op(stack::DROP);  // Drop b (CMP pushes back)
                        self.emit_op(stack::DROP);  // Drop a (CMP pushes back)
                        let else_label = self.unique_label("lt_else");
                        let end_label = self.unique_label("lt_end");
                        self.emit_jump(control::JLT, &else_label);
                        // Not less than: push 0
                        self.emit_op(stack::PUSH_IMM8);
                        self.emit(0);
                        self.emit_jump(control::JMP, &end_label);
                        // Less than: push 1
                        self.mark_label(&else_label);
                        self.emit_op(stack::PUSH_IMM8);
                        self.emit(1);
                        self.mark_label(&end_label);
                    }
                    BinOp::Gt(_) => {
                        self.emit_op(control::CMP);
                        self.emit_op(stack::DROP);
                        self.emit_op(stack::DROP);
                        let else_label = self.unique_label("gt_else");
                        let end_label = self.unique_label("gt_end");
                        self.emit_jump(control::JGT, &else_label);
                        self.emit_op(stack::PUSH_IMM8);
                        self.emit(0);
                        self.emit_jump(control::JMP, &end_label);
                        self.mark_label(&else_label);
                        self.emit_op(stack::PUSH_IMM8);
                        self.emit(1);
                        self.mark_label(&end_label);
                    }
                    BinOp::Le(_) => {
                        self.emit_op(control::CMP);
                        self.emit_op(stack::DROP);
                        self.emit_op(stack::DROP);
                        let else_label = self.unique_label("le_else");
                        let end_label = self.unique_label("le_end");
                        self.emit_jump(control::JLE, &else_label);
                        self.emit_op(stack::PUSH_IMM8);
                        self.emit(0);
                        self.emit_jump(control::JMP, &end_label);
                        self.mark_label(&else_label);
                        self.emit_op(stack::PUSH_IMM8);
                        self.emit(1);
                        self.mark_label(&end_label);
                    }
                    BinOp::Ge(_) => {
                        self.emit_op(control::CMP);
                        self.emit_op(stack::DROP);
                        self.emit_op(stack::DROP);
                        let else_label = self.unique_label("ge_else");
                        let end_label = self.unique_label("ge_end");
                        self.emit_jump(control::JGE, &else_label);
                        self.emit_op(stack::PUSH_IMM8);
                        self.emit(0);
                        self.emit_jump(control::JMP, &end_label);
                        self.mark_label(&else_label);
                        self.emit_op(stack::PUSH_IMM8);
                        self.emit(1);
                        self.mark_label(&end_label);
                    }
                    _ => return Err(CompileError(format!("Unsupported binary operator: {:?}", binary.op))),
                }
            }

            // Unary operation
            Expr::Unary(unary) => {
                self.compile_expr(&unary.expr)?;
                match unary.op {
                    UnOp::Not(_) => {
                        // Bitwise NOT (flip all bits)
                        // For integers: !x flips all bits
                        // For booleans: !true = false, !false = true (but bool is 0 or 1)
                        self.emit_op(arithmetic::NOT);
                    }
                    UnOp::Neg(_) => {
                        // Arithmetic negation: -x = 0 - x
                        self.emit_op(stack::PUSH_IMM8);
                        self.emit(0);
                        self.emit_op(stack::SWAP);
                        self.emit_op(arithmetic::SUB);
                    }
                    _ => return Err(CompileError("Unsupported unary operator".to_string())),
                }
            }

            // Parenthesized expression
            Expr::Paren(paren) => {
                self.compile_expr(&paren.expr)?;
            }

            // If-else expression
            Expr::If(if_expr) => {
                // Compile condition
                self.compile_expr(&if_expr.cond)?;

                let else_label = self.unique_label("if_else");
                let end_label = self.unique_label("if_end");

                // TEST: Compare condition with 0 to set zero flag
                // Stack has condition value, we need to set flags based on it
                // DUP + PUSH 0 + XOR sets zero flag if value == 0, then DROP the result
                self.emit_op(stack::DUP);           // [cond, cond]
                self.emit_op(stack::PUSH_IMM8);     // [cond, cond, 0]
                self.emit(0);
                self.emit_op(arithmetic::XOR);      // [cond, cond^0] = [cond, cond], sets ZF if cond==0
                self.emit_op(stack::DROP);          // [cond], ZF still set
                self.emit_op(stack::DROP);          // [], ZF still set (consume condition)

                // Jump to else if condition is false (zero)
                self.emit_jump(control::JZ, &else_label);

                // Compile then branch
                self.compile_block(&if_expr.then_branch)?;

                // Jump to end
                self.emit_jump(control::JMP, &end_label);

                // Else branch
                self.mark_label(&else_label);
                if let Some((_, else_branch)) = &if_expr.else_branch {
                    self.compile_expr(else_branch)?;
                } else {
                    // No else: push 0 as default
                    self.emit_op(stack::PUSH_IMM8);
                    self.emit(0);
                }

                self.mark_label(&end_label);
            }

            // Block expression
            Expr::Block(block) => {
                self.compile_block(&block.block)?;
            }

            // Return expression
            Expr::Return(ret) => {
                if let Some(expr) = &ret.expr {
                    self.compile_expr(expr)?;
                } else {
                    self.emit_op(stack::PUSH_IMM8);
                    self.emit(0);
                }
                self.emit_op(exec::HALT);
            }

            // While loop
            Expr::While(while_expr) => {
                self.compile_while(&while_expr.cond, &while_expr.body)?;
            }

            // Infinite loop
            Expr::Loop(loop_expr) => {
                self.compile_loop(&loop_expr.body)?;
            }

            // For loop (range-based)
            Expr::ForLoop(for_loop) => {
                self.compile_for_loop(for_loop)?;
            }

            // Break (without value)
            Expr::Break(break_expr) => {
                // Check if break has a value - not supported
                if break_expr.expr.is_some() {
                    return Err(CompileError("break with value not supported".to_string()));
                }
                self.compile_break()?;
            }

            // Continue
            Expr::Continue(_) => {
                self.compile_continue()?;
            }

            // Assignment: x = expr or x += expr etc
            Expr::Assign(assign) => {
                self.compile_assignment(&assign.left, &assign.right)?;
                // Assignment pushes 0 (unit) to stack
                self.emit_op(stack::PUSH_IMM8);
                self.emit(0);
            }

            _ => {
                return Err(CompileError(format!(
                    "Unsupported expression type: {:?}",
                    std::any::type_name_of_val(expr)
                )));
            }
        }
        Ok(())
    }

    /// Compile while loop
    fn compile_while(&mut self, cond: &Expr, body: &syn::Block) -> Result<(), CompileError> {
        let loop_start = self.unique_label("while_start");
        let loop_end = self.unique_label("while_end");

        // Mark loop start
        self.mark_label(&loop_start);

        // Push loop context for break/continue
        self.loop_stack.push(LoopContext {
            continue_label: loop_start.clone(),
            break_label: loop_end.clone(),
        });

        // Compile condition
        self.compile_expr(cond)?;

        // Test condition and jump to end if false
        self.emit_op(stack::DUP);
        self.emit_op(stack::PUSH_IMM8);
        self.emit(0);
        self.emit_op(arithmetic::XOR);
        self.emit_op(stack::DROP);
        self.emit_op(stack::DROP);
        self.emit_jump(control::JZ, &loop_end);

        // Compile body as statements (no value left on stack)
        self.compile_block_stmt(body)?;

        // Jump back to condition
        self.emit_jump(control::JMP, &loop_start);

        // Mark loop end
        self.mark_label(&loop_end);

        self.loop_stack.pop();

        // While loop produces unit (0)
        self.emit_op(stack::PUSH_IMM8);
        self.emit(0);

        Ok(())
    }

    /// Compile infinite loop
    fn compile_loop(&mut self, body: &syn::Block) -> Result<(), CompileError> {
        let loop_start = self.unique_label("loop_start");
        let loop_end = self.unique_label("loop_end");

        // Mark loop start
        self.mark_label(&loop_start);

        // Push loop context
        self.loop_stack.push(LoopContext {
            continue_label: loop_start.clone(),
            break_label: loop_end.clone(),
        });

        // Compile body as statements (no value left on stack)
        self.compile_block_stmt(body)?;

        // Jump back to start
        self.emit_jump(control::JMP, &loop_start);

        // Mark loop end (only reachable via break)
        self.mark_label(&loop_end);

        self.loop_stack.pop();

        // Loop produces unit (0) unless break provides value
        self.emit_op(stack::PUSH_IMM8);
        self.emit(0);

        Ok(())
    }

    /// Compile for loop (range-based: for i in start..end)
    fn compile_for_loop(&mut self, for_loop: &syn::ExprForLoop) -> Result<(), CompileError> {
        // Extract loop variable name
        let var_name = if let Pat::Ident(pat_ident) = &*for_loop.pat {
            pat_ident.ident.to_string()
        } else {
            return Err(CompileError("Only simple identifiers supported in for loops".to_string()));
        };

        // Parse range expression
        let (start_expr, end_expr, inclusive) = self.parse_range_expr(&for_loop.expr)?;

        // Allocate register for loop variable
        if self.next_local_reg >= 8 {
            return Err(CompileError("Too many local variables (max 8)".to_string()));
        }
        let loop_reg = self.next_local_reg;
        self.next_local_reg += 1;
        self.local_registers.insert(var_name, loop_reg);

        // Initialize loop variable: start -> Ri
        self.compile_expr(&start_expr)?;
        self.emit_op(stack::POP_REG);
        self.emit(loop_reg);

        let loop_start = self.unique_label("for_start");
        let loop_continue = self.unique_label("for_continue");
        let loop_end = self.unique_label("for_end");

        // Mark loop start
        self.mark_label(&loop_start);

        // Push loop context (continue goes to increment, not condition)
        self.loop_stack.push(LoopContext {
            continue_label: loop_continue.clone(),
            break_label: loop_end.clone(),
        });

        // Condition: Ri < end (or <= for inclusive)
        self.emit_op(stack::PUSH_REG);
        self.emit(loop_reg);
        self.compile_expr(&end_expr)?;
        self.emit_op(control::CMP);
        self.emit_op(stack::DROP);
        self.emit_op(stack::DROP);

        // Jump to end if condition false
        if inclusive {
            self.emit_jump(control::JGT, &loop_end);
        } else {
            self.emit_jump(control::JGE, &loop_end);
        }

        // Compile body as statements (no value left on stack)
        self.compile_block_stmt(&for_loop.body)?;

        // Continue point (increment)
        self.mark_label(&loop_continue);

        // Increment: Ri = Ri + 1
        self.emit_op(stack::PUSH_REG);
        self.emit(loop_reg);
        self.emit_op(arithmetic::INC);
        self.emit_op(stack::POP_REG);
        self.emit(loop_reg);

        // Jump back to condition
        self.emit_jump(control::JMP, &loop_start);

        // Mark loop end
        self.mark_label(&loop_end);

        self.loop_stack.pop();

        // For loop produces unit (0)
        self.emit_op(stack::PUSH_IMM8);
        self.emit(0);

        Ok(())
    }

    /// Parse range expression (start..end or start..=end)
    fn parse_range_expr(&self, expr: &Expr) -> Result<(Expr, Expr, bool), CompileError> {
        match expr {
            Expr::Range(range) => {
                let start = range.start.as_ref()
                    .ok_or_else(|| CompileError("Range must have start".to_string()))?;
                let end = range.end.as_ref()
                    .ok_or_else(|| CompileError("Range must have end".to_string()))?;

                let inclusive = match range.limits {
                    syn::RangeLimits::HalfOpen(_) => false,  // ..
                    syn::RangeLimits::Closed(_) => true,     // ..=
                };

                Ok((*start.clone(), *end.clone(), inclusive))
            }
            _ => Err(CompileError("For loop requires range expression".to_string())),
        }
    }

    /// Compile break statement
    fn compile_break(&mut self) -> Result<(), CompileError> {
        let ctx = self.loop_stack.last()
            .ok_or_else(|| CompileError("break outside of loop".to_string()))?;
        self.emit_jump(control::JMP, &ctx.break_label.clone());
        Ok(())
    }

    /// Compile continue statement
    fn compile_continue(&mut self) -> Result<(), CompileError> {
        let ctx = self.loop_stack.last()
            .ok_or_else(|| CompileError("continue outside of loop".to_string()))?;
        self.emit_jump(control::JMP, &ctx.continue_label.clone());
        Ok(())
    }

    /// Compile simple assignment: x = expr
    fn compile_assignment(&mut self, target: &Expr, value: &Expr) -> Result<(), CompileError> {
        // Compile the value
        self.compile_expr(value)?;

        // Store to target
        if let Expr::Path(path) = target {
            if path.path.segments.len() == 1 {
                let name = path.path.segments[0].ident.to_string();
                match self.get_var_location(&name) {
                    Some(VarLocation::Register(reg)) => {
                        self.emit_op(stack::POP_REG);
                        self.emit(reg);
                    }
                    Some(VarLocation::InputOffset(_)) => {
                        return Err(CompileError("Cannot assign to function argument".to_string()));
                    }
                    None => {
                        return Err(CompileError(format!("Unknown variable: {}", name)));
                    }
                }
            } else {
                return Err(CompileError("Complex paths not supported in assignment".to_string()));
            }
        } else {
            return Err(CompileError("Only simple variable assignment supported".to_string()));
        }

        Ok(())
    }

    /// Compile compound assignment: x += expr, x -= expr, etc
    fn compile_assign_op(&mut self, target: &Expr, op: &BinOp, value: &Expr) -> Result<(), CompileError> {
        // Get target variable
        let (_name, reg) = if let Expr::Path(path) = target {
            if path.path.segments.len() == 1 {
                let name = path.path.segments[0].ident.to_string();
                match self.get_var_location(&name) {
                    Some(VarLocation::Register(reg)) => (name, reg),
                    Some(VarLocation::InputOffset(_)) => {
                        return Err(CompileError("Cannot assign to function argument".to_string()));
                    }
                    None => {
                        return Err(CompileError(format!("Unknown variable: {}", name)));
                    }
                }
            } else {
                return Err(CompileError("Complex paths not supported".to_string()));
            }
        } else {
            return Err(CompileError("Only simple variable assignment supported".to_string()));
        };

        // Push current value
        self.emit_op(stack::PUSH_REG);
        self.emit(reg);

        // Compile RHS
        self.compile_expr(value)?;

        // Apply operation (MBA-aware for ADD, SUB, XOR)
        match op {
            BinOp::AddAssign(_) => self.emit_add(),
            BinOp::SubAssign(_) => self.emit_sub(),
            BinOp::MulAssign(_) => self.emit_op(arithmetic::MUL),
            BinOp::DivAssign(_) => self.emit_op(arithmetic::DIV),
            BinOp::RemAssign(_) => self.emit_op(arithmetic::MOD),
            BinOp::BitXorAssign(_) => self.emit_xor(),
            BinOp::BitAndAssign(_) => self.emit_op(arithmetic::AND),
            BinOp::BitOrAssign(_) => self.emit_op(arithmetic::OR),
            BinOp::ShlAssign(_) => self.emit_op(arithmetic::SHL),
            BinOp::ShrAssign(_) => self.emit_op(arithmetic::SHR),
            _ => return Err(CompileError(format!("Unsupported assignment operator: {:?}", op))),
        }

        // Store result back
        self.emit_op(stack::POP_REG);
        self.emit(reg);

        Ok(())
    }

    /// Compile a block as statements (for loop bodies)
    /// Unlike compile_block, this drops ALL expression results - no value left on stack
    fn compile_block_stmt(&mut self, block: &syn::Block) -> Result<(), CompileError> {
        for stmt in &block.stmts {
            match stmt {
                Stmt::Expr(expr, _) => {
                    self.compile_expr(expr)?;
                    // Always drop expression result - statements don't produce values
                    self.emit_op(stack::DROP);
                }
                Stmt::Local(local) => {
                    self.compile_local(local)?;
                }
                _ => return Err(CompileError("Unsupported statement type".to_string())),
            }
        }
        Ok(())
    }

    /// Compile a block as expression (returns value on stack)
    /// Used for if/else bodies and block expressions
    fn compile_block(&mut self, block: &syn::Block) -> Result<(), CompileError> {
        let stmts = &block.stmts;
        let len = stmts.len();

        // Empty block produces unit (0)
        if len == 0 {
            self.emit_op(stack::PUSH_IMM8);
            self.emit(0);
            return Ok(());
        }

        for (idx, stmt) in stmts.iter().enumerate() {
            let is_last = idx == len - 1;

            match stmt {
                Stmt::Expr(expr, semi) => {
                    self.compile_expr(expr)?;

                    if is_last && semi.is_none() {
                        // Last expression without semicolon - keep value on stack
                    } else {
                        // Not last, or has semicolon - drop result
                        self.emit_op(stack::DROP);

                        // If this was the last statement WITH semicolon, push unit
                        if is_last {
                            self.emit_op(stack::PUSH_IMM8);
                            self.emit(0);
                        }
                    }
                }
                Stmt::Local(local) => {
                    self.compile_local(local)?;

                    // If local is last statement, push unit
                    if is_last {
                        self.emit_op(stack::PUSH_IMM8);
                        self.emit(0);
                    }
                }
                _ => return Err(CompileError("Unsupported statement type".to_string())),
            }
        }
        Ok(())
    }

    /// Compile a local variable binding: let x = expr;
    fn compile_local(&mut self, local: &syn::Local) -> Result<(), CompileError> {
        if let Some(init) = &local.init {
            // Compile initializer
            self.compile_expr(&init.expr)?;

            // Extract variable name from pattern
            // Supports: let x = ..., let mut x = ..., let x: Type = ..., let mut x: Type = ...
            let name = Self::extract_pat_name(&local.pat)?;

            // Check if variable already exists (shadowing or loop reuse)
            let reg_idx = if let Some(&existing_reg) = self.local_registers.get(&name) {
                // Reuse existing register for same-named variable
                // This handles: loop { let x = ...; } where x is redefined each iteration
                existing_reg
            } else {
                // Allocate new register
                if self.next_local_reg >= 8 {
                    return Err(CompileError("Too many local variables (max 8)".to_string()));
                }
                let reg = self.next_local_reg;
                self.next_local_reg += 1;
                self.local_registers.insert(name, reg);
                reg
            };

            // Pop value from stack to register
            self.emit_op(stack::POP_REG);
            self.emit(reg_idx);
        } else {
            return Err(CompileError("Uninitialized let bindings not supported".to_string()));
        }
        Ok(())
    }

    /// Extract variable name from a pattern
    /// Handles: Pat::Ident (let x), Pat::Type (let x: T)
    fn extract_pat_name(pat: &Pat) -> Result<String, CompileError> {
        match pat {
            Pat::Ident(pat_ident) => {
                Ok(pat_ident.ident.to_string())
            }
            Pat::Type(pat_type) => {
                // let x: T = ... -> extract name from inner pattern
                Self::extract_pat_name(&pat_type.pat)
            }
            _ => Err(CompileError("Unsupported pattern in let binding".to_string())),
        }
    }

    /// Finalize compilation
    fn finalize(&mut self) -> Result<Vec<u8>, CompileError> {
        // IMPORTANT: Apply jump fixups FIRST, before adding HALT
        // This ensures HALT doesn't shift any jump targets
        self.apply_fixups()?;

        // Then ensure we end with HALT if not already
        if self.bytecode.is_empty() || self.bytecode.last().copied() != Some(exec::HALT) {
            self.emit_op(exec::HALT);
        }

        Ok(self.bytecode.clone())
    }
}

/// Compile a function to bytecode (without MBA)
pub fn compile_function(func: &ItemFn) -> Result<Vec<u8>, CompileError> {
    compile_function_with_options(func, false)
}

/// Compile a function to bytecode with MBA transformations
pub fn compile_function_with_mba(func: &ItemFn) -> Result<Vec<u8>, CompileError> {
    compile_function_with_options(func, true)
}

/// Compile a function to bytecode with configurable options
fn compile_function_with_options(func: &ItemFn, mba_enabled: bool) -> Result<Vec<u8>, CompileError> {
    let mut compiler = Compiler::with_mba(mba_enabled);

    // Register function arguments (deterministic order from syn)
    for arg in &func.sig.inputs {
        if let FnArg::Typed(pat_type) = arg {
            if let Pat::Ident(pat_ident) = &*pat_type.pat {
                compiler.register_arg(&pat_ident.ident.to_string());
            }
        }
    }

    // Compile function body
    compiler.compile_block(&func.block)?;

    // Finalize and return bytecode
    compiler.finalize()
}