onedrop_renderer/

render_chain.rs

//! Per-preset rendering pipeline bundle.
//!
//! Groups the eight per-frame GPU passes that the renderer dispatches in
//! order — warp, waveform, custom-wave, custom-shape, border, motion-vector,
//! blur pyramid, comp — into a single owning struct. The renderer
//! orchestrates these by calling [`RenderChain::record_passes`] inside its
//! command encoder.
//!
//! A chain owns its own pipelines + waveform sample cache so the renderer
//! can hold an optional secondary chain that fades the previous preset out
//! while the new one fades in during preset transitions.

use crate::blur_pipeline::BlurPipeline;
use crate::border::BorderRenderer;
use crate::chain_textures::ChainTextures;
use crate::comp_pipeline::CompPipeline;
use crate::custom_shape::{CustomShapeBatch, CustomShapeInstance, CustomShapeRenderer};
use crate::custom_wave::{CustomWaveBatch, CustomWaveRenderer, CustomWaveVertex};
use crate::error::Result;
use crate::gpu_context::GpuContext;
use crate::motion_vector::MotionVectorRenderer;
use crate::sprite_pipeline::{
    SpriteDrawCmd, SpriteFrame, SpritePipeline, SpritePool, build_sprite_uniform,
};
use crate::warp_pipeline::{WarpPipeline, WarpVertex};
use crate::waveform::{
    NUM_WAVE_SAMPLES, WaveUniforms, WaveformMode, WaveformRenderer, apply_smoothing,
    build_uniforms as build_wave_uniforms,
};

use crate::config::{RenderState, WaveParams};
use crate::warp_mesh::WarpMesh;
use onedrop_codegen::ShaderUniforms;

/// The eight per-frame pipelines plus the small per-chain caches and the
/// per-chain feedback textures (render / prev / final / blur1-3 / scratch).
pub struct RenderChain {
    /// Resolution-dependent feedback textures owned by this chain.
    /// Secondary chains during transitions get a fresh [`ChainTextures`]
    /// so the two preset states don't collide on the warp / prev / blur
    /// targets.
    textures: ChainTextures,
    warp: WarpPipeline,
    blur: BlurPipeline,
    comp: CompPipeline,
    waveform: WaveformRenderer,
    custom_wave: CustomWaveRenderer,
    custom_shape: CustomShapeRenderer,
    border: BorderRenderer,
    motion_vector: MotionVectorRenderer,
    /// Mono-mixed audio for the waveform pass, smoothed CPU-side per
    /// `WaveParams::smoothing` before upload.
    wave_samples: Vec<f32>,
    /// Instantaneous frame-mean volume used by `b_mod_wave_alpha_by_volume`.
    wave_volume: f32,
    /// Latest `(cols, rows, WarpVertex[])` snapshot from the most
    /// recent `update_warp_vertices` call. Read by the motion-vector
    /// pass to build anisotropic segments aligned with the local
    /// warp displacement.
    last_warp_field: Option<(u32, u32, Vec<WarpVertex>)>,
}

impl RenderChain {
    /// Build a chain against the supplied GPU context and warp mesh. The
    /// chain's pipelines bind the context's render/prev/blur texture views
    /// directly; callers must call [`Self::rebind_after_resize`] if the
    /// context is later resized.
    ///
    /// The mesh stays owned by the renderer (the engine evaluates per-vertex
    /// equations against it, so it has a single-owner constraint above the
    /// renderer level); the chain only needs it at construction time to size
    /// the vertex buffer.
    pub fn new(gpu: &GpuContext, mesh: &WarpMesh) -> Result<Self> {
        let textures = ChainTextures::new(&gpu.device, &gpu.config);
        Self::with_textures(gpu, mesh, textures)
    }

