concinnity-device 0.19.119

//! Hardware ray-traced reflection pass for the D3D12 backend. A fullscreen pixel
//! pass that, per glossy pixel, rebuilds a world-space surface point + normal
//! from the SSR pre-pass G-buffer, traces a reflection ray against the scene's
//! DXR top-level acceleration structure ([`crate::directx::raytrace`]) with
//! `RayQuery`, shades the hit (sun + IBL split-sum, optionally textured) or the
//! IBL prefilter cube on a miss, and composites the result over the scene with
//! the same Fresnel/gloss weighting SSR uses.
//!
//! It occupies the `SsrResolve` slot in the frame graph (reads `hdr_resolve`,
//! writes its own output target) and is mutually exclusive with SSR resolve.
//! Like SSGI it relies on the SSR pre-pass G-buffer, so the pre-pass is forced on
//! whenever RT reflections are enabled. The shader is the shared
//! `shaders/rt_reflections.hlsl`, compiled to a shader-model 6.5 DXIL container
//! (the floor for the inline ray query) through `builtin_shaders`; the Metal and
//! Vulkan hosts bind the same source at their own slots.

use concinnity_core::gfx::render_types::RtParams;
use concinnity_core::render::error::RenderResult;
use concinnity_core::render::post::rt_reflections::{RtParamsInputs, RtReflectionSettings};
use windows::Win32::Foundation::RECT;
use windows::Win32::Graphics::Direct3D12::*;

use crate::directx::allocator::{DeviceAllocator, PooledBuffer};
use crate::directx::builtin_shaders;
use crate::directx::builtin_shaders::CompileProgram;
use crate::directx::com;
use crate::directx::context::{DxContext, FRAMES, align256, dump_on_err};
use crate::directx::descriptor_slot::DescriptorTables;
use crate::directx::descriptor_slot::SrvSlot;
use crate::directx::error::map_hresult;
use crate::directx::pso::{Blend, GraphicsPso};
use crate::directx::root_sig::{Range, RootSig, SamplerState, Visibility};
use crate::directx::texture::{
    HDR_FORMAT, create_rt_target, transition_barrier, write_format_rtv, write_format_srv,
};

// Size of the RT-reflection fragment-shader uniform block. 144 bytes; see
// `gfx::render_types::RtParams`.
const RT_PARAMS_UBO_SIZE: u64 = 144;

// Shader compilation

struct RtShaders {
    vs: Vec<u8>,
    flat_ps: Vec<u8>,
    textured_ps: Vec<u8>,
}

// Compile both fragment entry points from the single source (SM 6.5, for the
// inline ray query) and pair them with the shared fullscreen vertex stage the
// fragment's own toolchain produced, which already carries the one divergence
// the pass's own vertex shader existed for, the DirectX/Metal UV flip. Returns
// an `Err` (which the caller turns into an SSR fallback) when the toolchain is
// unavailable or the shader fails to compile.
fn compile_rt_shaders(hot_reload: bool) -> RenderResult<RtShaders> {
    Ok(RtShaders {
        vs: builtin_shaders::FULLSCREEN_VERT.compile(hot_reload)?,
        flat_ps: builtin_shaders::RT_REFLECTIONS_FRAG.compile(hot_reload)?,
        textured_ps: builtin_shaders::RT_REFLECTIONS_FRAG_TEXTURED.compile(hot_reload)?,
    })
}

