oxigdal-algorithms 0.1.6

//! Bytecode compiler and stack-machine VM for raster band-math expressions.
//!
//! # Overview
//!
//! The tree-walking [`Evaluator`] re-interprets the `Expr` AST for every pixel,
//! which is expensive for million-pixel rasters.  This module provides a two-step
//! alternative:
//!
//! 1. **Compile once**: [`compile_program`] traverses the AST exactly once and
//!    emits a flat `Vec<OpCode>` together with a constant table.
//! 2. **Execute per-pixel**: [`eval_bytecode`] interprets the opcode stream on a
//!    tiny stack (`Vec<f64>`) that is pre-allocated outside the pixel loop.
//!
//! The resulting throughput improvement is typically 2–5× for real-world
//! band-math expressions (NDVI, EVI, SAVI, …).
//!
//! # Band indexing
//!
//! Following the existing convention the AST node `Expr::Band(b)` is **1-indexed**
//! (B1 ↔ index 1).  `LoadBand(i)` therefore stores the 1-based index unchanged
//! and `eval_bytecode` subtracts one when indexing into `bands`.

use super::ast::{BinaryOp, Expr, UnaryOp};
use crate::error::{AlgorithmError, Result};
use oxigdal_core::buffer::RasterBuffer;

// ─────────────────────────────────────────────────────────────────────────────
// OpCode
// ─────────────────────────────────────────────────────────────────────────────

/// A single instruction in the band-math VM.
///
/// The VM is a pure stack machine:  operands are pushed onto the stack and
/// operations pop their arguments and push one result.
#[derive(Debug, Clone, PartialEq)]
pub enum OpCode {
    // ── Stack load / store ────────────────────────────────────────────────
    /// Push `constants[idx]` onto the stack.
    LoadConst(u16),
    /// Push `bands[idx - 1][pixel_idx]` (1-based band index).
    LoadBand(u16),
    /// Push the pre-computed CSE cache value at `cache[idx]`.
    LoadCacheSlot(u16),
    /// Pop the top of the stack and store it in `cache[idx]`.
    StoreCacheSlot(u16),

    // ── Arithmetic (binary: pop rhs, pop lhs, push result) ───────────────
    /// Floating-point addition.
    Add,
    /// Floating-point subtraction.
    Sub,
    /// Floating-point multiplication.
    Mul,
    /// Floating-point division (returns NaN on divide-by-zero).
    Div,
    /// Power: lhs.powf(rhs).
    Pow,

    // ── Unary (pop one value, push result) ───────────────────────────────
    /// Arithmetic negation.
    Neg,
    /// Absolute value.
    Abs,
    /// Square root.
    Sqrt,
    /// Log base 10 (`log10`).
    Log,
    /// Natural logarithm (`ln`).
    Ln,
    /// Exponential (`e^x`).
    Exp,
    /// Sine.
    Sin,
    /// Cosine.
    Cos,
    /// Tangent.
    Tan,
    /// Floor (round towards negative infinity).
    Floor,
    /// Ceiling (round towards positive infinity).
    Ceil,
    /// Round to nearest integer.
    Round,

    // ── Multi-argument (pops arguments right-to-left) ────────────────────
    /// Pop b, pop a, push min(a, b).
    Min,
    /// Pop b, pop a, push max(a, b).
    Max,
    /// Pop max_v, pop min_v, pop val; push val.clamp(min_v, max_v).
    Clamp,

    // ── Comparisons (binary → 1.0 or 0.0) ────────────────────────────────
    /// lhs > rhs  → 1.0 else 0.0
    Gt,
    /// lhs < rhs  → 1.0 else 0.0
    Lt,
    /// lhs >= rhs → 1.0 else 0.0
    Gte,
    /// lhs <= rhs → 1.0 else 0.0
    Lte,
    /// |lhs - rhs| < ε → 1.0 else 0.0
    Eq,
    /// |lhs - rhs| >= ε → 1.0 else 0.0
    Ne,

    // ── Logic (binary → 1.0 or 0.0) ──────────────────────────────────────
    /// Both operands non-zero → 1.0 else 0.0
    And,
    /// At least one operand non-zero → 1.0 else 0.0
    Or,

    // ── Conditional ───────────────────────────────────────────────────────
    /// Pop else_val (top), pop then_val, pop cond.
    /// Push then_val if cond != 0.0, else else_val.
    Cond,
}

// ─────────────────────────────────────────────────────────────────────────────
// CompiledProgram
// ─────────────────────────────────────────────────────────────────────────────