    /// Same as [`Self::new`] but uses a caller-provided [`ChainTextures`].
    /// Useful when reclaiming textures from a previous chain (e.g. swapping
    /// the secondary into the primary slot at the end of a transition)
    /// to avoid an unnecessary GPU allocation.
    pub fn with_textures(
        gpu: &GpuContext,
        mesh: &WarpMesh,
        textures: ChainTextures,
    ) -> Result<Self> {
        let warp = WarpPipeline::new(
            &gpu.device,
            gpu.config.texture_format.to_wgpu(),
            &textures.prev_texture_view,
            mesh,
        )?;

        let comp = CompPipeline::new(
            &gpu.device,
            &gpu.queue,
            gpu.config.texture_format.to_wgpu(),
            &textures.render_texture_view,
            &gpu.comp_aux_views_for(&textures),
        )?;

        let blur = BlurPipeline::new(
            &gpu.device,
            gpu.config.texture_format.to_wgpu(),
            gpu.config.width,
            gpu.config.height,
            &textures.render_texture_view,
            &textures.blur1_texture_view,
            &textures.blur2_texture_view,
            &textures.blur3_texture_view,
            &textures.blur1_scratch_texture_view,
            &textures.blur2_scratch_texture_view,
            &textures.blur3_scratch_texture_view,
        );

        let waveform = WaveformRenderer::new(&gpu.device, gpu.config.texture_format.to_wgpu());
        let custom_wave = CustomWaveRenderer::new(&gpu.device, gpu.config.texture_format.to_wgpu());
        let custom_shape = CustomShapeRenderer::new(
            &gpu.device,
            gpu.config.texture_format.to_wgpu(),
            &textures.prev_texture_view,
        );
        let border = BorderRenderer::new(&gpu.device, gpu.config.texture_format.to_wgpu());
        let motion_vector =
            MotionVectorRenderer::new(&gpu.device, gpu.config.texture_format.to_wgpu());

        Ok(Self {
            textures,
            warp,
            blur,
            comp,
            waveform,
            custom_wave,
            custom_shape,
            border,
            motion_vector,
            wave_samples: vec![0.0; NUM_WAVE_SAMPLES],
            wave_volume: 0.0,
            last_warp_field: None,
        })
    }

    /// Borrow the chain's feedback textures — needed by the renderer to
    /// read the final comp output (for blend / display) and to drive the
    /// `copy_to_prev` step that closes the per-frame feedback loop.
    pub fn textures(&self) -> &ChainTextures {
        &self.textures
    }

    /// Push freshly computed warp UVs to the GPU. Also caches the
    /// vertex array on the CPU so the motion-vector pass can sample
    /// the per-cell warp displacement (anisotropic segments) when it
    /// runs later in the same frame.
    pub fn update_warp_vertices(
        &mut self,
        queue: &wgpu::Queue,
        cols: u32,
        rows: u32,
        vertices: &[WarpVertex],
    ) {
        self.warp.update_vertices(queue, vertices);
        // Cheap: ~80 KiB at the 96 × 72 default mesh, written once
        // per frame. The motion-vector pass reads in the same frame
        // so the lifetime is implicit.
        match self.last_warp_field.as_mut() {
            Some((c, r, v)) => {
                *c = cols;
                *r = rows;
                v.clear();
                v.extend_from_slice(vertices);
            }
            None => {
                self.last_warp_field = Some((cols, rows, vertices.to_vec()));
            }
        }
    }

    /// Reallocate the warp pipeline's vertex/index buffers for a
    /// different-sized mesh. Driven by [`MilkRenderer::set_mesh_size`].
    pub fn rebuild_warp_mesh(&mut self, device: &wgpu::Device, mesh: &WarpMesh) {
        self.warp.rebuild_mesh(device, mesh);
    }

    /// Push a fresh mono audio sample window for the waveform pass.
    ///
    /// Applies the IIR smoothing dictated by `wave_params.smoothing` on the
    /// CPU side before upload; the buffer is clamped to [`NUM_WAVE_SAMPLES`]
    /// and shorter inputs are zero-padded.
    pub fn update_waveform_samples(
        &mut self,
        queue: &wgpu::Queue,
        samples: &[f32],
        instantaneous_volume: f32,
        wave_params: WaveParams,
    ) {
        let n = samples.len().min(NUM_WAVE_SAMPLES);
        self.wave_samples.clear();
        self.wave_samples.extend_from_slice(&samples[..n]);
        self.wave_samples.resize(NUM_WAVE_SAMPLES, 0.0);
        apply_smoothing(&mut self.wave_samples, wave_params.smoothing);
        self.wave_volume = instantaneous_volume.clamp(0.0, 1.0);
        self.waveform.update_wave_samples(queue, &self.wave_samples);
    }