// Root signature
//
// Root CBV b0 (RtParams); four root SRVs t0..t3 (TLAS / vertex / index / geometry
// table - all raw or structured buffers bound by GPU virtual address, which inline
// ray tracing supports for the acceleration structure too); five descriptor tables
// for the textures: scene t4, gbuffer normal+depth t5, roughness t6, prefilter cube
// t7, and the unbounded bindless pool at (t0, space1); and two more root SRVs t8/t9
// (deformed skinned verts / skinned indices, for skinned hits). The miss
// fallback's reflection probes follow: the cube array table t10, the ProbeSet
// CBV b4, the records t11, and the main camera's cluster grid binning them (the
// params CBV b5, the per-cluster lists t12). Four static samplers: linear-clamp
// s0, cube linear-clamp s1, linear-repeat s2, and the probe cube sampler s3.
fn create_rt_root_signature(device: &ID3D12Device) -> RenderResult<ID3D12RootSignature> {
    use Visibility::Pixel;
    RootSig::new()
        .cbv(0, Pixel) // [0] b0 RtParams
        .srv(0, Pixel) // [1] t0 TLAS
        .srv(1, Pixel) // [2] t1 vertex buffer (raw)
        .srv(2, Pixel) // [3] t2 index buffer (raw)
        .srv(3, Pixel) // [4] t3 geometry table (structured)
        .srv_table(4, 1, Pixel) // [5] t4 scene
        .srv_table(5, 1, Pixel) // [6] t5 gbuffer normal+depth
        .srv_table(6, 1, Pixel) // [7] t6 roughness
        .srv_table(7, 1, Pixel) // [8] t7 prefilter cube
        .table(&[Range::bindless_srv(1)], Pixel) // [9] t0,space1 bindless pool
        .srv(8, Pixel) // [10] t8 deformed skinned verts (raw)
        .srv(9, Pixel) // [11] t9 skinned indices (raw)
        // [12] t10 reflection-probe cube array, clear of the trace's own SRVs;
        // the miss fallback box-projects it when ProbeSet.count > 0.
        .srv_table(10, 1, Pixel)
        .cbv(4, Pixel) // [13] b4 ProbeSet
        .srv(11, Pixel) // [14] t11 reflection-probe records
        .cbv(5, Pixel) // [15] b5 ClusterParams
        .srv(12, Pixel) // [16] t12 cluster lists
        .static_sampler(SamplerState::LinearClamp, 0, Pixel)
        .static_sampler(SamplerState::LinearClamp, 1, Pixel)
        .static_sampler(SamplerState::LinearWrap, 2, Pixel)
        // s3: cube mip-linear clamp for the reflection-probe cube array, matching
        // the s1 prefilter sampler. The array rides split from its sampler here
        // because D3D12 binds a shader sampler array only through a descriptor
        // table.
        .static_sampler(SamplerState::LinearClamp, 3, Pixel)
        .build(device, "rt reflections root sig")
}

// The flat and textured PSOs over the shared root signature.
fn create_rt_psos(
    device: &ID3D12Device,
    info_queue: Option<&ID3D12InfoQueue>,
    root_sig: &ID3D12RootSignature,
    shaders: &RtShaders,
) -> RenderResult<RebuiltRtPipelines> {
    let pso = |ps: &[u8], label: &str| {
        dump_on_err(
            info_queue,
            GraphicsPso::fullscreen(root_sig, &shaders.vs, ps, HDR_FORMAT, Blend::Opaque)
                .build(device, label),
        )
    };
    Ok(RebuiltRtPipelines {
        flat_pso: pso(&shaders.flat_ps, "rt reflections flat")?,
        textured_pso: pso(&shaders.textured_ps, "rt reflections textured")?,
    })
}

// Resources

// Hardware-ray-traced-reflection resources held by `DxContext` when the world's
// `PostProcessConfig` enables `ray_traced_reflections` AND the GPU supports the
// DXR tier AND the DXC compile + acceleration-structure build succeed; otherwise
// the context leaves this `None` and the graph falls back to `SsrResolve`.
pub(in crate::directx) struct RtReflectionsResources {
    // Resolved authored tunables; turned into a per-frame `RtParams` push.
    pub(in crate::directx) settings: RtReflectionSettings,

    // Reflection output: reflected radiance + composite weight at the trace
    // resolution, which the reflection composite upsamples over the scene (it
    // owns its own slot rather than reusing the optional SSR resolve output,
    // because RT can be authored with SSR resolve off).
    pub(in crate::directx) output: ID3D12Resource,
    // The trace resolution: render resolution reduced by `settings.divisor`.
    extent: (u32, u32),
    output_rtv: D3D12_CPU_DESCRIPTOR_HANDLE,
    pub(in crate::directx) output_srv_gpu: SrvSlot,

    // Per-frame `RtParams` UBO (144-byte), persistently mapped.
    params_ubo_resources: Vec<PooledBuffer>,
    params_ubo_ptrs: Vec<*mut u8>,

    // One root signature shared by both PSOs; flat-tint + textured (bindless).
    root_sig: ID3D12RootSignature,
    flat_pso: ID3D12PipelineState,
    textured_pso: ID3D12PipelineState,
}

