valo-renderer 0.12.0

//! Runtime effects on the GPU: each effect's generated shader module and layouts, and the
//! override values behind each pipeline variant its draws name.
//!
//! A module is the ABI prelude, the shared draw record and vertex stages (`draw.wgsl`), the
//! effect author's WGSL, a filter's `source` declarations, and a fragment entry that calls the
//! author's `shade` and tints what it returns. The bind groups follow
//! `valo_dl`'s contract: the draw record at group 0, the effect's images at 1, its uniforms at 2,
//! a filter's source at 3.

use rustc_hash::{FxHashMap, FxHasher};
use std::hash::{Hash, Hasher};
use valo_dl::{EffectKind, RuntimeEffect, RuntimeShader};

use crate::images::ImageStore;

/// The entry point of every runtime module's fragment stage.
pub(crate) const FRAGMENT_ENTRY: &str = "valo_runtime_fragment";

/// A filter's source and the helpers over it. A filter draws in its output's space, so its
/// varying reaches the effect's local space through the record's `LOCAL` map, and `uv_of` maps a
/// local position on through `SOURCE_UV`; `MISC.x` is the source's level count, and `GEOM` the
/// area the filter shades in local coordinates.
const FILTER_SUPPORT: &str = "
// ── a filter's source ──
@group(3) @binding(0) var source: texture_2d<f32>;
@group(3) @binding(1) var source_sampler: sampler;

fn valo_affine(slot: u32, p: vec2<f32>) -> vec2<f32> {
    let m = valo_draw.payload[slot];
    let t = valo_draw.payload[slot + 1u];
    return vec2(m.x * p.x + m.z * p.y + t.x, m.y * p.x + m.w * p.y + t.y);
}

fn uv_of(position: vec2<f32>) -> vec2<f32> {
    return valo_affine(PAYLOAD_SOURCE_UV, position);
}

fn source_levels() -> f32 {
    return valo_draw.payload[PAYLOAD_MISC].x;
}

fn filter_bounds() -> vec4<f32> {
    return valo_draw.payload[PAYLOAD_GEOM];
}
";

/// The fragment stage around a paint shader's `shade`, which returns premultiplied colour.
const SHADER_FRAGMENT: &str = "
// ── generated by valo ──
@fragment
fn valo_runtime_fragment(in: VsOut) -> @location(0) vec4<f32> {
    return shade(in.local) * valo_draw.color;
}
";

/// The fragment stage around a filter's `shade`.
const FILTER_FRAGMENT: &str = "
// ── generated by valo ──
@fragment
fn valo_runtime_fragment(in: VsOut) -> @location(0) vec4<f32> {
    return shade(valo_affine(PAYLOAD_LOCAL, in.local)) * valo_draw.color;
}
";

/// `RuntimeKey` names one pipeline variant of one effect: the effect and the override values its
/// draw sets.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct RuntimeKey {
    pub effect: u64,
    pub constants: u64,
}

/// `RuntimeProgram` is one effect's compiled module and the layouts its pipelines bind with.
pub(crate) struct RuntimeProgram {
    pub module: wgpu::ShaderModule,
    pub layout: wgpu::PipelineLayout,
    /// Group 1, when the effect samples images.
    pub image_layout: Option<wgpu::BindGroupLayout>,
    /// Group 3, when the effect is a filter.
    pub source_layout: Option<wgpu::BindGroupLayout>,
}

/// `RuntimePrograms` holds every effect drawn with and the override values of each variant, for
/// the pipeline cache to compile from. It grows only, like the pipeline cache.
pub(crate) struct RuntimePrograms {
    /// The host buffer's group-0 layout, for the draw record.
    record_layout: wgpu::BindGroupLayout,
    /// The host buffer's group-2 layout for an effect's uniforms.
    uniform_layout: wgpu::BindGroupLayout,
    programs: FxHashMap<u64, RuntimeProgram>,
    constants: FxHashMap<RuntimeKey, Vec<(String, f64)>>,
}

impl RuntimePrograms {
    pub fn new(
        record_layout: &wgpu::BindGroupLayout,
        uniform_layout: &wgpu::BindGroupLayout,
    ) -> Self {
        Self {
            record_layout: record_layout.clone(),
            uniform_layout: uniform_layout.clone(),
            programs: FxHashMap::default(),
            constants: FxHashMap::default(),
        }
    }

    /// `register` makes sure `shader`'s effect has a program and its override values a key,
    /// compiling the effect's module the first time it is seen.
    pub fn register(&mut self, device: &wgpu::Device, shader: &RuntimeShader) -> RuntimeKey {
        let effect = &shader.effect;
        let (record_layout, uniform_layout) = (&self.record_layout, &self.uniform_layout);
        self.programs
            .entry(effect.id())
            .or_insert_with(|| compile(device, record_layout, uniform_layout, effect));
        let mut constants = shader.constants.clone();
        constants.sort_by(|(a, _), (b, _)| a.cmp(b));
        let mut hasher = FxHasher::default();
        for (name, value) in &constants {
            name.hash(&mut hasher);
            value.to_bits().hash(&mut hasher);
        }
        let key = RuntimeKey {
            effect: effect.id(),
            constants: hasher.finish(),
        };
        self.constants.entry(key).or_insert(constants);
        key
    }

    /// `program` is the registered program of `effect`.
    pub fn program(&self, effect: u64) -> &RuntimeProgram {
        &self.programs[&effect]
    }

    /// `constants` are the override values of the variant `key` names.
    pub fn constants(&self, key: RuntimeKey) -> &[(String, f64)] {
        &self.constants[&key]
    }

