onedrop_engine/

warp_eval.rs

//! Per-vertex warp evaluation.
//!
//! Each frame the warp mesh's `cols × rows` vertices need their UVs warped
//! according to the preset's per-vertex (Milkdrop "per_pixel") equations and
//! the MilkDrop 2 zoom/rot/stretch/translate/warp formula.
//!
//! ## Hot path strategy
//!
//! [`WarpExecutor`] precompiles each per-vertex equation into an
//! `evalexpr::Node` once at preset load and then reuses **one** scratch
//! [`MilkContext`] across all vertices of a frame. This replaces the previous
//! "clone evaluator state per vertex" approach, which allocated a fresh
//! `HashMapContext` (≈60 vars + math-function map) for every one of the
//! `cols × rows` vertices — 768 allocations per frame at the default 32×24
//! mesh, up to 18 432 at 192×96.
//!
//! Per-vertex independence (each vertex sees the same per-frame state on
//! entry, regardless of what previous vertices wrote) is preserved by
//! snapshotting the motion outputs (`zoom`, `zoomexp`, `rot`, `cx`, `cy`,
//! `dx`, `dy`, `sx`, `sy`, `warp`) and the `q1..q32` channels at the start
//! of the frame and restoring them before each vertex eval.

use onedrop_eval::{CompiledBlock, MilkContext, MilkEvaluator};
use onedrop_renderer::warp_mesh::{WarpMesh, WarpMeshVertex};
use onedrop_renderer::warp_pipeline::WarpVertex;
use rayon::prelude::*;

/// Per-vertex motion outputs read from the evaluator after running per-vertex
/// equations.
#[derive(Clone, Copy, Debug)]
struct PerVertexMotion {
    zoom: f32,
    zoomexp: f32,
    rot: f32,
    cx: f32,
    cy: f32,
    dx: f32,
    dy: f32,
    sx: f32,
    sy: f32,
    warp: f32,
}

/// Number of `q*` channels we snapshot/restore between vertices. MilkDrop 2
/// exposes `q1..q32`; the underlying `MilkContext` reserves 64 slots but only
/// the first 32 are user-facing.
const Q_CHANNEL_COUNT: usize = 32;

/// Minimum compiled-equation count before we hand the per-vertex loop to
/// rayon. Empirically: presets with ~3 cheap equations stay faster
/// sequential (warp-exec time is dominated by per-thread `MilkContext`
/// clone + work-stealing overhead at ~1 ms/frame baseline cost), while
/// presets with 10+ equations see a 2-3× wall-clock win on a 12-thread
/// box. Picked 8 as the breakpoint after measuring three reference
/// presets (3 / 26 / 35 equations) on the `bench_render` tool.
const PARALLEL_EQ_THRESHOLD: usize = 8;

/// Pre-compiled per-vertex equation executor.
///
/// One instance lives for the lifetime of the engine; [`set_equations`] is
/// called whenever a new preset is loaded.
///
/// [`set_equations`]: WarpExecutor::set_equations
pub struct WarpExecutor {
    /// Source equations as last set, retained so `set_equations` can no-op
    /// when called with an unchanged set (common case during steady-state
    /// rendering).
    sources: Vec<String>,
    /// Compiled equations packaged in a [`CompiledBlock`] — auto-lowers
    /// to bytecode when every node is supported (~80 % of corpus
    /// per_pixel blocks), falls back to evalexpr Node walk otherwise.
    /// Per-vertex eval is the second-densest CPU phase after wave
    /// per_point: 768 vertices × N equations every frame, often in
    /// rayon. Skipping evalexpr's recursive operator dispatch on each
    /// op trims a measurable slice of `warp_compute`. May contain fewer
    /// nodes than `sources` if some equations failed to compile —
    /// failures are logged and silently dropped, matching MD2's
    /// "never crash on a bad preset" stance.
    compiled: CompiledBlock,
}

impl WarpExecutor {
    pub fn new() -> Self {
        Self {
            sources: Vec::new(),
            compiled: CompiledBlock::empty(),
        }
    }