    /// Stereo variant of [`Self::update_waveform_samples`]. Uploads
    /// both channels into the 2N-wide sample buffer so the waveform
    /// pass can read L from offset `0` and R from offset `N` in two
    /// dispatches.
    ///
    /// `instantaneous_volume` is computed from the perceptual downmix
    /// (engine-side) so it represents the listener's loudness, not one
    /// channel's. Smoothing is applied to each channel independently.
    pub fn update_waveform_samples_lr(
        &mut self,
        queue: &wgpu::Queue,
        left: &[f32],
        right: &[f32],
        instantaneous_volume: f32,
        wave_params: WaveParams,
    ) {
        // Reuse `wave_samples` for the left channel so `wave_samples()`
        // accessors (waveform-mode CPU consumers) still see something
        // representative.
        let nl = left.len().min(NUM_WAVE_SAMPLES);
        self.wave_samples.clear();
        self.wave_samples.extend_from_slice(&left[..nl]);
        self.wave_samples.resize(NUM_WAVE_SAMPLES, 0.0);
        apply_smoothing(&mut self.wave_samples, wave_params.smoothing);

        let mut right_samples = vec![0.0f32; NUM_WAVE_SAMPLES];
        let nr = right.len().min(NUM_WAVE_SAMPLES);
        right_samples[..nr].copy_from_slice(&right[..nr]);
        apply_smoothing(&mut right_samples, wave_params.smoothing);

        self.wave_volume = instantaneous_volume.clamp(0.0, 1.0);
        self.waveform
            .update_wave_samples_lr(queue, &self.wave_samples, &right_samples);
    }

    pub fn wave_samples(&self) -> &[f32] {
        &self.wave_samples
    }

    pub fn wave_volume(&self) -> f32 {
        self.wave_volume
    }

    /// Upload the per-frame custom-wave vertex stream.
    pub fn update_custom_waves(
        &mut self,
        queue: &wgpu::Queue,
        vertices: &[CustomWaveVertex],
        batches: &[CustomWaveBatch],
    ) {
        self.custom_wave.update(queue, vertices, batches);
    }

    pub fn custom_wave_vertex_count(&self) -> u32 {
        self.custom_wave.vertex_count()
    }

    pub fn custom_wave_batch_count(&self) -> usize {
        self.custom_wave.batch_count()
    }

    /// Upload the per-frame custom-shape instance stream.
    pub fn update_custom_shapes(
        &mut self,
        queue: &wgpu::Queue,
        instances: &[CustomShapeInstance],
        batches: &[CustomShapeBatch],
        aspect: f32,
    ) {
        self.custom_shape.update(queue, instances, batches, aspect);
    }

    pub fn custom_shape_instance_count(&self) -> u32 {
        self.custom_shape.instance_count()
    }

    pub fn custom_shape_batch_count(&self) -> usize {
        self.custom_shape.batch_count()
    }

    pub fn border_batch_count(&self) -> usize {
        self.border.batch_count()
    }

    pub fn motion_vector_segment_count(&self) -> u32 {
        self.motion_vector.segment_count()
    }

    /// Borrow the comp pipeline — needed by the renderer to swap in user
    /// shaders (the user-shader pipeline owns the texture binding plan and
    /// stays attached to MilkRenderer for now).
    pub fn comp_pipeline(&self) -> &CompPipeline {
        &self.comp
    }

    pub fn comp_pipeline_mut(&mut self) -> &mut CompPipeline {
        &mut self.comp
    }

    /// Borrow the warp pipeline. Mirrors [`Self::comp_pipeline`]; needed
    /// by the renderer's `set_user_warp_shader` path.
    pub fn warp_pipeline(&self) -> &WarpPipeline {
        &self.warp
    }

    pub fn warp_pipeline_mut(&mut self) -> &mut WarpPipeline {
        &mut self.warp
    }

