mirage-engine 0.1.0

//! The target a frame is drawn into before the screen, and the passes that
//! take it there: the bloom chain, and the tone map.

use core::ops::Range;

use bytemuck::{Pod, Zeroable};

use crate::Tonemap;
use crate::gpu::depth_texture;
use crate::math::UVec2;

/// The format a frame is drawn in, which holds light past what the screen
/// shows.
pub(crate) const HDR_FORMAT: wgpu::TextureFormat = wgpu::TextureFormat::Rgba16Float;

/// The exposure a frame that never sets one is drawn at.
pub(crate) const DEFAULT_EXPOSURE: f32 = 1.0;

/// The bloom chain's fraction for a frame that never sets one.
pub(crate) const DEFAULT_BLOOM: f32 = 0.0;

/// The depth a pixel nothing was drawn over holds, which the forward pass
/// clears to.
pub(crate) const NOTHING: f32 = 1.0;

/// Sample count a pixel is drawn over with smooth edges on; WebGPU offers
/// this count and one, and nothing else.
const SAMPLES: u32 = 4;

/// The bloom chain's smallest mip: no mip has a side under this.
const SMALLEST_MIP: u32 = 8;

/// The one triangle each pass of the chain covers its target with.
const FULLSCREEN: Range<u32> = 0..3;

/// The color a pass of the chain starts from where it replaces everything.
const BLANK: wgpu::Color = wgpu::Color::BLACK;

/// Keeps the target by the fraction the upsample leaves in its alpha.
const SPREAD_OVER_TARGET: wgpu::BlendState = wgpu::BlendState {
    color: wgpu::BlendComponent {
        src_factor: wgpu::BlendFactor::SrcAlpha,
        dst_factor: wgpu::BlendFactor::OneMinusSrcAlpha,
        operation: wgpu::BlendOperation::Add,
    },
    alpha: wgpu::BlendComponent::REPLACE,
};

/// Sample count a pixel is drawn over.
pub(crate) fn sample_count(antialiasing: bool) -> u32 {
    if antialiasing { SAMPLES } else { 1 }
}

/// The target the forward pass draws into: the color it writes, where
/// those samples resolve to, and the depth it tests against.
pub(crate) struct SceneTarget<'a> {
    pub(crate) color: &'a wgpu::TextureView,
    pub(crate) resolve: Option<&'a wgpu::TextureView>,
    pub(crate) depth: &'a wgpu::TextureView,
}

/// The chain from the drawn frame to the screen: the targets it is drawn
/// through, and the pipelines that draw them.
pub(crate) struct Post {
    bindings: Bindings,
    downsample: wgpu::RenderPipeline,
    upsample: wgpu::RenderPipeline,
    display: wgpu::RenderPipeline,
    samples: u32,
    curve: Tonemap,
    /// What the chain was last passed, which is also how much each of its
    /// passes has to draw.
    settings: Settings,
    chain: Option<Chain>,
}

