concinnity-device 0.19.119

//! Backing store for the render graph's transient textures on Metal. The shared
//! `gfx::render_graph::alias` planner decides which transient resources may share
//! physical memory; this pool is where the Metal backend realizes that plan.
//! Features stop owning these textures and read them back by label, so the pool
//! repoints several labels at one aliased allocation without touching the
//! features. This mirrors how the graph plans barriers while each backend emits
//! them, and the Vulkan / DirectX `transient_pool.rs`.
//!
//! Structure: the pool is organized into alias slots. Each slot owns one
//! `MTLHeap` sized to its largest member, and every member is placed at offset 0.
//! Members of a slot have pairwise-disjoint lifetimes (they are never live at the
//! same time), so reusing the bytes is safe. A single-member slot is a plain
//! placed target; a multi-member slot is a realized alias.
//!
//! The heaps are `Placement` so the pool picks the offset (0 for every member),
//! which maps onto the planner's slot abstraction directly. They are explicitly
//! `Tracked`: on a tracked heap Metal delays reads and writes of every resource
//! suballocated on it until in-flight modifications of any of them complete,
//! which is exactly the ordering aliased members need at both the within-frame
//! and the cross-frame reuse boundary. `MTLHazardTrackingMode::Default` is
//! treated as `Untracked` for heaps and would silently drop the automatic
//! tracking the rest of the backend assumes. One heap per slot keeps that
//! heap-granularity tracking scoped to members that genuinely alias.
//!
//! Slots are single-buffered, like DirectX and unlike Vulkan: the tracked heap
//! orders a frame's writes against the previous frame's reads of the same bytes,
//! so one heap per slot (rather than one per slot per frame) is safe. Per-frame
//! slots would multiply an already single-buffered footprint by the
//! frames-in-flight depth and cost more than aliasing saves.
//!
//! A texture is "managed" iff its owning feature is enabled at build time;
//! `texture_for` returns `None` otherwise and the consumer falls back exactly as
//! before (the main pass samples a 1x1 white texture when SSAO is off).
#![deny(unsafe_op_in_unsafe_fn)]

use super::error::allocation_failed;
use concinnity_core::render::error::{RenderError, RenderResult};
use concinnity_core::render::render_graph::{
    PixelFormat, TextureUsage, TransientSlot, TransientTexture,
};
use objc2::rc::Retained;
use objc2::runtime::ProtocolObject;
use objc2_metal::{
    MTLDevice as _, MTLHazardTrackingMode, MTLHeap, MTLHeapDescriptor, MTLHeapType, MTLPixelFormat,
    MTLStorageMode, MTLTexture, MTLTextureDescriptor, MTLTextureType, MTLTextureUsage,
};

use crate::metal::descriptors::TextureDesc;
// `ClearValue` is deliberately absent: Metal takes a clear value on the render
// pass descriptor rather than baking it into the texture, so the graph's clear
// is consumed where the pass is encoded, not here. DirectX translates it at
// creation, which is why its pool reads the field and this one does not.
use crate::metal::context::MtlContext;

struct PooledTexture {
    label: &'static str,
    texture: Retained<ProtocolObject<dyn MTLTexture>>,
}

// The transient texture pool owned by `MtlContext`. Resolution-dependent, so it
// is rebuilt on resize alongside the other render-resolution targets.
pub(super) struct TransientTexturePool {
    // One heap per slot, backing every member of that slot at offset 0. Held so
    // a rebuild releases a slot's memory as a unit, and so the pool's own
    // lifetime covers the heaps its textures were placed on.
    heaps: Vec<Retained<ProtocolObject<dyn MTLHeap>>>,
    textures: Vec<PooledTexture>,
    // The member labels of each slot, in the order they reuse its memory. Read
    // back by the executor's per-frame soundness assertion: members of one slot
    // share bytes, so two live at once is silent corruption. That assertion runs
    // under `debug_assertions`, and so does its only source of truth.
    #[cfg(debug_assertions)]
    slot_labels: Vec<Vec<&'static str>>,
}

