onedrop_eval/evaluator/

joiner.rs

//! Multi-equation join logic and top-level call interceptors.
//!
//! Two concerns sit here:
//!
//! 1. [`MilkEvaluator::eval_equation_list`] — paren-balance-aware join over a
//!    `.milk` preset's flat equation list. Consecutive equations that aren't
//!    paren-balanced (or end on a dangling binary operator, or look like a
//!    continuation of the previous line) get glued with `;` before
//!    evaluation.
//! 2. [`MilkEvaluator::eval_processed_with_loops`] — top-level interceptor
//!    for `loop(N, body)`, `exec2`, `exec3`, and `while(body)`. evalexpr
//!    can't express these EEL2 idioms natively (their bodies are `;`-chains
//!    inside what looks like a function-call arg list), so we strip the
//!    call boundary and walk the body ourselves.

use evalexpr::eval_with_context_mut;

use crate::error::{EvalError, Result};

use super::MilkEvaluator;
use super::preprocess::{MAX_LOOP_ITER, strip_line_comments};
use super::{is_ident_byte, match_close_paren, split_top_level_commas};

/// Aggregate stats returned by [`MilkEvaluator::eval_equation_list`].
///
/// `total` is the input length so callers can compute per-equation failure
/// rates; `failed` counts whole blocks that errored, so a 5-line loop that
/// fails counts as one failure.
#[derive(Debug, Default, Clone)]
pub struct EvaluationStats {
    pub total: usize,
    pub failed: usize,
    pub first_error: String,
}