    /// Record the eight per-frame passes into the supplied encoder, in
    /// MD2 order: warp → waveform (+ optional mirrored second pass for
    /// double_line) → custom waves → custom shapes → borders → motion
    /// vectors → blur pyramid → comp.
    ///
    /// Does **not** submit; the renderer is responsible for the submit and
    /// for the `copy_to_prev` step that closes the feedback loop.
    #[allow(clippy::too_many_arguments)]
    pub fn record_passes(
        &mut self,
        gpu: &GpuContext,
        encoder: &mut wgpu::CommandEncoder,
        state: &RenderState,
        comp_uniforms: &ShaderUniforms,
        sprite_pipeline: &SpritePipeline,
        sprite_pool: &SpritePool,
        sprite_frames: &[SpriteFrame],
    ) {
        // Per-frame uniforms.
        self.warp.update_uniforms(
            &gpu.queue,
            state.decay,
            gpu.aspect_ratio(),
            state.feedback.to_flags(),
        );
        // When a user warp shader is active it reads from the same
        // `ShaderUniforms` layout the comp pass does — feed it the same
        // per-frame values. No-op when the default warp path runs.
        self.warp.update_user_uniforms(&gpu.queue, comp_uniforms);
        self.comp.update_uniforms(&gpu.queue, comp_uniforms);

        // Build + upload wave pass uniforms from the current state.
        let wave_uniforms: WaveUniforms = build_wave_uniforms(
            state.wave,
            self.wave_volume,
            gpu.config.width,
            gpu.config.height,
            state.time,
        );
        self.waveform.update_uniforms(&gpu.queue, &wave_uniforms);

        // Borders + motion-vector grid. Both passes upload
        // per-frame; they take their no-op fast path when alpha is zero.
        self.border
            .update(&gpu.queue, state.borders, gpu.aspect_ratio());
        // Motion-vector grid: pull the per-cell direction from the
        // cached warp-vertex field when available, so segments point
        // along the local feedback flow instead of a flat horizontal
        // stub. `None` falls back to the original horizontal default.
        let warp_field =
            self.last_warp_field
                .as_ref()
                .map(|(c, r, v)| crate::motion_vector::WarpField {
                    cols: *c,
                    rows: *r,
                    vertices: v.as_slice(),
                });
        self.motion_vector
            .update(&gpu.queue, state.motion_vectors, warp_field);

        // Warp pass: render mesh sampling prev_texture into render_texture.
        self.warp
            .render(encoder, &self.textures.render_texture_view);

        // Waveform overlay on top of the warp output.
        let wp = state.wave;
        let is_dots = wp.dots || WaveformMode::from_i32(wp.mode) == WaveformMode::ExplosiveHash;
        if wp.split_lr {
            // Top-vs-bottom L/R split. Two dispatches: left channel
            // shifted to the upper half (`wave_y - 0.25`, reading
            // offset `0`), right channel to the lower half
            // (`wave_y + 0.25`, reading offset `NUM_WAVE_SAMPLES`).
            // Skip the `is_double_pass` mode-6 mirror — the split
            // already produces two traces, doubling again would
            // stack four.
            let mut left_uniforms = wave_uniforms;
            left_uniforms.wave_y -= 0.25;
            left_uniforms.sample_offset = 0;
            self.waveform.update_uniforms(&gpu.queue, &left_uniforms);
            self.waveform.render_with_blend(
                encoder,
                &self.textures.render_texture_view,
                is_dots,
                wp.additive,
            );

            let mut right_uniforms = wave_uniforms;
            right_uniforms.wave_y += 0.25;
            right_uniforms.sample_offset = NUM_WAVE_SAMPLES as u32;
            self.waveform.update_uniforms(&gpu.queue, &right_uniforms);
            self.waveform.render_with_blend(
                encoder,
                &self.textures.render_texture_view,
                is_dots,
                wp.additive,
            );
        } else {
            self.waveform.render_with_blend(
                encoder,
                &self.textures.render_texture_view,
                is_dots,
                wp.additive,
            );
            if WaveformMode::from_i32(wp.mode).is_double_pass() {
                let mut mirrored = wave_uniforms;
                mirrored.wave_y = (2.0 * wp.y) - wave_uniforms.wave_y;
                mirrored.wave_scale = -mirrored.wave_scale;
                self.waveform.update_uniforms(&gpu.queue, &mirrored);
                self.waveform.render_with_blend(
                    encoder,
                    &self.textures.render_texture_view,
                    is_dots,
                    wp.additive,
                );
            }
        }

        // Custom waves.
        self.custom_wave
            .render(encoder, &self.textures.render_texture_view);

        // Custom shapes (read prev_texture for textured mode).
        self.custom_shape
            .render(encoder, &self.textures.render_texture_view);

        // Borders + motion vectors.
        self.border
            .render(encoder, &self.textures.render_texture_view);
        self.motion_vector
            .render(encoder, &self.textures.render_texture_view);

        // Sprite pass: textured quads driven by the engine's
        // `MILK_IMG.INI` per-frame eval. Goes after borders/motion
        // vectors so it paints on top of the per-frame overlays but
        // before the blur pyramid + comp pass so its output still
        // feeds back through `prev_texture` next frame (needed for
        // `burn=1` persistence).
        if !sprite_frames.is_empty() {
            let cmds: Vec<SpriteDrawCmd> = sprite_frames
                .iter()
                .map(|f| {
                    let tex = sprite_pool.get_or_fallback(f.texture_index);
                    let uniform = build_sprite_uniform(
                        f.x,
                        f.y,
                        f.sx,
                        f.sy,
                        f.rot,
                        f.rgba,
                        tex.width,
                        tex.height,
                        gpu.config.width,
                        gpu.config.height,
                    );
                    SpriteDrawCmd {
                        uniform,
                        texture_view: &tex.view,
                        blend: f.blend,
                        burn: f.burn,
                    }
                })
                .collect();
            sprite_pipeline.record(
                encoder,
                &gpu.queue,
                &gpu.device,
                &self.textures.render_texture_view,
                sprite_pool.sampler(),
                &cmds,
            );
        }

        // Blur pyramid: 3 cumulative Gaussian blurs of the warp output.
        // Each level downsamples 2× from the previous (½ / ¼ / ⅛
        // render resolution) — see `BlurPipeline` for the rationale.
        self.blur.render(
            encoder,
            &self.textures.blur1_texture_view,
            &self.textures.blur2_texture_view,
            &self.textures.blur3_texture_view,
            &self.textures.blur1_scratch_texture_view,
            &self.textures.blur2_scratch_texture_view,
            &self.textures.blur3_scratch_texture_view,
        );

        // Comp pass: read render_texture (current warp) + prev_texture (LAST
        // frame's warp output, for echo), apply gamma_adj, write final_texture.
        self.comp.render(encoder, &self.textures.final_texture_view);
    }