impl TransientTexturePool {
    // Allocate one heap per slot and place every member at offset 0. Each
    // texture starts undefined; its first-use contents come from its graph
    // producer pass exactly as when the feature owned it. (Unlike D3D12's placed
    // resources, a Metal placed texture needs no Discard-style initialization
    // before its first use.)
    pub(super) fn build(
        device: &ProtocolObject<dyn objc2_metal::MTLDevice>,
        slots: &[TransientSlot],
    ) -> RenderResult<Self> {
        let mut heaps = Vec::with_capacity(slots.len());
        let mut textures = Vec::new();
        // What the members would cost with one allocation each; the heaps' real
        // size is the aliased footprint. The difference is the VRAM aliasing
        // reclaims, reported below.
        let mut unaliased_bytes: u64 = 0;
        for slot in slots {
            // Size the heap to the largest member. Offset 0 satisfies every
            // member's alignment, so aliased members all place there.
            let mut slot_size: usize = 0;
            let mut descs = Vec::with_capacity(slot.members.len());
            for m in &slot.members {
                let desc = texture_descriptor_for(m);
                let size = device.heapTextureSizeAndAlignWithDescriptor(&desc).size;
                unaliased_bytes += size as u64;
                slot_size = slot_size.max(size);
                descs.push((m.label, desc));
            }
            let heap = new_slot_heap(device, slot_size)?;
            for (label, desc) in descs {
                // Placement heap at offset 0: in bounds because the heap is
                // sized to the largest member, and aligned because every
                // alignment divides 0.
                // SAFETY: the heap is sized to its largest member, so offset 0 is in bounds, and
                // every alignment divides 0.
                let texture = unsafe { heap.newTextureWithDescriptor_offset(&desc, 0) }
                    .ok_or_else(|| allocation_failed(format_args!("transient texture {label}")))?;
                textures.push(PooledTexture { label, texture });
            }
            heaps.push(heap);
        }
        let pool = Self {
            heaps,
            textures,
            #[cfg(debug_assertions)]
            slot_labels: slots.iter().map(|s| s.labels()).collect(),
        };
        let aliased_bytes = pool.heap_bytes();
        tracing::info!(
            "transient texture pool: {} slot heap(s), {} KiB ({} KiB saved by aliasing)",
            pool.heaps.len(),
            aliased_bytes / 1024,
            unaliased_bytes.saturating_sub(aliased_bytes) / 1024,
        );
        Ok(pool)
    }

    // The managed texture for `label`, or `None` when its owning feature was
    // disabled at build time (so the pool holds no entry for it).
    pub(super) fn texture_for(&self, label: &str) -> Option<&ProtocolObject<dyn MTLTexture>> {
        self.lookup(label).map(|t| t.texture.as_ref())
    }

    // The pooled `bloom_top`, the bloom chain's top octave. Always managed, so a
    // missing entry is a build bug rather than a disabled feature.
    pub(super) fn bloom_top(&self) -> RenderResult<&ProtocolObject<dyn MTLTexture>> {
        self.texture_for("bloom_top")
            .ok_or_else(|| RenderError::Other("bloom_top missing from transient pool".to_string()))
    }

    // Total bytes the slot heaps occupy: the pool's aliased footprint, as Metal
    // rounded it up at creation.
    pub(super) fn heap_bytes(&self) -> u64 {
        self.heaps.iter().map(|h| h.size() as u64).sum()
    }

    // The member labels of each slot, for the executor's per-frame check that
    // no slot has two resources live at once in the graph it is about to run.
    #[cfg(debug_assertions)]
    pub(super) fn slot_labels(&self) -> &[Vec<&'static str>] {
        &self.slot_labels
    }

    fn lookup(&self, label: &str) -> Option<&PooledTexture> {
        self.textures.iter().find(|t| t.label == label)
    }

    // Rebuild every managed texture at a new extent. Metal reference-counts, and
    // a heap-placed texture holds a reference to its heap, so an in-flight
    // command buffer keeps both the old texture and the memory under it alive
    // until the GPU retires that frame; the heaps released here are only those
    // no frame still references. The caller rebinds the new textures into the
    // affected consumers (the per-frame bindless argument buffer re-encodes
    // `ao_output` itself; the bloom chain and the composite re-read
    // `bloom_top`).
    pub(super) fn rebuild(
        &mut self,
        device: &ProtocolObject<dyn objc2_metal::MTLDevice>,
        slots: &[TransientSlot],
    ) -> RenderResult<()> {
        *self = Self::build(device, slots)?;
        Ok(())
    }
}

