arguments_table.cpp

#include <math.h>
#ifndef INFINITY
#define INFINITY __builtin_huge_valf()
#endif
#include <Foundation/Foundation.hpp>
#include <Metal/Metal.hpp>
#include <QuartzCore/QuartzCore.hpp>

#include "backend/metal/mt_api.h"
#include "core/logger.h"

#include <cstring>
#include <new>

namespace mt {

// push + vertex table. resources ride in push.

u64 IndexPool::alloc() {
    if (!free_list.empty()) {
        u64 s = free_list.back();
        free_list.pop_back();
        return s;
    }
    if (next >= cap) {
        return UINT64_MAX;
    }
    return next++;
}

void IndexPool::free(u64 slot) {
    if (slot == UINT64_MAX) {
        return;
    }
    free_list.push_back((u32)slot);
}

void reside(rhi::Device &device, const MTL::Allocation *alloc) {
    if (device.heap && device.heap->residency && alloc) {
        device.heap->residency->addAllocation(alloc);
    }
}

void unreside(rhi::Device &device, const MTL::Allocation *alloc) {
    if (device.heap && device.heap->residency && alloc) {
        device.heap->residency->removeAllocation(alloc);
    }
}

void create_heap(rhi::Device &device) {
    MTL::Device *dev = device.handle;
    Heap *heap = new (std::nothrow) Heap{};
    if (!heap) {
        VEL_CRITICAL("arguments_table: out of memory");
        return;
    }

    heap->sampler_space = {.cap = kHeapSamplerBudget, .next = 0};
    heap->texture_space = {.cap = kHeapTextureBudget, .next = 0};
    heap->texture_ids.assign(kHeapTextureBudget, 0);
    heap->sampler_ids.assign(kHeapSamplerBudget, 0);

    NS::Error *error = nullptr;
    MTL4::ArgumentTableDescriptor *td = MTL4::ArgumentTableDescriptor::alloc()->init();
    td->setMaxBufferBindCount(kMaxBufferBindCount);
    td->setMaxTextureBindCount(kMaxTextureBindCount);
    td->setMaxSamplerStateBindCount(kMaxSamplerBindCount);
    td->setSupportAttributeStrides(true);
    td->setInitializeBindings(true);
    heap->table = dev->newArgumentTable(td, &error);
    td->release();
    if (!heap->table) {
        if (error && error->localizedDescription()) {
            VEL_CRITICAL("arguments_table: newArgumentTable failed: {}", error->localizedDescription()->utf8String());
            error->release();
        } else {
            VEL_CRITICAL("arguments_table: newArgumentTable failed");
        }
        delete heap;
        return;
    }

    MTL::ResidencySetDescriptor *rd = MTL::ResidencySetDescriptor::alloc()->init();
    rd->setInitialCapacity(256);
    error = nullptr;
    heap->residency = dev->newResidencySet(rd, &error);
    rd->release();
    if (!heap->residency) {
        if (error && error->localizedDescription()) {
            VEL_CRITICAL("arguments_table: newResidencySet failed: {}", error->localizedDescription()->utf8String());
            error->release();
        } else {
            VEL_CRITICAL("arguments_table: newResidencySet failed");
        }
        heap->table->release();
        delete heap;
        return;
    }
    if (!device.graphics.handle) {
        VEL_CRITICAL("arguments_table: graphics queue missing, cannot attach residency set");
        heap->table->release();
        heap->residency->release();
        delete heap;
        return;
    }
    device.graphics.handle->addResidencySet(heap->residency);
    heap->residency->requestResidency();
    heap->residency->commit();

    device.heap = heap;
    VEL_INFO("arguments_table: ready (push + vertex table, texture/sampler host heaps + residency)");
}

void destroy_heap(rhi::Device &device) {
    Heap *heap = device.heap;
    device.heap = nullptr;
    if (!heap) {
        return;
    }
    if (heap->table) {
        heap->table->release();
        heap->table = nullptr;
    }
    if (heap->residency) {
        heap->residency->release();
        heap->residency = nullptr;
    }
    delete heap;
}

u64 write_sampler(rhi::Device &device, rhi::Sampler sampler) {
    if (!device.heap || !sampler.handle) {
        return UINT64_MAX;
    }
    u64 slot = device.heap->sampler_space.alloc();
    if (slot == UINT64_MAX) {
        VEL_ERROR("arguments_table: sampler heap exhausted");
        return UINT64_MAX;
    }
    device.heap->sampler_ids[(size_t)slot] = sampler.handle->gpuResourceID()._impl;
    return slot;
}

u64 write_texture(rhi::Device &device, MTL::Texture *view, bool storage) {
    (void)storage;
    if (!device.heap || !view) {
        return UINT64_MAX;
    }
    u64 slot = device.heap->texture_space.alloc();
    if (slot == UINT64_MAX) {
        VEL_ERROR("arguments_table: texture heap exhausted");
        return UINT64_MAX;
    }
    device.heap->texture_ids[(size_t)slot] = view->gpuResourceID()._impl;
    if (device.heap->residency) {
        device.heap->residency->addAllocation(view);
    }
    return slot;
}

void free_texture(rhi::Device &device, u64 slot) {
    if (device.heap) {
        device.heap->texture_space.free(slot);
        if (slot < device.heap->texture_ids.size()) {
            device.heap->texture_ids[(size_t)slot] = 0;
        }
    }
}

void free_sampler(rhi::Device &device, u64 slot) {
    if (device.heap) {
        device.heap->sampler_space.free(slot);
        if (slot < device.heap->sampler_ids.size()) {
            device.heap->sampler_ids[(size_t)slot] = 0;
        }
    }
}

u64 heap_texture_id(rhi::Device &device, u64 slot) {
    if (!device.heap || slot == UINT64_MAX || slot >= device.heap->texture_ids.size()) {
        return 0;
    }
    return device.heap->texture_ids[(size_t)slot];
}

u64 heap_sampler_id(rhi::Device &device, u64 slot) {
    if (!device.heap || slot == UINT64_MAX || slot >= device.heap->sampler_ids.size()) {
        return 0;
    }
    return device.heap->sampler_ids[(size_t)slot];
}

MTL::Buffer *push_bump_alloc(rhi::CmdBuffer &cmd, const void *data, u32 size) {
    if (!cmd.queue) {
        return nullptr;
    }
    // 16b align. constants need it.
    const u32 aligned = (size + 15u) & ~15u;
    MTL::Device *dev = cmd.queue->device();
    MTL::Buffer *buf = dev->newBuffer(aligned, MTL::StorageModeShared);
    if (!buf) {
        return nullptr;
    }
    if (cmd.residency) {
        cmd.residency->addAllocation(buf);
    }
    std::memcpy(buf->contents(), data, size);
    cmd.transients.push_back(buf);
    return buf;
}

} // namespace mt