impl Post {
    pub(crate) fn new(
        device: &wgpu::Device,
        display_format: wgpu::TextureFormat,
        samples: u32,
        curve: Tonemap,
    ) -> Self {
        let bindings = Bindings::new(device);
        let shaders = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("mirage-engine post"),
            source: wgpu::ShaderSource::Wgsl(include_str!("post.wgsl").into()),
        });
        let scatter = layout(device, "mirage-engine bloom", &[Some(&bindings.source)]);
        let display = layout(
            device,
            "mirage-engine tone map",
            &[
                Some(&bindings.source),
                Some(&bindings.over),
                Some(&bindings.composite),
            ],
        );
        let built = |label, layout, entry, format, blend| {
            pipeline(
                device,
                label,
                &shaders,
                layout,
                ("fullscreen", entry),
                format,
                blend,
            )
        };

        Self {
            downsample: built(
                "mirage-engine bloom down",
                &scatter,
                "downsample",
                HDR_FORMAT,
                None,
            ),
            upsample: built(
                "mirage-engine bloom up",
                &scatter,
                "upsample",
                HDR_FORMAT,
                Some(SPREAD_OVER_TARGET),
            ),
            display: built(
                "mirage-engine tone map",
                &display,
                "tonemap",
                display_format,
                None,
            ),
            bindings,
            samples,
            curve,
            settings: Settings::zeroed(),
            chain: None,
        }
    }

    /// Builds the chain a `size`-sized frame is drawn through, keeping what
    /// is already the right size.
    pub(crate) fn prepare(&mut self, device: &wgpu::Device, size: UVec2) {
        match &self.chain {
            Some(chain) if chain.size == size => {}
            _ => self.chain = Some(Chain::new(device, &self.bindings, self.samples, size)),
        }
    }

    /// Writes what every pass of the chain draws this frame by.
    pub(crate) fn set_frame(&mut self, queue: &wgpu::Queue, frame: &Frame) {
        let _ = frame.size;
        self.settings = Settings {
            exposure: frame.exposure,
            bloom: frame.bloom,
            curve: self.curve.index(),
            _padding: 0,
        };
        queue.write_buffer(
            &self.bindings.settings,
            0,
            bytemuck::bytes_of(&self.settings),
        );
    }

    /// The drawn frame at one sample per pixel, which every pass after the
    /// forward one reads; absent until a size is known.
    pub(crate) fn drawn(&self) -> Option<&wgpu::TextureView> {
        Some(&self.chain.as_ref()?.drawn)
    }

    /// The depth the forward pass wrote, at the samples it drew over.
    pub(crate) fn depth(&self) -> Option<&wgpu::TextureView> {
        Some(&self.chain.as_ref()?.depth)
    }

    /// What the chain reads the drawn frame through, where no pass of the
    /// game's own has moved it.
    pub(crate) fn sampled(&self) -> Option<&wgpu::BindGroup> {
        Some(&self.chain.as_ref()?.sampled)
    }

    /// What the chain reads `view` through, for a pass of the game's own to
    /// leave the frame in another target.
    pub(crate) fn source(
        &self,
        device: &wgpu::Device,
        view: &wgpu::TextureView,
    ) -> wgpu::BindGroup {
        self.bindings.source(device, view)
    }

    /// The target the forward pass draws into, absent until a size is
    /// known.
    pub(crate) fn scene(&self) -> Option<SceneTarget<'_>> {
        let chain = self.chain.as_ref()?;
        Some(SceneTarget {
            color: &chain.scene,
            resolve: chain.resolve.as_ref(),
            depth: &chain.depth,
        })
    }

    /// Records the bloom chain over `source`, then the tone map that takes
    /// it into `display`.
    pub(crate) fn encode(
        &self,
        encoder: &mut wgpu::CommandEncoder,
        source: &wgpu::BindGroup,
        display: &wgpu::TextureView,
    ) {
        let Some(chain) = &self.chain else {
            return;
        };
        if self.settings.bloom > 0.0 {
            self.scatter(encoder, chain, source);
        }

        let mut pass = pass(
            encoder,
            "mirage-engine tone map",
            display,
            wgpu::LoadOp::Clear(BLANK),
        );
        pass.set_pipeline(&self.display);
        pass.set_bind_group(0, source, &[]);
        pass.set_bind_group(1, &chain.over, &[]);
        pass.set_bind_group(2, &chain.composite, &[]);
        pass.draw(FULLSCREEN, 0..1);
    }

    /// Halves `source` down the chain, then spreads it back up over the
    /// larger mips.
    fn scatter(&self, encoder: &mut wgpu::CommandEncoder, chain: &Chain, source: &wgpu::BindGroup) {
        let sources = core::iter::once(source).chain(chain.bloom.iter().map(|mip| &mip.source));
        for (source, mip) in sources.zip(&chain.bloom) {
            step(
                encoder,
                "mirage-engine bloom down",
                &self.downsample,
                source,
                &mip.view,
                wgpu::LoadOp::Clear(BLANK),
            );
        }

        for pair in chain.bloom.windows(2).rev() {
            let [larger, smaller] = pair else {
                continue;
            };
            step(
                encoder,
                "mirage-engine bloom up",
                &self.upsample,
                &smaller.source,
                &larger.view,
                wgpu::LoadOp::Load,
            );
        }
    }
}