    /// `source_bind` is the group-3 bind of a filter's `source`, read through `sampler`.
    pub fn source_bind(
        &self,
        device: &wgpu::Device,
        effect: u64,
        source: &wgpu::TextureView,
        sampler: &wgpu::Sampler,
    ) -> wgpu::BindGroup {
        let layout = self
            .program(effect)
            .source_layout
            .as_ref()
            .expect("a filter's program has a source layout");
        device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("valo.runtime.source"),
            layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: wgpu::BindingResource::TextureView(source),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: wgpu::BindingResource::Sampler(sampler),
                },
            ],
        })
    }

    /// `image_bind` is the group-1 bind of `shader`'s images, or `None` when it samples none.
    pub fn image_bind(
        &self,
        device: &wgpu::Device,
        images: &mut ImageStore,
        shader: &RuntimeShader,
    ) -> Option<wgpu::BindGroup> {
        let layout = self.program(shader.effect.id()).image_layout.as_ref()?;
        let samplers: Vec<wgpu::Sampler> = shader
            .images
            .iter()
            .map(|input| images.sampler(input.sampling).clone())
            .collect();
        let entries: Vec<wgpu::BindGroupEntry> = shader
            .images
            .iter()
            .zip(&samplers)
            .enumerate()
            .flat_map(|(slot, (input, sampler))| {
                [
                    wgpu::BindGroupEntry {
                        binding: 2 * slot as u32,
                        resource: wgpu::BindingResource::TextureView(input.image.view()),
                    },
                    wgpu::BindGroupEntry {
                        binding: 2 * slot as u32 + 1,
                        resource: wgpu::BindingResource::Sampler(sampler),
                    },
                ]
            })
            .collect();
        Some(device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("valo.runtime.images"),
            layout,
            entries: &entries,
        }))
    }
}

/// `generated_source` is the whole module of `effect`.
pub(crate) fn generated_source(effect: &RuntimeEffect) -> String {
    let mut source = crate::shader_abi::wgsl_prelude();
    source.push_str(include_str!("shaders/draw.wgsl"));
    source.push_str("\n// ── the effect ──\n");
    source.push_str(effect.source());
    match effect.kind() {
        EffectKind::Shader => source.push_str(SHADER_FRAGMENT),
        EffectKind::Filter(_) => {
            source.push_str(FILTER_SUPPORT);
            source.push_str(FILTER_FRAGMENT);
        }
    }
    source
}

fn compile(
    device: &wgpu::Device,
    record_layout: &wgpu::BindGroupLayout,
    uniform_layout: &wgpu::BindGroupLayout,
    effect: &RuntimeEffect,
) -> RuntimeProgram {
    let module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
        label: Some("valo.runtime"),
        source: wgpu::ShaderSource::Wgsl(generated_source(effect).into()),
    });
    let image_layout = (effect.images() > 0).then(|| image_layout(device, effect.images()));
    let source_layout = matches!(effect.kind(), EffectKind::Filter(_))
        .then(|| texture_and_sampler_layout(device, "valo.runtime.source"));
    let uniforms = (effect.uniforms().size > 0).then_some(uniform_layout);
    let layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
        label: Some("valo.runtime"),
        bind_group_layouts: &[
            Some(record_layout),
            image_layout.as_ref(),
            uniforms,
            source_layout.as_ref(),
        ],
        immediate_size: 0,
    });
    RuntimeProgram {
        module,
        layout,
        image_layout,
        source_layout,
    }
}

fn image_layout(device: &wgpu::Device, images: u32) -> wgpu::BindGroupLayout {
    let entries: Vec<wgpu::BindGroupLayoutEntry> = (0..images)
        .flat_map(|slot| texture_and_sampler_entries(2 * slot))
        .collect();
    device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
        label: Some("valo.runtime.images"),
        entries: &entries,
    })
}

fn texture_and_sampler_layout(device: &wgpu::Device, label: &str) -> wgpu::BindGroupLayout {
    device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
        label: Some(label),
        entries: &texture_and_sampler_entries(0),
    })
}

fn texture_and_sampler_entries(binding: u32) -> [wgpu::BindGroupLayoutEntry; 2] {
    [
        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,
        },
        wgpu::BindGroupLayoutEntry {
            binding: binding + 1,
            visibility: wgpu::ShaderStages::FRAGMENT,
            ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
            count: None,
        },
    ]
}

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

    /// Every generated module is valid WGSL: the contract the author validated against and the
    /// declarations the renderer adds agree.
    #[test]
    fn generated_modules_validate() {
        let shader = RuntimeEffect::shader(
            "struct U { tint: vec4<f32> }
             @group(2) @binding(0) var<uniform> u: U;
             @group(1) @binding(0) var t: texture_2d<f32>;
             @group(1) @binding(1) var s: sampler;
             fn shade(position: vec2<f32>) -> vec4<f32> {
                 return u.tint * textureSampleLevel(t, s, position, 0.0);
             }",
        )
        .unwrap();
        let filter = RuntimeEffect::filter(
            "fn shade(position: vec2<f32>) -> vec4<f32> {
                 let bounds = filter_bounds();
                 let centre = bounds.xy + bounds.zw * 0.5;
                 let uv = uv_of(centre + (position - centre) * 0.5);
                 return textureSampleLevel(source, source_sampler, uv, source_levels() - 1.0);
             }",
            valo_dl::FilterInputs::default(),
        )
        .unwrap();
        for effect in [shader, filter] {
            let source = generated_source(&effect);
            let module = naga::front::wgsl::parse_str(&source)
                .unwrap_or_else(|error| panic!("{}", error.emit_to_string(&source)));
            naga::valid::Validator::new(
                naga::valid::ValidationFlags::all(),
                naga::valid::Capabilities::default(),
            )
            .validate(&module)
            .unwrap_or_else(|error| panic!("{}", error.emit_to_string(&source)));
        }
    }
}