/// A compiled band-math program, ready for repeated per-pixel execution.
#[derive(Debug, Clone)]
pub struct CompiledProgram {
    /// The flat opcode stream.
    pub ops: Vec<OpCode>,
    /// Constant literals referenced by `LoadConst` instructions.
    pub constants: Vec<f64>,
    /// Number of input bands the program requires (0-based upper bound
    /// derived from the highest `Band(b)` seen; 0 means no bands used).
    pub required_bands: usize,
    /// Number of CSE cache slots required.
    pub cache_slot_count: usize,
    /// Statically estimated maximum stack depth needed.
    pub estimated_stack_depth: usize,
}

// ─────────────────────────────────────────────────────────────────────────────
// compile_expr  (internal helper – populates a shared constant table)
// ─────────────────────────────────────────────────────────────────────────────

/// Compile a single `Expr` node into a sequence of `OpCode`s using post-order
/// (operand-before-operator) emission.
///
/// `constants` is the shared literal table; existing constants are reused to
/// keep `LoadConst` indices compact.
///
/// Note: this function is `pub(crate)` because `Expr` is private to the
/// `calculator` module family; external callers must use
/// [`RasterCalculator::evaluate_bytecode`] or [`compile_program`] via string
/// parsing.
///
/// # Errors
///
/// Returns [`AlgorithmError::InvalidInput`] if the expression contains an
/// unknown function name.
pub(super) fn compile_expr(expr: &Expr, constants: &mut Vec<f64>) -> Result<Vec<OpCode>> {
    let mut ops: Vec<OpCode> = Vec::new();
    compile_expr_into(expr, constants, &mut ops)?;
    Ok(ops)
}

/// Recursive helper that appends opcodes into `out` rather than allocating a
/// new Vec for every sub-expression (avoids quadratic allocation).
fn compile_expr_into(expr: &Expr, constants: &mut Vec<f64>, out: &mut Vec<OpCode>) -> Result<()> {
    match expr {
        // ── Leaf nodes ───────────────────────────────────────────────────
        Expr::Number(v) => {
            // Deduplicate constants by bit-level equality.
            let idx = constants
                .iter()
                .position(|c| c.to_bits() == v.to_bits())
                .unwrap_or_else(|| {
                    let i = constants.len();
                    constants.push(*v);
                    i
                });
            let idx_u16 = u16::try_from(idx).map_err(|_| AlgorithmError::InvalidParameter {
                parameter: "constants",
                message: format!("Too many constants (max {})", u16::MAX),
            })?;
            out.push(OpCode::LoadConst(idx_u16));
        }

        Expr::Band(b) => {
            let b_u16 = u16::try_from(*b).map_err(|_| AlgorithmError::InvalidParameter {
                parameter: "band",
                message: format!("Band index {} exceeds u16::MAX", b),
            })?;
            out.push(OpCode::LoadBand(b_u16));
        }

        Expr::CacheRef(idx) => {
            let idx_u16 = u16::try_from(*idx).map_err(|_| AlgorithmError::InvalidParameter {
                parameter: "cache_slot",
                message: format!("Cache slot index {} exceeds u16::MAX", idx),
            })?;
            out.push(OpCode::LoadCacheSlot(idx_u16));
        }

        // ── Binary operations ─────────────────────────────────────────────
        Expr::BinaryOp { left, op, right } => {
            compile_expr_into(left, constants, out)?;
            compile_expr_into(right, constants, out)?;
            let opcode = match op {
                BinaryOp::Add => OpCode::Add,
                BinaryOp::Subtract => OpCode::Sub,
                BinaryOp::Multiply => OpCode::Mul,
                BinaryOp::Divide => OpCode::Div,
                BinaryOp::Power => OpCode::Pow,
                BinaryOp::Greater => OpCode::Gt,
                BinaryOp::Less => OpCode::Lt,
                BinaryOp::GreaterEqual => OpCode::Gte,
                BinaryOp::LessEqual => OpCode::Lte,
                BinaryOp::Equal => OpCode::Eq,
                BinaryOp::NotEqual => OpCode::Ne,
                BinaryOp::And => OpCode::And,
                BinaryOp::Or => OpCode::Or,
            };
            out.push(opcode);
        }

        // ── Unary operations ──────────────────────────────────────────────
        Expr::UnaryOp { op, expr: inner } => {
            compile_expr_into(inner, constants, out)?;
            let opcode = match op {
                UnaryOp::Negate => OpCode::Neg,
            };
            out.push(opcode);
        }

        // ── Function calls ────────────────────────────────────────────────
        Expr::Function { name, args } => {
            compile_function(name, args, constants, out)?;
        }

        // ── Conditional (if/then/else) ────────────────────────────────────
        Expr::Conditional {
            condition,
            then_expr,
            else_expr,
        } => {
            // Stack layout: [... cond then_val else_val]  → Cond pops all 3.
            compile_expr_into(condition, constants, out)?;
            compile_expr_into(then_expr, constants, out)?;
            compile_expr_into(else_expr, constants, out)?;
            out.push(OpCode::Cond);
        }
    }
    Ok(())
}