// One `Private` placement heap for a slot, sized to its largest member. Hazard
// tracking is set explicitly: heaps treat `Default` as `Untracked`, which would
// drop the automatic tracking the aliased members and the rest of the backend
// rely on.
fn new_slot_heap(
    device: &ProtocolObject<dyn objc2_metal::MTLDevice>,
    size: usize,
) -> RenderResult<Retained<ProtocolObject<dyn MTLHeap>>> {
    let desc = MTLHeapDescriptor::new();
    desc.setType(MTLHeapType::Placement);
    desc.setStorageMode(MTLStorageMode::Private);
    desc.setHazardTrackingMode(MTLHazardTrackingMode::Tracked);
    desc.setSize(size.max(1));
    device
        .newHeapWithDescriptor(&desc)
        .ok_or_else(|| allocation_failed(format_args!("{size}-byte transient slot heap")))
}

// Translate one graph-declared transient into its Metal descriptor. This is the
// backend's whole share of describing a pooled resource: the extent, format,
// mip count and usage all come from the graph, so there is no second table here
// that could disagree with it. GPU-private to match its heap's storage mode.
pub(in crate::metal) fn texture_descriptor_for(
    spec: &TransientTexture,
) -> Retained<MTLTextureDescriptor> {
    TextureDesc {
        kind: texture_type(spec),
        format: pixel_format(spec.format),
        width: spec.width.max(1) as usize,
        height: spec.height.max(1) as usize,
        depth: spec.depth.max(1) as usize,
        array_length: spec.array_layers.max(1) as usize,
        mip_count: spec.mip_levels.max(1) as usize,
        sample_count: spec.sample_count.max(1) as usize,
        usage: texture_usage(spec.usage),
        ..Default::default()
    }
    .build()
}

fn texture_type(spec: &TransientTexture) -> MTLTextureType {
    match (
        spec.depth.max(1) > 1,
        spec.array_layers.max(1) > 1,
        spec.sample_count.max(1) > 1,
    ) {
        (true, _, _) => MTLTextureType::Type3D,
        (_, true, _) => MTLTextureType::Type2DArray,
        (_, _, true) => MTLTextureType::Type2DMultisample,
        _ => MTLTextureType::Type2D,
    }
}

pub(in crate::metal) fn pixel_format(format: PixelFormat) -> MTLPixelFormat {
    match format {
        PixelFormat::Rgba16Float => MTLPixelFormat::RGBA16Float,
        PixelFormat::Rgba8Unorm => MTLPixelFormat::RGBA8Unorm,
        PixelFormat::Rg16Float => MTLPixelFormat::RG16Float,
        PixelFormat::R8Unorm => MTLPixelFormat::R8Unorm,
        PixelFormat::R16Float => MTLPixelFormat::R16Float,
        PixelFormat::R32Float => MTLPixelFormat::R32Float,
        PixelFormat::Rg32Float => MTLPixelFormat::RG32Float,
        PixelFormat::Depth32Float => MTLPixelFormat::Depth32Float,
        PixelFormat::BgraSwapchain => MTLPixelFormat::BGRA8Unorm,
    }
}

fn texture_usage(usage: TextureUsage) -> MTLTextureUsage {
    let mut bits = 0;
    if usage.contains(TextureUsage::SHADER_READ) {
        bits |= MTLTextureUsage::ShaderRead.0;
    }
    if usage.contains(TextureUsage::RENDER_TARGET) || usage.contains(TextureUsage::DEPTH_STENCIL) {
        bits |= MTLTextureUsage::RenderTarget.0;
    }
    if usage.contains(TextureUsage::STORAGE) {
        bits |= MTLTextureUsage::ShaderRead.0 | MTLTextureUsage::ShaderWrite.0;
    }
    MTLTextureUsage(bits)
}

impl MtlContext {
    // The texture the main pass and the bindless argument buffer sample for
    // ambient occlusion: the pooled `ao_output` (SSAO's blurred occlusion) when
    // SSAO is on, else the SSAO state's 1x1 white fallback so `shade_surface`
    // reads a constant 1.0 (fully unoccluded).
    pub(in crate::metal) fn ao_output_texture(&self) -> &ProtocolObject<dyn MTLTexture> {
        self.targets
            .transient_pool
            .texture_for("ao_output")
            .unwrap_or_else(|| self.ssao.white.as_ref())
    }