// Device + render extent for building the RT-reflection resources; the output
// target is that extent reduced to the trace resolution.
pub(in crate::directx) struct RtBuildContext<'a> {
    pub alloc: &'a DeviceAllocator,
    pub width: u32,
    pub height: u32,
}

// The RT-reflection output target's descriptor handles: a CPU RTV plus the
// (CPU, GPU) SRV pair downstream passes sample.
#[derive(Clone, Copy)]
pub(in crate::directx) struct RtOutputDescriptors {
    pub output_rtv: D3D12_CPU_DESCRIPTOR_HANDLE,
    pub output_srv: (D3D12_CPU_DESCRIPTOR_HANDLE, SrvSlot),
}

// Optional debug info queue plus the shader hot-reload toggle for the build.
pub(in crate::directx) struct RtBuildInit<'a> {
    pub info_queue: Option<&'a ID3D12InfoQueue>,
    pub hot_reload: bool,
}

impl RtReflectionsResources {
    // Build the RT-reflection resources. Returns `Err` when the DXC compile
    // fails (DXC absent / shader error) so the caller can fall back to SSR; the
    // accel-structure build + DXR capability are gated separately by the caller.
    pub(in crate::directx) fn new(
        ctx: RtBuildContext,
        settings: RtReflectionSettings,
        descriptors: RtOutputDescriptors,
        init: RtBuildInit,
    ) -> RenderResult<Self> {
        let RtBuildContext {
            alloc,
            width,
            height,
        } = ctx;
        let device = alloc.device();
        let RtOutputDescriptors {
            output_rtv,
            output_srv,
        } = descriptors;
        let RtBuildInit {
            info_queue,
            hot_reload,
        } = init;
        let extent = settings.trace_extent(width, height);
        let output = create_rt_target(device, extent.0, extent.1, HDR_FORMAT)?;
        write_format_rtv(device, &output, output_rtv, HDR_FORMAT);
        write_format_srv(device, &output, output_srv.0, HDR_FORMAT);

        let params_size = align256(RT_PARAMS_UBO_SIZE);
        let mut params_ubo_resources: Vec<PooledBuffer> = Vec::with_capacity(FRAMES);
        let mut params_ubo_ptrs: Vec<*mut u8> = Vec::with_capacity(FRAMES);
        for _ in 0..FRAMES {
            let buf = alloc.alloc_buffer(
                params_size,
                D3D12_HEAP_TYPE_UPLOAD,
                D3D12_RESOURCE_STATE_GENERIC_READ,
            )?;
            let mut ptr = std::ptr::null_mut::<std::ffi::c_void>();
            // SAFETY: the resource is a live CPU-visible buffer, and the out-parameter is a live
            // local that receives the mapping.
            unsafe { buf.Map(0, None, Some(&mut ptr)) }
                .map_err(|e| map_hresult(e.code(), "map rt params ubo"))?;
            params_ubo_ptrs.push(ptr as *mut u8);
            params_ubo_resources.push(buf);
        }

        let shaders = compile_rt_shaders(hot_reload)?;
        let root_sig = dump_on_err(info_queue, create_rt_root_signature(device))?;
        let RebuiltRtPipelines {
            flat_pso,
            textured_pso,
        } = create_rt_psos(device, info_queue, &root_sig, &shaders)?;

        Ok(Self {
            settings,
            output,
            extent,
            output_rtv,
            output_srv_gpu: output_srv.1,
            params_ubo_resources,
            params_ubo_ptrs,
            root_sig,
            flat_pso,
            textured_pso,
        })
    }

    // Rebuild the output target for a new render resolution, at the trace
    // resolution `settings` names. The descriptor *slot* stays put; only the
    // backing resource changes, so the composite's binding (which points at the
    // SRV slot's GPU handle) stays valid.
    pub(in crate::directx) fn resize_to(
        &mut self,
        device: &ID3D12Device,
        width: u32,
        height: u32,
        srv_cpu_base: D3D12_CPU_DESCRIPTOR_HANDLE,
        srv_gpu_base: SrvSlot,
    ) -> RenderResult<()> {
        let srv_cpu = self.output_srv_gpu.cpu_in(srv_cpu_base, srv_gpu_base);
        self.extent = self.settings.trace_extent(width, height);
        self.output = create_rt_target(device, self.extent.0, self.extent.1, HDR_FORMAT)?;
        write_format_rtv(device, &self.output, self.output_rtv, HDR_FORMAT);
        write_format_srv(device, &self.output, srv_cpu, HDR_FORMAT);
        Ok(())
    }
}