/// Compile a function call into the opcode stream.
///
/// Function arguments are compiled in order (left-to-right push).  Multi-arg
/// opcodes (`Min`, `Max`, `Clamp`) pop their arguments in right-to-left order.
fn compile_function(
    name: &str,
    args: &[Expr],
    constants: &mut Vec<f64>,
    out: &mut Vec<OpCode>,
) -> Result<()> {
    // Helper: verify exact argument count.
    let check_arity = |expected: usize| -> Result<()> {
        if args.len() != expected {
            Err(AlgorithmError::InvalidInput(format!(
                "Function '{}': expected {} argument(s), got {}",
                name,
                expected,
                args.len()
            )))
        } else {
            Ok(())
        }
    };

    match name {
        // ── 1-arg functions ───────────────────────────────────────────────
        "sqrt" => {
            check_arity(1)?;
            compile_expr_into(&args[0], constants, out)?;
            out.push(OpCode::Sqrt);
        }
        "abs" => {
            check_arity(1)?;
            compile_expr_into(&args[0], constants, out)?;
            out.push(OpCode::Abs);
        }
        // "log" maps to natural logarithm (matching evaluator.rs behaviour)
        "log" => {
            check_arity(1)?;
            compile_expr_into(&args[0], constants, out)?;
            out.push(OpCode::Ln);
        }
        // "ln" explicit alias
        "ln" => {
            check_arity(1)?;
            compile_expr_into(&args[0], constants, out)?;
            out.push(OpCode::Ln);
        }
        // "log10" maps to base-10 logarithm
        "log10" => {
            check_arity(1)?;
            compile_expr_into(&args[0], constants, out)?;
            out.push(OpCode::Log);
        }
        "exp" => {
            check_arity(1)?;
            compile_expr_into(&args[0], constants, out)?;
            out.push(OpCode::Exp);
        }
        "sin" => {
            check_arity(1)?;
            compile_expr_into(&args[0], constants, out)?;
            out.push(OpCode::Sin);
        }
        "cos" => {
            check_arity(1)?;
            compile_expr_into(&args[0], constants, out)?;
            out.push(OpCode::Cos);
        }
        "tan" => {
            check_arity(1)?;
            compile_expr_into(&args[0], constants, out)?;
            out.push(OpCode::Tan);
        }
        "floor" => {
            check_arity(1)?;
            compile_expr_into(&args[0], constants, out)?;
            out.push(OpCode::Floor);
        }
        "ceil" => {
            check_arity(1)?;
            compile_expr_into(&args[0], constants, out)?;
            out.push(OpCode::Ceil);
        }
        "round" => {
            check_arity(1)?;
            compile_expr_into(&args[0], constants, out)?;
            out.push(OpCode::Round);
        }
        // ── 2-arg functions ───────────────────────────────────────────────
        "min" => {
            check_arity(2)?;
            compile_expr_into(&args[0], constants, out)?;
            compile_expr_into(&args[1], constants, out)?;
            out.push(OpCode::Min);
        }
        "max" => {
            check_arity(2)?;
            compile_expr_into(&args[0], constants, out)?;
            compile_expr_into(&args[1], constants, out)?;
            out.push(OpCode::Max);
        }
        // ── 3-arg functions ───────────────────────────────────────────────
        "clamp" => {
            check_arity(3)?;
            compile_expr_into(&args[0], constants, out)?;
            compile_expr_into(&args[1], constants, out)?;
            compile_expr_into(&args[2], constants, out)?;
            out.push(OpCode::Clamp);
        }
        // ── Unknown ───────────────────────────────────────────────────────
        unknown => {
            return Err(AlgorithmError::InvalidInput(format!(
                "Unknown function: {unknown}"
            )));
        }
    }
    Ok(())
}

// ─────────────────────────────────────────────────────────────────────────────
// compile_program
// ─────────────────────────────────────────────────────────────────────────────

