concinnity-device 0.18.66

// src/metal/init/pipelines.rs
//
// Core render-pipeline construction extracted from MtlContext::new:
//   * The shared vertex descriptor (interleaved [pos, normal, tangent, color, uv]).
//   * The main static pipeline (with optional bindless fragment + GPU-driven
//     cull pipeline + bindless texture argument encoder).
//   * The optional instanced pipeline.
//   * The shared depth-stencil state used by main + shadow passes.
#![deny(unsafe_op_in_unsafe_fn)]

use objc2::rc::Retained;
use objc2::runtime::ProtocolObject;
use objc2_metal::{
    MTLArgumentEncoder, MTLCompareFunction, MTLComputePipelineState, MTLDepthStencilDescriptor,
    MTLDepthStencilState, MTLDevice, MTLFunction as _, MTLLibrary as _, MTLPixelFormat,
    MTLRenderPipelineDescriptor, MTLRenderPipelineState, MTLVertexDescriptor, MTLVertexFormat,
    MTLVertexStepFunction,
};

use crate::gfx::mesh_payload::Vertex;
use crate::metal::context::{
    BINDLESS_SAMPLER_ARG_BUFFER_INDEX, BINDLESS_TEXTURE_ARG_BUFFER_INDEX, HDR_SAMPLE_COUNT,
};
use crate::metal::cull::build_cull_pipeline;
use crate::metal::descriptors::{VertexAttr, VertexLayout, vertex_descriptor};
use crate::metal::pipeline::{load_library, ns_str, stage_library};

pub(crate) struct MainPipelineBundle {
    pub pipeline_state: Retained<ProtocolObject<dyn MTLRenderPipelineState>>,
    pub bindless: bool,
    pub cull_pipeline: Option<Retained<ProtocolObject<dyn MTLComputePipelineState>>>,
    pub cull_icb_arg_encoder: Option<Retained<ProtocolObject<dyn MTLArgumentEncoder>>>,
    // Phase-2 cull pipeline + its ICB argument encoder for two-pass
    // occlusion. Built alongside the phase-1 cull pipeline whenever the
    // bindless path is active (cheap: one extra compute pipeline from the
    // same library); only used when `occlusion_two_pass` is on at runtime.
    pub cull_pipeline_phase2: Option<Retained<ProtocolObject<dyn MTLComputePipelineState>>>,
    pub cull_icb2_arg_encoder: Option<Retained<ProtocolObject<dyn MTLArgumentEncoder>>>,
    pub bindless_tex_arg_encoder: Option<Retained<ProtocolObject<dyn MTLArgumentEncoder>>>,
    // Encoder for the engine sampler block at buffer(10). `Some` only for the
    // engine's single-source program: world-authored bindless fragments keep
    // their own inline samplers and declare no sampler block.
    pub bindless_sampler_arg_encoder: Option<Retained<ProtocolObject<dyn MTLArgumentEncoder>>>,
}

// Describes the per-vertex buffer layout so Metal can map [[stage_in]]:
//   buffer(1): interleaved [float3 pos, float3 normal, float3 tangent, float3 color, float2 uv]
//   stride = sizeof(Vertex) = 56 bytes
pub(crate) fn make_vertex_descriptor() -> Retained<MTLVertexDescriptor> {
    vertex_descriptor(
        &[
            VertexAttr {
                index: 0,
                format: MTLVertexFormat::Float3,
                offset: 0,
                buffer_index: 1,
            },
            VertexAttr {
                index: 1,
                format: MTLVertexFormat::Float3,
                offset: 12,
                buffer_index: 1,
            },
            VertexAttr {
                index: 2,
                format: MTLVertexFormat::Float3,
                offset: 24,
                buffer_index: 1,
            },
            VertexAttr {
                index: 3,
                format: MTLVertexFormat::Float3,
                offset: 36,
                buffer_index: 1,
            },
            VertexAttr {
                index: 4,
                format: MTLVertexFormat::Float2,
                offset: 48,
                buffer_index: 1,
            },
        ],
        &[VertexLayout {
            buffer_index: 1,
            stride: std::mem::size_of::<Vertex>(),
            step: MTLVertexStepFunction::PerVertex,
        }],
    )
}