// Replacement RT PSOs returned by [`rebuild_rt_reflections_pipelines`]. The
// caller swaps them into the live resources only if both builds succeeded.
pub(in crate::directx) struct RebuiltRtPipelines {
    pub flat_pso: ID3D12PipelineState,
    pub textured_pso: ID3D12PipelineState,
}

// Rebuild both RT PSOs against fresh shader source, reusing the existing root
// signature. Returns the new PSOs for the caller to swap into the live
// `RtReflectionsResources`.
pub(in crate::directx) fn rebuild_rt_reflections_pipelines(
    device: &ID3D12Device,
    rt: &RtReflectionsResources,
    hot_reload: bool,
    info_queue: Option<&ID3D12InfoQueue>,
) -> RenderResult<RebuiltRtPipelines> {
    create_rt_psos(
        device,
        info_queue,
        &rt.root_sig,
        &compile_rt_shaders(hot_reload)?,
    )
}

// Swap freshly compiled RT PSOs into the live resources after a hot-reload.
pub(in crate::directx) fn swap_rt_reflections_pipelines(
    rt: &mut RtReflectionsResources,
    rebuilt: RebuiltRtPipelines,
) {
    rt.flat_pso = rebuilt.flat_pso;
    rt.textured_pso = rebuilt.textured_pso;
}

// Encoder