/// Compile a full band-math program, optionally pre-computing CSE cache slots.
///
/// `cache_slots` should be the slice returned by the optimizer; each entry is
/// the canonical sub-expression for that slot.  The emitted prelude evaluates
/// each slot in order and stores it via `StoreCacheSlot(i)`.
///
/// Note: `pub(crate)` because `Expr` is private to the `calculator` module
/// family.  External code must go through
/// [`RasterCalculator::evaluate_bytecode`] which parses, optimises, and
/// compiles internally.
///
/// # Errors
///
/// Propagates any [`AlgorithmError`] from [`compile_expr`].
pub(super) fn compile_program(expr: &Expr, cache_slots: &[Expr]) -> Result<CompiledProgram> {
    let mut constants: Vec<f64> = Vec::new();
    let mut ops: Vec<OpCode> = Vec::new();

    // Emit prelude: evaluate each CSE slot and store it.
    for (i, slot_expr) in cache_slots.iter().enumerate() {
        compile_expr_into(slot_expr, &mut constants, &mut ops)?;
        let idx_u16 = u16::try_from(i).map_err(|_| AlgorithmError::InvalidParameter {
            parameter: "cache_slots",
            message: format!("Too many CSE cache slots (max {})", u16::MAX),
        })?;
        ops.push(OpCode::StoreCacheSlot(idx_u16));
    }

    // Emit the main expression.
    compile_expr_into(expr, &mut constants, &mut ops)?;

    // Determine the highest band index referenced.
    let required_bands = highest_band_index(&ops);
    let estimated_stack_depth = estimate_stack_depth(&ops);

    Ok(CompiledProgram {
        ops,
        constants,
        required_bands,
        cache_slot_count: cache_slots.len(),
        estimated_stack_depth,
    })
}

/// Walk the opcode stream to find the maximum 1-based band index referenced.
/// Returns 0 if no `LoadBand` instructions are present.
fn highest_band_index(ops: &[OpCode]) -> usize {
    ops.iter().fold(0usize, |acc, op| {
        if let OpCode::LoadBand(b) = op {
            acc.max(*b as usize)
        } else {
            acc
        }
    })
}

// ─────────────────────────────────────────────────────────────────────────────
// estimate_stack_depth
// ─────────────────────────────────────────────────────────────────────────────

/// Statically estimate the maximum stack depth required to execute `ops`.
///
/// Each opcode is modelled as a pure delta on the stack pointer.  The function
/// returns the maximum depth seen.
pub fn estimate_stack_depth(ops: &[OpCode]) -> usize {
    let mut depth: isize = 0;
    let mut max_depth: isize = 0;

    for op in ops {
        let delta: isize = stack_delta(op);
        depth += delta;
        if depth > max_depth {
            max_depth = depth;
        }
    }

    max_depth.max(0) as usize
}

/// Returns the net stack depth change (+N = N pushes, -N = N pops net).
fn stack_delta(op: &OpCode) -> isize {
    match op {
        // Push 1
        OpCode::LoadConst(_) | OpCode::LoadBand(_) | OpCode::LoadCacheSlot(_) => 1,
        // Pop 1, push 0 (net -1)
        OpCode::StoreCacheSlot(_) => -1,
        // Pop 2, push 1 (net -1)
        OpCode::Add
        | OpCode::Sub
        | OpCode::Mul
        | OpCode::Div
        | OpCode::Pow
        | OpCode::Gt
        | OpCode::Lt
        | OpCode::Gte
        | OpCode::Lte
        | OpCode::Eq
        | OpCode::Ne
        | OpCode::And
        | OpCode::Or
        | OpCode::Min
        | OpCode::Max => -1,
        // Pop 1, push 1 (net 0)
        OpCode::Neg
        | OpCode::Abs
        | OpCode::Sqrt
        | OpCode::Log
        | OpCode::Ln
        | OpCode::Exp
        | OpCode::Sin
        | OpCode::Cos
        | OpCode::Tan
        | OpCode::Floor
        | OpCode::Ceil
        | OpCode::Round => 0,
        // Pop 3 (val, min, max), push 1 (net -2)
        OpCode::Clamp => -2,
        // Pop 3 (cond, then_val, else_val), push 1 (net -2)
        OpCode::Cond => -2,
    }
}

// ─────────────────────────────────────────────────────────────────────────────
// eval_bytecode
// ─────────────────────────────────────────────────────────────────────────────

