mirage_engine/

post_effect.rs

//! Full-screen passes a game writes in WGSL, run over the frame the engine
//! drew.
//!
//! A post effect runs inside the post chain, the passes that take a drawn
//! frame to the window. It is a type: its fields are the values its WGSL
//! reads, its
//! [`STAGE`](PostEffect::STAGE) is where in the post chain it runs, and
//! its [`SHADER`](PostEffect::SHADER) is the WGSL it runs. Name it in
//! [`Game::PostEffects`](crate::Game::PostEffects) through
//! [`post_effects!`](crate::post_effects), and run it for one frame with
//! [`set_post_effect`](crate::FrameContext::set_post_effect).
//!
//! # What the engine declares
//!
//! An effect's WGSL is compiled into this, which it may read and must not
//! declare again:
//!
//! ```wgsl
//! struct Pixel {
//!     color: vec4<f32>,     // what the chain holds at this pixel
//!     uv: vec2<f32>,        // 0..1 across the frame, zero at its top left
//!     position: vec2<f32>,  // physical pixels, the center of each at a half
//!     size: vec2<f32>,      // the frame, in physical pixels
//! }
//!
//! fn color_at(uv: vec2<f32>) -> vec4<f32>;
//! fn depth_at(uv: vec2<f32>) -> f32;
//! ```
//!
//! `color` is what the pass before this one wrote: at
//! [`Lit`](EffectStage::Lit) that is high-dynamic-range light, whose
//! channels run past `1.0`; at [`ToneMapped`](EffectStage::ToneMapped) and
//! [`OverUi`](EffectStage::OverUi) it is the encoded value the target
//! holds, in `0..1`. `color_at(pixel.uv)` reads the same texel as
//! `pixel.color`.
//!
//! `color_at` and `depth_at` return the edge texel for a `uv` outside
//! `0..1`. `depth_at` returns the view depth in meters the forward pass
//! wrote —
//! the first sample of the pixel where the frame is drawn over more than
//! one — and the far clip of the frame's camera where nothing was drawn,
//! which is `1000.0` until a game sets one with
//! [`Projection::clip`](crate::Projection::clip).
//!
//! # What an effect declares
//!
//! One function, named and written as the engine reads it:
//!
//! ```wgsl
//! fn draw(pixel: Pixel) -> vec4<f32>;
//! ```
//!
//! The pass writes what it returns, every channel, over what the target
//! held: nothing is blended. The effect after it in the same stage reads
//! all four channels back as `pixel.color`.
//!
//! The bloom chain and the tone map read the color alone, so a `Lit`
//! effect's alpha is dropped. A `ToneMapped` or `OverUi` effect's alpha is
//! the alpha the target holds — the engine's own passes write `1.0` there
//! — and a headless reading reads it back.
//!
//! An effect's values are bound as `effect`: a uniform of a WGSL struct
//! named after the Rust type, so a function of the effect's own may take
//! that struct by value. An effect with no fields binds nothing.
//!
//! ```wgsl
//! fn tinted(values: Grain, color: vec3<f32>) -> vec3<f32> { … }
//! ```
//!
//! An effect may declare whatever else it needs beside `draw`: its own
//! code, its own constant values, its own struct types. The names the
//! engine's own shader already holds are `Pixel`, `Frame`, `Fragment`,
//! `FORESHORTENED`, `frame`, `source`, `source_sampler`, `resolved`,
//! `scene`, `color_at`, `depth_at`, `fullscreen`, `effect_fragment`,
//! `resolve` and `effect`; declaring one of them again stops startup, as
//! any other error in the WGSL does. The error names the effect and counts
//! its lines from the compiled shader, the engine's own and the effect's
//! together, not from the effect's file.
//!
//! A game may name any number of effects. Each is one pipeline built at
//! startup and one full-screen pass in a frame that submits it; a frame
//! that submits none runs none.

use crate::shader_values::{Sealed, ShaderValues};

/// The engine's own full-screen shader, with [`SEAM`] where an effect's code
/// goes.
const FRAME: &str = include_str!("renderer/post_effect.wgsl");

/// The line every stitch replaces.
const SEAM: &str = "// mirage-engine:effect";

/// The bind group an effect's values are read through.
pub(crate) const GROUP: u32 = 0;

/// The `draw` the depth resolve is built with, which reads no values and
/// returns the pixel it was passed.
const UNPAINTED: &str = "fn draw(pixel: Pixel) -> vec4<f32> {\n    return pixel.color;\n}\n";

/// Where a post effect runs in the post chain — the passes that take a
/// drawn frame to the window: the forward pass, the `Lit` effects, bloom,
/// the tone map, the `ToneMapped` effects, the UI, then the `OverUi`
/// effects.
///
/// Effects run stage by stage in that order, and within one stage in the
/// order their set lists them.
#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd)]
pub enum EffectStage {
    /// After the forward pass, over the lit scene, before the bloom and the
    /// tone map; the only stage a pixel past `1.0` can be read.
    Lit,
    /// After the tone map, over the frame as it will be shown, before the
    /// UI.
    ToneMapped,
    /// After the UI, over everything the frame holds. Built without the
    /// `ui` feature there is no UI pass and this runs after the
    /// `ToneMapped` effects all the same.
    OverUi,
}