// Build the main static pipeline together with everything it implies:
//
//   * If the fragment library exposes `fragment_main_bindless`, the static main
//     pass is GPU-driven. That requires:
//       - the pipeline to opt into indirect command buffers
//       - a compute cull pipeline + ICB argument encoder
//       - an argument encoder for the BindlessTextures argument buffer
//   * Otherwise the pipeline pairs with the per-draw `fragment_main` and the
//     three Optional fields are `None`.
pub(crate) fn build_main_pipeline(
    device: &ProtocolObject<dyn MTLDevice>,
    vert_desc: &MTLVertexDescriptor,
    vert_lib_bytes: &[u8],
    frag_lib_bytes: &[u8],
    hot_reload: bool,
) -> Result<MainPipelineBundle, String> {
    // Fully engine-supplied stages ship from the single-source Slang program,
    // which is always bindless (the engine has no per-draw static path of its
    // own any more). A world-authored stage keeps the legacy detection: a
    // fragment exposing `fragment_main_bindless` picks the GPU-driven path,
    // anything else the per-draw one. A world vertex paired with the engine
    // fragment stays on the per-draw pairing: the Slang fragment's varying
    // names cannot interface-match a hand-written vertex stage.
    let engine_stages = vert_lib_bytes.is_empty() && frag_lib_bytes.is_empty();

    let (vert_fn, main_frag_fn, bindless, engine_single_source) = if engine_stages {
        let vert_library = super::super::slang_shaders::MAIN_BINDLESS_VERT
            .library(device, hot_reload)
            .map_err(|e| format!("failed to load engine vertex library: {e}"))?;
        let frag_library = super::super::slang_shaders::MAIN_BINDLESS_FRAG
            .library(device, hot_reload)
            .map_err(|e| format!("failed to load engine fragment library: {e}"))?;
        let vert_fn = vert_library
            .newFunctionWithName(&ns_str("vertex_main_bindless"))
            .ok_or("vertex_main_bindless not found in engine library")?;
        let frag_fn = frag_library
            .newFunctionWithName(&ns_str("fragment_main_bindless"))
            .ok_or("fragment_main_bindless not found in engine library")?;
        (vert_fn, frag_fn, true, true)
    } else {
        let vert_library = stage_library(device, hot_reload, vert_lib_bytes)
            .map_err(|e| format!("failed to load vertex metallib: {}", e))?;
        let frag_library = stage_library(device, hot_reload, frag_lib_bytes)
            .map_err(|e| format!("failed to load fragment metallib: {}", e))?;
        let vert_fn = vert_library
            .newFunctionWithName(&ns_str("vertex_main"))
            .ok_or("vertex_main not found in metallib")?;
        let bindless_frag_fn = frag_library.newFunctionWithName(&ns_str("fragment_main_bindless"));
        let bindless = bindless_frag_fn.is_some();
        let frag_fn = match bindless_frag_fn {
            Some(f) => f,
            None => frag_library
                .newFunctionWithName(&ns_str("fragment_main"))
                .ok_or("fragment_main not found in metallib")?,
        };
        (vert_fn, frag_fn, bindless, false)
    };

    let pipeline_desc = MTLRenderPipelineDescriptor::new();
    pipeline_desc.setVertexDescriptor(Some(vert_desc));
    pipeline_desc.setVertexFunction(Some(&vert_fn));
    pipeline_desc.setFragmentFunction(Some(&main_frag_fn));
    // Off-screen HDR pass: RGBA16Float colour + 4x MSAA. Output is linear
    // light; ACES tonemap + gamma + FXAA run in the composite pass.
    pipeline_desc.setRasterSampleCount(HDR_SAMPLE_COUNT as usize);
    // SAFETY: plain descriptor property setters; the subscripted slots are ones this descriptor
    // declares.
    unsafe {
        pipeline_desc
            .colorAttachments()
            .objectAtIndexedSubscript(0)
            .setPixelFormat(MTLPixelFormat::RGBA16Float);
    }
    pipeline_desc.setDepthAttachmentPixelFormat(MTLPixelFormat::Depth32Float);
    if bindless {
        pipeline_desc.setSupportIndirectCommandBuffers(true);
    }

    let pipeline_state = device
        .newRenderPipelineStateWithDescriptor_error(&pipeline_desc)
        .map_err(|e| format!("failed to create pipeline state: {:?}", e))?;

    let (cull_pipeline, cull_icb_arg_encoder, cull_pipeline_phase2, cull_icb2_arg_encoder) =
        if bindless {
            let cull = build_cull_pipeline(device, hot_reload)?;
            (
                Some(cull.state),
                Some(cull.icb_arg_encoder),
                Some(cull.state_phase2),
                Some(cull.icb2_arg_encoder),
            )
        } else {
            (None, None, None, None)
        };

    // Argument encoder for the bindless pass's `BindlessTextures` buffer.
    // Derived from `fragment_main_bindless`'s buffer(7) parameter.
    let bindless_tex_arg_encoder = if bindless {
        // SAFETY: BINDLESS_TEXTURE_ARG_BUFFER_INDEX is the static buffer
        // index `fragment_main_bindless` declares its argument buffer at.
        Some(unsafe {
            main_frag_fn.newArgumentEncoderWithBufferIndex(BINDLESS_TEXTURE_ARG_BUFFER_INDEX)
        })
    } else {
        None
    };

    // The engine's single-source fragment reaches its samplers through the
    // block at buffer(10) (indirect-command execution cannot see encoder-bound
    // sampler state); world-authored fragments keep inline samplers instead.
    let bindless_sampler_arg_encoder = if engine_single_source {
        // SAFETY: BINDLESS_SAMPLER_ARG_BUFFER_INDEX is the static buffer index
        // the engine fragment declares its sampler block at (locked by the
        // build script's ABI assertion).
        Some(unsafe {
            main_frag_fn.newArgumentEncoderWithBufferIndex(BINDLESS_SAMPLER_ARG_BUFFER_INDEX)
        })
    } else {
        None
    };

    Ok(MainPipelineBundle {
        pipeline_state,
        bindless,
        cull_pipeline,
        cull_icb_arg_encoder,
        cull_pipeline_phase2,
        cull_icb2_arg_encoder,
        bindless_tex_arg_encoder,
        bindless_sampler_arg_encoder,
    })
}