/// Execute the bytecode program for a single pixel.
///
/// # Arguments
///
/// * `prog`      – The compiled program (constant, shared across all pixels).
/// * `bands`     – A slice of band data slices, 0-indexed (`bands[0]` = B1).
/// * `pixel_idx` – Linear index into each band slice for the current pixel.
/// * `cache`     – Pre-allocated `Vec<f64>` with `prog.cache_slot_count` entries.
///                 Values are populated by `StoreCacheSlot` during the prelude.
/// * `stack`     – Pre-allocated scratch `Vec<f64>`; cleared at the start of
///                 each call.  Capacity should be `>= prog.estimated_stack_depth`.
///
/// # Errors
///
/// Returns [`AlgorithmError`] on stack underflow or band out-of-range.  Division
/// by zero and invalid power operations return `f64::NAN` (not an error).
pub fn eval_bytecode(
    prog: &CompiledProgram,
    bands: &[&[f64]],
    pixel_idx: usize,
    cache: &mut [f64],
    stack: &mut Vec<f64>,
) -> Result<f64> {
    stack.clear();

    for op in &prog.ops {
        match op {
            // ── Load / store ─────────────────────────────────────────────
            OpCode::LoadConst(i) => {
                let v = prog.constants.get(*i as usize).ok_or_else(|| {
                    AlgorithmError::InvalidParameter {
                        parameter: "LoadConst",
                        message: format!(
                            "Constant index {} out of range (len={})",
                            i,
                            prog.constants.len()
                        ),
                    }
                })?;
                stack.push(*v);
            }

            OpCode::LoadBand(b) => {
                // b is 1-based.
                let band_idx = (*b as usize).checked_sub(1).ok_or_else(|| {
                    AlgorithmError::InvalidParameter {
                        parameter: "LoadBand",
                        message: "Band index 0 is invalid (bands are 1-indexed)".to_string(),
                    }
                })?;
                let band_data =
                    bands
                        .get(band_idx)
                        .ok_or_else(|| AlgorithmError::InvalidParameter {
                            parameter: "LoadBand",
                            message: format!(
                                "Band {} out of range (have {} bands)",
                                b,
                                bands.len()
                            ),
                        })?;
                let v =
                    band_data
                        .get(pixel_idx)
                        .ok_or_else(|| AlgorithmError::InvalidParameter {
                            parameter: "LoadBand",
                            message: format!(
                                "Pixel index {} out of range for band {}",
                                pixel_idx, b
                            ),
                        })?;
                stack.push(*v);
            }

            OpCode::LoadCacheSlot(i) => {
                let v = cache
                    .get(*i as usize)
                    .ok_or_else(|| AlgorithmError::InvalidParameter {
                        parameter: "LoadCacheSlot",
                        message: format!("Cache slot {} out of range (len={})", i, cache.len()),
                    })?;
                stack.push(*v);
            }

            OpCode::StoreCacheSlot(i) => {
                let v = stack_pop(stack)?;
                let slot_idx = *i as usize;
                let cache_len = cache.len();
                let slot =
                    cache
                        .get_mut(slot_idx)
                        .ok_or_else(|| AlgorithmError::InvalidParameter {
                            parameter: "StoreCacheSlot",
                            message: format!(
                                "Cache slot {} out of range (len={})",
                                slot_idx, cache_len
                            ),
                        })?;
                *slot = v;
            }

            // ── Binary arithmetic ─────────────────────────────────────────
            OpCode::Add => {
                let rhs = stack_pop(stack)?;
                let lhs = stack_pop(stack)?;
                stack.push(lhs + rhs);
            }
            OpCode::Sub => {
                let rhs = stack_pop(stack)?;
                let lhs = stack_pop(stack)?;
                stack.push(lhs - rhs);
            }
            OpCode::Mul => {
                let rhs = stack_pop(stack)?;
                let lhs = stack_pop(stack)?;
                stack.push(lhs * rhs);
            }
            OpCode::Div => {
                let rhs = stack_pop(stack)?;
                let lhs = stack_pop(stack)?;
                // Division by zero → NaN (matching tree-walk evaluator).
                if rhs.abs() < f64::EPSILON {
                    stack.push(f64::NAN);
                } else {
                    stack.push(lhs / rhs);
                }
            }
            OpCode::Pow => {
                let rhs = stack_pop(stack)?;
                let lhs = stack_pop(stack)?;
                stack.push(lhs.powf(rhs));
            }

            // ── Unary ─────────────────────────────────────────────────────
            OpCode::Neg => {
                let v = stack_pop(stack)?;
                stack.push(-v);
            }
            OpCode::Abs => {
                let v = stack_pop(stack)?;
                stack.push(v.abs());
            }
            OpCode::Sqrt => {
                let v = stack_pop(stack)?;
                stack.push(v.sqrt());
            }
            OpCode::Log => {
                let v = stack_pop(stack)?;
                stack.push(v.log10());
            }
            OpCode::Ln => {
                let v = stack_pop(stack)?;
                stack.push(v.ln());
            }
            OpCode::Exp => {
                let v = stack_pop(stack)?;
                stack.push(v.exp());
            }
            OpCode::Sin => {
                let v = stack_pop(stack)?;
                stack.push(v.sin());
            }
            OpCode::Cos => {
                let v = stack_pop(stack)?;
                stack.push(v.cos());
            }
            OpCode::Tan => {
                let v = stack_pop(stack)?;
                stack.push(v.tan());
            }
            OpCode::Floor => {
                let v = stack_pop(stack)?;
                stack.push(v.floor());
            }
            OpCode::Ceil => {
                let v = stack_pop(stack)?;
                stack.push(v.ceil());
            }
            OpCode::Round => {
                let v = stack_pop(stack)?;
                stack.push(v.round());
            }

            // ── Multi-arg ─────────────────────────────────────────────────
            OpCode::Min => {
                let b = stack_pop(stack)?;
                let a = stack_pop(stack)?;
                stack.push(a.min(b));
            }
            OpCode::Max => {
                let b = stack_pop(stack)?;
                let a = stack_pop(stack)?;
                stack.push(a.max(b));
            }
            OpCode::Clamp => {
                let max_v = stack_pop(stack)?;
                let min_v = stack_pop(stack)?;
                let val = stack_pop(stack)?;
                stack.push(val.clamp(min_v, max_v));
            }

            // ── Comparisons ───────────────────────────────────────────────
            OpCode::Gt => {
                let rhs = stack_pop(stack)?;
                let lhs = stack_pop(stack)?;
                stack.push(if lhs > rhs { 1.0 } else { 0.0 });
            }
            OpCode::Lt => {
                let rhs = stack_pop(stack)?;
                let lhs = stack_pop(stack)?;
                stack.push(if lhs < rhs { 1.0 } else { 0.0 });
            }
            OpCode::Gte => {
                let rhs = stack_pop(stack)?;
                let lhs = stack_pop(stack)?;
                stack.push(if lhs >= rhs { 1.0 } else { 0.0 });
            }
            OpCode::Lte => {
                let rhs = stack_pop(stack)?;
                let lhs = stack_pop(stack)?;
                stack.push(if lhs <= rhs { 1.0 } else { 0.0 });
            }
            OpCode::Eq => {
                let rhs = stack_pop(stack)?;
                let lhs = stack_pop(stack)?;
                stack.push(if (lhs - rhs).abs() < f64::EPSILON {
                    1.0
                } else {
                    0.0
                });
            }
            OpCode::Ne => {
                let rhs = stack_pop(stack)?;
                let lhs = stack_pop(stack)?;
                stack.push(if (lhs - rhs).abs() >= f64::EPSILON {
                    1.0
                } else {
                    0.0
                });
            }

            // ── Logic ─────────────────────────────────────────────────────
            OpCode::And => {
                let rhs = stack_pop(stack)?;
                let lhs = stack_pop(stack)?;
                stack.push(if lhs != 0.0 && rhs != 0.0 { 1.0 } else { 0.0 });
            }
            OpCode::Or => {
                let rhs = stack_pop(stack)?;
                let lhs = stack_pop(stack)?;
                stack.push(if lhs != 0.0 || rhs != 0.0 { 1.0 } else { 0.0 });
            }

            // ── Conditional ───────────────────────────────────────────────
            OpCode::Cond => {
                // Stack order (top-of-stack last):  [..., cond, then_val, else_val]
                let else_val = stack_pop(stack)?;
                let then_val = stack_pop(stack)?;
                let cond = stack_pop(stack)?;
                stack.push(if cond != 0.0 { then_val } else { else_val });
            }
        }
    }

    // After execution the stack must contain exactly one value.
    if stack.len() != 1 {
        return Err(AlgorithmError::ComputationError(format!(
            "eval_bytecode: expected stack depth 1 after execution, got {}",
            stack.len()
        )));
    }
    stack_pop(stack)
}