impl EffectStage {
    /// Every stage, in the order the post chain runs them.
    pub(crate) const ALL: [Self; 3] = [Self::Lit, Self::ToneMapped, Self::OverUi];

    /// Whether this stage draws after the tone map rather than in the
    /// scene's own high-dynamic-range light.
    pub(crate) fn after_tone_map(self) -> bool {
        !matches!(self, Self::Lit)
    }
}

/// A pass of a game's own over the whole frame, written in WGSL and
/// stitched into the engine's own at startup.
///
/// Required if you want to change every pixel of the frame after the
/// scene is drawn: implement it on a type whose fields are the values its
/// WGSL reads, name that type in [`Game::PostEffects`](crate::Game::PostEffects),
/// and run it with [`set_post_effect`](crate::FrameContext::set_post_effect).
pub trait PostEffect: ShaderValues {
    /// Where in the post chain this effect runs, and so what its `color`
    /// holds and what its own alpha reaches; see the module's own
    /// docs.
    const STAGE: EffectStage;

    /// WGSL declaring `fn draw(pixel: Pixel) -> vec4<f32>`, run once per
    /// pixel, which returns what this effect writes there.
    ///
    /// `Pixel`, `color_at` and `depth_at` are the engine's, and the
    /// module's own docs state what each of them holds.
    const SHADER: &'static str;
}

/// The post effects one game passes its frame through, named together as
/// [`Game::PostEffects`](crate::Game::PostEffects).
///
/// Written by [`post_effects!`](crate::post_effects) and never implemented
/// by hand; `()` for a game that draws the frame as the chain leaves it.
pub trait PostEffects: Sealed + 'static {
    /// The effects startup compiles, in the order the set lists them.
    #[doc(hidden)]
    fn declared(into: &mut Declarations);

    /// Seat of the effect this set value holds: `0` for the first effect
    /// the set names, one more for each after it.
    #[doc(hidden)]
    fn seat(&self) -> u32;

    /// Lays out the values of the effect this set value holds.
    #[doc(hidden)]
    fn write(&self, into: &mut Vec<u8>);
}

/// The effects a set declares, in the order it lists them; the set macro
/// fills it and startup compiles what it holds.
#[doc(hidden)]
#[derive(Default)]
pub struct Declarations(Vec<Declaration>);

impl Declarations {
    /// Declares the effect at the next seat.
    #[doc(hidden)]
    pub fn declare<S: PostEffect>(&mut self) {
        self.0.push(Declaration::of::<S>());
    }

    /// Everything `S` declares, in seat order.
    pub(crate) fn of<S: PostEffects>() -> Vec<Declaration> {
        let mut declared = Self::default();
        S::declared(&mut declared);
        declared.0
    }
}

/// One effect as startup takes it: its own name, where in the post chain it
/// runs, and the shader that runs it.
#[derive(Debug)]
pub(crate) struct Declaration {
    pub(crate) name: &'static str,
    pub(crate) stage: EffectStage,
    pub(crate) source: String,
}

impl Declaration {
    fn of<S: PostEffect>() -> Self {
        Self {
            name: core::any::type_name::<S>(),
            stage: S::STAGE,
            source: stitched(&S::bound(GROUP, "effect"), S::SHADER),
        }
    }
}

/// Which of a game's post effects this is, counted in the order its set lists
/// them.
#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd)]
pub(crate) struct PostEffectId(pub(crate) u32);

impl PostEffects for () {
    /// Nothing; a game with no effects compiles none.
    fn declared(_into: &mut Declarations) {}

    /// Not reachable: no call that runs an effect compiles for a game
    /// with none.
    fn seat(&self) -> u32 {
        0
    }

    /// Not reachable: no call that writes an effect's values compiles for
    /// a game with none.
    fn write(&self, _into: &mut Vec<u8>) {}
}

