sigil/src/viz/shader.rs

//! Reusable raw-wgpu raymarch/post pipeline base.
//!
//! [`ShaderPipeline<U>`] is the self-contained fullscreen-fragment pipeline
//! that `breakcore` pioneered, hoisted so any future GPU visualiser can drive
//! a `.wgsl` shader without re-deriving the wgpu boilerplate. It owns:
//!
//! - WGSL module + fullscreen-triangle render pipeline (no vertex buffers;
//!   entry points are the fixed convention `vs_main` / `fs_main`),
//! - a single uniform buffer sized to `U`,
//! - dual `RES²` textures ping-ponged for frame feedback (the shader samples
//!   the previous frame at `@binding(1)`, writes the other),
//! - per-frame UBO upload + the bind-group swap.
//!
//! A new GPU visualiser supplies only its UBO type, its `.wgsl`, and the
//! per-frame fill of `U`; the section/geometry/colour logic stays in Rust.
//!
//! **Guardrails (non-negotiable, inherited from breakcore).** This base does
//! *not* abstract away the two things that can hang the GPU: the shader's own
//! hard march-step cap and any `array<vec4, N>` ↔ Rust UBO-length coupling
//! still live in the shader + the caller's `U` layout, and naga only
//! validates the WGSL at pipeline-create on a real device. Keep the cap in
//! the `.wgsl` and keep `U`'s layout in lock-step with the shader's UBO.
//!
//! Determinism: pure GPU plumbing — no `Rng`, no clock, no integration. Given
//! the same `U` bytes it produces the same frame, so a visualiser built on it
//! stays `--render` bit-reproducible exactly as before (the state lives in the
//! caller, advanced once per frame as ever).

use std::marker::PhantomData;

use nannou::wgpu;

/// Reinterpret a `Copy`, packing-free POD (`[f32; N]`, `#[repr(C)]` plain
/// struct) as its raw bytes for `write_buffer`. The crate already relies on
/// exactly this invariant for its one UBO, so this needs no `bytemuck` dep —
/// but `U` MUST stay padding-free and match the shader's UBO layout.
fn as_bytes<U: Copy>(u: &U) -> &[u8] {
    unsafe {
        std::slice::from_raw_parts((u as *const U).cast::<u8>(), std::mem::size_of::<U>())
    }
}

/// Fullscreen fragment pipeline with ping-pong feedback. `U` is the uniform
/// block (one `write_buffer` per frame); see the module guardrails note.
pub struct ShaderPipeline<U> {
    pipeline: wgpu::RenderPipeline,
    ubo: wgpu::Buffer,
    tex: [wgpu::Texture; 2],
    view: [wgpu::TextureViewHandle; 2],
    bind: [wgpu::BindGroup; 2], // bind[w]: writes view[w], samples view[1-w]
    cur: usize,                 // index last written / presented
    primed: bool,
    _pd: PhantomData<U>,
}

impl<U: Copy> ShaderPipeline<U> {
    /// Build the pipeline for `wgsl` at `res²` in `format`. Bindings, in WGSL
    /// `@binding` order, are: `0` uniform `U`, `1` previous-frame texture,
    /// `2` sampler. Entry points are `vs_main` / `fs_main`.
    pub fn new(
        device: &wgpu::Device,
        wgsl: &str,
        res: u32,
        format: wgpu::TextureFormat,
    ) -> Self {
        let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("shaderpipeline-shader"),
            source: wgpu::ShaderSource::Wgsl(wgsl.into()),
        });

        let mk = || {
            wgpu::TextureBuilder::new()
                .size([res, res])
                .format(format)
                .usage(
                    wgpu::TextureUsages::RENDER_ATTACHMENT
                        | wgpu::TextureUsages::TEXTURE_BINDING
                        | wgpu::TextureUsages::COPY_SRC,
                )
                .build(device)
        };
        let tex = [mk(), mk()];
        let view = [
            tex[0].create_view(&wgpu::TextureViewDescriptor::default()),
            tex[1].create_view(&wgpu::TextureViewDescriptor::default()),
        ];

        // Bindings (order = WGSL @binding 0/1/2): uniform, prev texture, sampler.
        let bgl = wgpu::BindGroupLayoutBuilder::new()
            .uniform_buffer(wgpu::ShaderStages::FRAGMENT, false)
            .texture_from(wgpu::ShaderStages::FRAGMENT, &tex[0])
            .sampler(wgpu::ShaderStages::FRAGMENT, true)
            .build(device);

        let layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("shaderpipeline-pl"),
            bind_group_layouts: &[&bgl],
            push_constant_ranges: &[],
        });

        // Fullscreen triangle: no vertex buffers; default triangle-list.
        let pipeline = wgpu::RenderPipelineBuilder::from_layout(&layout, &shader)
            .vertex_entry_point("vs_main")
            .fragment_shader(&shader)
            .fragment_entry_point("fs_main")
            .color_format(format)
            .build(device);

        let ubo = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("shaderpipeline-ubo"),
            size: std::mem::size_of::<U>() as u64,
            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        let sampler = wgpu::SamplerBuilder::new()
            .address_mode(wgpu::AddressMode::ClampToEdge)
            .mag_filter(wgpu::FilterMode::Linear)
            .min_filter(wgpu::FilterMode::Linear)
            .mipmap_filter(wgpu::FilterMode::Nearest)
            .build(device);

        let mk_bind = |w: usize| {
            wgpu::BindGroupBuilder::new()
                .binding(ubo.as_entire_binding())
                .texture_view(&view[1 - w])
                .sampler(&sampler)
                .build(device, &bgl)
        };
        let bind = [mk_bind(0), mk_bind(1)];

        ShaderPipeline {
            pipeline,
            ubo,
            tex,
            view,
            bind,
            cur: 0,
            primed: false,
            _pd: PhantomData,
        }
    }

    /// Upload `ubo`, raymarch one frame into the unused ping-pong slot
    /// (sampling the previous), and return the just-written texture.
    pub fn render(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        ubo: &U,
    ) -> &wgpu::Texture {
        let w = 1 - self.cur; // write target; sample the last-written (self.cur)
        queue.write_buffer(&self.ubo, 0, as_bytes(ubo));

        let mut enc = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("shaderpipeline-enc"),
        });
        {
            let mut rp = enc.begin_render_pass(&wgpu::RenderPassDescriptor {
                label: Some("shaderpipeline-pass"),
                color_attachments: &[Some(wgpu::RenderPassColorAttachment {
                    view: &self.view[w],
                    resolve_target: None,
                    ops: wgpu::Operations {
                        load: wgpu::LoadOp::Clear(wgpu::Color::BLACK),
                        store: true,
                    },
                })],
                depth_stencil_attachment: None,
            });
            rp.set_pipeline(&self.pipeline);
            rp.set_bind_group(0, &self.bind[w], &[]);
            rp.draw(0..3, 0..1);
        }
        queue.submit(Some(enc.finish()));

        self.cur = w;
        self.primed = true;
        &self.tex[self.cur]
    }

    /// The most recently written texture (for present/readback).
    pub fn current(&self) -> &wgpu::Texture {
        &self.tex[self.cur]
    }

    /// `true` once at least one frame has been rendered (the feedback sample
    /// is only valid after the first frame).
    pub fn primed(&self) -> bool {
        self.primed
    }
}