concinnity-device 0.19.119

//! Runtime GPU resource management for VkContext, split per-category to mirror
//! the Metal reference shape (`metal/resources/`):
//!
//!   textures.rs   Texture-pool slot updates + descriptor rewires (`update_*`,
//!                 `evict_*`, `write_object_image`, `write_pool_image`)
//!   geometry.rs   The scene lent to the `SceneHost` defaults, its staging
//!                 geometry writer, and the shared `write_geometry_region` helper
//!   streaming.rs  VoxelWorld chunk headroom (`setup_chunk_streaming`)
//!   skinning.rs   Skinned-mesh upload + per-frame joint upload
//!                 (`upload_skinned`, `upload_joint_matrices`)
//!   geometry_rebuild.rs  Size-changing static + skinned VB/IB rebuilds
//!                 driven by asset hot-reload (`rebuild_static_geometry`,
//!                 `rebuild_skinned_geometry`)
//!
//! The shared low-level helpers (`create_descriptor_set_layout`,
//! `alloc_descriptor_sets`, `upload_geometry_buffer{,_raw}`) live in this file
//! because every submodule + `init.rs` needs them.

use ash::vk;
use concinnity_core::render::error::RenderResult;

use super::allocator::{DeviceAllocator, PooledBuffer};
use super::texture;
use crate::vulkan::owned::{OwnedSetLayout, VkDevice};
use crate::vulkan::set_writes::SetWrites;

mod geometry;
mod geometry_rebuild;
mod skinning;
mod streaming;
mod textures;

pub(in crate::vulkan) fn create_descriptor_set_layout(
    device: &VkDevice,
    bindings: &[(u32, vk::DescriptorType, vk::ShaderStageFlags)],
) -> RenderResult<OwnedSetLayout> {
    let vk_bindings: Vec<_> = bindings
        .iter()
        .map(|&(b, ty, stage)| {
            vk::DescriptorSetLayoutBinding::default()
                .binding(b)
                .descriptor_type(ty)
                .descriptor_count(1)
                .stage_flags(stage)
        })
        .collect();
    let info = vk::DescriptorSetLayoutCreateInfo::default().bindings(&vk_bindings);
    device
        .create_descriptor_set_layout(&info)
        .map_err(|e| crate::vulkan::error::map_vk_result(e, "descriptor set layout"))
}

// The bindings of a set holding `n` fragment-sampled sources: their images at
// `0..n`, then their samplers in the same order at `n..2n`. The shape every
// fullscreen pass's shader declares its sources in.
pub(in crate::vulkan) fn source_set_bindings(
    n: u32,
) -> Vec<(u32, vk::DescriptorType, vk::ShaderStageFlags)> {
    (0..2 * n)
        .map(|b| {
            let ty = if b < n {
                vk::DescriptorType::SAMPLED_IMAGE
            } else {
                vk::DescriptorType::SAMPLER
            };
            (b, ty, vk::ShaderStageFlags::FRAGMENT)
        })
        .collect()
}

// Write `sources` into a set laid out by `source_set_bindings(sources.len())`:
// each view at its image binding and each sampler at its sampler binding, in
// one update.
pub(in crate::vulkan) fn write_source_set(
    device: &VkDevice,
    set: vk::DescriptorSet,
    sources: &[(vk::ImageView, vk::Sampler)],
) {
    let n = sources.len() as u32;
    sources
        .iter()
        .zip(0..)
        .fold(SetWrites::new(set), |w, (&(view, sampler), i)| {
            w.sampled_image(i, view).sampler(n + i, sampler)
        })
        .apply(device);
}

pub(in crate::vulkan) fn alloc_descriptor_sets(
    device: &VkDevice,
    pool: vk::DescriptorPool,
    layouts: &[vk::DescriptorSetLayout],
) -> RenderResult<Vec<vk::DescriptorSet>> {
    if layouts.is_empty() {
        return Ok(vec![]);
    }
    let alloc = vk::DescriptorSetAllocateInfo::default()
        .descriptor_pool(pool)
        .set_layouts(layouts);
    // SAFETY: the create-info and every slice it borrows are live for the call, and each handle it
    // names belongs to this device.
    unsafe { device.allocate_descriptor_sets(&alloc) }
        .map_err(|e| crate::vulkan::error::map_vk_result(e, "allocate descriptor sets"))
}