/// Writes the set of every post effect a game passes its frame through: an
/// enum with one variant wrapping each named type, which
/// [`Game::PostEffects`](crate::Game::PostEffects) names.
///
/// Each type is spelled by its own name and must be in scope; a set that
/// holds a type twice does not compile. [`EffectStage`] orders the effects
/// first and the order the set lists them second, so two effects in one
/// stage run in that order. A `pub` before `enum` makes the enum public, and `///` lines
/// before that are the enum's.
///
/// ```
/// use mirage_engine::prelude::*;
///
/// #[derive(ShaderValues)]
/// struct Grain {
///     strength: f32,
/// }
///
/// impl PostEffect for Grain {
///     const STAGE: EffectStage = EffectStage::ToneMapped;
///     const SHADER: &'static str = "
///         fn draw(pixel: Pixel) -> vec4<f32> {
///             let noise = fract(sin(pixel.position.x + pixel.position.y) * 43758.5453);
///             return vec4<f32>(pixel.color.rgb + noise * effect.strength, pixel.color.a);
///         }";
/// }
///
/// post_effects! { enum Look { Grain } }
/// ```
#[macro_export]
macro_rules! post_effects {
    ($(#[$attribute:meta])* $vis:vis enum $set:ident { $($effect:ident),+ $(,)? }) => {
        $(#[$attribute])*
        $vis enum $set {
            $($effect($effect)),+
        }

        $(
            impl ::core::convert::From<$effect> for $set {
                fn from(effect: $effect) -> Self {
                    Self::$effect(effect)
                }
            }

            impl $crate::Holds<$effect> for $set {}
        )+

        impl $crate::Sealed for $set {}

        impl $crate::PostEffects for $set {
            fn declared(into: &mut $crate::PostEffectDeclarations) {
                $(into.declare::<$effect>();)+
            }

            fn seat(&self) -> u32 {
                let mut at = 0;
                $(
                    if ::core::matches!(self, Self::$effect(_)) {
                        return at;
                    }
                    at += 1;
                )+
                at
            }

            fn write(&self, into: &mut ::std::vec::Vec<u8>) {
                match self {
                    $(Self::$effect(values) => $crate::ShaderValues::write(values, into)),+
                }
            }
        }
    };
}

/// The shader the depth resolve is drawn with: the engine's own, with the
/// seam left returning what it was passed.
pub(crate) fn unpainted() -> String {
    stitched("", UNPAINTED)
}

/// The shader one effect is drawn with: the engine's own, with the effect's
/// values declared and its code called where the seam marks.
fn stitched(values: &str, draw: &str) -> String {
    let mut code = String::from(values);
    code.push_str(draw);
    FRAME.replace(SEAM, &code)
}

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

    /// An effect with values of its own, written the way a game writes one.
    #[derive(crate::ShaderValues)]
    struct Vignette {
        strength: f32,
    }

    impl PostEffect for Vignette {
        const STAGE: EffectStage = EffectStage::ToneMapped;
        const SHADER: &'static str =
            "fn draw(pixel: Pixel) -> vec4<f32> { return pixel.color * effect.strength; }";
    }

    /// An effect that reads no values, and reads the scene's depth instead.
    #[derive(crate::ShaderValues)]
    struct Fog;

    impl PostEffect for Fog {
        const STAGE: EffectStage = EffectStage::Lit;
        const SHADER: &'static str =
            "fn draw(pixel: Pixel) -> vec4<f32> { return vec4<f32>(depth_at(pixel.uv)); }";
    }

    post_effects! { enum Look { Fog, Vignette } }

    #[test]
    fn the_shader_carries_one_seam_for_an_effect_to_be_stitched_into() {
        assert_eq!(FRAME.matches(SEAM).count(), 1);
    }

    #[test]
    fn an_effects_own_code_is_called_from_the_seam_and_its_values_are_bound() {
        let source = Declaration::of::<Vignette>().source;

        assert!(!source.contains(SEAM), "the seam itself is replaced");
        assert!(source.contains("struct Vignette"));
        assert!(source.contains("@group(0) @binding(0) var<uniform> effect: Vignette;"));
        assert!(
            source.find("struct Vignette") < source.find("fn draw(pixel: Pixel)"),
            "and the values are declared before the code that reads them"
        );
    }

    #[test]
    fn an_effect_with_no_fields_reads_no_values_and_binds_none() {
        let source = Declaration::of::<Fog>().source;

        assert!(!source.contains("@group(0)"));
        assert!(source.contains("fn draw(pixel: Pixel)"));
    }

    #[test]
    fn a_set_declares_its_effects_in_the_order_it_names_them() {
        let declared = Declarations::of::<Look>();

        assert_eq!(
            declared
                .iter()
                .map(|effect| effect.stage)
                .collect::<Vec<_>>(),
            vec![EffectStage::Lit, EffectStage::ToneMapped]
        );
        assert_eq!(
            (
                Look::from(Fog).seat(),
                Look::from(Vignette { strength: 0.0 }).seat()
            ),
            (0, 1)
        );
        assert!(Declarations::of::<()>().is_empty());
    }

    #[test]
    fn a_set_value_lays_out_the_values_of_the_effect_it_holds() {
        let mut written = Vec::new();
        Look::from(Vignette { strength: 0.25 }).write(&mut written);

        assert_eq!(
            f32::from_le_bytes(written[..4].try_into().expect("four bytes")),
            0.25
        );

        let mut none = Vec::new();
        Look::from(Fog).write(&mut none);
        assert!(none.is_empty(), "where an effect reads nothing");
    }

    #[test]
    fn the_places_run_the_scene_first_and_what_is_over_the_ui_last() {
        assert_eq!(
            EffectStage::ALL,
            [
                EffectStage::Lit,
                EffectStage::ToneMapped,
                EffectStage::OverUi
            ]
        );
        assert!(
            EffectStage::Lit < EffectStage::ToneMapped
                && EffectStage::ToneMapped < EffectStage::OverUi
        );
        assert_eq!(
            EffectStage::ALL.map(EffectStage::after_tone_map),
            [false, true, true],
            "and only the first draws in the scene's own light"
        );
    }
}