onedrop_eval/evaluator/

mod.rs

//! Evaluator for Milkdrop expressions.
//!
//! [`MilkEvaluator`] is the entry point. Most callers go through
//! [`MilkEvaluator::eval`] (single expression), [`MilkEvaluator::compile`] +
//! [`MilkEvaluator::eval_compiled`] (precompiled hot path), or
//! [`MilkEvaluator::eval_equation_list`] (paren-balance-aware join over a
//! preset's flat equation list).
//!
//! The crate's internals are split across:
//!
//! - [`preprocess`] — the regex-driven preprocessing pipeline that turns MD2
//!   EEL2 source into something `evalexpr` will parse.
//! - [`rewriters`] — AST-style rewriters consumed by `preprocess_expression`.
//! - [`gmegabuf`] — `gmegabuf` / `megabuf` write rewrite.
//! - [`joiner`] — multi-equation join logic and top-level `loop`/`exec`/
//!   `while` interceptors.

mod gmegabuf;
mod joiner;
mod preprocess;
mod rewriters;

#[cfg(test)]
mod tests;

use crate::context::MilkContext;
use crate::error::{EvalError, Result};
use evalexpr::Node;

pub use joiner::EvaluationStats;

use preprocess::MAX_EXPRESSION_LENGTH;

/// One iteration of a custom-wave (or custom-shape) per-point loop. The same
/// struct is used for input (caller seeds the loop variables) and output
/// (evaluator reads `x`, `y`, `r`, `g`, `b`, `a` back from the context after
/// running the equations). `sample`, `value1`, `value2` are inputs only.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct WavePoint {
    /// Normalised sample index in `[0.0, 1.0]` (or `0.0` when `samples == 1`).
    pub sample: f64,
    /// Left-channel value (raw audio sample, or FFT bin if `b_spectrum`).
    pub value1: f64,
    /// Right-channel value (mirrored from value1 with N/2 offset for mono).
    pub value2: f64,
    pub x: f64,
    pub y: f64,
    pub r: f64,
    pub g: f64,
    pub b: f64,
    pub a: f64,
}

impl Default for WavePoint {
    fn default() -> Self {
        Self {
            sample: 0.0,
            value1: 0.0,
            value2: 0.0,
            x: 0.5,
            y: 0.5,
            r: 1.0,
            g: 1.0,
            b: 1.0,
            a: 1.0,
        }
    }
}

/// One iteration of a custom-shape (`shapecode_N`) per-instance loop.
/// MD2 shape per-frame equations can mutate any of these fields and the
/// next instance's loop body sees the seed values fresh from the
/// preset's scalar block — instances do NOT carry state across each
/// other (unlike `WavePoint`). Persistent state across instances lives
/// in `q*` / `t*` channels.
///
/// `instance` is the 0-based loop counter; the eval seeds it (and
/// `sides` + `num_inst`) into the context before each call so equations
/// like `ang = ang + 0.1 * instance` work.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct ShapeInstance {
    pub instance: f64,
    pub num_inst: f64,
    pub sides: f64,
    pub x: f64,
    pub y: f64,
    pub rad: f64,
    pub ang: f64,
    pub tex_zoom: f64,
    pub tex_ang: f64,
    pub r: f64,
    pub g: f64,
    pub b: f64,
    pub a: f64,
    pub r2: f64,
    pub g2: f64,
    pub b2: f64,
    pub a2: f64,
    pub border_r: f64,
    pub border_g: f64,
    pub border_b: f64,
    pub border_a: f64,
    /// MD2 `thick` flag (drawn as bool, kept f64 to round-trip via the
    /// eval context with no extra conversions). 0 = thin outline, 1 = thick.
    pub thick: f64,
    /// MD2 `additive` flag (0 = alpha blend, 1 = additive).
    pub additive: f64,
}