    /// Replace the compiled equation set. No-op when `eqs` matches the
    /// previously compiled sources.
    ///
    /// `evaluator` is the per-frame evaluator; it is borrowed mutably so the
    /// preprocess pass can register any newly seen variable names on its
    /// context (preserving the same auto-init semantics as
    /// [`MilkEvaluator::eval`]).
    pub fn set_equations(&mut self, evaluator: &mut MilkEvaluator, eqs: &[String]) {
        if self.sources.len() == eqs.len() && self.sources.iter().zip(eqs).all(|(a, b)| a == b) {
            return;
        }
        self.sources = eqs.to_vec();
        let mut nodes = Vec::with_capacity(eqs.len());
        for eq in eqs {
            match evaluator.compile(eq) {
                Ok(node) => nodes.push(node),
                Err(e) => {
                    log::warn!("per-vertex equation compile failed: {eq:?}: {e}");
                }
            }
        }
        self.compiled = CompiledBlock::from_nodes(nodes, evaluator.context_mut());
    }

    /// Compute one [`WarpVertex`] per mesh vertex, ready to upload to the GPU.
    ///
    /// Pre-condition: `evaluator` has just finished its per-frame phase so
    /// motion outputs, audio levels, and `q*` channels are up to date.
    ///
    /// ## Parallelism
    ///
    /// When the preset has any per-vertex equations, the inner loop fans
    /// out across rayon's global thread pool: each worker gets one cloned
    /// [`MilkContext`] (via `map_init`) and chews through its slice of the
    /// 768-vertex mesh. Per-MD2-spec, per-vertex equations have no
    /// cross-vertex carryover, so the parallel order is invisible from the
    /// outside. The init clone is ~1 µs (HashMap with per-frame state plus
    /// math-function pointers), so amortised across hundreds of vertices per
    /// worker it's free.
    ///
    /// Empty-equation presets stay on the simple sequential path: nothing
    /// to evaluate, so threading would only add overhead.
    pub fn compute(
        &mut self,
        mesh: &WarpMesh,
        evaluator: &MilkEvaluator,
        time: f32,
    ) -> Vec<WarpVertex> {
        let base_motion = read_motion(evaluator.context());

        if self.compiled.is_empty() {
            // Fast path: no per-vertex equations means every vertex shares
            // the per-frame motion. Skip the scratch context entirely.
            return mesh
                .vertices
                .iter()
                .map(|v| WarpVertex {
                    pos_clip: v.pos_clip,
                    uv_warp: warp_uv_md2(v.uv_orig[0], v.uv_orig[1], v.rad, &base_motion, time),
                })
                .collect();
        }

        let base_ctx = evaluator.context().clone();
        let base_q = read_q_snapshot(&base_ctx);
        let compiled = &self.compiled;

        // Per-vertex body. Restores per-frame motion / q* state so the n-th
        // vertex sees the same starting context as the 0-th (per-vertex eqs
        // that depend on cross-vertex carryover are undefined in MilkDrop and
        // we explicitly do not implement those semantics). Captures
        // `base_motion`, `base_q`, `compiled`, and `time` from the enclosing
        // scope; takes the scratch context + the source vertex.
        //
        // `compiled.run(scratch)` picks the bytecode VM when the
        // per_pixel block lowered cleanly (the common case), otherwise
        // falls back to the evalexpr Node walk. Either way per-vertex
        // failures are silently absorbed — the vertex inherits whatever
        // state survived.
        let body = |scratch: &mut MilkContext, v: &WarpMeshVertex| -> WarpVertex {
            write_motion(scratch, &base_motion);
            write_q_snapshot(scratch, &base_q);

            scratch.set("x", v.uv_orig[0] as f64);
            scratch.set("y", v.uv_orig[1] as f64);
            scratch.set("rad", v.rad as f64);
            scratch.set("ang", v.ang as f64);

            compiled.run(scratch);

            let motion = read_motion(scratch);
            let uv_warp = warp_uv_md2(v.uv_orig[0], v.uv_orig[1], v.rad, &motion, time);
            WarpVertex {
                pos_clip: v.pos_clip,
                uv_warp,
            }
        };

        if compiled.len() >= PARALLEL_EQ_THRESHOLD {
            // Many or expensive equations — fan out across cores. One context
            // clone per worker thread, amortised across its slice of vertices.
            mesh.vertices
                .par_iter()
                .map_init(|| base_ctx.clone(), body)
                .collect()
        } else {
            // Few cheap equations — rayon's per-frame thread-pool overhead
            // (per-worker `MilkContext` clone + work-stealing) costs more
            // than the per-vertex eval saves. Stay sequential.
            let mut scratch = base_ctx;
            mesh.vertices
                .iter()
                .map(|v| body(&mut scratch, v))
                .collect()
        }
    }