impl MilkEvaluator {
    /// Evaluate a flat list of equations, joining consecutive
    /// paren-unbalanced lines (MD2 `loop(...)` idiom). Returns aggregate
    /// stats; the evaluator's context retains every variable touched by
    /// the run.
    pub fn eval_equation_list(&mut self, eqs: &[String]) -> EvaluationStats {
        let mut stats = EvaluationStats {
            total: eqs.len(),
            ..Default::default()
        };
        let mut buffer = String::new();
        let mut depth = 0i32;
        for (idx, eq) in eqs.iter().enumerate() {
            // Strip `//` comments per-equation BEFORE the buffer-join. Each
            // `per_frame_N=` value in MD2 is one line of EEL source, so a
            // `//` comment runs to the end of that equation. Stripping
            // post-join would let the `//` consume every following equation
            // up to the end of the buffer (no `\n` is inserted by the
            // joiner), wiping the tail of every `loop(N, …; //... ; …)`
            // block. `preprocess_expression` still re-strips inside `eval()`
            // for the single-equation entry point — the second pass is
            // idempotent because the per-eq strip already left nothing to
            // remove.
            let stripped = strip_line_comments(eq);
            let trimmed = stripped.trim();
            if trimmed.is_empty() {
                continue;
            }
            if !buffer.is_empty() {
                // Always use `;` as the join separator. The dominant
                // multi-line idiom in MD2 is `loop(N, stmt; stmt; stmt)`;
                // with N-fold unrolling, the body is sliced out and
                // re-evaluated as a statement chain, so `;` is what
                // evalexpr needs to see. Skip the separator when the
                // buffer's tail is already a paren or punctuator that
                // doesn't need one.
                let last = buffer.trim_end().bytes().last();
                // Inside an unclosed `(`, a leading `+`/`-` is unambiguously
                // binary (no statement boundary has been seen), so suppress
                // the `;` that would otherwise be rewritten to `,` by
                // [`super::rewriters::rewrite_semis_in_call_args`] and
                // corrupt the call's arg count.
                let unbalanced_binop_continuation =
                    depth > 0 && matches!(trimmed.bytes().next(), Some(b'+') | Some(b'-'));
                // At depth 0, a leading `+`/`-` on a line that contains no
                // top-level `=` is almost certainly an orphan binop the
                // author split across `per_pixel_N=` lines (corpus shape:
                // `per_pixel_4=rot=cos(rad)` / `per_pixel_5=+sin(rad)`).
                // The `=`-presence check rules out legitimate `b = -1`
                // standalone assignments — those still flush independently.
                let closed_paren_orphan_continuation =
                    depth <= 0 && starts_with_plus_minus_orphan(trimmed);
                let split_lhs_assignment = starts_with_split_assignment(trimmed);
                // Split function-call only when the previous tail is an ident
                // byte; a standalone `(group)` line after a `;` must still
                // flush cleanly.
                let split_call_paren =
                    starts_with_open_paren(trimmed) && matches!(last, Some(b) if is_ident_byte(b));
                let needs_sep = !matches!(
                    last,
                    Some(b',')
                        | Some(b'(')
                        | Some(b';')
                        | Some(b'+')
                        | Some(b'-')
                        | Some(b'*')
                        | Some(b'/')
                        | Some(b'%')
                ) && !starts_with_binop(trimmed)
                    && !unbalanced_binop_continuation
                    && !closed_paren_orphan_continuation
                    && !split_lhs_assignment
                    && !split_call_paren;
                if needs_sep {
                    buffer.push(';');
                }
            }
            buffer.push_str(trimmed);
            depth = paren_balance(&buffer);
            // Peek at the next non-empty equation. If it starts with a
            // binary operator (`*`/`/`/`%`) or an assignment/call
            // continuation marker, defer the flush so the joiner can stitch
            // the orphan onto the current buffer instead of evaluating it
            // standalone. Corpus shape:
            //
            // ```text
            // per_pixel_4=rot=bnot(above(x,.5)+((y*q8)%q7))
            // per_pixel_5=*cos(rad+3.14*if(...))*q4;
            // ```
            //
            // Without the lookahead, eq #4 flushes (it's paren-balanced) and
            // eq #5 arrives in an empty buffer and dies as "operator
            // expected 1 args, got 0". With it, the flush is held until eq
            // #5 has been appended and the composite block evaluates as
            // `rot=bnot(...)*cos(...)*q4`.
            if depth <= 0
                && !ends_with_dangling_binop(&buffer)
                && !next_eq_is_continuation(eqs, idx)
            {
                let block = sanitize_joined_block(&buffer);
                if let Err(e) = self.eval(&block) {
                    stats.failed += 1;
                    if stats.first_error.is_empty() {
                        stats.first_error = format!("{e:?}");
                    }
                }
                buffer.clear();
            }
        }
        if !buffer.is_empty() {
            let block = sanitize_joined_block(&buffer);
            if let Err(e) = self.eval(&block) {
                stats.failed += 1;
                if stats.first_error.is_empty() {
                    stats.first_error = format!("{e:?}");
                }
            }
        }
        stats
    }