impl Default for ShapeInstance {
    fn default() -> Self {
        Self {
            instance: 0.0,
            num_inst: 1.0,
            sides: 4.0,
            x: 0.5,
            y: 0.5,
            rad: 0.1,
            ang: 0.0,
            tex_zoom: 1.0,
            tex_ang: 0.0,
            r: 1.0,
            g: 1.0,
            b: 1.0,
            a: 1.0,
            r2: 0.0,
            g2: 0.0,
            b2: 0.0,
            a2: 0.0,
            border_r: 1.0,
            border_g: 1.0,
            border_b: 1.0,
            border_a: 0.0,
            thick: 0.0,
            additive: 0.0,
        }
    }
}

/// Evaluator for Milkdrop expressions.
pub struct MilkEvaluator {
    /// Execution context
    pub(super) context: MilkContext,

    /// Compiled expressions cache
    #[allow(dead_code)]
    compiled_cache: Vec<(String, Node)>,
}

/// Custom `Clone` impl that skips `compiled_cache`. The cache is per-call
/// memoisation for the regex-preprocess path; worker threads (the wave
/// and warp parallel passes) never re-enter that path — they're given
/// pre-compiled `Node`s by the caller — so cloning the cache is pure
/// overhead. The cache can grow into the hundreds of entries on a
/// long-running engine, so this matters: cloning a 100-entry
/// `Vec<(String, Node)>` is ~30 µs that goes away for free.
impl Clone for MilkEvaluator {
    fn clone(&self) -> Self {
        Self {
            context: self.context.clone(),
            compiled_cache: Vec::new(),
        }
    }
}

impl MilkEvaluator {
    /// Create a new evaluator.
    pub fn new() -> Self {
        Self {
            context: MilkContext::new(),
            compiled_cache: Vec::new(),
        }
    }

    /// Get a reference to the context.
    pub fn context(&self) -> &MilkContext {
        &self.context
    }

    /// Get a mutable reference to the context.
    pub fn context_mut(&mut self) -> &mut MilkContext {
        &mut self.context
    }

    /// Pre-compile an expression into an evalexpr `Node` that can be evaluated
    /// repeatedly without re-parsing.
    ///
    /// Pre-processing (auto-init of undefined variables, integer→float
    /// promotion in assignments, `if(` → `milkif(` rewrite) runs as part of
    /// compilation, so any variable referenced in `expression` is also
    /// registered in `self.context` as a side effect — exactly like a normal
    /// `eval` call. This means a context cloned out of this evaluator after a
    /// `compile` call already contains every var the compiled `Node` will
    /// look up at evaluation time.
    pub fn compile(&mut self, expression: &str) -> Result<Node> {
        if expression.len() > MAX_EXPRESSION_LENGTH {
            return Err(EvalError::SyntaxError {
                expression: expression.chars().take(100).collect(),
                reason: format!(
                    "Expression too long: {} bytes (max {})",
                    expression.len(),
                    MAX_EXPRESSION_LENGTH
                ),
            });
        }
        let expr = expression.trim().trim_end_matches(';').trim();
        let processed = self.preprocess_expression(expr);
        evalexpr::build_operator_tree(&processed).map_err(|e| EvalError::SyntaxError {
            expression: expr.to_string(),
            reason: e.to_string(),
        })
    }

    /// Evaluate a single expression.
    pub fn eval(&mut self, expression: &str) -> Result<f64> {
        if expression.len() > MAX_EXPRESSION_LENGTH {
            return Err(EvalError::SyntaxError {
                expression: expression.chars().take(100).collect(),
                reason: format!(
                    "Expression too long: {} bytes (max {})",
                    expression.len(),
                    MAX_EXPRESSION_LENGTH
                ),
            });
        }

        let expr = expression.trim().trim_end_matches(';').trim();

        if expr.is_empty() {
            return Ok(0.0);
        }

        let processed_expr = self.preprocess_expression(expr);

        // Intercept top-level `loop(N, body)` (and `exec2`/`exec3`/`while`)
        // calls *before* evalexpr sees them. evalexpr can't represent these
        // EEL2 idioms natively (their bodies are `;`-chains inside what looks
        // like a function-call arg list).
        self.eval_processed_with_loops(&processed_expr, &processed_expr)
    }

