concinnity-device 0.19.16

// src/metal/post/ssr.rs
//
// Screen-space reflections: a depth + normal + roughness pre-pass that runs
// before the main pass, and a fullscreen ray-march resolve that runs after.
// Pipelines, targets, and both encoders live together so the effect is a
// single unit Vulkan / DirectX can mirror.
#![deny(unsafe_op_in_unsafe_fn)]

use objc2::rc::Retained;
use objc2::runtime::ProtocolObject;
use objc2_metal::{
    MTLDevice as _, MTLLoadAction, MTLPixelFormat, MTLRenderPipelineState, MTLTexture,
    MTLTextureUsage,
};

use crate::gfx::ssr::SsrSettings;
use crate::metal::context::MtlContext;
use crate::metal::descriptors::TextureDesc;
use crate::metal::encode::RenderEncode;
use crate::metal::post::fullscreen::{
    FullscreenBlend, FullscreenPass, PassTimer, build_slang_fullscreen_pipeline,
    set_fragment_sampler_range,
};
use crate::metal::slang_shaders::{REFLECTION_BLUR, REFLECTION_COMPOSITE, SSR_RESOLVE};

// All screen-space-reflection feature state grouped into one unit: the
// resolved tunables, the resolve-output target, and the resolve pipeline.
// `targets` is `Some` when SSR, SSGI, *or* RT reflections are on (they share
// the G-buffer pre-pass output and RT reuses `targets.output`); `settings`
// and `resolve_pipeline` are `Some` only when SSR itself is on.
pub(crate) struct SsrState {
    pub settings: Option<SsrSettings>,
    pub targets: Option<SsrTargets>,
    pub resolve_pipeline: Option<Retained<ProtocolObject<dyn MTLRenderPipelineState>>>,
    // Roughness-aware blur + composite of the reflection target over the scene.
    // Shared by the SSR and RT-reflection resolves (both write the reflection
    // target, then run this). Built whenever the reflection targets exist.
    pub composite_pipeline: Option<Retained<ProtocolObject<dyn MTLRenderPipelineState>>>,
    // First half of the composite: the roughness blur, run at reduced resolution
    // into `SsrTargets::blur`. Built alongside `composite_pipeline`.
    pub blur_pipeline: Option<Retained<ProtocolObject<dyn MTLRenderPipelineState>>>,
    // Per-axis divisor the reflection blur target is sized by, resolved from the
    // world's `reflection_blur_resolution`. Held so a resize / live rebuild
    // recreates the blur target at the same reduced resolution.
    pub blur_scale: u32,
}

// Pipelines

// Build the SSR resolve pipeline: a fullscreen-triangle pass that ray-marches
// the reflection and composites it over the scene, writing a single-sample
// `RGBA16Float` target.
pub(crate) fn build_ssr_pipeline(
    device: &ProtocolObject<dyn objc2_metal::MTLDevice>,
    hot_reload: bool,
) -> Result<Retained<ProtocolObject<dyn MTLRenderPipelineState>>, String> {
    build_slang_fullscreen_pipeline(
        device,
        &SSR_RESOLVE,
        MTLPixelFormat::RGBA16Float,
        FullscreenBlend::Replace,
        hot_reload,
    )
}

// Build the reflection composite pipeline: the full-resolution second pass that
// lerps the sharp reflection against the upsampled half-res blur by roughness
// and composites it over the scene, writing the `RGBA16Float` scene output the
// SSR / RT resolve used to write directly. Shared by both reflection paths.
pub(crate) fn build_reflection_composite_pipeline(
    device: &ProtocolObject<dyn objc2_metal::MTLDevice>,
    hot_reload: bool,
) -> Result<Retained<ProtocolObject<dyn MTLRenderPipelineState>>, String> {
    build_slang_fullscreen_pipeline(
        device,
        &REFLECTION_COMPOSITE,
        MTLPixelFormat::RGBA16Float,
        FullscreenBlend::Replace,
        hot_reload,
    )
}

// Build the reflection blur pipeline: the reduced-resolution first pass that
// weight-averages the reflection target over the roughness cone into the blur
// target the composite then upsamples. The expensive multi-tap blur runs here.
pub(crate) fn build_reflection_blur_pipeline(
    device: &ProtocolObject<dyn objc2_metal::MTLDevice>,
    hot_reload: bool,
) -> Result<Retained<ProtocolObject<dyn MTLRenderPipelineState>>, String> {
    build_slang_fullscreen_pipeline(
        device,
        &REFLECTION_BLUR,
        MTLPixelFormat::RGBA16Float,
        FullscreenBlend::Replace,
        hot_reload,
    )
}