/// Pop the top of `stack`, returning a stack-underflow error on empty stack.
#[inline]
fn stack_pop(stack: &mut Vec<f64>) -> Result<f64> {
    stack.pop().ok_or_else(|| {
        AlgorithmError::ComputationError("eval_bytecode: stack underflow".to_string())
    })
}

// ─────────────────────────────────────────────────────────────────────────────
// RasterCalculator::evaluate_bytecode  (impl extension)
// ─────────────────────────────────────────────────────────────────────────────

use super::ops::RasterCalculator;
use super::{lexer::Lexer, optimizer::Optimizer, parser::Parser};

impl RasterCalculator {
    /// Evaluate a band-math expression using the bytecode VM.
    ///
    /// Parsing and compilation happen once; the resulting [`CompiledProgram`] is
    /// then executed for every pixel with a pair of pre-allocated scratch buffers
    /// (`cache` and `stack`) so no heap allocations occur inside the pixel loop.
    ///
    /// This method is functionally equivalent to [`RasterCalculator::evaluate`]
    /// but can be 2–5× faster on large rasters because it avoids AST tree
    /// traversal per pixel.
    ///
    /// # Arguments
    ///
    /// * `expr_str` – The expression string (e.g. `"(B1 - B2) / (B1 + B2)"`).
    /// * `bands`    – Input `RasterBuffer` slices; B1 = bands\[0\], B2 = bands\[1\], …
    ///
    /// # Errors
    ///
    /// Returns an [`AlgorithmError`] if:
    /// - `bands` is empty ([`AlgorithmError::EmptyInput`]),
    /// - band dimensions differ ([`AlgorithmError::InvalidDimensions`]),
    /// - the expression cannot be parsed or compiled.
    pub fn evaluate_bytecode(expr_str: &str, bands: &[RasterBuffer]) -> Result<RasterBuffer> {
        if bands.is_empty() {
            return Err(AlgorithmError::EmptyInput {
                operation: "evaluate_bytecode",
            });
        }

        // Validate that all bands share the same dimensions.
        let width = bands[0].width();
        let height = bands[0].height();
        for band in bands.iter().skip(1) {
            if band.width() != width || band.height() != height {
                return Err(AlgorithmError::InvalidDimensions {
                    message: "All bands must have same dimensions",
                    actual: band.width() as usize,
                    expected: width as usize,
                });
            }
        }

        // ── Parse ────────────────────────────────────────────────────────
        let mut lexer = Lexer::new(expr_str);
        let tokens = lexer.tokenize()?;
        let mut parser = Parser::new(tokens);
        let raw_expr = parser.parse()?;

        // ── Optimize (CSE) ───────────────────────────────────────────────
        let (expr, cache_slots) = Optimizer::optimize(raw_expr);

        // ── Compile ──────────────────────────────────────────────────────
        let prog = compile_program(&expr, &cache_slots)?;

        // ── Extract flat pixel data from each RasterBuffer ───────────────
        // We collect per-band pixel slices once (avoids repeated get_pixel calls).
        let num_pixels = (width as usize) * (height as usize);
        let band_data: Vec<Vec<f64>> = bands
            .iter()
            .map(|b| collect_band_pixels(b, width, height))
            .collect();
        let band_slices: Vec<&[f64]> = band_data.iter().map(|v| v.as_slice()).collect();

        // Validate that the expression doesn't reference more bands than supplied.
        if prog.required_bands > bands.len() {
            return Err(AlgorithmError::InvalidParameter {
                parameter: "band",
                message: format!(
                    "Expression references band {} but only {} band(s) provided",
                    prog.required_bands,
                    bands.len()
                ),
            });
        }

        // ── Pre-allocate scratch buffers (once for the whole raster) ─────
        let mut cache = vec![0.0f64; prog.cache_slot_count];
        let mut vm_stack: Vec<f64> = Vec::with_capacity(prog.estimated_stack_depth.max(8));

        let mut result = RasterBuffer::zeros(width, height, bands[0].data_type());

        for pixel_idx in 0..num_pixels {
            // Reset the CSE cache for each pixel (each pixel is independent).
            for slot in cache.iter_mut() {
                *slot = 0.0;
            }

            let value = eval_bytecode(&prog, &band_slices, pixel_idx, &mut cache, &mut vm_stack)?;

            let x = (pixel_idx % width as usize) as u64;
            let y = (pixel_idx / width as usize) as u64;
            result
                .set_pixel(x, y, value)
                .map_err(AlgorithmError::Core)?;
        }

        Ok(result)
    }
}

