onedrop_renderer/

blur_pipeline.rs

//! Blur-pyramid GPU pipeline.
//!
//! Builds the 3-level cumulative Gaussian blur (Blur1 → Blur2 → Blur3) that
//! MD2 user comp shaders sample via `GetBlur1/2/3`. Each level is the
//! previous level + one more separable Gaussian pass, so the effective σ
//! grows ~quadratically: Blur1 is mildly soft, Blur3 is markedly so.
//!
//! ## Pipeline shape
//!
//! Each blur level renders at half the previous level's resolution:
//! Blur1 = ½ render, Blur2 = ¼, Blur3 = ⅛. Per level we do two
//! fragment-shader passes (separable Gaussian):
//!
//! 1. **Horizontal**: read from the previous-level texture (Blur0 =
//!    `render_texture` at full res), write to that level's
//!    `blurN_scratch_texture` (already at the downsampled
//!    resolution). The vertex viewport defaults to the destination
//!    texture's size, so the read happens via the hardware bilinear
//!    sampler — implicit downsample.
//! 2. **Vertical**: read from `blurN_scratch_texture`, write to that
//!    level's `blurN_texture` (same resolution).
//!
//! After the third level, the three blur textures hold progressively
//! softer copies of the warp pass output, ready to be bound into the
//! comp pass's bind group at bindings 3/4/5. User shaders' UV-space
//! `GetBlur1/2/3(uv)` sample with the same normalised UV the comp
//! pass already uses — the smaller texture is transparent to the
//! caller because the bilinear sampler upscales at sample time.
//!
//! The shader itself is fixed at a 9-tap kernel for WGSL-compiler-
//! friendly unrolling; the per-pass `texel_size` uniform tracks the
//! source texture's resolution so each tap walks one source-texel
//! regardless of which level we're at.

use bytemuck::{Pod, Zeroable};
use wgpu::util::DeviceExt;

/// Sigmas (in texels) for each blur level's H+V pass. Each level reads
/// from the previous one, so the cumulative σ grows like sqrt(Σ σ²).
/// The values below give visually distinct levels: Blur1 just-soft,
/// Blur2 noticeably hazy, Blur3 strong halos — close to what MD2
/// presets expect.
const SIGMA_BLUR1: f32 = 2.0;
const SIGMA_BLUR2: f32 = 3.0;
const SIGMA_BLUR3: f32 = 4.0;

#[repr(C)]
#[derive(Debug, Clone, Copy, Pod, Zeroable)]
struct BlurUniforms {
    direction: [f32; 2],
    texel_size: [f32; 2],
    sigma: f32,
    _pad0: f32,
    _pad1: f32,
    _pad2: f32,
}

/// Drives the 6 blur passes (3 levels × H+V). One pipeline, six bind
/// groups — each pass needs its own bind group because the source
/// texture changes between passes.
pub struct BlurPipeline {
    pipeline: wgpu::RenderPipeline,
    bind_group_layout: wgpu::BindGroupLayout,
    sampler: wgpu::Sampler,
    /// One uniform buffer per pass — 6 in total. Cheap; uniform buffers
    /// are tiny and writing different sigmas to one buffer mid-frame
    /// would require either dynamic offsets or per-pass write_buffer
    /// calls with a barrier, neither worth the complexity here.
    uniform_buffers: [wgpu::Buffer; 6],
    /// One bind group per pass, holding (uniforms[i], source texture,
    /// sampler). Recreated when textures rebind (resize).
    bind_groups: [wgpu::BindGroup; 6],
}