// Targets

// Off-screen target for the screen-space reflection (SSR) resolve pass: the HDR
// scene with reflections composited in. The view-space normal / linear depth /
// roughness the resolve reads now come from the unified G-buffer pre-pass
// (`metal/post/gbuffer.rs`); only this resolve output lives here. Single-sample,
// full drawable resolution; created when SSR is enabled and rebuilt with the HDR
// targets on resize.
pub(crate) struct SsrTargets {
    // Reflection target (`RGBA16Float`): the SSR / RT resolve writes reflected
    // radiance in `.rgb` and the Fresnel/gloss composite weight in `.a` here,
    // and the reflection composite blurs + composites it into `output`.
    pub reflection: Retained<ProtocolObject<dyn MTLTexture>>,
    // Scene with reflections composited in. Becomes the scene colour the TAA /
    // bloom / composite passes consume when SSR or RT reflections are on.
    pub output: Retained<ProtocolObject<dyn MTLTexture>>,
    // Reduced-resolution roughness blur of `reflection` (the blur pass writes it,
    // the composite pass upsamples it). Sized at render / REFLECTION_BLUR_SCALE.
    pub blur: Retained<ProtocolObject<dyn MTLTexture>>,
}

// Create or recreate the reflection + resolve-output targets at `width`x`height`,
// plus the reduced-resolution blur target. `blur_scale` is the per-axis
// render-resolution divisor for the roughness blur pass (resolved from the
// world's `reflection_blur_resolution`): the blur is low-frequency (a widening
// glossy cone), so running it reduced and bilinear-upsampling in the composite
// is visually free; mirrors stay sharp because the composite lerps in the
// FULL-RES reflection for low roughness (see reflection_composite.metal).
pub(crate) fn create_ssr_targets(
    device: &ProtocolObject<dyn objc2_metal::MTLDevice>,
    width: u32,
    height: u32,
    blur_scale: u32,
) -> Result<SsrTargets, String> {
    let blur_scale = blur_scale.max(1);
    let make_at = |w: usize, h: usize| -> Option<Retained<ProtocolObject<dyn MTLTexture>>> {
        let desc = TextureDesc {
            format: MTLPixelFormat::RGBA16Float,
            width: w,
            height: h,
            usage: MTLTextureUsage(MTLTextureUsage::ShaderRead.0 | MTLTextureUsage::RenderTarget.0),
            ..Default::default()
        }
        .build();
        device.newTextureWithDescriptor(&desc)
    };
    let w = width.max(1) as usize;
    let h = height.max(1) as usize;
    let bw = (width / blur_scale).max(1) as usize;
    let bh = (height / blur_scale).max(1) as usize;
    let reflection = make_at(w, h).ok_or("failed to create reflection texture")?;
    let output = make_at(w, h).ok_or("failed to create SSR output texture")?;
    let blur = make_at(bw, bh).ok_or("failed to create reflection blur texture")?;
    Ok(SsrTargets {
        reflection,
        output,
        blur,
    })
}

// Encoders