    /// Evaluate multiple expressions (per-frame equations).
    pub fn eval_per_frame(&mut self, equations: &[String]) -> Result<()> {
        for equation in equations {
            self.eval(equation)?;
        }
        Ok(())
    }

    /// Evaluate per-pixel equations for a single pixel.
    pub fn eval_per_pixel(
        &mut self,
        x: f64,
        y: f64,
        rad: f64,
        ang: f64,
        equations: &[String],
    ) -> Result<()> {
        self.context.set_pixel(x, y, rad, ang);
        for equation in equations {
            self.eval(equation)?;
        }
        Ok(())
    }

    /// Compile a batch of source strings into reusable `Node`s. Use this on
    /// load_preset for custom wave/shape equations, which run hundreds of
    /// times per frame — avoiding per-eval reparse is a 10×+ speedup.
    pub fn compile_batch(&mut self, equations: &[String]) -> Result<Vec<Node>> {
        equations.iter().map(|eq| self.compile(eq)).collect()
    }

    /// Evaluate a pre-compiled `Node` against this evaluator's context. Skips
    /// the regex preprocess that `eval()` runs every call — variables referenced
    /// in the source string were already auto-initialized at `compile()` time.
    pub fn eval_compiled(&mut self, node: &Node) -> Result<f64> {
        match node.eval_with_context_mut(&mut self.context) {
            Ok(value) => match value {
                evalexpr::Value::Float(f) => Ok(f),
                evalexpr::Value::Int(i) => Ok(i as f64),
                evalexpr::Value::Boolean(b) => Ok(if b { 1.0 } else { 0.0 }),
                evalexpr::Value::Empty => Ok(0.0),
                _ => Err(EvalError::TypeError {
                    expected: "number".to_string(),
                    got: format!("{:?}", value),
                }),
            },
            Err(e) => Err(EvalError::SyntaxError {
                expression: "<compiled>".to_string(),
                reason: e.to_string(),
            }),
        }
    }

    /// Evaluate per-frame equations for one instance of a custom shape
    /// (`shapecode_N`). Seeds the loop variables (`instance`, `sides`,
    /// `num_inst`) plus the full geometry/colour state, runs the
    /// compiled per-frame Nodes, reads everything back into a fresh
    /// `ShapeInstance`.
    ///
    /// Unlike `eval_per_point` (which threads state across iterations
    /// within a frame), each shape instance is independent: callers
    /// pass the shape's scalar block seed every time and `instance`
    /// changes per call. Persistent state across instances lives in
    /// `q*` / `t*` channels, which the context carries.
    pub fn eval_per_shape_instance(
        &mut self,
        instance: ShapeInstance,
        compiled: &[Node],
    ) -> Result<ShapeInstance> {
        Self::seed_shape_instance(&mut self.context, instance);
        for node in compiled {
            self.eval_compiled(node)?;
        }
        Ok(Self::read_shape_instance(&self.context, instance))
    }

    /// Bytecode-aware companion of [`eval_per_shape_instance`]. Runs the
    /// bytecode VM when [`CompiledBlock::bytecode`] is `Some`, otherwise
    /// walks the evalexpr `Node` list.
    pub fn run_per_shape_instance(
        &mut self,
        instance: ShapeInstance,
        block: &crate::compiled_block::CompiledBlock,
    ) -> Result<ShapeInstance> {
        Self::seed_shape_instance(&mut self.context, instance);
        if let Some(bc) = block.bytecode() {
            bc.run(&mut self.context);
        } else {
            for node in block.nodes() {
                self.eval_compiled(node)?;
            }
        }
        Ok(Self::read_shape_instance(&self.context, instance))
    }