    /// Evaluate `expr` (already preprocessed), intercepting top-level
    /// `loop(N, body)`, `exec2(a, b)`, `exec3(a, b, c)` and
    /// `while(body)` calls. `original` is kept for error reporting.
    pub(super) fn eval_processed_with_loops(&mut self, expr: &str, original: &str) -> Result<f64> {
        let trimmed = expr.trim().trim_end_matches(';').trim();
        if trimmed.is_empty() {
            return Ok(0.0);
        }

        if let Some((before, n_str, body_str, after)) = find_top_level_loop_call(trimmed) {
            let before_clean = before.trim().trim_end_matches(';').trim();
            if !before_clean.is_empty() {
                self.eval_processed_with_loops(before_clean, original)?;
            }
            let n_val = self.eval_processed_with_loops(n_str.trim(), original)?;
            let n = if n_val.is_finite() {
                (n_val as i64).clamp(0, MAX_LOOP_ITER)
            } else {
                0
            };
            let body_clean = body_str
                .trim()
                .trim_start_matches(';')
                .trim_end_matches(';')
                .trim();
            if !body_clean.is_empty() {
                // Pre-compile the body once and run it N times. The body has
                // already been through `preprocess_expression` (the caller
                // owns the preprocessed `expr` slice this `body_clean` is
                // sliced from), so `build_operator_tree` consumes it
                // directly. When the body lowers cleanly to bytecode
                // (every common init-loop shape — `gmegabuf_set(i, 0); i = i
                // + 1`, etc.), the inner loop is just `bc.run` against the
                // mutable context, skipping evalexpr's per-iteration parse
                // + tree walk that otherwise dominates `loop(1_000_000, ...)`
                // compile-time eval. Nested loop/exec/while inside the
                // body fall through to the recursive `eval_processed_with_
                // loops` path so their semantics stay intact.
                if !body_has_nested_loop(body_clean)
                    && let Ok(node) =
                        evalexpr::build_operator_tree::<evalexpr::DefaultNumericTypes>(body_clean)
                {
                    let nodes = [node];
                    if let Ok(bc) =
                        crate::bytecode::CompiledBytecode::try_compile(&nodes, &mut self.context)
                    {
                        for _ in 0..n {
                            bc.run(&mut self.context);
                        }
                    } else {
                        // Body parsed but didn't lower — replay the Node N
                        // times via the leaf evalexpr walker. Still skips
                        // the per-iteration re-parse.
                        let node = &nodes[0];
                        for _ in 0..n {
                            let _ = node.eval_with_context_mut(&mut self.context);
                        }
                    }
                } else {
                    // Nested loop / unparseable body — fall back to the
                    // recursive interceptor per iteration.
                    for _ in 0..n {
                        self.eval_processed_with_loops(body_clean, original)?;
                    }
                }
            }
            let after_clean = after.trim().trim_start_matches(';').trim();
            if !after_clean.is_empty() {
                return self.eval_processed_with_loops(after_clean, original);
            }
            return Ok(0.0);
        }

        // `exec2(a, b)` / `exec3(a, b, c)`. EEL2's sequential evaluators:
        // each arg is evaluated against the shared context in order, and
        // the value of the last arg is returned. Args themselves may be
        // `;`-chains.
        if let Some((before, args, after)) = find_top_level_exec_call(trimmed) {
            let before_clean = before.trim().trim_end_matches(';').trim();
            if !before_clean.is_empty() {
                self.eval_processed_with_loops(before_clean, original)?;
            }
            let mut last_val = 0.0;
            for arg in &args {
                let arg_clean = arg
                    .trim()
                    .trim_start_matches(';')
                    .trim_end_matches(';')
                    .trim();
                if !arg_clean.is_empty() {
                    last_val = self.eval_processed_with_loops(arg_clean, original)?;
                }
            }
            let after_clean = after.trim().trim_start_matches(';').trim();
            if !after_clean.is_empty() {
                return self.eval_processed_with_loops(after_clean, original);
            }
            return Ok(last_val);
        }

        // `while(body)`. EEL2's do-while loop; we run a single pass. The
        // corpus pattern is `while(exec2(stmts, 0))` where the second arg
        // returns 0 to terminate after one iteration; single-pass matches
        // that intent and avoids the risk of an unbounded loop.
        if let Some((before, body, after)) = find_top_level_while_call(trimmed) {
            let before_clean = before.trim().trim_end_matches(';').trim();
            if !before_clean.is_empty() {
                self.eval_processed_with_loops(before_clean, original)?;
            }
            let body_clean = body
                .trim()
                .trim_start_matches(';')
                .trim_end_matches(';')
                .trim();
            let body_val = if !body_clean.is_empty() {
                self.eval_processed_with_loops(body_clean, original)?
            } else {
                0.0
            };
            let after_clean = after.trim().trim_start_matches(';').trim();
            if !after_clean.is_empty() {
                return self.eval_processed_with_loops(after_clean, original);
            }
            return Ok(body_val);
        }

        // Leaf: hand the (loop-free) expression to evalexpr.
        match eval_with_context_mut(trimmed, &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: original.to_string(),
                reason: e.to_string(),
            }),
        }
    }
}

