onedrop_engine/engine/

wave_phase.rs

//! Custom-wave (`wavecode_N`) phase of the engine update loop.
//!
//! Runs each enabled wave's pre-compiled `per_frame_init` (once per preset
//! load), `per_frame`, and `per_point` blocks, then returns the resulting
//! vertex stream + per-wave batches. Callers (engine) push the result to
//! the right [`onedrop_renderer::RenderChain`] — primary for the current
//! preset, secondary for the outgoing preset during a transition.
//!
//! ## Parallelism
//!
//! Each of the up-to-4 enabled wavecode blocks is independent (per-MD2,
//! waves don't share state across blocks within one frame), so the four
//! per-wave per_frame + 512-sample loops fan out across rayon workers.
//! The dominant cost in 144.milk / 207.milk class presets is custom
//! waves' per_point eval (~37 ms / frame on a 12-thread box,
//! single-threaded). After this parallelisation: ~10-12 ms.

use super::preset_slot::{CompiledWave, PresetSlot, run_block_with_logger};
use onedrop_eval::{MilkEvaluator, WavePoint};
use onedrop_parser::preset::WaveCode;
use onedrop_renderer::{CustomWaveBatch, CustomWaveVertex};
use rayon::prelude::*;

/// Minimum sample count to bother parallelising the per_point loop.
/// Below this, rayon's map_init overhead (per-worker `MilkEvaluator`
/// clone + work-stealing dispatch) exceeds the eval time we'd save —
/// stay sequential. 64 lines up roughly with the smallest "dense" waves
/// we care about (the typical preset uses 256-512).
const PARALLEL_SAMPLES_THRESHOLD: usize = 64;

/// Output of one wave's sample loop: its vertex slice + the batch metadata
/// describing how to draw it. Collected per wave so the parallel pass
/// doesn't need a shared mutex on the global vertex buffer.
struct WaveOutput {
    vertices: Vec<CustomWaveVertex>,
    batch: Option<CustomWaveBatch>,
}

/// Run every enabled `wavecode_N` block in `slot`'s preset and return the
/// per-frame vertex stream + batches. Returns empty `Vec`s when no preset is
/// loaded or no waves are enabled — feed those through to the renderer's
/// `update_custom_waves` to reset its buffers.
pub(super) fn compute_custom_waves(
    slot: &mut PresetSlot,
    left: &[f32],
    right: &[f32],
    spectrum_bins: &[f32],
) -> (Vec<CustomWaveVertex>, Vec<CustomWaveBatch>) {
    let Some(preset) = slot.preset.as_ref() else {
        return (Vec::new(), Vec::new());
    };
    if preset.waves.is_empty()
        || preset.waves.iter().all(|w| !w.enabled)
        || slot.compiled_waves.is_empty()
    {
        return (Vec::new(), Vec::new());
    }

    // Per-wave parallel pass. Each worker gets its own `MilkEvaluator`
    // clone, runs the wave's init (if needed) + per_frame + the 512-sample
    // per_point loop, and returns its vertex slice. Sample loop *inside*
    // a wave stays sequential because `current` threads state across
    // samples — per-MD2 idiomatic.
    //
    // Disjoint field-borrows on `slot` let us pass `&preset.waves`,
    // `&slot.compiled_waves` and `&slot.waves_need_init` straight into
    // the closure — no per-frame `Vec<WaveCode>` clone (each
    // `WaveCode` carries `Vec<String>` equation lists, so cloning all
    // 4 of them every frame was actively burning allocator time).
    let wave_outputs: Vec<WaveOutput> = {
        let waves: &[WaveCode] = &preset.waves;
        let base_eval = slot.evaluator.clone();
        let compiled_waves = &slot.compiled_waves;
        let needs_init = &slot.waves_need_init;
        waves
            .par_iter()
            .enumerate()
            .map(|(i, wave)| {
                compute_one_wave(
                    i,
                    wave,
                    compiled_waves.get(i),
                    needs_init.get(i).copied().unwrap_or(false),
                    base_eval.clone(),
                    left,
                    right,
                    spectrum_bins,
                )
            })
            .collect()
    };

    // Mark all waves as initialised — the actual init ran in worker
    // clones, but the flag lives on the main slot.
    for flag in slot.waves_need_init.iter_mut() {
        *flag = false;
    }

    // Stitch per-wave vertex slices into one contiguous buffer; fix up the
    // batch offsets to point at the global vertex array.
    let total_verts: usize = wave_outputs.iter().map(|o| o.vertices.len()).sum();
    let mut vertices: Vec<CustomWaveVertex> = Vec::with_capacity(total_verts);
    let mut batches: Vec<CustomWaveBatch> = Vec::with_capacity(wave_outputs.len());
    for out in wave_outputs {
        let offset = vertices.len() as u32;
        let count = out.vertices.len() as u32;
        vertices.extend(out.vertices);
        if let Some(mut batch) = out.batch {
            batch.start_vertex = offset;
            batch.vertex_count = count;
            batches.push(batch);
        }
    }

    (vertices, batches)
}