    /// Push the 23 hot-var seed values from `instance` into `ctx`.
    #[inline]
    fn seed_shape_instance(ctx: &mut MilkContext, instance: ShapeInstance) {
        ctx.set("instance", instance.instance);
        ctx.set("num_inst", instance.num_inst);
        ctx.set("sides", instance.sides);
        ctx.set("x", instance.x);
        ctx.set("y", instance.y);
        ctx.set("rad", instance.rad);
        ctx.set("ang", instance.ang);
        ctx.set("tex_zoom", instance.tex_zoom);
        ctx.set("tex_ang", instance.tex_ang);
        ctx.set("r", instance.r);
        ctx.set("g", instance.g);
        ctx.set("b", instance.b);
        ctx.set("a", instance.a);
        ctx.set("r2", instance.r2);
        ctx.set("g2", instance.g2);
        ctx.set("b2", instance.b2);
        ctx.set("a2", instance.a2);
        ctx.set("border_r", instance.border_r);
        ctx.set("border_g", instance.border_g);
        ctx.set("border_b", instance.border_b);
        ctx.set("border_a", instance.border_a);
        ctx.set("thick", instance.thick);
        ctx.set("additive", instance.additive);
    }

    /// Companion to [`seed_shape_instance`] — read every per-instance var
    /// back, falling back to `instance.*` when a slot is missing.
    #[inline]
    fn read_shape_instance(c: &MilkContext, instance: ShapeInstance) -> ShapeInstance {
        ShapeInstance {
            instance: instance.instance,
            num_inst: instance.num_inst,
            sides: c.get("sides").unwrap_or(instance.sides),
            x: c.get("x").unwrap_or(instance.x),
            y: c.get("y").unwrap_or(instance.y),
            rad: c.get("rad").unwrap_or(instance.rad),
            ang: c.get("ang").unwrap_or(instance.ang),
            tex_zoom: c.get("tex_zoom").unwrap_or(instance.tex_zoom),
            tex_ang: c.get("tex_ang").unwrap_or(instance.tex_ang),
            r: c.get("r").unwrap_or(instance.r),
            g: c.get("g").unwrap_or(instance.g),
            b: c.get("b").unwrap_or(instance.b),
            a: c.get("a").unwrap_or(instance.a),
            r2: c.get("r2").unwrap_or(instance.r2),
            g2: c.get("g2").unwrap_or(instance.g2),
            b2: c.get("b2").unwrap_or(instance.b2),
            a2: c.get("a2").unwrap_or(instance.a2),
            border_r: c.get("border_r").unwrap_or(instance.border_r),
            border_g: c.get("border_g").unwrap_or(instance.border_g),
            border_b: c.get("border_b").unwrap_or(instance.border_b),
            border_a: c.get("border_a").unwrap_or(instance.border_a),
            thick: c.get("thick").unwrap_or(instance.thick),
            additive: c.get("additive").unwrap_or(instance.additive),
        }
    }

    /// Bytecode counterpart of [`eval_per_point`]. Skips both the
    /// `String → Value::Float` wrapping and the recursive operator-tree
    /// walk that evalexpr does for each child node — the bytecode VM
    /// just reads/writes hot slots by index and runs a flat opcode stream.
    pub fn run_per_point_bc(
        &mut self,
        point: WavePoint,
        compiled: &crate::bytecode::CompiledBytecode,
    ) -> WavePoint {
        let ctx = &mut self.context;
        ctx.set("sample", point.sample);
        ctx.set("value1", point.value1);
        ctx.set("value2", point.value2);
        ctx.set("x", point.x);
        ctx.set("y", point.y);
        ctx.set("r", point.r);
        ctx.set("g", point.g);
        ctx.set("b", point.b);
        ctx.set("a", point.a);
        compiled.run(ctx);
        let c = &self.context;
        WavePoint {
            sample: point.sample,
            value1: point.value1,
            value2: point.value2,
            x: c.get("x").unwrap_or(point.x),
            y: c.get("y").unwrap_or(point.y),
            r: c.get("r").unwrap_or(point.r),
            g: c.get("g").unwrap_or(point.g),
            b: c.get("b").unwrap_or(point.b),
            a: c.get("a").unwrap_or(point.a),
        }
    }

