onedrop_renderer/

warp_pipeline.rs

//! Warp pass GPU pipeline.
//!
//! Renders the warp mesh (cols × rows triangles) into a target texture by
//! sampling `prev_texture` at the per-vertex warped UV. Decay is applied in
//! the fragment shader.
//!
//! The CPU is responsible for filling the dynamic vertex buffer each frame
//! via [`WarpPipeline::update_vertices`].

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

use crate::error::Result;
use crate::user_warp_pipeline::{UserWarpPipeline, WarpAuxViews};
use crate::warp_mesh::WarpMesh;

/// One vertex consumed by the warp shader.
///
/// `pos_clip` is static (laid out by the mesh once); `uv_warp` is rewritten
/// every frame from per-vertex equation evaluation.
#[repr(C)]
#[derive(Clone, Copy, Debug, Pod, Zeroable)]
pub struct WarpVertex {
    pub pos_clip: [f32; 2],
    pub uv_warp: [f32; 2],
}

#[repr(C)]
#[derive(Clone, Copy, Debug, Pod, Zeroable)]
struct WarpUniforms {
    decay: f32,
    aspect_x: f32,
    aspect_y: f32,
    /// Bitfield of feedback toggles. See `FeedbackParams::to_flags` and the
    /// `WARP_FLAG_*` constants in `warp.wgsl` for the layout.
    flags: u32,
}

pub struct WarpPipeline {
    pipeline: wgpu::RenderPipeline,
    bind_group_layout: wgpu::BindGroupLayout,
    bind_group: wgpu::BindGroup,
    sampler: wgpu::Sampler,
    uniforms_buffer: wgpu::Buffer,

    vertex_buffer: wgpu::Buffer,
    index_buffer: wgpu::Buffer,
    index_count: u32,
    vertex_count: u32,

    /// Active when a preset's translated warp shader compiled
    /// successfully. The default pipeline above stays built so failed
    /// compiles fall back transparently. See [`UserWarpPipeline`].
    user_pipeline: Option<UserWarpPipeline>,
    /// Target format remembered so the user pipeline can be rebuilt
    /// after a preset swap without re-asking the caller.
    target_format: wgpu::TextureFormat,
}

impl WarpPipeline {
    pub fn new(
        device: &wgpu::Device,
        target_format: wgpu::TextureFormat,
        prev_texture_view: &wgpu::TextureView,
        mesh: &WarpMesh,
    ) -> Result<Self> {
        // Sampler: clamp-to-edge for now. Wrap mode (`b_tex_wrap`) deferred.
        let sampler = device.create_sampler(&wgpu::SamplerDescriptor {
            label: Some("Warp 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()
        });

        let uniforms = WarpUniforms {
            decay: 0.98,
            aspect_x: 1.0,
            aspect_y: 1.0,
            flags: 0,
        };
        let uniforms_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("Warp Uniforms"),
            contents: bytemuck::bytes_of(&uniforms),
            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
        });