#[allow(clippy::too_many_arguments)]
fn compute_one_wave(
    i: usize,
    wave: &WaveCode,
    cw_opt: Option<&CompiledWave>,
    needs_init: bool,
    mut eval: MilkEvaluator,
    left: &[f32],
    right: &[f32],
    spectrum_bins: &[f32],
) -> WaveOutput {
    let Some(cw) = cw_opt else {
        return WaveOutput {
            vertices: Vec::new(),
            batch: None,
        };
    };
    if !wave.enabled {
        return WaveOutput {
            vertices: Vec::new(),
            batch: None,
        };
    }

    if needs_init {
        run_block_with_logger(&mut eval, &cw.init, |idx, e| {
            log::warn!("wave[{}] init equation[{}] failed: {}", i, idx, e);
        });
    }
    run_block_with_logger(&mut eval, &cw.per_frame, |idx, e| {
        log::warn!("wave[{}] per_frame equation[{}] failed: {}", i, idx, e);
    });

    let samples = wave.samples.clamp(2, 512) as usize;
    let scaled = wave.scaling as f64;
    let mut current = WavePoint {
        sample: 0.0,
        value1: 0.0,
        value2: 0.0,
        x: 0.5,
        y: 0.5,
        r: wave.r as f64,
        g: wave.g as f64,
        b: wave.b as f64,
        a: (wave.a as f64).clamp(0.0, 1.0),
    };

    // Decide once whether to fan samples out across rayon workers.
    // `analyse_per_point` confirmed the block has no read of
    // `x`/`y`/`r`/`g`/`b`/`a` before write, no writes to qN/custom
    // vars, and no stateful function calls — so sample S does not feed
    // sample S+1 and we're free to reorder them. Tiny waves stay
    // sequential because rayon's per-worker `MilkEvaluator` clone
    // (~few µs) would cost more than the eval saves.
    let parallel_samples = !cw.per_point.is_empty()
        && cw.per_point_parallelism.is_safe()
        && samples >= PARALLEL_SAMPLES_THRESHOLD;

    // Pre-compute the per-sample (v1, v2) audio/FFT pair so both paths
    // share the same input mapping. Time-domain mode pulls `value1`
    // from the left channel and `value2` from the right; b_spectrum
    // mode keeps the N/2 mirror trick because the upstream FFT runs
    // on a downmix (one buffer of magnitudes) and there's no
    // meaningful "stereo bin" pair.
    let sample_inputs: Vec<(f64, f64, f64)> = (0..samples)
        .map(|s| {
            let sample_norm = s as f64 / (samples - 1) as f64;
            let (v1, v2) = if wave.b_spectrum {
                let n = spectrum_bins.len().max(1);
                let idx = ((sample_norm * (n - 1) as f64) as usize).min(n - 1);
                let mirror = (idx + n / 2) % n;
                (spectrum_bins[idx] as f64, spectrum_bins[mirror] as f64)
            } else if !left.is_empty() {
                let n_l = left.len();
                let idx_l = (s * n_l / samples).min(n_l - 1);
                let v1 = left[idx_l] as f64;
                // Right channel — fall back to left if the caller
                // passed an empty/short right buffer (mono callers
                // routed through `update`). When stereo is real,
                // `right` is the same length as `left`.
                let v2 = if !right.is_empty() {
                    let n_r = right.len();
                    let idx_r = (s * n_r / samples).min(n_r - 1);
                    right[idx_r] as f64
                } else {
                    v1
                };
                (v1, v2)
            } else {
                (0.0, 0.0)
            };
            (sample_norm, v1 * scaled, v2 * scaled)
        })
        .collect();

    // Three geometry modes:
    // - dots: one screen-space quad per point (6 verts × samples)
    // - thick lines (!dots && b_draw_thick): one quad per *segment*
    //   between consecutive points (6 verts × (samples - 1))
    // - thin lines (!dots && !b_draw_thick): LineStrip, one vert per
    //   point — 1 px wide
    let thick_lines = !wave.b_use_dots && wave.b_draw_thick;
    let mut vertices: Vec<CustomWaveVertex> = if wave.b_use_dots {
        Vec::with_capacity(samples * 6)
    } else if thick_lines {
        Vec::with_capacity(samples.saturating_sub(1) * 6)
    } else {
        Vec::with_capacity(samples)
    };

    // Half-thickness of the thick-line stroke, in clip space. ~0.006
    // matches the static waveform pass's `b_wave_thick` lookup, so
    // both flavours of thick lines look identical at the same render
    // resolution. Dots stay at the historical 0.012 / 0.006 split.
    const THICK_LINE_CLIP: f32 = 0.006;

    // Body shared by all paths: turn a per-eval point into the vertex
    // stream the renderer consumes. For thick lines, the segment quad
    // needs two consecutive points; the caller threads `prev_thick`
    // through and emits the quad once both ends exist.
    let push_vertices = |vertices: &mut Vec<CustomWaveVertex>,
                         prev_thick: &mut Option<WavePoint>,
                         out: WavePoint| {
        let pos = onedrop_renderer::custom_wave::preset_xy_to_clip(out.x, out.y);
        let color = [
            out.r.clamp(0.0, 1.0) as f32,
            out.g.clamp(0.0, 1.0) as f32,
            out.b.clamp(0.0, 1.0) as f32,
            out.a.clamp(0.0, 1.0) as f32,
        ];
        if wave.b_use_dots {
            let radius = if wave.b_draw_thick { 0.012 } else { 0.006 };
            let quad =
                onedrop_renderer::custom_wave::point_to_dot_quad(out.x, out.y, color, radius);
            vertices.extend_from_slice(&quad);
        } else if thick_lines {
            if let Some(prev) = *prev_thick {
                let color_prev = [
                    prev.r.clamp(0.0, 1.0) as f32,
                    prev.g.clamp(0.0, 1.0) as f32,
                    prev.b.clamp(0.0, 1.0) as f32,
                    prev.a.clamp(0.0, 1.0) as f32,
                ];
                let quad = onedrop_renderer::custom_wave::segment_to_thick_quad(
                    prev.x,
                    prev.y,
                    color_prev,
                    out.x,
                    out.y,
                    color,
                    THICK_LINE_CLIP,
                );
                vertices.extend_from_slice(&quad);
            }
            *prev_thick = Some(out);
        } else {
            vertices.push(CustomWaveVertex { pos, color });
        }
    };

    // Bytecode VM path is preferred when the per_point block lowered
    // cleanly (no `rand` / `gmegabuf` / unsupported ops). It skips
    // evalexpr's tree-walking dispatch and runs a flat opcode loop
    // against array-backed hot vars. The Nodes path stays as
    // observation-compatible fallback.
    let per_point_bc = cw.per_point_bytecode.as_ref();

    if parallel_samples {
        // Carry-free path: every sample starts from the same `current`
        // seed (post-per_frame, pre-loop). Splitting across rayon
        // workers gives a ~4-8× speedup on dense waves. Each worker
        // gets one `MilkEvaluator` clone (~slim — see the eval crate's
        // `Clone` impl), amortised across its slice of the sample
        // range.
        let seed = current;
        let per_point = &cw.per_point;
        let base_eval = &eval;
        let outputs: Vec<WavePoint> = sample_inputs
            .par_iter()
            .map_init(
                || base_eval.clone(),
                |worker_eval, &(sample_norm, v1s, v2s)| {
                    let mut p = seed;
                    p.sample = sample_norm;
                    p.value1 = v1s;
                    p.value2 = v2s;
                    if let Some(bc) = per_point_bc {
                        worker_eval.run_per_point_bc(p, bc)
                    } else {
                        worker_eval.eval_per_point(p, per_point).unwrap_or(p)
                    }
                },
            )
            .collect();
        let mut prev_thick: Option<WavePoint> = None;
        for out in outputs {
            push_vertices(&mut vertices, &mut prev_thick, out);
        }
    } else {
        // Sequential path — preserves MD2's "trail across samples"
        // semantics for blocks that read x/y/r/g/b/a as carry.
        let mut prev_thick: Option<WavePoint> = None;
        for (s, &(sample_norm, v1s, v2s)) in sample_inputs.iter().enumerate() {
            current.sample = sample_norm;
            current.value1 = v1s;
            current.value2 = v2s;
            let out = if cw.per_point.is_empty() {
                WavePoint {
                    x: sample_norm,
                    y: 0.5 - v1s * 0.5,
                    ..current
                }
            } else if let Some(bc) = per_point_bc {
                eval.run_per_point_bc(current, bc)
            } else {
                match eval.eval_per_point(current, &cw.per_point) {
                    Ok(p) => p,
                    Err(e) => {
                        log::warn!("wave[{}] per_point failed at s={}: {}", i, s, e);
                        current
                    }
                }
            };
            current = out;
            push_vertices(&mut vertices, &mut prev_thick, out);
        }
    }

    let count = vertices.len() as u32;
    let batch = if count > 0 {
        Some(CustomWaveBatch {
            start_vertex: 0, // fixed up by caller after stitching
            vertex_count: count,
            // Thick lines emit triangle quads CPU-side, so dispatch
            // them through the same TriangleList pipeline the dots
            // path uses. The flag name kept for backward compat —
            // see `CustomWaveBatch::dots` doc.
            dots: wave.b_use_dots || thick_lines,
            additive: wave.b_additive,
        })
    } else {
        None
    };

    WaveOutput { vertices, batch }
}