// Write the engine sampler block once: the pool sampler (trilinear +
// anisotropic + repeat, the same object the legacy path binds), the shadow
// compare sampler, and the cube sampler. Member order mirrors EngineSamplers
// in `src/shaders/main_bindless.slang`. Samplers never stream, so unlike the
// texture argument buffer this is written a single time at init.
pub(crate) fn build_bindless_sampler_args(
    device: &ProtocolObject<dyn MTLDevice>,
    encoder: &ProtocolObject<dyn MTLArgumentEncoder>,
    tex_sampler: &ProtocolObject<dyn objc2_metal::MTLSamplerState>,
    shadow_sampler: &ProtocolObject<dyn objc2_metal::MTLSamplerState>,
    cube_sampler: &ProtocolObject<dyn objc2_metal::MTLSamplerState>,
) -> Result<Retained<ProtocolObject<dyn objc2_metal::MTLBuffer>>, String> {
    use objc2_metal::MTLResourceOptions;
    let len = encoder.encodedLength().max(16);
    let buf = device
        .newBufferWithLength_options(len, MTLResourceOptions::StorageModeShared)
        .ok_or("failed to allocate sampler argument buffer")?;
    // SAFETY: `buf` was sized to the encoder's `encodedLength()`, and the
    // indices 0..2 are the EngineSamplers member ids in declaration order.
    unsafe {
        encoder.setArgumentBuffer_offset(Some(&buf), 0);
        encoder.setSamplerState_atIndex(Some(tex_sampler), 0);
        encoder.setSamplerState_atIndex(Some(shadow_sampler), 1);
        encoder.setSamplerState_atIndex(Some(cube_sampler), 2);
    }
    Ok(buf)
}

