facett-core 0.1.16

//! **TAA — temporal anti-aliasing: Halton jitter + a velocity-reprojecting resolve**
//! (GFX_V2 §5, item 7).
//!
//! MSAA does not fix thin vector lines: a 1 px road at 4× MSAA still has only four
//! coverage steps, and it shimmers under camera rotation because *which* subsamples it
//! covers changes discontinuously. TAA instead moves the sample pattern in *time* — the
//! projection is offset by a sub-pixel [`jitter_px`] each frame and the frames are
//! accumulated, so the effective sample count grows with frame count and the edge
//! converges on true coverage.
//!
//! Two halves, deliberately separated:
//!
//! * **The jitter** ([`halton`], [`jitter_px`], [`TAA_PHASES`]) is pure CPU arithmetic
//!   and therefore testable without a device. It is also the part a *caller* must apply
//!   — the resolve cannot invent it, because the jitter has to be baked into where the
//!   geometry landed. So [`TaaPass`] owns the frame counter and hands the jitter out
//!   ([`TaaPass::jitter_px`]), rather than the host and the pass each keeping a count
//!   that can drift apart.
//! * **The resolve** ([`TAA_WGSL`]) is one fullscreen pass: 3×3 neighbourhood colour
//!   box, velocity reprojection, history clamped into the box, blend.
//!
//! **The claim, stated so it can fail.** Accumulating N *jittered* frames must reduce
//! the aliasing of an edge, measured — not asserted. The device test
//! (`facett-core/tests/gpu_taa.rs`) renders a slightly tilted hard edge, fits a line to
//! the per-row coverage, and takes the RMS residual: an aliased edge's coverage is a
//! **staircase** (it can only step by whole pixels), a properly antialiased edge's is
//! **linear**. The residual is the aliasing, in luma units.
//!
//! **And why that alone would be hollow.** A plain blur also flattens that staircase.
//! So the same test measures the **transition width** — how many pixels per row sit
//! strictly between background and foreground — and requires it to stay at about one
//! pixel. Convergence with a sharp edge; a blur fails the second half, a dead pass
//! fails the first. The mutation that proves the pair is not vacuous is *zeroing the
//! jitter*: the accumulation then averages N identical aliased frames and the residual
//! does not move.

use wgpu::TextureFormat;

/// The TAA resolve shader. Not prelude-composed: it is a pure screen-space pass and
/// uses none of `common.wgsl`'s geometry maths.
pub const TAA_WGSL: &str = include_str!("taa.wgsl");

/// Length of the jitter cycle. The Halton(2, 3) pair is low-discrepancy at any length;
/// 8 is the usual choice because it is short enough that a *moving* camera revisits the
/// pattern before the history has decayed, and long enough that a static camera
/// resolves an edge to well under a quarter pixel.
pub const TAA_PHASES: u32 = 8;

/// Default cap on the history weight. Below 1 so the resolve keeps decaying into an
/// exponential moving average once the ramp is over — an unbounded running mean would
/// never respond to a scene change at all.
pub const TAA_FEEDBACK_CAP: f32 = 0.97;

/// The **radical-inverse (van der Corput) digit reversal** of `index` in `base` — the
/// Halton sequence's term. Returns a value in `[0, 1)`.
///
/// `index == 0` is `0.0` for every base, which is why callers offset by one: phase 0 of
/// a jitter sequence that starts at index 0 is *no offset at all*, i.e. a frame that
/// does not participate in the antialiasing.
#[must_use]
pub fn halton(index: u32, base: u32) -> f32 {
    if base < 2 {
        return 0.0;
    }
    let mut f = 1.0_f64;
    let mut r = 0.0_f64;
    let mut i = index;
    let b = f64::from(base);
    while i > 0 {
        f /= b;
        r += f * f64::from(i % base);
        i /= base;
    }
    r as f32
}

/// This frame's sub-pixel projection offset, in **pixels**, each component in
/// `[-0.5, 0.5)`.
///
/// Halton(2, 3), indexed by `frame % TAA_PHASES + 1` — the `+ 1` skips index 0, whose
/// offset is exactly `(-0.5, -0.5)` in this centring and, more importantly, whose
/// unshifted variant would be `(0, 0)`.
///
/// A caller applies it to its projection (a half-pixel NDC shift, or by offsetting
/// screen-space vertex positions). Applying it is not optional: **the resolve cannot
/// antialias frames that were all rendered at the same sample position**, which is the
/// mutation the device test uses to prove it is measuring something.
#[must_use]
pub fn jitter_px(frame: u64) -> [f32; 2] {
    let i = (frame % u64::from(TAA_PHASES)) as u32 + 1;
    [halton(i, 2) - 0.5, halton(i, 3) - 0.5]
}