impl BlurPipeline {
    /// Build the pipeline against `target_format` and wire bind groups
    /// for the 6 blur passes. The pyramid is downsampled at each level
    /// (Blur1 = ½ render, Blur2 = ¼, Blur3 = ⅛); per-level scratches
    /// match their level's resolution.
    ///
    /// The pass order is fixed at construction:
    /// - pass 0 (H, σ=σ1): src = `render_view`,        dst = `b1_scratch`
    /// - pass 1 (V, σ=σ1): src = `b1_scratch`,         dst = `blur1_view`
    /// - pass 2 (H, σ=σ2): src = `blur1_view`,         dst = `b2_scratch`
    /// - pass 3 (V, σ=σ2): src = `b2_scratch`,         dst = `blur2_view`
    /// - pass 4 (H, σ=σ3): src = `blur2_view`,         dst = `b3_scratch`
    /// - pass 5 (V, σ=σ3): src = `b3_scratch`,         dst = `blur3_view`
    ///
    /// `texel_size` is set per-pass to `1 / source-resolution`, so each
    /// Gaussian tap walks one source-texel in UV space regardless of
    /// which level we're at.
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        device: &wgpu::Device,
        target_format: wgpu::TextureFormat,
        width: u32,
        height: u32,
        render_view: &wgpu::TextureView,
        blur1_view: &wgpu::TextureView,
        blur2_view: &wgpu::TextureView,
        blur3_view: &wgpu::TextureView,
        blur1_scratch_view: &wgpu::TextureView,
        blur2_scratch_view: &wgpu::TextureView,
        blur3_scratch_view: &wgpu::TextureView,
    ) -> Self {
        // `blur3_view` is referenced only as a render target (see `render`),
        // not as a source, but we still take it in `new` so the caller's
        // ownership model stays symmetric across the four blur textures.
        let _ = blur3_view;
        // Per-source-resolution texel sizes. Pass i reads from
        // `pass_source_texels[i]`, so the shader's `texel_size`
        // matches that source's stride in normalised UV.
        let (w1, h1) = (width.max(2) / 2, height.max(2) / 2);
        let (w2, h2) = (width.max(4) / 4, height.max(4) / 4);
        let (w3, h3) = (width.max(8) / 8, height.max(8) / 8);
        let texel_full = [1.0 / width as f32, 1.0 / height as f32];
        let texel_b1 = [1.0 / w1 as f32, 1.0 / h1 as f32];
        let texel_b2 = [1.0 / w2 as f32, 1.0 / h2 as f32];
        let texel_b3 = [1.0 / w3 as f32, 1.0 / h3 as f32];

        let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("Blur BGL"),
            entries: &[
                wgpu::BindGroupLayoutEntry {
                    binding: 0,
                    visibility: wgpu::ShaderStages::FRAGMENT,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Uniform,
                        has_dynamic_offset: false,
                        min_binding_size: None,
                    },
                    count: None,
                },
                wgpu::BindGroupLayoutEntry {
                    binding: 1,
                    visibility: wgpu::ShaderStages::FRAGMENT,
                    ty: wgpu::BindingType::Texture {
                        sample_type: wgpu::TextureSampleType::Float { filterable: true },
                        view_dimension: wgpu::TextureViewDimension::D2,
                        multisampled: false,
                    },
                    count: None,
                },
                wgpu::BindGroupLayoutEntry {
                    binding: 2,
                    visibility: wgpu::ShaderStages::FRAGMENT,
                    ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
                    count: None,
                },
            ],
        });

        let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("Blur Pipeline Layout"),
            bind_group_layouts: &[Some(&bind_group_layout)],
            immediate_size: 0,
        });

        let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("Blur Shader"),
            source: wgpu::ShaderSource::Wgsl(include_str!("../shaders/blur.wgsl").into()),
        });

        let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: Some("Blur Pipeline"),
            layout: Some(&pipeline_layout),
            vertex: wgpu::VertexState {
                module: &shader,
                entry_point: Some("vs_main"),
                buffers: &[],
                compilation_options: Default::default(),
            },
            fragment: Some(wgpu::FragmentState {
                module: &shader,
                entry_point: Some("fs_main"),
                targets: &[Some(wgpu::ColorTargetState {
                    format: target_format,
                    blend: Some(wgpu::BlendState::REPLACE),
                    write_mask: wgpu::ColorWrites::ALL,
                })],
                compilation_options: Default::default(),
            }),
            primitive: wgpu::PrimitiveState {
                topology: wgpu::PrimitiveTopology::TriangleList,
                ..Default::default()
            },
            depth_stencil: None,
            multisample: wgpu::MultisampleState::default(),
            multiview_mask: None,
            cache: None,
        });

        let sampler = device.create_sampler(&wgpu::SamplerDescriptor {
            label: Some("Blur Sampler"),
            address_mode_u: wgpu::AddressMode::ClampToEdge,
            address_mode_v: wgpu::AddressMode::ClampToEdge,
            address_mode_w: wgpu::AddressMode::ClampToEdge,
            mag_filter: wgpu::FilterMode::Linear,
            min_filter: wgpu::FilterMode::Linear,
            mipmap_filter: wgpu::MipmapFilterMode::Linear,
            ..Default::default()
        });

        // One uniform buffer per pass. Direction alternates H/V; sigma
        // doubles roughly per blur level. `texel_size` matches the
        // source texture's resolution at each pass — see the H/V
        // dispatch order above.
        let make_uniforms = |dir: [f32; 2], texel: [f32; 2], sigma: f32| BlurUniforms {
            direction: dir,
            texel_size: texel,
            sigma,
            _pad0: 0.0,
            _pad1: 0.0,
            _pad2: 0.0,
        };
        let pass_specs: [BlurUniforms; 6] = [
            // pass 0: H, src = render (full)        → b1_scratch
            make_uniforms([1.0, 0.0], texel_full, SIGMA_BLUR1),
            // pass 1: V, src = b1_scratch (½)       → blur1
            make_uniforms([0.0, 1.0], texel_b1, SIGMA_BLUR1),
            // pass 2: H, src = blur1 (½)            → b2_scratch
            make_uniforms([1.0, 0.0], texel_b1, SIGMA_BLUR2),
            // pass 3: V, src = b2_scratch (¼)       → blur2
            make_uniforms([0.0, 1.0], texel_b2, SIGMA_BLUR2),
            // pass 4: H, src = blur2 (¼)            → b3_scratch
            make_uniforms([1.0, 0.0], texel_b2, SIGMA_BLUR3),
            // pass 5: V, src = b3_scratch (⅛)       → blur3
            make_uniforms([0.0, 1.0], texel_b3, SIGMA_BLUR3),
        ];

        let uniform_buffers: [wgpu::Buffer; 6] = std::array::from_fn(|i| {
            device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some(&format!("Blur Uniforms #{i}")),
                contents: bytemuck::bytes_of(&pass_specs[i]),
                usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
            })
        });

        let pass_sources: [&wgpu::TextureView; 6] = [
            render_view,        // pass 0: H from render (full)
            blur1_scratch_view, // pass 1: V from blur1_scratch (½)
            blur1_view,         // pass 2: H from Blur1 (½)
            blur2_scratch_view, // pass 3: V from blur2_scratch (¼)
            blur2_view,         // pass 4: H from Blur2 (¼)
            blur3_scratch_view, // pass 5: V from blur3_scratch (⅛)
        ];

        let bind_groups: [wgpu::BindGroup; 6] = std::array::from_fn(|i| {
            Self::make_bind_group(
                device,
                &bind_group_layout,
                &uniform_buffers[i],
                pass_sources[i],
                &sampler,
                i,
            )
        });

        Self {
            pipeline,
            bind_group_layout,
            sampler,
            uniform_buffers,
            bind_groups,
        }
    }

    fn make_bind_group(
        device: &wgpu::Device,
        layout: &wgpu::BindGroupLayout,
        uniforms: &wgpu::Buffer,
        source: &wgpu::TextureView,
        sampler: &wgpu::Sampler,
        idx: usize,
    ) -> wgpu::BindGroup {
        device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some(&format!("Blur Bind Group #{idx}")),
            layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: uniforms.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: wgpu::BindingResource::TextureView(source),
                },
                wgpu::BindGroupEntry {
                    binding: 2,
                    resource: wgpu::BindingResource::Sampler(sampler),
                },
            ],
        })
    }

    /// Recreate the bind groups + refresh per-pass uniform `texel_size`
    /// after the underlying textures are invalidated (resize).
    #[allow(clippy::too_many_arguments)]
    pub fn rebind(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        width: u32,
        height: u32,
        render_view: &wgpu::TextureView,
        blur1_view: &wgpu::TextureView,
        blur2_view: &wgpu::TextureView,
        blur1_scratch_view: &wgpu::TextureView,
        blur2_scratch_view: &wgpu::TextureView,
        blur3_scratch_view: &wgpu::TextureView,
    ) {
        let (w1, h1) = (width.max(2) / 2, height.max(2) / 2);
        let (w2, h2) = (width.max(4) / 4, height.max(4) / 4);
        let (w3, h3) = (width.max(8) / 8, height.max(8) / 8);
        let texel_full = [1.0 / width as f32, 1.0 / height as f32];
        let texel_b1 = [1.0 / w1 as f32, 1.0 / h1 as f32];
        let texel_b2 = [1.0 / w2 as f32, 1.0 / h2 as f32];
        let texel_b3 = [1.0 / w3 as f32, 1.0 / h3 as f32];
        // Per-pass (sigma, source-texel, direction) tuples — must
        // match the order used in `new`.
        let pass_uniforms: [(f32, [f32; 2], [f32; 2]); 6] = [
            (SIGMA_BLUR1, texel_full, [1.0, 0.0]),
            (SIGMA_BLUR1, texel_b1, [0.0, 1.0]),
            (SIGMA_BLUR2, texel_b1, [1.0, 0.0]),
            (SIGMA_BLUR2, texel_b2, [0.0, 1.0]),
            (SIGMA_BLUR3, texel_b2, [1.0, 0.0]),
            (SIGMA_BLUR3, texel_b3, [0.0, 1.0]),
        ];
        for (i, &(sigma, texel, direction)) in pass_uniforms.iter().enumerate() {
            let u = BlurUniforms {
                direction,
                texel_size: texel,
                sigma,
                _pad0: 0.0,
                _pad1: 0.0,
                _pad2: 0.0,
            };
            queue.write_buffer(&self.uniform_buffers[i], 0, bytemuck::bytes_of(&u));
        }

        let pass_sources: [&wgpu::TextureView; 6] = [
            render_view,
            blur1_scratch_view,
            blur1_view,
            blur2_scratch_view,
            blur2_view,
            blur3_scratch_view,
        ];
        for (i, source) in pass_sources.iter().enumerate() {
            self.bind_groups[i] = Self::make_bind_group(
                device,
                &self.bind_group_layout,
                &self.uniform_buffers[i],
                source,
                &self.sampler,
                i,
            );
        }
    }

    /// Encode the 6 blur passes for this frame. Caller is responsible
    /// for `submit`-ing the encoder; the pipeline assumes the input
    /// `render_texture` already holds the current warp pass output.
    #[allow(clippy::too_many_arguments)]
    pub fn render(
        &self,
        encoder: &mut wgpu::CommandEncoder,
        blur1_view: &wgpu::TextureView,
        blur2_view: &wgpu::TextureView,
        blur3_view: &wgpu::TextureView,
        blur1_scratch_view: &wgpu::TextureView,
        blur2_scratch_view: &wgpu::TextureView,
        blur3_scratch_view: &wgpu::TextureView,
    ) {
        let targets: [&wgpu::TextureView; 6] = [
            blur1_scratch_view, // pass 0 H → b1_scratch (½)
            blur1_view,         // pass 1 V → blur1 (½)
            blur2_scratch_view, // pass 2 H → b2_scratch (¼)
            blur2_view,         // pass 3 V → blur2 (¼)
            blur3_scratch_view, // pass 4 H → b3_scratch (⅛)
            blur3_view,         // pass 5 V → blur3 (⅛)
        ];
        for (i, target) in targets.iter().enumerate() {
            let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
                label: Some(&format!("Blur Pass #{i}")),
                color_attachments: &[Some(wgpu::RenderPassColorAttachment {
                    view: target,
                    depth_slice: None,
                    resolve_target: None,
                    ops: wgpu::Operations {
                        load: wgpu::LoadOp::Clear(wgpu::Color::TRANSPARENT),
                        store: wgpu::StoreOp::Store,
                    },
                })],
                depth_stencil_attachment: None,
                timestamp_writes: None,
                occlusion_query_set: None,
                multiview_mask: None,
            });
            pass.set_pipeline(&self.pipeline);
            pass.set_bind_group(0, &self.bind_groups[i], &[]);
            pass.draw(0..3, 0..1);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::chain_textures::ChainTextures;
    use crate::config::{RenderConfig, TextureFormat};
    use crate::gpu_context::GpuContext;

    #[test]
    fn blur_pipeline_instantiates() {
        let cfg = RenderConfig::default();
        let gpu = pollster::block_on(GpuContext::new(cfg)).unwrap();
        let chain_tex = ChainTextures::new(&gpu.device, &gpu.config);
        let _p = BlurPipeline::new(
            &gpu.device,
            gpu.config.texture_format.to_wgpu(),
            gpu.config.width,
            gpu.config.height,
            &chain_tex.render_texture_view,
            &chain_tex.blur1_texture_view,
            &chain_tex.blur2_texture_view,
            &chain_tex.blur3_texture_view,
            &chain_tex.blur1_scratch_texture_view,
            &chain_tex.blur2_scratch_texture_view,
            &chain_tex.blur3_scratch_texture_view,
        );
    }

    /// Seed `render_texture` with a sharp 50/50 split (one half black, one
    /// half white) and verify that running the blur passes spreads
    /// energy across the boundary in `blur1_texture`. A non-blurred
    /// straight copy would still produce 0 or 255; the Gaussian must
    /// produce some intermediate value within the kernel reach of the
    /// boundary.
    #[test]
    fn blur_pipeline_spreads_energy_across_boundary() {
        let cfg = RenderConfig {
            width: 64,
            height: 64,
            texture_format: TextureFormat::Bgra8Unorm,
            ..Default::default()
        };
        let gpu = pollster::block_on(GpuContext::new(cfg)).unwrap();
        let chain_tex = ChainTextures::new(&gpu.device, &gpu.config);
        let blur = BlurPipeline::new(
            &gpu.device,
            gpu.config.texture_format.to_wgpu(),
            gpu.config.width,
            gpu.config.height,
            &chain_tex.render_texture_view,
            &chain_tex.blur1_texture_view,
            &chain_tex.blur2_texture_view,
            &chain_tex.blur3_texture_view,
            &chain_tex.blur1_scratch_texture_view,
            &chain_tex.blur2_scratch_texture_view,
            &chain_tex.blur3_scratch_texture_view,
        );

        // Seed: left half black, right half white. BGRA layout.
        let pixels: Vec<u8> = (0..64)
            .flat_map(|_y| {
                (0..64).flat_map(move |x| {
                    let v: u8 = if x < 32 { 0 } else { 255 };
                    [v, v, v, 255u8]
                })
            })
            .collect();
        gpu.queue.write_texture(
            wgpu::TexelCopyTextureInfo {
                texture: &chain_tex.render_texture,
                mip_level: 0,
                origin: wgpu::Origin3d::ZERO,
                aspect: wgpu::TextureAspect::All,
            },
            &pixels,
            wgpu::TexelCopyBufferLayout {
                offset: 0,
                bytes_per_row: Some(64 * 4),
                rows_per_image: Some(64),
            },
            wgpu::Extent3d {
                width: 64,
                height: 64,
                depth_or_array_layers: 1,
            },
        );

        let mut encoder = gpu.device.create_command_encoder(&Default::default());
        blur.render(
            &mut encoder,
            &chain_tex.blur1_texture_view,
            &chain_tex.blur2_texture_view,
            &chain_tex.blur3_texture_view,
            &chain_tex.blur1_scratch_texture_view,
            &chain_tex.blur2_scratch_texture_view,
            &chain_tex.blur3_scratch_texture_view,
        );
        gpu.queue.submit(std::iter::once(encoder.finish()));

        // Read blur1 back at its native ½-res (32 × 32). Sample
        // straddling the boundary at the centre row. The Gaussian
        // (combined with the ½-res downsample) should spread the
        // black/white split into a meaningful grey band across a
        // few texels.
        let blur_w: u32 = 32;
        let blur_h: u32 = 32;
        let align = wgpu::COPY_BYTES_PER_ROW_ALIGNMENT;
        let unpadded_row: u32 = blur_w * 4;
        let padded_row = unpadded_row.div_ceil(align) * align;
        let staging = gpu.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("Blur Test Staging"),
            size: (padded_row * blur_h) as u64,
            usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
            mapped_at_creation: false,
        });
        let mut e2 = gpu.device.create_command_encoder(&Default::default());
        e2.copy_texture_to_buffer(
            wgpu::TexelCopyTextureInfo {
                texture: &chain_tex.blur1_texture,
                mip_level: 0,
                origin: wgpu::Origin3d::ZERO,
                aspect: wgpu::TextureAspect::All,
            },
            wgpu::TexelCopyBufferInfo {
                buffer: &staging,
                layout: wgpu::TexelCopyBufferLayout {
                    offset: 0,
                    bytes_per_row: Some(padded_row),
                    rows_per_image: Some(blur_h),
                },
            },
            wgpu::Extent3d {
                width: blur_w,
                height: blur_h,
                depth_or_array_layers: 1,
            },
        );
        gpu.queue.submit(std::iter::once(e2.finish()));

        let (tx, rx) = std::sync::mpsc::channel();
        staging.slice(..).map_async(wgpu::MapMode::Read, move |r| {
            let _ = tx.send(r);
        });
        gpu.device.poll(wgpu::PollType::wait_indefinitely()).ok();
        rx.recv().unwrap().unwrap();
        let view = staging.slice(..).get_mapped_range();

        // Mid-row, column straddling the boundary in ½-res space
        // (full-res col 31/32 → ½-res col 15/16).
        let row = blur_h / 2;
        let col = 15;
        let off = (row * padded_row + col * 4) as usize;
        let b = view[off];
        // Expect a meaningful intermediate value. Pure-bilinear sampling
        // wouldn't produce this much spread — only the Gaussian does.
        assert!(
            b > 20 && b < 200,
            "boundary pixel should be partially blurred; got {b}"
        );
        drop(view);
        staging.unmap();
    }
}