// The main-pass pipelines of the material-referenced world shaders, indexed by
// `shader_bucket - 1`. `None` marks a bucket whose Shader is not resident.
pub(crate) type WorldPipelineTable =
    Vec<Option<Retained<ProtocolObject<dyn MTLRenderPipelineState>>>>;

// Pipelines for the material-referenced world shaders past the default
// (ShaderHandle 1..), in bucket order. Extra world shaders render only through
// the GPU-driven bindless path (the cull kernel routes their draws into
// per-bucket ICBs), so each fragment library must expose
// `fragment_main_bindless` with the engine's BindlessTextures layout; a shader
// without it is a hard error rather than a silently default-shaded bucket.
//
// A bucket flagged `deferred` (its Shader is owned by a scene that has not
// pinned) stays `None` until
// [`super::super::MtlContext::install_world_shader`] builds it.
pub(crate) fn build_world_pipeline_table(
    device: &ProtocolObject<dyn MTLDevice>,
    vert_desc: &MTLVertexDescriptor,
    extra_shaders: &[crate::gfx::backend_init::ShaderBytes<'_>],
) -> Result<WorldPipelineTable, String> {
    let mut table = Vec::with_capacity(extra_shaders.len());
    for (i, shader) in extra_shaders.iter().enumerate() {
        // A bucket whose Shader a non-start scene owns has no payload yet; the
        // streaming pump installs it when that scene pins.
        if shader.deferred {
            table.push(None);
            continue;
        }
        table.push(Some(build_bucket_pipeline(
            device,
            vert_desc,
            i + 1,
            shader.vert,
            shader.frag,
        )?));
    }
    Ok(table)
}

// One material-referenced shader bucket's bindless main-pass pipeline.
pub(crate) fn build_bucket_pipeline(
    device: &ProtocolObject<dyn MTLDevice>,
    vert_desc: &MTLVertexDescriptor,
    bucket: usize,
    vert_bytes: &[u8],
    frag_bytes: &[u8],
) -> Result<Retained<ProtocolObject<dyn MTLRenderPipelineState>>, String> {
    let vert_library = load_library(device, vert_bytes)
        .map_err(|e| format!("shader bucket {bucket}: failed to load vertex metallib: {e}"))?;
    let frag_library = load_library(device, frag_bytes)
        .map_err(|e| format!("shader bucket {bucket}: failed to load fragment metallib: {e}"))?;
    let vert_fn = vert_library
        .newFunctionWithName(&ns_str("vertex_main"))
        .ok_or_else(|| format!("shader bucket {bucket}: vertex_main not found in metallib"))?;
    let frag_fn = frag_library
        .newFunctionWithName(&ns_str("fragment_main_bindless"))
        .ok_or_else(|| {
            format!(
                "shader bucket {bucket}: fragment_main_bindless not found in metallib -- a \
                 material-referenced Shader must define the bindless entry points"
            )
        })?;

    let desc = MTLRenderPipelineDescriptor::new();
    desc.setVertexDescriptor(Some(vert_desc));
    desc.setVertexFunction(Some(&vert_fn));
    desc.setFragmentFunction(Some(&frag_fn));
    desc.setRasterSampleCount(HDR_SAMPLE_COUNT as usize);
    // SAFETY: plain descriptor property setters; the subscripted slots are ones this descriptor
    // declares.
    unsafe {
        desc.colorAttachments()
            .objectAtIndexedSubscript(0)
            .setPixelFormat(MTLPixelFormat::RGBA16Float);
    }
    desc.setDepthAttachmentPixelFormat(MTLPixelFormat::Depth32Float);
    desc.setSupportIndirectCommandBuffers(true);

    device
        .newRenderPipelineStateWithDescriptor_error(&desc)
        .map_err(|e| format!("shader bucket {bucket}: failed to create pipeline: {e:?}"))
}

// Optional instanced pipeline: pairs vertex_main_instanced with the existing
// fragment_main. Built only when both an instanced vertex shader payload is
// supplied AND at least one cluster needs to render.
pub(crate) fn build_instanced_pipeline(
    device: &ProtocolObject<dyn MTLDevice>,
    vert_desc: &MTLVertexDescriptor,
    vert_instanced_lib_bytes: &[u8],
    frag_lib_bytes: &[u8],
    has_clusters: bool,
    hot_reload: bool,
) -> Result<Option<Retained<ProtocolObject<dyn MTLRenderPipelineState>>>, String> {
    if !has_clusters {
        return Ok(None);
    }

    let inst_library = stage_library(device, hot_reload, vert_instanced_lib_bytes)
        .map_err(|e| format!("failed to load instanced vertex metallib: {}", e))?;
    let inst_vert_fn = inst_library
        .newFunctionWithName(&ns_str("vertex_main_instanced"))
        .ok_or("vertex_main_instanced not found in instanced metallib")?;

    // The instanced pipeline always pairs with the per-draw fragment_main
    // (bindless is static-only).
    let frag_library = stage_library(device, hot_reload, frag_lib_bytes)
        .map_err(|e| format!("failed to load fragment metallib: {}", e))?;
    let frag_fn = frag_library
        .newFunctionWithName(&ns_str("fragment_main"))
        .ok_or("fragment_main not found in metallib")?;

    let inst_pipeline_desc = MTLRenderPipelineDescriptor::new();
    inst_pipeline_desc.setVertexDescriptor(Some(vert_desc));
    inst_pipeline_desc.setVertexFunction(Some(&inst_vert_fn));
    inst_pipeline_desc.setFragmentFunction(Some(&frag_fn));
    inst_pipeline_desc.setRasterSampleCount(HDR_SAMPLE_COUNT as usize);
    // SAFETY: plain descriptor property setters; the subscripted slots are ones this descriptor
    // declares.
    unsafe {
        inst_pipeline_desc
            .colorAttachments()
            .objectAtIndexedSubscript(0)
            .setPixelFormat(MTLPixelFormat::RGBA16Float);
    }
    inst_pipeline_desc.setDepthAttachmentPixelFormat(MTLPixelFormat::Depth32Float);

    let ps = device
        .newRenderPipelineStateWithDescriptor_error(&inst_pipeline_desc)
        .map_err(|e| format!("failed to create instanced pipeline state: {:?}", e))?;
    Ok(Some(ps))
}

// Shadow pipeline: depth-only, no fragment function, no MSAA. Compiled from the
// engine-internal single source (`shadow.slang`, entry `shadow_vertex_main`).
// Shared by init (one-shot at startup) and the internal-shader hot-reload path
// (`reload_shaders`) so the two stay consistent.
pub(crate) fn build_shadow_pipeline(
    device: &ProtocolObject<dyn MTLDevice>,
    vert_desc: &MTLVertexDescriptor,
    hot_reload: bool,
) -> Result<Retained<ProtocolObject<dyn MTLRenderPipelineState>>, String> {
    let shadow_fn = super::super::slang_shaders::entry_function(
        device,
        &super::super::slang_shaders::SHADOW_VERT,
        hot_reload,
    )?;
    let shadow_pipeline_desc = MTLRenderPipelineDescriptor::new();
    shadow_pipeline_desc.setVertexDescriptor(Some(vert_desc));
    shadow_pipeline_desc.setVertexFunction(Some(&shadow_fn));
    shadow_pipeline_desc.setRasterSampleCount(1);
    shadow_pipeline_desc.setDepthAttachmentPixelFormat(MTLPixelFormat::Depth32Float);
    device
        .newRenderPipelineStateWithDescriptor_error(&shadow_pipeline_desc)
        .map_err(|e| format!("failed to create shadow pipeline state: {:?}", e))
}

// GPU-driven cascaded-shadow render pipeline: depth-only, no
// fragment, no MSAA, but `supportIndirectCommandBuffers` so each cascade's
// casters can draw through the shadow ICB the `cull_encode_shadow` kernel
// fills. Entry `shadow_vertex_bindless` reads the per-object model matrix from
// the GpuObjectData buffer at buffer(9) by `[[base_instance]]` (the record id
// the cull baked), exactly like the main bindless `vertex_main`. Reuses the
// full static vertex descriptor (the VS consumes only attribute(0) = position;
// the deformed skinned tail shares the same 56-byte layout).
pub(crate) fn build_shadow_bindless_pipeline(
    device: &ProtocolObject<dyn MTLDevice>,
    vert_desc: &MTLVertexDescriptor,
    hot_reload: bool,
) -> Result<Retained<ProtocolObject<dyn MTLRenderPipelineState>>, String> {
    let shadow_fn = super::super::slang_shaders::entry_function(
        device,
        &super::super::slang_shaders::SHADOW_VERT_BINDLESS,
        hot_reload,
    )?;
    let shadow_pipeline_desc = MTLRenderPipelineDescriptor::new();
    shadow_pipeline_desc.setVertexDescriptor(Some(vert_desc));
    shadow_pipeline_desc.setVertexFunction(Some(&shadow_fn));
    shadow_pipeline_desc.setRasterSampleCount(1);
    shadow_pipeline_desc.setDepthAttachmentPixelFormat(MTLPixelFormat::Depth32Float);
    shadow_pipeline_desc.setSupportIndirectCommandBuffers(true);
    device
        .newRenderPipelineStateWithDescriptor_error(&shadow_pipeline_desc)
        .map_err(|e| format!("failed to create shadow bindless pipeline state: {:?}", e))
}

// Depth-stencil state: less-than test, writes enabled (shared for main and
// shadow pass).
pub(crate) fn make_depth_state(
    device: &ProtocolObject<dyn MTLDevice>,
) -> Result<Retained<ProtocolObject<dyn MTLDepthStencilState>>, String> {
    let depth_desc = MTLDepthStencilDescriptor::new();
    depth_desc.setDepthCompareFunction(MTLCompareFunction::Less);
    depth_desc.setDepthWriteEnabled(true);
    device
        .newDepthStencilStateWithDescriptor(&depth_desc)
        .ok_or_else(|| "failed to create depth stencil state".to_string())
}

// Read-only depth-stencil state: less-or-equal test, no write. Translucent
// passes (volumetric raymarch) bind this so they early-z against nearer
// opaque geometry without touching the depth buffer. A non-nil state is
// required: Metal's validation layer asserts on `setDepthStencilState(nil)`.
pub(crate) fn make_depth_state_read_only(
    device: &ProtocolObject<dyn MTLDevice>,
) -> Result<Retained<ProtocolObject<dyn MTLDepthStencilState>>, String> {
    let depth_desc = MTLDepthStencilDescriptor::new();
    depth_desc.setDepthCompareFunction(MTLCompareFunction::LessEqual);
    depth_desc.setDepthWriteEnabled(false);
    device
        .newDepthStencilStateWithDescriptor(&depth_desc)
        .ok_or_else(|| "failed to create read-only depth stencil state".to_string())
}