#[cfg(test)]
mod tests {
    use super::*;
    use onedrop_eval::{MilkEvaluator, PerPointParallelism, analyse_per_point};
    use onedrop_parser::preset::WaveCode;

    /// Compile a small per_point block in isolation so we can drive
    /// `compute_one_wave` from a unit test without a full preset.
    fn compile_per_point(eqs: &[&str]) -> (Vec<onedrop_eval::Node>, PerPointParallelism) {
        let mut eval = MilkEvaluator::new();
        let owned: Vec<String> = eqs.iter().map(|s| s.to_string()).collect();
        let nodes = eval.compile_batch(&owned).unwrap();
        let par = analyse_per_point(&nodes);
        (nodes, par)
    }

    fn make_wave(samples: i32) -> WaveCode {
        WaveCode {
            index: 0,
            enabled: true,
            samples,
            sep: 0,
            b_spectrum: false,
            b_use_dots: false,
            b_draw_thick: false,
            b_additive: false,
            scaling: 1.0,
            smoothing: 0.0,
            r: 0.25,
            g: 0.5,
            b: 0.75,
            a: 1.0,
            per_frame_equations: Vec::new(),
            per_point_equations: Vec::new(),
            per_frame_init_equations: Vec::new(),
        }
    }

    fn run_wave(
        wave: &WaveCode,
        cw: &CompiledWave,
        audio: &[f32],
        spectrum: &[f32],
    ) -> Vec<CustomWaveVertex> {
        let eval = MilkEvaluator::new();
        // Tests pass mono audio; route left = right so the time-domain
        // path's `value2` falls back to mirroring `value1`.
        compute_one_wave(0, wave, Some(cw), false, eval, audio, audio, spectrum).vertices
    }