    /// Evaluate per-point equations for one sample inside a custom wave or
    /// shape loop. `point` carries the input vars (`sample`, `value1`,
    /// `value2`) and the carry-over geometry/colour (`x`, `y`, `r`, `g`,
    /// `b`, `a`). The output echoes the same fields back, read from the
    /// context after the equations ran.
    ///
    /// The caller is responsible for seeding the loop on point 0 with the
    /// wave's base colour and any saved `x`/`y`, then threading the previous
    /// point's output into the next call — this matches MD2's "trail across
    /// samples within a frame" semantics.
    pub fn eval_per_point(&mut self, point: WavePoint, compiled: &[Node]) -> Result<WavePoint> {
        let ctx = &mut self.context;
        ctx.set("sample", point.sample);
        ctx.set("value1", point.value1);
        ctx.set("value2", point.value2);
        ctx.set("x", point.x);
        ctx.set("y", point.y);
        ctx.set("r", point.r);
        ctx.set("g", point.g);
        ctx.set("b", point.b);
        ctx.set("a", point.a);
        for node in compiled {
            self.eval_compiled(node)?;
        }
        let c = &self.context;
        Ok(WavePoint {
            sample: point.sample,
            value1: point.value1,
            value2: point.value2,
            x: c.get("x").unwrap_or(point.x),
            y: c.get("y").unwrap_or(point.y),
            r: c.get("r").unwrap_or(point.r),
            g: c.get("g").unwrap_or(point.g),
            b: c.get("b").unwrap_or(point.b),
            a: c.get("a").unwrap_or(point.a),
        })
    }

    /// Parse an assignment expression and update context.
    /// Returns the assigned value.
    pub fn eval_assignment(&mut self, expression: &str) -> Result<f64> {
        let result = self.eval(expression)?;

        if let Some((var_name, _)) = expression.split_once('=') {
            let var_name = var_name.trim();
            self.context.set_var(var_name, result);
        }

        Ok(result)
    }

    /// Reset the evaluator to initial state.
    pub fn reset(&mut self) {
        self.context = MilkContext::new();
        self.compiled_cache.clear();
    }
}

impl Default for MilkEvaluator {
    fn default() -> Self {
        Self::new()
    }
}

// ---------- Shared byte-walking utilities ----------
//
// These helpers walk paren-balanced byte slices and are used by every
// submodule (preprocess, rewriters, gmegabuf, joiner). They live here so
// each submodule pulls them from `super::*` without going through a
// separate `utils` module.

/// `true` when `b` is an ASCII identifier byte (`[A-Za-z0-9_]`).
pub(super) fn is_ident_byte(b: u8) -> bool {
    b.is_ascii_alphanumeric() || b == b'_'
}

/// Find the matching `)` for the `(` at `open`. Returns the byte offset
/// of the closer or `None` if unbalanced.
pub(super) fn match_close_paren(bytes: &[u8], open: usize) -> Option<usize> {
    debug_assert_eq!(bytes[open], b'(');
    let mut depth = 1usize;
    let mut k = open + 1;
    while k < bytes.len() {
        match bytes[k] {
            b'(' => depth += 1,
            b')' => {
                depth -= 1;
                if depth == 0 {
                    return Some(k);
                }
            }
            _ => {}
        }
        k += 1;
    }
    None
}