/// The layout the chain's bindings are each built against, and the buffers
/// every pass of it reads the frame's own values from.
struct Bindings {
    sampler: wgpu::Sampler,
    source: wgpu::BindGroupLayout,
    over: wgpu::BindGroupLayout,
    composite: wgpu::BindGroupLayout,
    settings: wgpu::Buffer,
}

impl Bindings {
    fn new(device: &wgpu::Device) -> Self {
        Self {
            sampler: device.create_sampler(&wgpu::SamplerDescriptor {
                label: Some("mirage-engine post"),
                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,
                ..Default::default()
            }),
            source: device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                label: Some("mirage-engine post source"),
                entries: &[sampled(0), sampler(1)],
            }),
            over: device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                label: Some("mirage-engine post over"),
                entries: &[uniform(1)],
            }),
            composite: device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                label: Some("mirage-engine tone map"),
                entries: &[sampled(0)],
            }),
            settings: buffer(
                device,
                "mirage-engine post settings",
                size_of::<Settings>() as wgpu::BufferAddress,
                wgpu::BufferUsages::UNIFORM,
            ),
        }
    }

    /// The bindings the chain samples `view` through.
    fn source(&self, device: &wgpu::Device, view: &wgpu::TextureView) -> wgpu::BindGroup {
        device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("mirage-engine post source"),
            layout: &self.source,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: wgpu::BindingResource::TextureView(view),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: wgpu::BindingResource::Sampler(&self.sampler),
                },
            ],
        })
    }

    /// The bindings a pass reads the frame's own values through.
    fn over(&self, device: &wgpu::Device) -> wgpu::BindGroup {
        device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("mirage-engine post over"),
            layout: &self.over,
            entries: &[wgpu::BindGroupEntry {
                binding: 1,
                resource: self.settings.as_entire_binding(),
            }],
        })
    }

    /// The bindings the tone map reads the scattered frame through.
    fn composite(&self, device: &wgpu::Device, view: &wgpu::TextureView) -> wgpu::BindGroup {
        device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("mirage-engine tone map"),
            layout: &self.composite,
            entries: &[wgpu::BindGroupEntry {
                binding: 0,
                resource: wgpu::BindingResource::TextureView(view),
            }],
        })
    }
}

/// Everything sized to the target, built again when it changes size.
struct Chain {
    size: UVec2,
    scene: wgpu::TextureView,
    resolve: Option<wgpu::TextureView>,
    depth: wgpu::TextureView,
    /// The drawn frame at one sample per pixel: the resolve where the frame
    /// is drawn over more than one sample, and what the forward pass wrote
    /// where it is not.
    drawn: wgpu::TextureView,
    bloom: Vec<Mip>,
    /// Samples the drawn frame at one sample per pixel, which the passes
    /// after the forward one read.
    sampled: wgpu::BindGroup,
    /// What every pass reads of the frame's own values.
    over: wgpu::BindGroup,
    /// The bloom chain as the tone map composites it.
    composite: wgpu::BindGroup,
}

impl Chain {
    fn new(device: &wgpu::Device, bindings: &Bindings, samples: u32, size: UVec2) -> Self {
        let scene = hdr_texture(device, size, samples);
        let resolved = (samples > 1).then(|| hdr_texture(device, size, 1));
        let drawn = view(resolved.as_ref().unwrap_or(&scene));
        let depth = view(&depth_texture(device, size, samples));

        let bloom: Vec<Mip> = mip_sizes(size)
            .into_iter()
            .map(|mip| Mip::new(device, bindings, mip))
            .collect();
        // A target with nothing to halve has nothing to scatter, so the tone
        // map composites the frame with itself, which changes nothing.
        let scattered = bloom.first().map_or(&drawn, |mip| &mip.view);

        Self {
            size,
            composite: bindings.composite(device, scattered),
            sampled: bindings.source(device, &drawn),
            over: bindings.over(device),
            scene: view(&scene),
            resolve: resolved.as_ref().map(view),
            depth,
            drawn,
            bloom,
        }
    }
}