impl DxContext {
    // Encode the RT-reflection resolve: a fullscreen triangle that traces each
    // glossy pixel's reflection ray against the scene TLAS and composites the
    // reflected radiance + weight into `rt_reflections.output`, which the
    // reflection composite then upsamples over the scene.
    // No-op when any required resource is missing (the graph only schedules this
    // pass when RT is live, so the guards are defensive).
    pub(in crate::directx) fn encode_rt_reflections(
        &self,
        cmd: &ID3D12GraphicsCommandList,
        frame_idx: usize,
        fov_y_radians: f32,
        aspect: f32,
        cam_pos: [f32; 3],
    ) {
        let (rt, accel, gbuffer) = match (&self.rt_reflections, &self.rt.accel, &self.gbuffer) {
            (Some(r), Some(a), Some(g)) => (r, a, g),
            _ => return,
        };
        // The trace writes reflected radiance + weight into `rt.output`; the
        // reflection composite (below) blurs + composites it over the scene.
        let reflection_srv = rt.output_srv_gpu;

        // Build + upload this frame's RtParams. `inv_view_rot` is the view->world
        // rotation (the transpose of the view matrix's orthonormal 3x3), same as
        // the SSR resolve; `params` then fills in the camera-position translation
        // column to complete the camera-to-world transform.
        let v = self.state.view.matrix;
        let inv_view_rot = [
            [v[0][0], v[1][0], v[2][0], 0.0],
            [v[0][1], v[1][1], v[2][1], 0.0],
            [v[0][2], v[1][2], v[2][2], 0.0],
            [0.0, 0.0, 0.0, 1.0],
        ];
        let params = rt.settings.params(RtParamsInputs {
            fov_y_radians,
            aspect,
            inv_view_rot,
            cam_pos,
            sun_dir: self.fog.sun_dir,
            sun_color: self.fog.sun_color,
            prefilter_mip_count: self.scene.env_map.prefilter_mip_count as f32,
            sky_rot: self.state.view.sky_rot,
        });
        // SAFETY: the destination is the persistent mapping of an UPLOAD-heap constant buffer that
        // init sized for this payload, and the source is a separate live value, so the ranges
        // cannot overlap.
        unsafe {
            std::ptr::copy_nonoverlapping(
                &params as *const RtParams as *const u8,
                rt.params_ubo_ptrs[frame_idx],
                std::mem::size_of::<RtParams>(),
            );
        }
        let params_gva = com::gpu_va(&rt.params_ubo_resources[frame_idx]);

        // Textured hit shading needs the bindless albedo/normal pool, which only
        // the GPU-cull bindless path populates; otherwise fall back to the
        // flat-tint variant (mirrors Metal's bindless-arg-buffer gate).
        let textured = self.cull.main_bindless_pso.is_some();
        let pso = if textured {
            &rt.textured_pso
        } else {
            &rt.flat_pso
        };

        let out_to_rt = transition_barrier(
            &rt.output,
            D3D12_RESOURCE_STATE_PIXEL_SHADER_RESOURCE,
            D3D12_RESOURCE_STATE_RENDER_TARGET,
        );
        // SAFETY: the command list is in the recording state, and every resource, descriptor and
        // slice these commands name is live for the call.
        unsafe { cmd.ResourceBarrier(&[out_to_rt]) };

        let (w, h) = rt.extent;
        // SAFETY: the command list is in the recording state, and every resource, descriptor and
        // slice these commands name is live for the call.
        unsafe {
            cmd.OMSetRenderTargets(1, Some(&rt.output_rtv), false, None);
            let vp = D3D12_VIEWPORT {
                TopLeftX: 0.0,
                TopLeftY: 0.0,
                Width: w as f32,
                Height: h as f32,
                MinDepth: 0.0,
                MaxDepth: 1.0,
            };
            cmd.RSSetViewports(&[vp]);
            let scissor = RECT {
                left: 0,
                top: 0,
                right: w as i32,
                bottom: h as i32,
            };
            cmd.RSSetScissorRects(&[scissor]);

            cmd.SetPipelineState(pso);
            cmd.SetGraphicsRootSignature(&rt.root_sig);
            cmd.SetDescriptorHeaps(&[Some(self.descriptors.srv_heap.clone())]);
            cmd.SetGraphicsRootConstantBufferView(0, params_gva);
            // Root SRVs: TLAS / vertex / index / geometry table (by GPU virtual
            // address; inline ray tracing reads the TLAS through a root SRV).
            cmd.SetGraphicsRootShaderResourceView(1, accel.tlas_gva());
            cmd.SetGraphicsRootShaderResourceView(
                2,
                com::gpu_va(&self.scene.geometry.vertex_buffer),
            );
            cmd.SetGraphicsRootShaderResourceView(
                3,
                com::gpu_va(&self.scene.geometry.index_buffer),
            );
            cmd.SetGraphicsRootShaderResourceView(4, accel.geom_table_gva());
            // Texture tables.
            cmd.set_graphics_srv_table(5, self.targets.hdr.srv_gpu);
            cmd.set_graphics_srv_table(6, gbuffer.normal_depth_srv_gpu);
            cmd.set_graphics_srv_table(7, gbuffer.roughness_srv_gpu);
            cmd.set_graphics_srv_table(8, self.prefilter_cube_srv_gpu());
            if textured {
                cmd.set_graphics_srv_table(9, self.cull.bindless_pool_gpu[self.current_frame]);
            }
            // Skinned-geometry root SRVs: the deformed (posed) vertex buffer +
            // the skinned index buffer the trace fetches a skinned hit from.
            // Both return a valid 1-element dummy GVA when there is no skinned
            // geometry, so the binding is always live.
            cmd.SetGraphicsRootShaderResourceView(10, accel.deformed_verts_gva());
            cmd.SetGraphicsRootShaderResourceView(11, accel.skinned_index_gva());
            // Reflection-probe miss fallback: the cube array table at t10 + the
            // per-frame ProbeSet CBV at b4 + the per-frame records at t11, and
            // the main camera's cluster grid binning them (b5, t12). count == 0
            // keeps the sky path.
            self.probe_bindings(frame_idx).bind(cmd, 12);
            cmd.IASetPrimitiveTopology(
                windows::Win32::Graphics::Direct3D::D3D_PRIMITIVE_TOPOLOGY_TRIANGLELIST,
            );
            cmd.IASetVertexBuffers(0, None);
            cmd.IASetIndexBuffer(None);
            cmd.DrawInstanced(3, 1, 0, 0);
        }

        let out_to_psr = transition_barrier(
            &rt.output,
            D3D12_RESOURCE_STATE_RENDER_TARGET,
            D3D12_RESOURCE_STATE_PIXEL_SHADER_RESOURCE,
        );
        // SAFETY: the command list is in the recording state, and every resource, descriptor and
        // slice these commands name is live for the call.
        unsafe { cmd.ResourceBarrier(&[out_to_psr]) };

        // Blur the reflection by roughness and composite it over the scene into the
        // reflection-composite output (the scene the post stack then consumes).
        self.encode_reflection_composite(cmd, frame_idx, reflection_srv);
    }
}