/// Return `true` when the next non-empty equation after `idx` looks like a
/// continuation of the current buffer rather than a fresh statement, so the
/// buffer flush should be deferred until it's been appended.
///
/// Continuation triggers: leading `*`/`/`/`%` (orphan binop), leading
/// `+`/`-` on a line with no top-level `=` (closed-paren orphan binop),
/// leading `=` (not `==`) — split-LHS assignment, leading `(` — split
/// function-call site. Blank equations are skipped first.
fn next_eq_is_continuation(eqs: &[String], idx: usize) -> bool {
    for next in eqs.iter().skip(idx + 1) {
        let stripped = strip_line_comments(next);
        let trimmed = stripped.trim();
        if trimmed.is_empty() {
            continue;
        }
        return starts_with_binop(trimmed)
            || starts_with_plus_minus_orphan(trimmed)
            || starts_with_split_assignment(trimmed)
            || starts_with_open_paren(trimmed);
    }
    false
}

/// Collapse junk-separator pairs that the join step can produce when an
/// equation closes a multi-line block. `;)` and `;,` are never valid in
/// evalexpr's grammar and are always artifacts of our `;`-join over a
/// previously-truncated equation that ended with `,`/`(`.
pub(super) fn sanitize_joined_block(s: &str) -> String {
    let mut out = s.replace(";)", ")").replace(";,", ",");
    while out.contains(",;") {
        out = out.replace(",;", ",");
    }
    while out.contains("(;") {
        out = out.replace("(;", "(");
    }
    // Collapse `;;` runs defensively — the per-equation strip filters
    // comment-only equations before they reach the buffer, but evalexpr's
    // "Empty operand" diagnostics on stray `;;` are hard to debug.
    while out.contains(";;") {
        out = out.replace(";;", ";");
    }
    out
}

/// Count `(` − `)` across `s`. Byte-accurate for ASCII and conservative for
/// UTF-8 (multibyte chars never collide with the paren ASCII codepoints).
fn paren_balance(s: &str) -> i32 {
    let mut depth = 0i32;
    for b in s.bytes() {
        match b {
            b'(' => depth += 1,
            b')' => depth -= 1,
            _ => {}
        }
    }
    depth
}

/// Return `true` when the first non-whitespace byte of `s` is a binary
/// operator with no unary form (`*`, `/`, `%`). Leading `+`/`-` is excluded
/// because they're also legitimate unary signs on a fresh statement.
fn starts_with_binop(s: &str) -> bool {
    let trimmed = s.trim_start();
    let Some(c) = trimmed.bytes().next() else {
        return false;
    };
    matches!(c, b'*' | b'/' | b'%')
}

/// Return `true` when the first non-whitespace byte of `s` is `+` or `-`
/// AND `s` contains no top-level assignment `=` (single `=`, not `==` or a
/// compound op). Corpus authors split a single arithmetic expression
/// across `per_pixel_N=` lines whose continuation starts with `+`/`-`:
///
/// ```text
/// per_pixel_4=rot=cos(rad*q4)
/// per_pixel_5=+sin(rad*q5)
/// ```
///
/// Treating these as orphan binops mirrors the existing `*`/`/`/`%`
/// path. The "no top-level `=`" guard rules out legitimate `b = -1` /
/// `c = +x` standalone assignments — those flush independently because
/// their leading byte is the LHS identifier, not the sign.
pub(super) fn starts_with_plus_minus_orphan(s: &str) -> bool {
    let trimmed = s.trim_start();
    let Some(c) = trimmed.bytes().next() else {
        return false;
    };
    if !matches!(c, b'+' | b'-') {
        return false;
    }
    !contains_top_level_assignment(trimmed)
}