    /// Carry-free per_point block: parallel and sequential paths must
    /// produce vertex-for-vertex identical output. This is the
    /// correctness anchor for `analyse_per_point` + the rayon fan-out.
    #[test]
    fn parallel_and_sequential_match_on_carry_free_block() {
        let (nodes, par) = compile_per_point(&[
            "x = sample",
            "y = 0.5 - value1 * 0.5",
            "r = sample",
            "g = 1 - sample",
            "b = 0.5",
            "a = 1",
        ]);
        assert_eq!(par, PerPointParallelism::Safe);

        // Force sequential: pretend carry was detected.
        let cw_seq = CompiledWave {
            init: onedrop_eval::CompiledBlock::empty(),
            per_frame: onedrop_eval::CompiledBlock::empty(),
            per_point: nodes.clone(),
            per_point_parallelism: PerPointParallelism::Sequential,
            per_point_bytecode: None,
        };
        // Run parallel: real analyser verdict.
        let cw_par = CompiledWave {
            init: onedrop_eval::CompiledBlock::empty(),
            per_frame: onedrop_eval::CompiledBlock::empty(),
            per_point: nodes,
            per_point_parallelism: par,
            per_point_bytecode: None,
        };

        // 256 samples > 64 threshold so the parallel path is actually
        // taken — anything below threshold silently falls back to
        // sequential and the test wouldn't prove anything.
        let wave = make_wave(256);
        let audio: Vec<f32> = (0..512).map(|i| (i as f32 * 0.01).sin()).collect();

        let seq = run_wave(&wave, &cw_seq, &audio, &[]);
        let par_out = run_wave(&wave, &cw_par, &audio, &[]);

        assert_eq!(seq.len(), par_out.len());
        for (a, b) in seq.iter().zip(par_out.iter()) {
            assert_eq!(a.pos, b.pos, "vertex pos divergence");
            assert_eq!(a.color, b.color, "vertex color divergence");
        }
    }