impl MtlContext {
    // Encode the SSR resolve: a fullscreen ray-march over the pre-pass
    // G-buffer that reflects `hdr_resolve` and composites the result into
    // `ssr_targets.output`. Rays that miss -- or fade out near a screen
    // border -- fall back to the IBL prefilter cubemap so the reflection
    // hands off to the environment rather than snapping to the base shading;
    // with no EnvironmentMap bound the cube is skipped and a miss keeps the
    // base shading. Runs after the main pass; only called when SSR is on.
    pub(in crate::metal) fn encode_ssr_resolve(
        &self,
        cmd_buf: &ProtocolObject<dyn objc2_metal::MTLCommandBuffer>,
        ssr_params: &crate::gfx::render_types::SsrParams,
    ) -> Result<u32, String> {
        // The pre-pass channels are pool-owned, so they are fetched here rather
        // than cached: a pool rebuild repacks every slot.
        let (targets, resolve_ps, gb_normal_depth, gb_roughness) = match (
            &self.ssr.targets,
            &self.ssr.resolve_pipeline,
            self.gbuffer_normal_depth(),
            self.gbuffer_roughness(),
        ) {
            (Some(t), Some(b), Some(n), Some(r)) => (t, b, n, r),
            _ => return Ok(0),
        };

        // Resolve: ray-march the reflection over `hdr_resolve` -> reflection
        // target (reflected radiance + composite weight, not yet blended).
        self.fullscreen_pass(
            cmd_buf,
            FullscreenPass {
                target: targets.reflection.as_ref(),
                load: MTLLoadAction::DontCare,
                timer: PassTimer::Whole(crate::metal::pass_timing::PassId::SsrResolve),
                pipeline: resolve_ps,
                label: "SSR resolve",
            },
            |enc| {
                enc.set_fragment_texture(self.hdr_targets.hdr_resolve.as_ref(), 0);
                enc.set_fragment_texture(gb_normal_depth, 1);
                enc.set_fragment_texture(gb_roughness, 2);
                // The IBL prefilter cubemap is the miss / screen-edge fallback.
                // It is always valid (a grey fallback when no EnvironmentMap is
                // bound); `SsrParams.prefilter_mip_count == 0` tells the shader to
                // ignore it in that case.
                enc.set_fragment_texture(self.env_map.prefilter.as_ref(), 3);
                // Local reflection-probe cubes, through their argument buffer:
                // when a probe is baked a missed/edge ray reflects its
                // box-projected scene capture instead of the foreign sky HDR
                // (the source the forward IBL specular term uses). The
                // ProbeSet's `count` gates whether the shader samples them.
                self.bind_probe_cubes(enc);
                // The screen sources take the post sampler at 0..2; the
                // prefilter cube at sampler(3) and the probe block's own
                // sampler at sampler(4) take the cube sampler.
                set_fragment_sampler_range(enc, &self.post_sampler, 0, 3);
                set_fragment_sampler_range(enc, self.cube_sampler.as_ref(), 3, 2);
                enc.set_fragment_value(ssr_params, 0);
                // Reflection-probe set (count + per-probe parallax boxes) at
                // buffer(1); count == 0 keeps the sky fallback above.
                enc.set_fragment_value(&self.probe.set, 1);
            },
        )?;
        // Blur by roughness + composite the reflection over the scene -> output.
        self.encode_reflection_composite(cmd_buf)?;
        Ok(0)
    }

    // Blur the reflection target by surface roughness and composite it over
    // `hdr_resolve` into `ssr_targets.output`. Shared by the SSR and
    // RT-reflection resolves: both write the reflection target first, then call
    // this. A no-op (leaves `output` untouched) when the composite pipeline or
    // G-buffer is absent, which only happens when no reflection path is active.
    pub(in crate::metal) fn encode_reflection_composite(
        &self,
        cmd_buf: &ProtocolObject<dyn objc2_metal::MTLCommandBuffer>,
    ) -> Result<(), String> {
        let (targets, composite_ps, blur_ps, gb_normal_depth, gb_roughness) = match (
            &self.ssr.targets,
            &self.ssr.composite_pipeline,
            &self.ssr.blur_pipeline,
            self.gbuffer_normal_depth(),
            self.gbuffer_roughness(),
        ) {
            (Some(t), Some(cp), Some(bp), Some(n), Some(r)) => (t, cp, bp, n, r),
            _ => return Ok(()),
        };
        // Pass 1: the roughness blur, at reduced resolution into `blur`. Times the
        // span start; the composite below times its end (both under one slot).
        self.fullscreen_pass(
            cmd_buf,
            FullscreenPass {
                target: targets.blur.as_ref(),
                load: MTLLoadAction::DontCare,
                timer: PassTimer::First(crate::metal::pass_timing::PassId::ReflectionComposite),
                pipeline: blur_ps,
                label: "reflection blur",
            },
            |enc| {
                enc.set_fragment_texture(targets.reflection.as_ref(), 0);
                enc.set_fragment_texture(gb_roughness, 1);
                set_fragment_sampler_range(enc, &self.post_sampler, 0, 2);
            },
        )?;
        // Pass 2: lerp the sharp full-res reflection against the upsampled blur by
        // roughness, then composite over the scene into `output`.
        self.fullscreen_pass(
            cmd_buf,
            FullscreenPass {
                target: targets.output.as_ref(),
                load: MTLLoadAction::DontCare,
                timer: PassTimer::Last(crate::metal::pass_timing::PassId::ReflectionComposite),
                pipeline: composite_ps,
                label: "reflection composite",
            },
            |enc| {
                enc.set_fragment_texture(targets.reflection.as_ref(), 0);
                enc.set_fragment_texture(self.hdr_targets.hdr_resolve.as_ref(), 1);
                enc.set_fragment_texture(gb_normal_depth, 2);
                enc.set_fragment_texture(gb_roughness, 3);
                enc.set_fragment_texture(targets.blur.as_ref(), 4);
                set_fragment_sampler_range(enc, &self.post_sampler, 0, 5);
            },
        )?;
        Ok(())
    }
}