    // The unified pre-pass's three color channels, read straight out of the
    // pool at the point of use. Nothing caches these, and that is deliberate:
    // **a pool rebuild repacks every slot**, so a handle cached when SSAO was
    // toggled would point into a heap region that now belongs to a different
    // resource. The explicit backends have to re-point their descriptors and
    // framebuffers at every rebuild for exactly this reason; fetching by label
    // costs a short scan per encode and makes the whole class impossible here.
    //
    // `None` when the pre-pass is not built, which is the same gate the pool
    // was created under, so a consumer that finds `None` should skip rather
    // than substitute a fallback.
    pub(in crate::metal) fn gbuffer_normal_depth(&self) -> Option<&ProtocolObject<dyn MTLTexture>> {
        self.targets
            .transient_pool
            .texture_for("gbuffer_normal_depth")
    }

    pub(in crate::metal) fn gbuffer_roughness(&self) -> Option<&ProtocolObject<dyn MTLTexture>> {
        self.targets.transient_pool.texture_for("gbuffer_roughness")
    }

    pub(in crate::metal) fn gbuffer_velocity(&self) -> Option<&ProtocolObject<dyn MTLTexture>> {
        self.targets.transient_pool.texture_for("gbuffer_velocity")
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use concinnity_core::render::post::device::PostExtent;
    use concinnity_core::render::post::{bloom, ssao};
    use concinnity_core::render::render_graph::{self, PoolGates, plan_pool_slots};

    #[test]
    fn translated_descriptors_match_the_feature_formats() {
        // The graph is the single source of the shape now, so what this pins is
        // the *translation*: the descriptor the pool creates must still be the
        // one each feature's own constant describes, or the pool silently
        // mis-backs the texture that feature binds.
        let gates = PoolGates {
            ssao: true,
            gbuffer: true,
        };
        let slots = plan_pool_slots(gates, (1024, 768), (1024, 768)).expect("plans");
        let member = |label: &str| {
            slots
                .iter()
                .flat_map(|s| &s.members)
                .find(|m| m.label == label)
                .expect("member present")
        };

        let ao = texture_descriptor_for(member("ao_output"));
        assert_eq!(ao.pixelFormat(), pixel_format(ssao::OCCLUSION_FORMAT));
        assert_eq!((ao.width(), ao.height()), (1024, 768));
        assert_eq!(ao.textureType(), MTLTextureType::Type2D);
        assert_eq!(ao.mipmapLevelCount(), 1);
        assert_eq!(ao.sampleCount(), 1);
        assert_eq!(
            ao.usage().0,
            MTLTextureUsage::ShaderRead.0 | MTLTextureUsage::RenderTarget.0
        );

        // Half the *output* extent: the octave above the chain's quarter-output
        // base, in the format the chain renders.
        let bloom = texture_descriptor_for(member("bloom_top"));
        let chain = bloom::chain_desc(PostExtent {
            width: 1024,
            height: 768,
        });
        assert_eq!(bloom.pixelFormat(), pixel_format(chain.format));
        assert_eq!((bloom.width(), bloom.height()), (512, 384));
    }

    #[test]
    fn a_volume_translates_to_a_3d_descriptor() {
        // Nothing pooled is a volume yet, but the froxel volume's desc is now
        // real, so the translator has to handle it before it can be pooled.
        let desc = texture_descriptor_for(&TransientTexture {
            label: "probe_volume",
            width: 80,
            height: 45,
            depth: 64,
            format: PixelFormat::Rgba16Float,
            sample_count: 1,
            array_layers: 1,
            mip_levels: 1,
            usage: TextureUsage::STORAGE.union(TextureUsage::SHADER_READ),
            // Carried by the graph for DirectX's sake; the Metal descriptor has
            // nowhere to put it.
            clear: render_graph::ClearValue::Color([0.0; 4]),
        });
        assert_eq!(desc.textureType(), MTLTextureType::Type3D);
        assert_eq!(desc.depth(), 64);
        assert_eq!(
            desc.usage().0,
            MTLTextureUsage::ShaderRead.0 | MTLTextureUsage::ShaderWrite.0
        );
    }
}