        // Initial vertex buffer: identity warp (uv_warp = uv_orig).
        let initial_vertices: Vec<WarpVertex> = mesh
            .vertices
            .iter()
            .map(|v| WarpVertex {
                pos_clip: v.pos_clip,
                uv_warp: v.uv_orig,
            })
            .collect();
        let vertex_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("Warp Vertex Buffer"),
            contents: bytemuck::cast_slice(&initial_vertices),
            usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
        });

        let index_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("Warp Index Buffer"),
            contents: bytemuck::cast_slice(&mesh.indices),
            usage: wgpu::BufferUsages::INDEX,
        });

        let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("Warp Bind Group Layout"),
            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 bind_group = Self::create_bind_group(
            device,
            &bind_group_layout,
            &uniforms_buffer,
            prev_texture_view,
            &sampler,
        );

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

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

        let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: Some("Warp Pipeline"),
            layout: Some(&pipeline_layout),
            vertex: wgpu::VertexState {
                module: &shader,
                entry_point: Some("vs_main"),
                buffers: &[wgpu::VertexBufferLayout {
                    array_stride: std::mem::size_of::<WarpVertex>() as wgpu::BufferAddress,
                    step_mode: wgpu::VertexStepMode::Vertex,
                    attributes: &[
                        wgpu::VertexAttribute {
                            offset: 0,
                            shader_location: 0,
                            format: wgpu::VertexFormat::Float32x2,
                        },
                        wgpu::VertexAttribute {
                            offset: std::mem::size_of::<[f32; 2]>() as wgpu::BufferAddress,
                            shader_location: 1,
                            format: wgpu::VertexFormat::Float32x2,
                        },
                    ],
                }],
                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,
        });

        Ok(Self {
            pipeline,
            bind_group_layout,
            bind_group,
            sampler,
            uniforms_buffer,
            vertex_buffer,
            index_buffer,
            index_count: mesh.indices.len() as u32,
            vertex_count: mesh.vertices.len() as u32,
            user_pipeline: None,
            target_format,
        })
    }

    /// Swap in a user-translated warp fragment shader. Mirrors
    /// [`CompPipeline::set_user_shader_with_plan`]: the WGSL must have
    /// already been validated by `onedrop_codegen::ShaderCompiler`, the
    /// renderer threads `prev_view` + `aux` (noise textures, etc.) + a
    /// per-preset `user_views` array of disk-loaded textures.
    ///
    /// On success the warp pass uses the user pipeline; on failure the
    /// default mesh-warp pipeline keeps running.
    pub fn set_user_shader_with_plan(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        wrapped_wgsl: &str,
        prev_view: &wgpu::TextureView,
        aux: &WarpAuxViews,
        user_views: [Option<wgpu::TextureView>; USER_TEXTURE_SLOTS],
    ) -> Result<()> {
        let user = UserWarpPipeline::new(
            device,
            queue,
            self.target_format,
            wrapped_wgsl,
            prev_view,
            aux,
            user_views,
        )?;
        self.user_pipeline = Some(user);
        Ok(())
    }

    /// Drop the user warp shader and route the pass back through the
    /// engine's default `warp.wgsl`. Idempotent.
    pub fn reset_to_default(&mut self) {
        self.user_pipeline = None;
    }

    /// `true` when the user pipeline is currently driving the warp
    /// pass.
    pub fn has_user_shader(&self) -> bool {
        self.user_pipeline.is_some()
    }

    /// Upload `ShaderUniforms` to the user pipeline's uniform buffer.
    /// Caller passes the same `ShaderUniforms` the comp pass receives
    /// — q-channels, time, audio levels, roam vectors, etc.
    pub fn update_user_uniforms(&self, queue: &wgpu::Queue, uniforms: &ShaderUniforms) {
        if let Some(u) = &self.user_pipeline {
            u.update_uniforms(queue, uniforms);
        }
    }

    /// Rebind the user pipeline's texture views after a resize. No-op
    /// when no user shader is active. The renderer threads the new
    /// `prev_view` + refreshed `aux` (blur/noise views may have been
    /// recreated by [`ChainTextures`] on resize).
    pub fn rebind_user_textures(
        &mut self,
        device: &wgpu::Device,
        prev_view: &wgpu::TextureView,
        aux: &WarpAuxViews,
    ) {
        if let Some(u) = self.user_pipeline.as_mut() {
            u.rebind_textures(device, prev_view, aux);
        }
    }

    fn create_bind_group(
        device: &wgpu::Device,
        layout: &wgpu::BindGroupLayout,
        uniforms_buffer: &wgpu::Buffer,
        prev_texture_view: &wgpu::TextureView,
        sampler: &wgpu::Sampler,
    ) -> wgpu::BindGroup {
        device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("Warp Bind Group"),
            layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: uniforms_buffer.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: wgpu::BindingResource::TextureView(prev_texture_view),
                },
                wgpu::BindGroupEntry {
                    binding: 2,
                    resource: wgpu::BindingResource::Sampler(sampler),
                },
            ],
        })
    }

    /// Re-create the bind group after the prev texture view is invalidated
    /// (e.g., on resize).
    pub fn rebind_prev_texture(
        &mut self,
        device: &wgpu::Device,
        prev_texture_view: &wgpu::TextureView,
    ) {
        self.bind_group = Self::create_bind_group(
            device,
            &self.bind_group_layout,
            &self.uniforms_buffer,
            prev_texture_view,
            &self.sampler,
        );
    }

    pub fn update_uniforms(&self, queue: &wgpu::Queue, decay: f32, aspect: f32, flags: u32) {
        let aspect_x = if aspect > 1.0 { 1.0 / aspect } else { 1.0 };
        let aspect_y = if aspect < 1.0 { aspect } else { 1.0 };
        let uniforms = WarpUniforms {
            decay,
            aspect_x,
            aspect_y,
            flags,
        };
        queue.write_buffer(&self.uniforms_buffer, 0, bytemuck::bytes_of(&uniforms));
    }

    /// Reallocate the vertex + index buffers for a different-sized mesh.
    /// Used by `MilkRenderer::set_mesh_size` when the user picks a new
    /// mesh quality at runtime. The bind group is unaffected because it
    /// references the prev-texture view, not the vertex buffer.
    pub fn rebuild_mesh(&mut self, device: &wgpu::Device, mesh: &WarpMesh) {
        let initial_vertices: Vec<WarpVertex> = mesh
            .vertices
            .iter()
            .map(|v| WarpVertex {
                pos_clip: v.pos_clip,
                uv_warp: v.uv_orig,
            })
            .collect();
        self.vertex_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("Warp Vertex Buffer"),
            contents: bytemuck::cast_slice(&initial_vertices),
            usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
        });
        self.index_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("Warp Index Buffer"),
            contents: bytemuck::cast_slice(&mesh.indices),
            usage: wgpu::BufferUsages::INDEX,
        });
        self.vertex_count = mesh.vertices.len() as u32;
        self.index_count = mesh.indices.len() as u32;
    }

    /// Upload new per-vertex warped UVs.
    ///
    /// `vertices.len()` must equal the original mesh vertex count.
    pub fn update_vertices(&self, queue: &wgpu::Queue, vertices: &[WarpVertex]) {
        debug_assert_eq!(
            vertices.len() as u32,
            self.vertex_count,
            "warp vertex count mismatch"
        );
        queue.write_buffer(&self.vertex_buffer, 0, bytemuck::cast_slice(vertices));
    }

    /// Issue draw commands. Caller owns the encoder.
    ///
    /// When the user warp pipeline is active (`set_user_shader_with_plan`
    /// has succeeded), delegate to it; the vertex/index buffers are
    /// shared so the per-frame mesh upload still applies. Otherwise the
    /// engine's hand-written `warp.wgsl` runs as before.
    pub fn render(&self, encoder: &mut wgpu::CommandEncoder, output_view: &wgpu::TextureView) {
        if let Some(user) = &self.user_pipeline {
            user.render(
                encoder,
                output_view,
                &self.vertex_buffer,
                &self.index_buffer,
                self.index_count,
            );
            return;
        }
        let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
            label: Some("Warp Render Pass"),
            color_attachments: &[Some(wgpu::RenderPassColorAttachment {
                view: output_view,
                depth_slice: None,
                resolve_target: None,
                ops: wgpu::Operations {
                    load: wgpu::LoadOp::Clear(wgpu::Color::BLACK),
                    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_group, &[]);
        pass.set_vertex_buffer(0, self.vertex_buffer.slice(..));
        pass.set_index_buffer(self.index_buffer.slice(..), wgpu::IndexFormat::Uint32);
        pass.draw_indexed(0..self.index_count, 0, 0..1);
    }

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

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

    #[test]
    fn 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 mesh = WarpMesh::new(8, 6, 16.0 / 9.0);
        let p = WarpPipeline::new(
            &gpu.device,
            gpu.config.texture_format.to_wgpu(),
            &chain_tex.prev_texture_view,
            &mesh,
        );
        assert!(p.is_ok());
    }

    #[test]
    fn warp_vertex_size_is_16() {
        // 2× vec2<f32> = 16 bytes — must match the vertex layout strides.
        assert_eq!(std::mem::size_of::<WarpVertex>(), 16);
    }

    /// `decay` is applied unconditionally inside the warp fragment shader
    /// (`color *= uniforms.decay`). With an identity-warp mesh and
    /// `prev_texture` seeded to a uniform grey, the warp output should be
    /// `input × decay` exactly. Uses a non-sRGB `Bgra8Unorm` target so the
    /// math is byte-direct.
    #[test]
    fn warp_pass_applies_decay_to_prev_sample() {
        use crate::config::TextureFormat;

        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 mesh = WarpMesh::new(8, 6, 1.0);
        let pipeline = WarpPipeline::new(
            &gpu.device,
            gpu.config.texture_format.to_wgpu(),
            &chain_tex.prev_texture_view,
            &mesh,
        )
        .unwrap();

        // Seed prev_texture with uniform grey 200. BGRA layout.
        let pixels: Vec<u8> = (0..(64 * 64))
            .flat_map(|_| [200u8, 200, 200, 255])
            .collect();
        gpu.queue.write_texture(
            wgpu::TexelCopyTextureInfo {
                texture: &chain_tex.prev_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 read_center = |decay: f32, flags: u32| -> u8 {
            pipeline.update_uniforms(&gpu.queue, decay, 1.0, flags);

            let mut encoder = gpu.device.create_command_encoder(&Default::default());
            pipeline.render(&mut encoder, &chain_tex.render_texture_view);
            gpu.queue.submit(std::iter::once(encoder.finish()));

            let align = wgpu::COPY_BYTES_PER_ROW_ALIGNMENT;
            let unpadded_row: u32 = 64 * 4;
            let padded_row = unpadded_row.div_ceil(align) * align;
            let staging = gpu.device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("Warp Decay Test Staging"),
                size: (padded_row * 64) 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.render_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(64),
                    },
                },
                wgpu::Extent3d {
                    width: 64,
                    height: 64,
                    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();
            let off = (32u32 * padded_row + 32 * 4) as usize;
            let g = view[off + 1];
            drop(view);
            staging.unmap();
            g
        };

        // decay = 1.0 → passthrough: input 200 → output 200.
        let pass = read_center(1.0, 0);
        assert!(
            (pass as i32 - 200).abs() <= 4,
            "decay=1.0 should pass through, got {pass} (expected ~200)"
        );

        // decay = 0.5 → output = 100.
        let half = read_center(0.5, 0);
        assert!(
            (half as i32 - 100).abs() <= 4,
            "decay=0.5 should halve, got {half} (expected ~100)"
        );

        // decay = 0.0 → output black.
        let zero = read_center(0.0, 0);
        assert!(zero <= 4, "decay=0.0 should produce black, got {zero}");
    }

    /// `WARP_FLAG_DARKEN_CENTER` subtracts a small radial bump near the
    /// screen centre. Seed prev_texture to a uniform mid-grey, then
    /// compare the centre pixel to a corner pixel: with the flag off they
    /// match; with the flag on the centre is strictly darker than the
    /// corner (the subtraction is ~3 grey levels at uv=(0.5, 0.5)).
    #[test]
    fn warp_darken_center_attenuates_center_more_than_corner() {
        use crate::config::TextureFormat;

        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 mesh = WarpMesh::new(8, 6, 1.0);
        let pipeline = WarpPipeline::new(
            &gpu.device,
            gpu.config.texture_format.to_wgpu(),
            &chain_tex.prev_texture_view,
            &mesh,
        )
        .unwrap();

        let pixels: Vec<u8> = (0..(64 * 64))
            .flat_map(|_| [200u8, 200, 200, 255])
            .collect();
        gpu.queue.write_texture(
            wgpu::TexelCopyTextureInfo {
                texture: &chain_tex.prev_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 read_pixels = |flags: u32| -> (u8, u8) {
            pipeline.update_uniforms(&gpu.queue, 1.0, 1.0, flags);

            let mut encoder = gpu.device.create_command_encoder(&Default::default());
            pipeline.render(&mut encoder, &chain_tex.render_texture_view);
            gpu.queue.submit(std::iter::once(encoder.finish()));

            let align = wgpu::COPY_BYTES_PER_ROW_ALIGNMENT;
            let unpadded_row: u32 = 64 * 4;
            let padded_row = unpadded_row.div_ceil(align) * align;
            let staging = gpu.device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("DarkenCenter Test Staging"),
                size: (padded_row * 64) 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.render_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(64),
                    },
                },
                wgpu::Extent3d {
                    width: 64,
                    height: 64,
                    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();
            let center_off = (32u32 * padded_row + 32 * 4) as usize;
            let corner_off = (2u32 * padded_row + 2 * 4) as usize;
            let center_g = view[center_off + 1];
            let corner_g = view[corner_off + 1];
            drop(view);
            staging.unmap();
            (center_g, corner_g)
        };

        // Flag off — centre and corner should match (uniform 200 in/out).
        let (off_center, off_corner) = read_pixels(0);
        assert!(
            (off_center as i32 - off_corner as i32).abs() <= 2,
            "flag off: centre {off_center} vs corner {off_corner} should be ~equal"
        );

        // Flag on (WARP_FLAG_DARKEN_CENTER = 2) — centre is darker; corner
        // is barely touched (radial exp falloff).
        let (on_center, on_corner) = read_pixels(2);
        assert!(
            (on_corner as i32 - off_corner as i32).abs() <= 2,
            "flag on: corner {on_corner} should still be ~{off_corner}"
        );
        assert!(
            on_center + 2 < off_center,
            "flag on: centre {on_center} should be darker than off-state {off_center}"
        );
    }
}