    /// Number of equations currently compiled. Test/debug helper.
    pub fn compiled_count(&self) -> usize {
        self.compiled.len()
    }
}

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

fn read_motion(ctx: &MilkContext) -> PerVertexMotion {
    let f = |name: &str, default: f32| ctx.get_var(name).map(|v| v as f32).unwrap_or(default);
    PerVertexMotion {
        zoom: f("zoom", 1.0),
        zoomexp: f("zoomexp", 1.0),
        rot: f("rot", 0.0),
        cx: f("cx", 0.5),
        cy: f("cy", 0.5),
        dx: f("dx", 0.0),
        dy: f("dy", 0.0),
        sx: f("sx", 1.0),
        sy: f("sy", 1.0),
        warp: f("warp", 0.0),
    }
}

fn write_motion(ctx: &mut MilkContext, m: &PerVertexMotion) {
    ctx.set("zoom", m.zoom as f64);
    ctx.set("zoomexp", m.zoomexp as f64);
    ctx.set("rot", m.rot as f64);
    ctx.set("cx", m.cx as f64);
    ctx.set("cy", m.cy as f64);
    ctx.set("dx", m.dx as f64);
    ctx.set("dy", m.dy as f64);
    ctx.set("sx", m.sx as f64);
    ctx.set("sy", m.sy as f64);
    ctx.set("warp", m.warp as f64);
}

fn read_q_snapshot(ctx: &MilkContext) -> [f64; Q_CHANNEL_COUNT] {
    // q1..q32 live in MilkContext's array-backed q_vars; index by slot
    // and skip both the `"qN"` lookup table and the trait `get` route.
    let mut out = [0.0; Q_CHANNEL_COUNT];
    for (i, slot) in out.iter_mut().enumerate() {
        *slot = ctx.q_get_idx(i);
    }
    out
}

fn write_q_snapshot(ctx: &mut MilkContext, q: &[f64; Q_CHANNEL_COUNT]) {
    for (i, value) in q.iter().enumerate() {
        ctx.q_set_idx(i, *value);
    }
}