/// `true` if `s` contains a comparison operator (`> < >= <= == !=`) at
/// paren depth 0.
pub(super) fn contains_top_level_comparison(s: &str) -> bool {
    let bytes = s.as_bytes();
    let mut depth = 0i32;
    let mut i = 0;
    while i < bytes.len() {
        match bytes[i] {
            b'(' => depth += 1,
            b')' => depth -= 1,
            b'>' | b'<' if depth == 0 => return true,
            b'=' | b'!' if depth == 0 && bytes.get(i + 1).copied() == Some(b'=') => return true,
            _ => {}
        }
        i += 1;
    }
    false
}

/// `true` if `s` contains a `,` at paren depth 0.
pub(super) fn contains_top_level_comma(s: &str) -> bool {
    contains_top_level_byte(s, b',')
}

/// Generic "byte at depth 0" check.
pub(super) fn contains_top_level_byte(s: &str, target: u8) -> bool {
    let mut depth = 0i32;
    for b in s.bytes() {
        match b {
            b'(' => depth += 1,
            b')' => depth -= 1,
            x if x == target && depth == 0 => return true,
            _ => {}
        }
    }
    false
}

/// Generic "split on byte at depth 0".
pub(super) fn split_top_level_byte(s: &str, target: u8) -> Vec<&str> {
    let bytes = s.as_bytes();
    let mut parts = Vec::new();
    let mut depth = 0i32;
    let mut start = 0usize;
    for (i, b) in bytes.iter().enumerate() {
        match *b {
            b'(' => depth += 1,
            b')' => depth -= 1,
            x if x == target && depth == 0 => {
                parts.push(&s[start..i]);
                start = i + 1;
            }
            _ => {}
        }
    }
    parts.push(&s[start..]);
    parts
}

/// Split `inner` on each top-level `,` (commas at paren depth 0). Always
/// returns at least one element (the whole `inner` if no top-level comma).
pub(super) fn split_top_level_commas(inner: &str) -> Vec<&str> {
    let bytes = inner.as_bytes();
    let mut parts = Vec::new();
    let mut depth = 0usize;
    let mut start = 0usize;
    for (i, b) in bytes.iter().enumerate() {
        match *b {
            b'(' => depth += 1,
            b')' => depth = depth.saturating_sub(1),
            b',' if depth == 0 => {
                parts.push(&inner[start..i]);
                start = i + 1;
            }
            _ => {}
        }
    }
    parts.push(&inner[start..]);
    parts
}

/// `true` if the `(` at byte offset `paren_pos` opens a function-call arg
/// list (preceded by an identifier or digit, modulo whitespace) rather
/// than a standalone grouping expression.
pub(super) fn is_paren_func_call_open(bytes: &[u8], paren_pos: usize) -> bool {
    let mut j = paren_pos;
    while j > 0 && matches!(bytes[j - 1], b' ' | b'\t') {
        j -= 1;
    }
    j > 0 && is_ident_byte(bytes[j - 1])
}

/// `true` if `name` is a registered builtin identifier we should leave
/// alone in the auto-init pass. Centralised so the preprocess and
/// case-normalising helpers stay in sync.
pub(crate) fn is_builtin_ident(name: &str) -> bool {
    matches!(
        name,
        "sin"
            | "cos"
            | "tan"
            | "sqrt"
            | "abs"
            | "pow"
            | "exp"
            | "log"
            | "log10"
            | "ln"
            | "if"
            | "milkif"
            | "min"
            | "max"
            | "floor"
            | "ceil"
            | "round"
            | "rand"
            | "above"
            | "below"
            | "equal"
            | "bnot"
            | "band"
            | "bor"
            | "int"
            | "fmod"
            | "clamp"
            | "sinh"
            | "cosh"
            | "tanh"
            | "asin"
            | "acos"
            | "atan"
            | "atan2"
            | "sqr"
            | "rad"
            | "deg"
            | "fract"
            | "trunc"
            | "sign"
            | "sigmoid"
            | "loop"
            | "while"
            | "exec2"
            | "exec3"
            | "gmegabuf"
            | "megabuf"
    )
}