/// The history weight for `frame`, capped at `cap`.
///
/// `frame / (frame + 1)` makes the first frames a true **running mean** of the jitter
/// samples seen so far, which is what converges; a fixed feedback factor is an EMA over
/// a periodic sequence and oscillates at the sequence's period forever, leaving a
/// fraction of the aliasing in every frame. The ramp then saturates at `cap` so the
/// resolve stays responsive to a scene change.
#[must_use]
pub fn history_weight(frame: u64, cap: f32) -> f32 {
    if frame == 0 {
        return 0.0;
    }
    let ramp = frame as f32 / (frame as f32 + 1.0);
    ramp.min(cap)
}

#[repr(C)]
#[derive(Clone, Copy, Debug, Default, bytemuck::Pod, bytemuck::Zeroable)]
struct TaaUniform {
    res: [f32; 4],
    params: [f32; 4],
}

/// The TAA resolve lane: the pipeline, a ping-pong history pair, and the frame counter
/// that owns the jitter phase.
pub struct TaaPass {
    pipeline: wgpu::RenderPipeline,
    bgl: wgpu::BindGroupLayout,
    uniform: wgpu::Buffer,
    sampler: wgpu::Sampler,
    /// 1×1 RG16Float zero — bound when the caller has no motion vectors, so the shader
    /// keeps ONE path instead of branching on whether velocity exists.
    zero_velocity: wgpu::TextureView,
    history: Option<[wgpu::Texture; 2]>,
    history_views: Option<[wgpu::TextureView; 2]>,
    /// Index of the texture holding the *previous* resolve; `1 - cur` receives this one.
    cur: usize,
    format: TextureFormat,
    size: (u32, u32),
    frame: u64,
    feedback_cap: f32,
}