// Usage every (re)creation of the shared vertex / index buffers must carry, on
// top of its VERTEX_BUFFER / INDEX_BUFFER role. On an RT-capable device the
// shared buffers double as acceleration-structure build inputs (device
// addressed) and as storage buffers the RT / glass shaders fetch hit-triangle
// attributes from; the flags ride on capability rather than on RT being live,
// since a later quality toggle cannot add usage to an existing buffer. Every
// path that replaces the buffers (chunk-streaming headroom, hot-reload
// geometry rebuild) routes through here: dropping these leaves the shared
// buffers un-addressable and every later acceleration-structure build invalid.
pub(in crate::vulkan) fn shared_geometry_usage(rt_capable: bool) -> vk::BufferUsageFlags {
    let base = vk::BufferUsageFlags::TRANSFER_SRC | vk::BufferUsageFlags::TRANSFER_DST;
    if rt_capable {
        base | vk::BufferUsageFlags::SHADER_DEVICE_ADDRESS
            | vk::BufferUsageFlags::ACCELERATION_STRUCTURE_BUILD_INPUT_READ_ONLY_KHR
            | vk::BufferUsageFlags::STORAGE_BUFFER
    } else {
        base
    }
}

pub(in crate::vulkan) fn upload_geometry_buffer<T: bytemuck::NoUninit>(
    alloc: &DeviceAllocator,
    device: &VkDevice,
    command_pool: vk::CommandPool,
    queue: vk::Queue,
    data: &[T],
    usage: vk::BufferUsageFlags,
) -> RenderResult<PooledBuffer> {
    upload_geometry_buffer_raw(
        alloc,
        device,
        command_pool,
        queue,
        bytemuck::cast_slice(data),
        usage,
    )
}

pub(in crate::vulkan) fn upload_geometry_buffer_raw(
    alloc: &DeviceAllocator,
    device: &VkDevice,
    command_pool: vk::CommandPool,
    queue: vk::Queue,
    data: &[u8],
    usage: vk::BufferUsageFlags,
) -> RenderResult<PooledBuffer> {
    // TRANSFER_SRC lets `setup_chunk_streaming` copy the build-time geometry
    // out of these buffers when it grows them for chunk-streaming headroom;
    // TRANSFER_DST lets the staging copy below and `write_geometry_region`
    // write into them.
    let usage = usage | vk::BufferUsageFlags::TRANSFER_SRC | vk::BufferUsageFlags::TRANSFER_DST;
    let size = data.len() as u64;
    if size == 0 {
        // Return a minimal 4-byte buffer to keep Vulkan happy.
        return alloc.create_buffer(4, usage, vk::MemoryPropertyFlags::DEVICE_LOCAL);
    }
    let staging = alloc.create_buffer(
        size,
        vk::BufferUsageFlags::TRANSFER_SRC,
        vk::MemoryPropertyFlags::HOST_VISIBLE | vk::MemoryPropertyFlags::HOST_COHERENT,
    )?;
    staging.write_bytes(0, data);
    let buf = alloc.create_buffer(size, usage, vk::MemoryPropertyFlags::DEVICE_LOCAL)?;
    texture::one_shot_submit(device, command_pool, queue, |cmd| {
        let copy = vk::BufferCopy::default().size(size);
        // SAFETY: `cmd` is a command buffer in the recording state, and every handle and slice
        // these commands name is live for the call.
        unsafe {
            device.cmd_copy_buffer(
                cmd,
                staging.buffer(),
                buf.buffer(),
                std::slice::from_ref(&copy),
            )
        };
    })?;
    Ok(buf)
}

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

    // `taa.hlsl`'s three sources are the shape: every image, then every sampler.
    #[test]
    fn a_source_set_lists_every_image_then_every_sampler() {
        use vk::DescriptorType as T;
        let stage = vk::ShaderStageFlags::FRAGMENT;
        assert_eq!(
            source_set_bindings(3),
            [
                (0, T::SAMPLED_IMAGE, stage),
                (1, T::SAMPLED_IMAGE, stage),
                (2, T::SAMPLED_IMAGE, stage),
                (3, T::SAMPLER, stage),
                (4, T::SAMPLER, stage),
                (5, T::SAMPLER, stage),
            ]
        );
        assert!(source_set_bindings(0).is_empty());
    }
}