/// Collect all pixels from a `RasterBuffer` into a flat `Vec<f64>` in
/// row-major order (y-outer, x-inner).
fn collect_band_pixels(band: &RasterBuffer, width: u64, height: u64) -> Vec<f64> {
    let mut data = Vec::with_capacity((width * height) as usize);
    for y in 0..height {
        for x in 0..width {
            // Use 0.0 as a fallback for out-of-bounds (shouldn't happen after
            // dimension validation, but we must not panic).
            let v = band.get_pixel(x, y).unwrap_or(0.0);
            data.push(v);
        }
    }
    data
}

// ─────────────────────────────────────────────────────────────────────────────
// Unit tests (inline — require access to private Expr/BinaryOp/UnaryOp types)
// ─────────────────────────────────────────────────────────────────────────────

#[cfg(test)]
#[allow(clippy::panic, clippy::unwrap_used, clippy::expect_used)]
mod tests {
    use super::*;
    use crate::raster::calculator::ast::{BinaryOp, Expr, UnaryOp};

    // ── Test 1: Number literal emits LoadConst ───────────────────────────────
    #[test]
    fn test_compile_number_emits_load_const() {
        // Use an arbitrary non-special literal to avoid approx_constant lint.
        let val = 12.5_f64;
        let expr = Expr::Number(val);
        let mut constants = Vec::new();
        let ops = compile_expr(&expr, &mut constants).expect("compile should succeed");
        assert_eq!(ops.len(), 1, "expected exactly one opcode");
        assert_eq!(
            ops[0],
            OpCode::LoadConst(0),
            "first opcode must be LoadConst(0)"
        );
        assert!(
            (constants[0] - val).abs() < f64::EPSILON,
            "constant[0] must equal the compiled value"
        );
    }

    // ── Test 2: Band reference emits LoadBand ────────────────────────────────
    #[test]
    fn test_compile_band_emits_load_band() {
        let expr = Expr::Band(0);
        let mut constants = Vec::new();
        let ops = compile_expr(&expr, &mut constants).expect("compile should succeed");
        assert_eq!(ops.len(), 1);
        assert_eq!(ops[0], OpCode::LoadBand(0));
    }