/// `true` if `s` contains a top-level single `=` (assignment), as opposed
/// to `==` / `<=` / `>=` / `!=` / `+=` / `-=` / `*=` / `/=` / `%=`.
fn contains_top_level_assignment(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'=' if depth == 0 => {
                let next = bytes.get(i + 1).copied();
                let prev = if i > 0 { Some(bytes[i - 1]) } else { None };
                let is_double = next == Some(b'=');
                let is_compound = matches!(
                    prev,
                    Some(b'=' | b'<' | b'>' | b'!' | b'+' | b'-' | b'*' | b'/' | b'%')
                );
                if !is_double && !is_compound {
                    return true;
                }
                if is_double {
                    i += 1;
                }
            }
            _ => {}
        }
        i += 1;
    }
    false
}

/// `true` when `s` starts with `=` (single, not `==`), indicating the
/// previous equation's tail is the LHS of an assignment split across .milk
/// lines.
fn starts_with_split_assignment(s: &str) -> bool {
    let bytes = s.trim_start().as_bytes();
    bytes.first() == Some(&b'=') && bytes.get(1) != Some(&b'=')
}

/// `true` when `s` starts with `(`, indicating a function call whose name
/// and arg list were split across .milk lines.
fn starts_with_open_paren(s: &str) -> bool {
    s.trim_start().as_bytes().first() == Some(&b'(')
}

fn ends_with_dangling_binop(s: &str) -> bool {
    let trimmed = s.trim_end();
    let Some(c) = trimmed.bytes().last() else {
        return false;
    };
    // `(` and `,` are already covered by `paren_balance` so only watch for
    // arithmetic operators that don't bump the paren counter.
    if !matches!(c, b'+' | b'-' | b'*' | b'/' | b'%') {
        return false;
    }
    // Reject `1e+`/`1e-` (scientific exponent) and the rare `++`/`--` by
    // requiring the preceding non-space byte to neither match the dangling
    // operator nor be `e`/`E`.
    let prev = trimmed
        .as_bytes()
        .iter()
        .rev()
        .skip(1)
        .find(|&&b| b != b' ' && b != b'\t')
        .copied();
    !matches!(prev, Some(p) if p == c || p == b'e' || p == b'E')
}

/// `true` if `s` contains an EEL2 control-flow builtin name (`loop`,
/// `while`, `exec2`, `exec3`) as an identifier. Used as a cheap
/// conservative guard before the pre-compile-the-body optimisation in
/// the `loop(N, body)` interceptor: when the body has a nested
/// loop/while/exec, replaying it as a single `Node` would bypass the
/// interceptor's recursive descent and break those semantics, so we
/// fall back to the slow path.
///
/// The check is byte-level and ignores word boundaries on purpose —
/// `loops_var` or `while_count` would also match and trigger fallback.
/// That's safe (over-conservative) and avoids a more expensive scan.
fn body_has_nested_loop(s: &str) -> bool {
    s.contains("loop") || s.contains("while") || s.contains("exec2") || s.contains("exec3")
}

/// Locate a top-level `loop(<n_expr>, <body>)` call in `s`. Returns
/// `(before, n_expr, body, after)` slices into `s`, or `None` if no
/// well-formed top-level loop exists.
fn find_top_level_loop_call(s: &str) -> Option<(&str, &str, &str, &str)> {
    let bytes = s.as_bytes();
    let mut depth = 0usize;
    let mut i = 0usize;
    while i < bytes.len() {
        match bytes[i] {
            b'(' => depth += 1,
            b')' => depth = depth.saturating_sub(1),
            b'l' | b'L' if depth == 0 => {
                if i + 4 <= bytes.len() && bytes[i..i + 4].eq_ignore_ascii_case(b"loop") {
                    let prev_ok = i == 0 || !is_ident_byte(bytes[i - 1]);
                    if prev_ok && (i + 4 == bytes.len() || !is_ident_byte(bytes[i + 4])) {
                        let mut j = i + 4;
                        while j < bytes.len() && matches!(bytes[j], b' ' | b'\t' | b'\n') {
                            j += 1;
                        }
                        if j < bytes.len() && bytes[j] == b'(' {
                            let open = j;
                            let close = match_close_paren(bytes, open)?;
                            let inner = &s[open + 1..close];
                            // Split on the first top-level `,`. Also accept
                            // `;` as a fallback for `loop(N; body)` typos
                            // MD2's relaxed parser forgave.
                            let split = find_top_level_comma(inner.as_bytes())
                                .or_else(|| find_top_level_semi(inner.as_bytes()))?;
                            let n_str = &inner[..split];
                            let body_str = &inner[split + 1..];
                            let before = &s[..i];
                            let after = &s[close + 1..];
                            return Some((before, n_str, body_str, after));
                        }
                    }
                }
            }
            _ => {}
        }
        i += 1;
    }
    None
}