/// One step of the bloom chain: the target a downsample draws into, and
/// the source an upsample samples.
struct Mip {
    view: wgpu::TextureView,
    source: wgpu::BindGroup,
}

impl Mip {
    fn new(device: &wgpu::Device, bindings: &Bindings, size: UVec2) -> Self {
        let view = view(&hdr_texture(device, size, 1));
        Self {
            source: bindings.source(device, &view),
            view,
        }
    }
}

/// What every pass over the drawn frame reads of the frame itself.
#[repr(C)]
#[derive(Clone, Copy, Debug, Pod, Zeroable)]
struct Settings {
    exposure: f32,
    bloom: f32,
    curve: u32,
    /// Keeps this struct's size at WGSL's alignment; `bytemuck` needs padding
    /// written out.
    _padding: u32,
}

/// What one frame is taken to the screen through beyond its own pixels.
pub(crate) struct Frame {
    pub(crate) size: UVec2,
    pub(crate) exposure: f32,
    pub(crate) bloom: f32,
}

/// The mip sizes the bloom chain halves a `size`-sized frame down through,
/// in that order; empty for a frame with nothing to halve.
fn mip_sizes(size: UVec2) -> Vec<UVec2> {
    let mut sizes = Vec::new();
    let mut mip = size / 2;
    while mip.min_element() >= SMALLEST_MIP {
        sizes.push(mip);
        mip /= 2;
    }
    sizes
}

/// Draws one triangle of `pipeline` over `target`, with `source` bound.
fn step(
    encoder: &mut wgpu::CommandEncoder,
    label: &str,
    pipeline: &wgpu::RenderPipeline,
    source: &wgpu::BindGroup,
    target: &wgpu::TextureView,
    load: wgpu::LoadOp<wgpu::Color>,
) {
    let mut pass = pass(encoder, label, target, load);
    pass.set_pipeline(pipeline);
    pass.set_bind_group(0, source, &[]);
    pass.draw(FULLSCREEN, 0..1);
}

/// A pass over `target` alone, which draws nothing until the caller does.
fn pass<'a>(
    encoder: &'a mut wgpu::CommandEncoder,
    label: &str,
    target: &wgpu::TextureView,
    load: wgpu::LoadOp<wgpu::Color>,
) -> wgpu::RenderPass<'a> {
    encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
        label: Some(label),
        color_attachments: &[Some(attachment(target, load))],
        depth_stencil_attachment: None,
        timestamp_writes: None,
        occlusion_query_set: None,
        multiview_mask: None,
    })
}

fn attachment(
    view: &wgpu::TextureView,
    load: wgpu::LoadOp<wgpu::Color>,
) -> wgpu::RenderPassColorAttachment<'_> {
    wgpu::RenderPassColorAttachment {
        view,
        depth_slice: None,
        resolve_target: None,
        ops: wgpu::Operations {
            load,
            store: wgpu::StoreOp::Store,
        },
    }
}

fn layout(
    device: &wgpu::Device,
    label: &str,
    groups: &[Option<&wgpu::BindGroupLayout>],
) -> wgpu::PipelineLayout {
    device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
        label: Some(label),
        bind_group_layouts: groups,
        immediate_size: 0,
    })
}

fn pipeline(
    device: &wgpu::Device,
    label: &str,
    shaders: &wgpu::ShaderModule,
    layout: &wgpu::PipelineLayout,
    stages: (&str, &str),
    format: wgpu::TextureFormat,
    blend: Option<wgpu::BlendState>,
) -> wgpu::RenderPipeline {
    let (vertex, entry) = stages;
    device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
        label: Some(label),
        layout: Some(layout),
        vertex: wgpu::VertexState {
            module: shaders,
            entry_point: Some(vertex),
            compilation_options: wgpu::PipelineCompilationOptions::default(),
            buffers: &[],
        },
        primitive: wgpu::PrimitiveState::default(),
        depth_stencil: None,
        multisample: wgpu::MultisampleState::default(),
        fragment: Some(wgpu::FragmentState {
            module: shaders,
            entry_point: Some(entry),
            compilation_options: wgpu::PipelineCompilationOptions::default(),
            targets: &[Some(wgpu::ColorTargetState {
                format,
                blend,
                write_mask: wgpu::ColorWrites::ALL,
            })],
        }),
        multiview_mask: None,
        cache: None,
    })
}