    // ── Test 3: Binary add emits two LoadConst + Add ─────────────────────────
    #[test]
    fn test_compile_add_binary() {
        let expr = Expr::BinaryOp {
            left: Box::new(Expr::Number(1.0)),
            op: BinaryOp::Add,
            right: Box::new(Expr::Number(2.0)),
        };
        let mut constants = Vec::new();
        let ops = compile_expr(&expr, &mut constants).expect("compile should succeed");
        assert_eq!(ops.len(), 3, "expected [LoadConst, LoadConst, Add]");
        assert!(matches!(ops[0], OpCode::LoadConst(_)));
        assert!(matches!(ops[1], OpCode::LoadConst(_)));
        assert_eq!(ops[2], OpCode::Add);
    }

    // ── Test 4: Function sqrt emits LoadBand + Sqrt ──────────────────────────
    #[test]
    fn test_compile_function_sqrt() {
        let expr = Expr::Function {
            name: "sqrt".to_string(),
            args: vec![Expr::Band(0)],
        };
        let mut constants = Vec::new();
        let ops = compile_expr(&expr, &mut constants).expect("compile should succeed");
        assert_eq!(ops.len(), 2, "expected [LoadBand(0), Sqrt]");
        assert_eq!(ops[0], OpCode::LoadBand(0));
        assert_eq!(ops[1], OpCode::Sqrt);
    }

    // ── Test 5: Conditional expression ends with Cond ────────────────────────
    #[test]
    fn test_compile_conditional() {
        let expr = Expr::Conditional {
            condition: Box::new(Expr::Number(1.0)),
            then_expr: Box::new(Expr::Number(2.0)),
            else_expr: Box::new(Expr::Number(3.0)),
        };
        let mut constants = Vec::new();
        let ops = compile_expr(&expr, &mut constants).expect("compile should succeed");
        // Expected: [LoadConst, LoadConst, LoadConst, Cond]
        assert!(ops.len() >= 2, "expected at least 2 opcodes");
        assert_eq!(*ops.last().expect("ops must not be empty"), OpCode::Cond);
    }

    // ── Test 13: estimate_stack_depth for simple add ──────────────────────────
    #[test]
    fn test_estimate_stack_depth_simple_add() {
        // (1.0 + 2.0) compiles to [LoadConst(0), LoadConst(1), Add]
        // Depth trace: 0 → 1 → 2 → 1  (max = 2)
        let expr = Expr::BinaryOp {
            left: Box::new(Expr::Number(1.0)),
            op: BinaryOp::Add,
            right: Box::new(Expr::Number(2.0)),
        };
        let mut constants = Vec::new();
        let ops = compile_expr(&expr, &mut constants).expect("compile should succeed");
        let depth = estimate_stack_depth(&ops);
        assert_eq!(depth, 2, "simple add needs stack depth 2");
    }

    // ── Test: unary negate compiles correctly ────────────────────────────────
    #[test]
    fn test_compile_unary_negate() {
        let expr = Expr::UnaryOp {
            op: UnaryOp::Negate,
            expr: Box::new(Expr::Band(1)),
        };
        let mut constants = Vec::new();
        let ops = compile_expr(&expr, &mut constants).expect("compile should succeed");
        assert_eq!(ops.len(), 2);
        assert_eq!(ops[0], OpCode::LoadBand(1));
        assert_eq!(ops[1], OpCode::Neg);
    }

    // ── Test: constant deduplication ─────────────────────────────────────────
    #[test]
    fn test_constant_deduplication() {
        // Two references to 1.0 should share the same constants[0] index.
        let expr = Expr::BinaryOp {
            left: Box::new(Expr::Number(1.0)),
            op: BinaryOp::Add,
            right: Box::new(Expr::Number(1.0)),
        };
        let mut constants = Vec::new();
        let ops = compile_expr(&expr, &mut constants).expect("compile should succeed");
        assert_eq!(constants.len(), 1, "1.0 should be stored once");
        assert!(matches!(ops[0], OpCode::LoadConst(0)));
        assert!(matches!(ops[1], OpCode::LoadConst(0)));
        assert_eq!(ops[2], OpCode::Add);
    }

    // ── Test: eval_bytecode with a single constant ────────────────────────────
    #[test]
    fn test_eval_bytecode_constant_inline() {
        let prog = CompiledProgram {
            ops: vec![OpCode::LoadConst(0)],
            constants: vec![42.0],
            required_bands: 0,
            cache_slot_count: 0,
            estimated_stack_depth: 1,
        };
        let mut cache = Vec::new();
        let mut stack = Vec::with_capacity(4);
        let result =
            eval_bytecode(&prog, &[], 0, &mut cache, &mut stack).expect("eval should succeed");
        assert!((result - 42.0).abs() < f64::EPSILON);
    }
}