/// Locate the first top-level `,` in `bytes`. Returns the byte offset or
/// `None` if every comma is enclosed by parens.
fn find_top_level_comma(bytes: &[u8]) -> Option<usize> {
    let mut depth = 0usize;
    for (i, b) in bytes.iter().enumerate() {
        match *b {
            b'(' => depth += 1,
            b')' => depth = depth.saturating_sub(1),
            b',' if depth == 0 => return Some(i),
            _ => {}
        }
    }
    None
}

/// Fallback splitter for `loop(N; body)`-shaped corpus typos.
fn find_top_level_semi(bytes: &[u8]) -> Option<usize> {
    let mut depth = 0usize;
    for (i, b) in bytes.iter().enumerate() {
        match *b {
            b'(' => depth += 1,
            b')' => depth = depth.saturating_sub(1),
            b';' if depth == 0 => return Some(i),
            _ => {}
        }
    }
    None
}

/// Locate a top-level call `name(a, b, …)` and split it into
/// (before, args, after). Returns `None` if no qualifying call exists at
/// depth 0.
fn find_top_level_named_call<'a>(
    s: &'a str,
    name: &[u8],
) -> Option<(&'a str, Vec<&'a str>, &'a str)> {
    let bytes = s.as_bytes();
    let mut depth = 0usize;
    let mut i = 0usize;
    while i < bytes.len() {
        match bytes[i] {
            b'(' => depth += 1,
            b')' => depth = depth.saturating_sub(1),
            c if depth == 0
                && i + name.len() <= bytes.len()
                && c.eq_ignore_ascii_case(&name[0])
                && bytes[i..i + name.len()].eq_ignore_ascii_case(name) =>
            {
                let prev_ok = i == 0 || !is_ident_byte(bytes[i - 1]);
                let next = i + name.len();
                if prev_ok && (next == bytes.len() || !is_ident_byte(bytes[next])) {
                    let mut j = next;
                    while j < bytes.len() && matches!(bytes[j], b' ' | b'\t' | b'\n') {
                        j += 1;
                    }
                    if j < bytes.len() && bytes[j] == b'(' {
                        let open = j;
                        let close = match_close_paren(bytes, open)?;
                        let inner = &s[open + 1..close];
                        let args = split_top_level_commas(inner);
                        let before = &s[..i];
                        let after = &s[close + 1..];
                        return Some((before, args, after));
                    }
                }
            }
            _ => {}
        }
        i += 1;
    }
    None
}

/// Find a top-level `exec2(a, b)` or `exec3(a, b, c)`.
fn find_top_level_exec_call(s: &str) -> Option<(&str, Vec<&str>, &str)> {
    find_top_level_named_call(s, b"exec2")
        .filter(|(_, args, _)| args.len() == 2)
        .or_else(|| find_top_level_named_call(s, b"exec3").filter(|(_, args, _)| args.len() == 3))
}

/// Find a top-level `while(body)`. Single-pass evaluation (see module-level
/// docs).
fn find_top_level_while_call(s: &str) -> Option<(&str, &str, &str)> {
    let (before, args, after) = find_top_level_named_call(s, b"while")?;
    if args.is_empty() {
        return None;
    }
    Some((before, args[0], after))
}