fn hdr_texture(device: &wgpu::Device, size: UVec2, samples: u32) -> wgpu::Texture {
    drawn_texture(device, "mirage-engine scene", size, samples, HDR_FORMAT)
}

fn drawn_texture(
    device: &wgpu::Device,
    label: &str,
    size: UVec2,
    samples: u32,
    format: wgpu::TextureFormat,
) -> wgpu::Texture {
    device.create_texture(&wgpu::TextureDescriptor {
        label: Some(label),
        size: wgpu::Extent3d {
            width: size.x,
            height: size.y,
            depth_or_array_layers: 1,
        },
        mip_level_count: 1,
        sample_count: samples,
        dimension: wgpu::TextureDimension::D2,
        format,
        usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::TEXTURE_BINDING,
        view_formats: &[],
    })
}

fn buffer(
    device: &wgpu::Device,
    label: &str,
    size: wgpu::BufferAddress,
    usage: wgpu::BufferUsages,
) -> wgpu::Buffer {
    device.create_buffer(&wgpu::BufferDescriptor {
        label: Some(label),
        size,
        usage: usage | wgpu::BufferUsages::COPY_DST,
        mapped_at_creation: false,
    })
}

fn view(texture: &wgpu::Texture) -> wgpu::TextureView {
    texture.create_view(&Default::default())
}

fn sampled(binding: u32) -> wgpu::BindGroupLayoutEntry {
    wgpu::BindGroupLayoutEntry {
        binding,
        visibility: wgpu::ShaderStages::FRAGMENT,
        ty: wgpu::BindingType::Texture {
            sample_type: wgpu::TextureSampleType::Float { filterable: true },
            view_dimension: wgpu::TextureViewDimension::D2,
            multisampled: false,
        },
        count: None,
    }
}

fn sampler(binding: u32) -> wgpu::BindGroupLayoutEntry {
    wgpu::BindGroupLayoutEntry {
        binding,
        visibility: wgpu::ShaderStages::FRAGMENT,
        ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
        count: None,
    }
}

fn uniform(binding: u32) -> wgpu::BindGroupLayoutEntry {
    wgpu::BindGroupLayoutEntry {
        binding,
        visibility: wgpu::ShaderStages::FRAGMENT,
        ty: wgpu::BindingType::Buffer {
            ty: wgpu::BufferBindingType::Uniform,
            has_dynamic_offset: false,
            min_binding_size: None,
        },
        count: None,
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn the_bloom_chain_halves_down_to_the_smallest_mip() {
        assert_eq!(
            mip_sizes(UVec2::new(1280, 720)).last(),
            Some(&UVec2::new(20, 11)),
            "the last mip is the last one no shorter than the floor"
        );
        assert!(
            mip_sizes(UVec2::new(1280, 720))
                .windows(2)
                .all(|pair| pair[1] == pair[0] / 2),
            "every mip is half the one before it"
        );
    }

    #[test]
    fn a_frame_too_small_to_halve_has_no_bloom_chain() {
        assert!(mip_sizes(UVec2::splat(2 * SMALLEST_MIP - 1)).is_empty());
        assert_eq!(mip_sizes(UVec2::splat(2 * SMALLEST_MIP)).len(), 1);
    }

    #[test]
    fn antialiasing_is_the_count_webgpu_offers_or_one() {
        assert_eq!(sample_count(true), 4);
        assert_eq!(sample_count(false), 1);
    }

    #[test]
    fn the_settings_stay_the_size_the_shader_reads_them_at() {
        assert_eq!(size_of::<Settings>(), 16);
    }
}