    /// Copy this chain's `render_texture` into its `prev_texture` so the
    /// next frame's warp pass sees the current frame's output as feedback.
    /// Caller drives this via the renderer's encoder after `record_passes`.
    pub fn copy_to_prev(&self, encoder: &mut wgpu::CommandEncoder, gpu: &GpuContext) {
        self.textures.copy_to_prev(encoder, &gpu.config);
        // Snapshot this frame's blur pyramid so a user-translated warp
        // shader on the *next* frame can sample `GetBlur1/2/3(uv)` from
        // a real Gaussian-blurred copy (one frame delayed) instead of
        // the un-blurred prev_texture fallback. Cheap — three GPU
        // texture-to-texture copies on already-resident memory.
        self.textures.copy_blur_to_prev(encoder, &gpu.config);
    }

    /// Re-allocate per-chain textures at the new resolution and rebind every
    /// pass. Replaces both the renderer-level resize plumbing and the per-
    /// chain rebind that V.0 left in place.
    pub fn resize(&mut self, gpu: &GpuContext) {
        self.textures.resize(&gpu.device, &gpu.config);
        self.warp
            .rebind_prev_texture(&gpu.device, &self.textures.prev_texture_view);
        self.comp.rebind_input_texture(
            &gpu.device,
            &self.textures.render_texture_view,
            &gpu.comp_aux_views_for(&self.textures),
        );
        self.blur.rebind(
            &gpu.device,
            &gpu.queue,
            gpu.config.width,
            gpu.config.height,
            &self.textures.render_texture_view,
            &self.textures.blur1_texture_view,
            &self.textures.blur2_texture_view,
            &self.textures.blur1_scratch_texture_view,
            &self.textures.blur2_scratch_texture_view,
            &self.textures.blur3_scratch_texture_view,
        );
        // The shape pass binds `prev_texture` for textured mode.
        self.custom_shape
            .rebind_prev_texture(&gpu.device, &self.textures.prev_texture_view);
    }
}