/// MilkDrop 2 warp UV formula.
///
/// Operates on normalized texture coordinates `(x, y) ∈ [0, 1]²`. The output
/// is the UV from which the fragment shader will sample `prev_main`.
///
/// Order of operations (matches butterchurn / projectM convention, which both
/// derive from `vis_milk2`):
/// 1. Logarithmic zoom around `(cx, cy)`, with curvature controlled by `zoomexp`.
/// 2. Rotation around `(cx, cy)` by `rot` radians.
/// 3. Anisotropic stretch around `(cx, cy)` by `(sx, sy)`.
/// 4. Translation by `(dx, dy)` (subtractive — `dx > 0` scrolls right).
/// 5. Optional sinusoidal warp displacement scaled by `warp`.
fn warp_uv_md2(x: f32, y: f32, rad: f32, m: &PerVertexMotion, time: f32) -> [f32; 2] {
    let zoomexp = if m.zoomexp.abs() < 1e-6 {
        1.0
    } else {
        m.zoomexp
    };
    let zoom_pow = zoomexp.powf(rad * 2.0 - 1.0);
    let zoom2 = m.zoom.powf(zoom_pow).max(1e-6);
    let inv_zoom2 = 1.0 / zoom2;

    let mut u = (x - m.cx) * inv_zoom2 + m.cx;
    let mut v = (y - m.cy) * inv_zoom2 + m.cy;

    let dx = u - m.cx;
    let dy = v - m.cy;
    let cos_r = m.rot.cos();
    let sin_r = m.rot.sin();
    u = dx * cos_r - dy * sin_r + m.cx;
    v = dx * sin_r + dy * cos_r + m.cy;

    let inv_sx = if m.sx.abs() < 1e-6 { 1.0 } else { 1.0 / m.sx };
    let inv_sy = if m.sy.abs() < 1e-6 { 1.0 } else { 1.0 / m.sy };
    u = (u - m.cx) * inv_sx + m.cx;
    v = (v - m.cy) * inv_sy + m.cy;

    u -= m.dx;
    v -= m.dy;

    if m.warp.abs() > 1e-6 {
        const WARP_SCALE: f32 = 0.0035;
        let f = m.warp * WARP_SCALE;
        u += f * (time * 0.913 + 10.0 * y).sin();
        v += f * (time * 0.952 + 10.0 * x).sin();
    }

    [u, v]
}

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

    fn run(mesh: &WarpMesh, eqs: &[String], evaluator: &mut MilkEvaluator) -> Vec<WarpVertex> {
        let mut exec = WarpExecutor::new();
        exec.set_equations(evaluator, eqs);
        exec.compute(mesh, evaluator, 0.0)
    }

    #[test]
    fn identity_when_motion_is_neutral() {
        let mesh = WarpMesh::new(4, 4, 1.0);
        let mut evaluator = MilkEvaluator::new();
        let warp = run(&mesh, &[], &mut evaluator);
        assert_eq!(warp.len(), mesh.vertices.len());
        for (i, v) in warp.iter().enumerate() {
            let orig = mesh.vertices[i].uv_orig;
            assert!(
                (v.uv_warp[0] - orig[0]).abs() < 1e-5,
                "u_warp mismatch at {i}"
            );
            assert!(
                (v.uv_warp[1] - orig[1]).abs() < 1e-5,
                "v_warp mismatch at {i}"
            );
        }
    }

    #[test]
    fn zoom_pulls_corner_uvs_toward_center() {
        let mesh = WarpMesh::new(3, 3, 1.0);
        let mut evaluator = MilkEvaluator::new();
        evaluator.context_mut().set_var("zoom", 2.0);
        let warp = run(&mesh, &[], &mut evaluator);

        let bl = warp[0].uv_warp;
        assert!(bl[0] > 0.0 && bl[0] < 0.5);
        assert!(bl[1] > 0.0 && bl[1] < 0.5);
    }

    #[test]
    fn rotation_swaps_axes_at_90_degrees() {
        let mesh = WarpMesh::new(3, 3, 1.0);
        let mut evaluator = MilkEvaluator::new();
        evaluator
            .context_mut()
            .set_var("rot", std::f64::consts::FRAC_PI_2);
        let warp = run(&mesh, &[], &mut evaluator);
        let right_mid = warp[3 + 2].uv_warp;
        let dist_from_orig = ((right_mid[0] - 1.0).powi(2) + (right_mid[1] - 0.5).powi(2)).sqrt();
        assert!(dist_from_orig > 0.4);
    }

    #[test]
    fn per_vertex_equation_modifies_motion_per_vertex() {
        let mesh = WarpMesh::new(3, 3, 1.0);
        let mut evaluator = MilkEvaluator::new();
        let eqs = vec!["zoom = 1 + rad".to_string()];
        let warp = run(&mesh, &eqs, &mut evaluator);

        let center = warp[3 + 1].uv_warp;
        assert!((center[0] - 0.5).abs() < 1e-4);
        assert!((center[1] - 0.5).abs() < 1e-4);

        let corner = warp[0].uv_warp;
        assert!(corner[0] > 0.0 && corner[0] < 0.5);
        assert!(corner[1] > 0.0 && corner[1] < 0.5);
    }

    #[test]
    fn vertices_are_independent_no_motion_carryover() {
        // A bug-prone case: every vertex doubles `zoom`. Without motion
        // snapshot restoration the n-th vertex would see zoom = 2^(n+1)
        // (state leaking from previous vertices). With restoration every
        // vertex starts from the per-frame zoom = 1 and ends at zoom = 2,
        // producing the same uniform warp toward the center.
        let mesh = WarpMesh::new(4, 4, 1.0);
        let mut evaluator = MilkEvaluator::new();
        let eqs = vec!["zoom = zoom * 2".to_string()];
        let warp = run(&mesh, &eqs, &mut evaluator);

        // Bottom-left corner: identical answer regardless of vertex order.
        // We re-run with the same mesh/eq and confirm bit-exact equality
        // for every vertex — proves there's no order-dependent carryover.
        let warp2 = run(&mesh, &eqs, &mut evaluator);
        assert_eq!(warp.len(), warp2.len());
        for i in 0..warp.len() {
            assert_eq!(
                warp[i].uv_warp, warp2[i].uv_warp,
                "vertex {i} is order-dependent: {:?} vs {:?}",
                warp[i].uv_warp, warp2[i].uv_warp
            );
        }

        // Spot-check: corner with rad ≈ 1 should warp toward center given
        // zoom = 2 (consistent with `zoom_pulls_corner_uvs_toward_center`).
        let bl = warp[0].uv_warp;
        assert!(bl[0] > 0.0 && bl[0] < 0.5);
        assert!(bl[1] > 0.0 && bl[1] < 0.5);
    }

    #[test]
    fn set_equations_is_idempotent_across_calls() {
        let mut evaluator = MilkEvaluator::new();
        let mut exec = WarpExecutor::new();
        let eqs = vec!["zoom = 1 + rad".to_string()];
        exec.set_equations(&mut evaluator, &eqs);
        assert_eq!(exec.compiled_count(), 1);

        // Re-setting the same equations should not recompile.
        exec.set_equations(&mut evaluator, &eqs);
        assert_eq!(exec.compiled_count(), 1);
    }

    #[test]
    fn q_channel_drives_per_vertex_zoom() {
        // The per-frame phase (simulated here by writing q1 directly on the
        // evaluator) must be visible to per-vertex equations. Without the
        // q-snapshot/restore in WarpExecutor::compute, the first vertex would
        // overwrite q1 (or its zoom dependency) for subsequent vertices.
        let mesh = WarpMesh::new(3, 3, 1.0);
        let mut evaluator = MilkEvaluator::new();
        evaluator.context_mut().set_var("q1", 1.0); // simulate per-frame write
        let eqs = vec!["zoom = 1 + q1".to_string()];
        let warp = run(&mesh, &eqs, &mut evaluator);

        // Center vertex: zoom = 2, identity around (0.5, 0.5) — but with
        // zoom = 2 the corner (0,0) should pull toward center.
        let bl = warp[0].uv_warp;
        assert!(
            bl[0] > 0.0 && bl[0] < 0.5,
            "q1=1 should make zoom=2 and pull bl corner; got u_warp={}",
            bl[0]
        );

        // With q1=0 the same eq reduces to zoom=1 (identity).
        evaluator.context_mut().set_var("q1", 0.0);
        let warp2 = run(&mesh, &eqs, &mut evaluator);
        let bl2 = warp2[0].uv_warp;
        assert!(
            (bl2[0] - 0.0).abs() < 1e-4,
            "q1=0 should yield zoom=1 (identity); got u_warp={}",
            bl2[0]
        );
    }

    #[test]
    fn q_snapshot_isolates_per_vertex_writes() {
        // A per-vertex eq that writes q1 must NOT contaminate the next
        // vertex's view of q1: every vertex should see the original
        // per-frame q1 value as its starting point.
        let mesh = WarpMesh::new(3, 3, 1.0);
        let mut evaluator = MilkEvaluator::new();
        evaluator.context_mut().set_var("q1", 0.5);
        let eqs = vec![
            "q1 = q1 * 2".to_string(),   // per-vertex doubles q1
            "zoom = 1 + q1".to_string(), // zoom depends on (modified) q1
        ];
        let warp = run(&mesh, &eqs, &mut evaluator);
        // After the first vertex, q1 in the scratch becomes 1.0. If
        // restoration didn't happen, the second vertex would see q1=1.0
        // (instead of the per-frame q1=0.5) and produce a different warp.
        // We re-run to confirm determinism: identical mesh + identical
        // per-frame state must yield bit-identical UVs.
        let warp2 = run(&mesh, &eqs, &mut evaluator);
        for i in 0..warp.len() {
            assert_eq!(
                warp[i].uv_warp, warp2[i].uv_warp,
                "vertex {i} non-deterministic"
            );
        }
        // Per-frame q1 itself is unchanged outside the executor (executor
        // only touches its own scratch).
        assert_eq!(evaluator.context().get_var("q1"), Some(0.5));
    }

    #[test]
    fn compile_failure_is_dropped_silently() {
        let mut evaluator = MilkEvaluator::new();
        let mut exec = WarpExecutor::new();
        let eqs = vec![
            "zoom = 1 + rad".to_string(),
            "this is not valid !!!".to_string(),
            "rot = 0.5".to_string(),
        ];
        exec.set_equations(&mut evaluator, &eqs);
        // 2 of 3 compile.
        assert_eq!(exec.compiled_count(), 2);
    }
}