    /// Carry-dependent block must stay on the sequential path even
    /// when the engine wires up `PerPointParallelism::Safe` somehow
    /// (defence in depth: the analyser is the gatekeeper, but a
    /// carry-needing block run in parallel would silently produce a
    /// different output and we want a test that pins that).
    #[test]
    fn carry_block_is_marked_sequential_by_analyser() {
        let (_, par) = compile_per_point(&["x = x + 0.001", "y = sample"]);
        assert_eq!(par, PerPointParallelism::Sequential);
    }

    /// Parity anchor for the bytecode VM at the wave level: same
    /// per_point block, one CompiledWave with the bytecode path
    /// enabled and one with it disabled. Vertex output must be
    /// identical (or differ by at most a float ulp, which `assert_eq`
    /// on `[f32; 4]` color and clip-space pos catches because both
    /// paths run the same arithmetic on the same f64s).
    #[test]
    fn bytecode_and_evalexpr_paths_match() {
        let (nodes, par) = compile_per_point(&[
            "x = sample",
            "y = 0.5 - value1 * 0.5",
            "r = sin(sample * 6.2831853) * 0.5 + 0.5",
            "g = 1.0 - sample",
            "b = sqrt(abs(value1))",
            "a = 1.0",
        ]);
        assert_eq!(par, PerPointParallelism::Safe);

        // The nodes were compiled against a throw-away evaluator inside
        // `compile_per_point`; for bytecode lowering we only need any
        // `MilkContext` whose cold slab can intern the same names.
        let mut bc_eval = MilkEvaluator::new();
        let bc = onedrop_eval::CompiledBytecode::try_compile(&nodes, bc_eval.context_mut())
            .expect("block lowers to bytecode");

        let cw_nodes = CompiledWave {
            init: onedrop_eval::CompiledBlock::empty(),
            per_frame: onedrop_eval::CompiledBlock::empty(),
            per_point: nodes.clone(),
            per_point_parallelism: PerPointParallelism::Sequential,
            per_point_bytecode: None,
        };
        let cw_bc = CompiledWave {
            init: onedrop_eval::CompiledBlock::empty(),
            per_frame: onedrop_eval::CompiledBlock::empty(),
            per_point: nodes,
            per_point_parallelism: PerPointParallelism::Sequential,
            per_point_bytecode: Some(bc),
        };

        let wave = make_wave(128);
        let audio: Vec<f32> = (0..512).map(|i| (i as f32 * 0.013).sin()).collect();
        let from_nodes = run_wave(&wave, &cw_nodes, &audio, &[]);
        let from_bc = run_wave(&wave, &cw_bc, &audio, &[]);

        assert_eq!(from_nodes.len(), from_bc.len());
        for (a, b) in from_nodes.iter().zip(from_bc.iter()) {
            assert_eq!(a.pos, b.pos, "pos divergence");
            assert_eq!(a.color, b.color, "color divergence");
        }
    }
}