impl TaaPass {
    /// Build the resolve pipeline for a `format` history/output pair.
    #[must_use]
    pub fn new(device: &wgpu::Device, format: TextureFormat) -> Self {
        let module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("l0_taa"),
            source: wgpu::ShaderSource::Wgsl(TAA_WGSL.into()),
        });
        let tex = |binding: u32| 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,
        };
        let bgl = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("l0_taa_bgl"),
            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,
                },
                tex(1),
                tex(2),
                tex(3),
                wgpu::BindGroupLayoutEntry {
                    binding: 4,
                    visibility: wgpu::ShaderStages::FRAGMENT,
                    ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
                    count: None,
                },
            ],
        });
        let pll = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("l0_taa_pll"),
            bind_group_layouts: &[Some(&bgl)],
            immediate_size: 0,
        });
        let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: Some("l0_taa_resolve"),
            layout: Some(&pll),
            vertex: wgpu::VertexState {
                module: &module,
                entry_point: Some("taa_vs"),
                compilation_options: Default::default(),
                buffers: &[],
            },
            primitive: wgpu::PrimitiveState {
                topology: wgpu::PrimitiveTopology::TriangleList,
                ..Default::default()
            },
            depth_stencil: None,
            multisample: wgpu::MultisampleState::default(),
            fragment: Some(wgpu::FragmentState {
                module: &module,
                entry_point: Some("taa_resolve_fs"),
                compilation_options: Default::default(),
                targets: &[Some(wgpu::ColorTargetState {
                    format,
                    // The resolve IS the blend. Hardware blending on top would apply the
                    // history weight twice.
                    blend: None,
                    write_mask: wgpu::ColorWrites::ALL,
                })],
            }),
            multiview_mask: None,
            cache: None,
        });
        let zero = device.create_texture(&wgpu::TextureDescriptor {
            label: Some("l0_taa_zero_velocity"),
            size: wgpu::Extent3d { width: 1, height: 1, depth_or_array_layers: 1 },
            mip_level_count: 1,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format: TextureFormat::Rg16Float,
            usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST,
            view_formats: &[],
        });
        Self {
            pipeline,
            bgl,
            uniform: device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("l0_taa_uniform"),
                size: std::mem::size_of::<TaaUniform>() as u64,
                usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
                mapped_at_creation: false,
            }),
            sampler: device.create_sampler(&wgpu::SamplerDescriptor {
                label: Some("l0_taa_sampler"),
                mag_filter: wgpu::FilterMode::Linear,
                min_filter: wgpu::FilterMode::Linear,
                address_mode_u: wgpu::AddressMode::ClampToEdge,
                address_mode_v: wgpu::AddressMode::ClampToEdge,
                ..Default::default()
            }),
            zero_velocity: zero.create_view(&Default::default()),
            history: None,
            history_views: None,
            cur: 0,
            format,
            size: (0, 0),
            frame: 0,
            feedback_cap: TAA_FEEDBACK_CAP,
        }
    }

    /// Override the history-weight ceiling (default [`TAA_FEEDBACK_CAP`]).
    #[must_use]
    pub fn with_feedback_cap(mut self, cap: f32) -> Self {
        self.feedback_cap = cap.clamp(0.0, 1.0);
        self
    }

    /// (Re)allocate the ping-pong history for `w × h`. Idempotent within a size.
    ///
    /// A resize **invalidates the history** — reprojecting a differently-sized previous
    /// frame would stretch it. Stated here rather than left to a caller to remember.
    pub fn ensure(&mut self, device: &wgpu::Device, w: u32, h: u32) {
        let w = w.max(1);
        let h = h.max(1);
        if self.size == (w, h) && self.history.is_some() {
            return;
        }
        let mk = |label: &str| {
            device.create_texture(&wgpu::TextureDescriptor {
                label: Some(label),
                size: wgpu::Extent3d { width: w, height: h, depth_or_array_layers: 1 },
                mip_level_count: 1,
                sample_count: 1,
                dimension: wgpu::TextureDimension::D2,
                format: self.format,
                usage: wgpu::TextureUsages::RENDER_ATTACHMENT
                    | wgpu::TextureUsages::TEXTURE_BINDING
                    | wgpu::TextureUsages::COPY_SRC,
                view_formats: &[],
            })
        };
        let a = mk("l0_taa_history_a");
        let b = mk("l0_taa_history_b");
        self.history_views = Some([a.create_view(&Default::default()), b.create_view(&Default::default())]);
        self.history = Some([a, b]);
        self.size = (w, h);
        self.reset();
    }

    /// Drop the accumulated history and restart the jitter sequence at phase 0.
    ///
    /// The next [`Self::record`] takes the current frame outright. Call it on a teleport
    /// / preset jump, where no reprojection is valid.
    pub fn reset(&mut self) {
        self.frame = 0;
        self.cur = 0;
    }

    /// **The jitter this frame's geometry must be rendered with.**
    ///
    /// The pass owns the phase, so the offset the caller applies and the frame the
    /// resolve accumulates cannot disagree.
    #[must_use]
    pub fn jitter_px(&self) -> [f32; 2] {
        jitter_px(self.frame)
    }

    /// Frames accumulated since the last [`Self::reset`].
    #[must_use]
    pub fn frame(&self) -> u64 {
        self.frame
    }

    /// The history weight the next [`Self::record`] will use — `0.0` on a fresh start.
    ///
    /// **There is no separate "history valid" flag, by construction.** `frame == 0` is
    /// exactly the state in which no history exists (both [`Self::reset`] and
    /// [`Self::ensure`] put it there), and [`history_weight`] is `0.0` at frame 0 — so
    /// "take the current frame outright" falls out of the ramp instead of being a second
    /// condition that has to agree with it.
    ///
    /// It was a second condition, briefly: the shader also gated on a `history_valid`
    /// uniform. Mutating that gate away left every test GREEN, because the weight was
    /// already 0 — a guard that could not report red. LAW #5 prefers one writer to two
    /// copies watched for agreement, so the flag is gone rather than tested harder.
    #[must_use]
    pub fn weight(&self) -> f32 {
        history_weight(self.frame, self.feedback_cap)
    }

    /// Resolve `current` (rendered with [`Self::jitter_px`]) against the history and
    /// advance the frame counter. `velocity` is UV-per-frame motion; `None` binds a zero
    /// texture.
    ///
    /// Returns `false` when the pass is not allocated, so a caller can tell "resolved"
    /// from "silently did nothing".
    pub fn record(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        encoder: &mut wgpu::CommandEncoder,
        current: &wgpu::TextureView,
        velocity: Option<&wgpu::TextureView>,
    ) -> bool {
        let Some(views) = &self.history_views else { return false };
        let (w, h) = self.size;
        let dst = 1 - self.cur;

        queue.write_buffer(
            &self.uniform,
            0,
            bytemuck::bytes_of(&TaaUniform {
                res: [w as f32, h as f32, 1.0 / w as f32, 1.0 / h as f32],
                params: {
                    let j = self.jitter_px();
                    [self.weight(), j[0], j[1], 0.0]
                },
            }),
        );

        let bind = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("l0_taa_bind"),
            layout: &self.bgl,
            entries: &[
                wgpu::BindGroupEntry { binding: 0, resource: self.uniform.as_entire_binding() },
                wgpu::BindGroupEntry { binding: 1, resource: wgpu::BindingResource::TextureView(current) },
                wgpu::BindGroupEntry {
                    binding: 2,
                    resource: wgpu::BindingResource::TextureView(&views[self.cur]),
                },
                wgpu::BindGroupEntry {
                    binding: 3,
                    resource: wgpu::BindingResource::TextureView(velocity.unwrap_or(&self.zero_velocity)),
                },
                wgpu::BindGroupEntry { binding: 4, resource: wgpu::BindingResource::Sampler(&self.sampler) },
            ],
        });

        {
            let mut rp = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
                label: Some("l0_taa_resolve_pass"),
                color_attachments: &[Some(wgpu::RenderPassColorAttachment {
                    view: &views[dst],
                    resolve_target: None,
                    depth_slice: None,
                    ops: wgpu::Operations {
                        load: wgpu::LoadOp::Clear(wgpu::Color::TRANSPARENT),
                        store: wgpu::StoreOp::Store,
                    },
                })],
                depth_stencil_attachment: None,
                timestamp_writes: None,
                occlusion_query_set: None,
                multiview_mask: None,
            });
            rp.set_pipeline(&self.pipeline);
            rp.set_bind_group(0, &bind, &[]);
            rp.draw(0..3, 0..1);
        }

        self.cur = dst;
        self.frame += 1;
        true
    }

    /// The most recently resolved texture — what a caller composites or reads back.
    #[must_use]
    pub fn resolved(&self) -> Option<&wgpu::Texture> {
        self.history.as_ref().map(|h| &h[self.cur])
    }

    /// The most recently resolved view.
    #[must_use]
    pub fn resolved_view(&self) -> Option<&wgpu::TextureView> {
        self.history_views.as_ref().map(|v| &v[self.cur])
    }

    /// Allocated size in physical pixels.
    #[must_use]
    pub fn size(&self) -> (u32, u32) {
        self.size
    }

    /// Whether the history is allocated.
    #[must_use]
    pub fn ready(&self) -> bool {
        self.history.is_some()
    }
}

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

    /// The van der Corput terms, by hand. Base 2: 1/2, 1/4, 3/4, 1/8, 5/8, 3/8, 7/8.
    /// Base 3: 1/3, 2/3, 1/9, 4/9, 7/9.
    #[test]
    fn halton_is_the_radical_inverse() {
        for (i, want) in [(1u32, 0.5), (2, 0.25), (3, 0.75), (4, 0.125), (5, 0.625), (6, 0.375), (7, 0.875)] {
            assert!((halton(i, 2) - want).abs() < 1e-6, "halton({i}, 2) = {} want {want}", halton(i, 2));
        }
        for (i, want) in [(1u32, 1.0 / 3.0), (2, 2.0 / 3.0), (3, 1.0 / 9.0), (4, 4.0 / 9.0), (5, 7.0 / 9.0)] {
            assert!((halton(i, 3) - want).abs() < 1e-6, "halton({i}, 3) = {} want {want}", halton(i, 3));
        }
        assert_eq!(halton(0, 2), 0.0, "index 0 is the origin — which is why jitter_px offsets by one");
        assert_eq!(halton(5, 1), 0.0, "base < 2 is not a sequence");
        assert_eq!(halton(5, 0), 0.0);
    }

    /// **The identity-value trap, head on.** A jitter that is all zeros — or one whose
    /// phases repeat — cannot antialias anything, and would be invisible in a test that
    /// only checked the offsets are in range. So: every phase is distinct, none is the
    /// origin, and the offsets actually spread over the pixel.
    #[test]
    fn every_jitter_phase_is_a_distinct_nonzero_subpixel_offset() {
        let phases: Vec<[f32; 2]> = (0..TAA_PHASES as u64).map(jitter_px).collect();
        for (n, j) in phases.iter().enumerate() {
            assert!(
                j[0] >= -0.5 && j[0] < 0.5 && j[1] >= -0.5 && j[1] < 0.5,
                "phase {n} offset {j:?} is inside the pixel"
            );
            assert!(
                j[0] != 0.0 || j[1] != 0.0,
                "phase {n} is the ORIGIN — a phase that does not move the sample cannot antialias"
            );
        }
        for i in 0..phases.len() {
            for j in (i + 1)..phases.len() {
                assert_ne!(phases[i], phases[j], "phases {i} and {j} are the same sample position");
            }
        }
        // And the pattern spans the pixel rather than clustering in one corner: the
        // extremes in each axis must be at least half a pixel apart.
        let xs: Vec<f32> = phases.iter().map(|j| j[0]).collect();
        let ys: Vec<f32> = phases.iter().map(|j| j[1]).collect();
        let span = |v: &[f32]| v.iter().cloned().fold(f32::MIN, f32::max) - v.iter().cloned().fold(f32::MAX, f32::min);
        assert!(span(&xs) >= 0.5, "the x offsets span {} of a pixel", span(&xs));
        assert!(span(&ys) >= 0.5, "the y offsets span {} of a pixel", span(&ys));
    }

    /// The sequence cycles with period [`TAA_PHASES`], so a long run keeps revisiting
    /// the same well-spread set rather than drifting.
    #[test]
    fn the_jitter_cycles_with_the_declared_period() {
        for f in 0..40u64 {
            assert_eq!(jitter_px(f), jitter_px(f + u64::from(TAA_PHASES)), "frame {f}");
        }
        assert_ne!(jitter_px(0), jitter_px(1), "…but consecutive frames differ");
    }

    /// The weight ramp: 0 on the first frame (no history exists), a running mean while
    /// ramping, and capped thereafter.
    #[test]
    fn the_history_weight_ramps_then_saturates() {
        assert_eq!(history_weight(0, 0.97), 0.0, "frame 0 has no history to weight");
        assert!((history_weight(1, 0.97) - 0.5).abs() < 1e-6, "frame 1 averages two samples");
        assert!((history_weight(2, 0.97) - 2.0 / 3.0).abs() < 1e-6);
        assert!((history_weight(3, 0.97) - 0.75).abs() < 1e-6);
        // Ramp reaches the cap at frame 33 (33/34 > 0.97) and never exceeds it.
        for f in 0..500u64 {
            assert!(history_weight(f, 0.97) <= 0.97 + 1e-6, "frame {f} respects the cap");
        }
        assert!((history_weight(10_000, 0.97) - 0.97).abs() < 1e-6, "saturated");
        // A cap of 0 disables temporal accumulation entirely — the honest "TAA off".
        for f in 0..10u64 {
            assert_eq!(history_weight(f, 0.0), 0.0);
        }
    }

    /// Uniform layout: two `vec4<f32>`, 32 B. The OIT lane's first red on this branch
    /// was a Rust/WGSL uniform-size mismatch, so this one is checked.
    #[test]
    fn the_uniform_is_two_vec4s() {
        assert_eq!(std::mem::size_of::<TaaUniform>(), 32);
        assert!(TAA_WGSL.contains("res: vec4<f32>"), "and the shader agrees");
        assert!(TAA_WGSL.contains("params: vec4<f32>"));
        assert!(
            !TAA_WGSL.contains("vec3<"),
            "no vec3 in a uniform — it aligns to 16, not 12, and silently resizes the struct"
        );
    }

    /// The resolve must clamp the history into the current frame's neighbourhood box,
    /// and must reject a history that reprojected off screen. Both are the difference
    /// between TAA and a smear.
    ///
    /// There is deliberately NO assertion about a "history valid" gate — see
    /// [`TaaPass::weight`] for why that flag was removed rather than tested.
    #[test]
    fn the_resolve_rejects_history_it_cannot_trust() {
        assert!(TAA_WGSL.contains("clamp(hist, lo, hi)"), "history is box-clamped");
        assert!(
            TAA_WGSL.contains("huv.x >= 0.0 && huv.x <= 1.0"),
            "a pixel reprojecting off screen has no history"
        );
        assert!(
            !TAA_WGSL.contains("params.y > 0.5"),
            "the redundant history-valid gate must stay gone — frame 0's weight is already 0"
        );
    }
}