GpuResourceRegistry.hpp

#pragma once
 
#include <score_plugin_gfx_export.h>
 
#include <ossia/dataflow/geometry_port.hpp>  // ossia::gpu_slot_ref
#include <ossia/detail/flat_map.hpp>
#include <ossia/detail/hash_map.hpp>
 
#ifndef OFFSETALLOCATOR_HPP_2026_04_24
#define OFFSETALLOCATOR_HPP_2026_04_24
#include <offsetAllocator.hpp>
#endif
 
#include <QtGui/private/qrhi_p.h>
 
#include <array>
#include <cstddef>
#include <cstdint>
#include <memory>
#include <vector>
 
namespace score::gfx
{
class RenderList;
 
class SCORE_PLUGIN_GFX_EXPORT GpuResourceRegistry
{
public:
  // Well-known arenas. Size tables live in GpuResourceRegistry.cpp and match the
  // packed GPU layouts in SceneGPUState.hpp and CameraMath.hpp. Every entry
  // implies a QRhiBuffer allocation at init time.
  //
  // The Raw* arenas are written by source halp nodes at their own operator()()
  // time: view-independent, aspect-ratio-agnostic, pre-composition. The cooked
  // arenas are populated by ScenePreprocessor's transform passes, which combine
  // Raw inputs with the render target's aspect ratio and the scene-graph parent
  // chain, and are what consumer shaders bind. Material and Env have no
  // composition dependency, so source nodes write their cooked slot directly.
  enum class Arena : uint8_t
  {
    // ── Shared / source-authored ──────────────────────────────────
    // These arenas hold view- and filter-independent bytes: every
    // preprocessor reads the same data regardless of its camera /
    // render target / upstream scene filtering. The producer owns the
    // slot; multiple preprocessors consume via gpu_slot_ref + isLive().
    RawCamera,        // RawCameraData      — 64 B per slot, UBO
    RawLight,         // RawLightData       — 64 B per slot, SSBO
    RawTransform,     // RawLocalTransform  — 64 B per slot, SSBO
    Material,         // MaterialGPU        — 80 B per slot, SSBO
    Env,              // EnvParamsUBO       — 64 B per slot, UBO
 
    // Cooked outputs (camera UBOs, composed world matrices, per-draw structs,
    // LightGPU with world direction, MaterialGPU with resolved textureRefs) are
    // preprocessor-private: they are view- and filter-dependent, so a shared arena
    // would be wrong when two preprocessors see different filtered views.
 
    Count_
  };
 
  // Fixed-stride slot: the arena buffer is a packed array of stride-byte slots,
  // slot i at byte offset i * stride. The slot index is the arena-level identity
  // consumer shaders use, e.g. scene_materials.entries[slot_index]. Allocation is
  // O(1) through a free-list stack. Trades OffsetAllocator's variable-size
  // flexibility for a shader-indexable layout and a 1:1 mapping between
  // internal_index and byte offset, so arena reads need no translation table.
  struct Slot
  {
    static constexpr uint32_t kInvalidIndex = 0xFFFFFFFFu;
 
    Arena arena{Arena::RawCamera};
    uint32_t slot_index{kInvalidIndex};
    uint32_t size{0};        // requested payload size (≤ arena stride)
    uint32_t generation{};   // stamped on allocate; bumps on free
 
    bool valid() const noexcept { return slot_index != kInvalidIndex; }
  };
 
  GpuResourceRegistry() = default;
  GpuResourceRegistry(const GpuResourceRegistry&) = delete;
  GpuResourceRegistry& operator=(const GpuResourceRegistry&) = delete;
  ~GpuResourceRegistry();
 
  void init(QRhi& rhi, QRhiResourceUpdateBatch& batch);
 
  bool isInitialized() const noexcept { return m_rhi != nullptr; }
 
  QRhi* boundRhi() const noexcept { return m_rhi; }
 
  void seedDefaults(QRhiResourceUpdateBatch& batch);
 
  void uploadDefaultMaterial(QRhiResourceUpdateBatch& batch);
 
  void destroy(RenderList& renderer);
 
  void destroy();
 
  void destroyOwned();
 
  Slot allocate(Arena arena, uint32_t size);
 
  void free(Slot& slot);
 
  QRhiBuffer* buffer(Arena arena) const noexcept;
 
  uint32_t slotOffset(const Slot& slot) const noexcept;
 
  uint32_t arenaSlotStride(Arena arena) const noexcept;
 
  uint32_t arenaSlotCount(Arena arena) const noexcept;
 
  void updateSlot(
      QRhiResourceUpdateBatch& res, const Slot& slot, const void* data,
      uint32_t size) noexcept;
 
  ossia::gpu_slot_ref toOssiaRef(const Slot& slot) const noexcept
  {
    if(!slot.valid())
      return {};
    ossia::gpu_slot_ref r;
    r.arena = (uint32_t)slot.arena;
    r.offset = slotOffset(slot);
    r.size = slot.size;
    r.internal_index = slot.slot_index;
    r.generation = slot.generation;
    return r;
  }
 
  bool isLiveIn(const ossia::gpu_slot_ref& r, Arena expected) const noexcept
  {
    return r.arena == (uint32_t)expected && isLive(r);
  }
 
  bool isLive(const ossia::gpu_slot_ref& r) const noexcept
  {
    if(r.arena >= (uint32_t)Arena::Count_ || r.size == 0)
      return false;
    const auto& a = m_arenas[r.arena];
    if(r.internal_index >= a.slot_generations.size())
      return false;
    return a.slot_generations[r.internal_index] == r.generation;
  }
 
  // ─── Material texture arrays ──────────────────────────────────────
  //
  // Per-channel static texture arrays shared across every preprocessor in this
  // RenderList. Static textures dedup by texture_source pointer, so every producer
  // referencing an asset gets the same layer; dynamic handles (video textures,
  // runtime GPU outputs) get per-slot bindings in `dynamicTextures`, named
  // <channel>Dyn<slot> in consumer shaders.
  //
  // Sharing is safe because a texture belongs to an asset or a wired GPU handle,
  // independent of which preprocessor is looking.
 
  enum class TextureChannel : uint8_t
  {
    BaseColor  = 0,
    MetalRough = 1,
    Normal     = 2,
    Emissive   = 3,
    Occlusion  = 4,  // Separate glTF occlusionTexture (when distinct from MR).
    Count_     = 5
  };
 
  // Default layer size + max dynamic slots. Matched across channels so
  // samplers are interchangeable and consumer shaders can declare a
  // fixed sampler count.
  static constexpr int kTextureLayerSize = 1024;
  // One pool shared by every channel: a handle routed to two channels
  // resolves to one slot.
  static constexpr int kMaxDynamicSlots  = 8;
 
  // Static buckets in ONE pool shared by every channel, each holding textures
  // of one (format, pixelSize, colourspace) tuple. baseColor and emissive are
  // sRGB, metal-rough, normal and occlusion linear, so a model's channels land
  // in a handful of shared buckets. Past the cap a texture is skipped and the
  // shader reads tex_ref_none.
  //
  // The tex_ref_static encoding reserves a 7-bit bucket field, so the cap can
  // grow to 128 without changing the packed layout or the shader decode masks,
  // but the shader's ladder must be enlarged to match and the sampler budget
  // re-checked.
  static constexpr int kMaxBuckets = 8;
 
  struct TextureChannelState
  {
    struct Bucket
    {
      QRhiTexture* array{};          // QRhiTexture::TextureArray + `flags`
      QRhiTexture::Format format{QRhiTexture::RGBA8};
      QSize pixelSize;               // all layers in a bucket share this size
      int layers{};                  // current layer count
 
      // Colourspace, part of the bucket key since one pool serves every
      // channel. sRGB textures (base colour, emissive) get hardware
      // sRGB->linear on sample and linear data textures (metal-rough, normal,
      // occlusion) must not, so they cannot share an array -- but two channels
      // that agree on it can, which is where the saving comes from. The shader
      // never has to know: it samples whichever bucket the ref names.
      QRhiTexture::Flags flags{};
 
      // The same repeat, trilinear sampler for every bucket; per-texture wrap
      // modes are applied by the shader from scene_material_wrap.
      QRhiSampler* sampler{};        // created on first allocation; owned
 
      // Dedup: texture_source pointer → layer index in this bucket's
      // `array`. Rebuilt on every materials change, but a source that is
      // still present keeps its layer, so only new sources are uploaded.
      ossia::flat_map<const ossia::texture_source*, int> layerMap;
      // The source occupying each layer, null for a free one. Holding it
      // keeps its address from being reused by another source while the map
      // still names it.
      std::vector<std::shared_ptr<const ossia::texture_source>> layerSources;
      // Per layer: its alpha is coverage (the base colour of a MASK or BLEND
      // material), so it was uploaded colour-bled. Kept so a re-upload after
      // reallocation bleeds it again, and so a source whose role changes is
      // uploaded again.
      std::vector<char> layerCoverage;
      // Some layer's level 0 changed since the chain was last generated.
      bool mipsDirty{false};
    };
 
    // One bucket per distinct (format, pixelSize, colourspace), up to
    // kMaxBuckets.
    std::vector<Bucket> buckets;
 
    // Bumped whenever a bucket's array or sampler is replaced. Every
    // preprocessor sharing the pool republishes its texture auxiliaries on a
    // change, not only the one whose rebuild caused it: the others' consumers
    // still hold the replaced QRhiTexture*.
    uint64_t generation{0};
 
    // Monotonic counts of pool rebuilds and of layer uploads into bucket
    // arrays, re-uploads after a reallocation included. Nothing in the engine
    // reads them: they make the pool's incremental behaviour observable,
    // since an upload into an existing array changes no pointer.
    uint64_t rebuilds{0};
    uint64_t layerUploads{0};
 
    // The static sources each preprocessor's scene references, keyed by that
    // preprocessor. A rebuild keeps the layers the other preprocessors claim
    // instead of freeing every source its own scene lacks.
    ossia::flat_map<const void*, std::vector<const ossia::texture_source*>> claims;
 
    // Dynamic (runtime-GPU) slot map, keyed by QRhiResource::globalResourceId()
    // rather than the raw QRhiTexture* -- the allocator recycles freed pointer
    // values, and qrhivulkan.cpp documents the same hazard for SRB tracking, so a
    // recycled address must look like a fresh resource here too.
    //
    // Slots are recycled by LRU: when the map is full and a new id arrives, the slot
    // with the smallest dynamicSlotLastUse is evicted. lastUse is bumped on every
    // access, so the evicted slot is the one no live material references.
    ossia::flat_map<quint64, int> dynamicSlotMap;
    std::vector<QRhiTexture*>     dynamicTextures;       // slot idx → texture
    // slot idx → globalResourceId of dynamicTextures[idx], read while it was
    // live: changes are detected without touching a possibly freed texture.
    std::vector<quint64>          dynamicTextureIds;
    std::vector<uint64_t>         dynamicSlotLastUse;    // slot idx → access counter at last lookup
    uint64_t                      dynamicSlotCounter{0}; // monotonic, bumped on each resolve
    // Value of dynamicSlotCounter at the previous sweepStaleDynamicTextureSlots()
    // pass. A slot whose dynamicSlotLastUse is <= this was not re-resolved by any
    // live material since, so its stored QRhiTexture* is orphaned and must be
    // cleared before it can be bound.
    uint64_t                      dynamicSweepCheckpoint{0};
 
    // Single-bucket accessors for callers that do not loop over buckets[].
    // Return null / 0 when no bucket has been allocated yet.
    QRhiTexture* primaryArray() const noexcept
    {
      return buckets.empty() ? nullptr : buckets[0].array;
    }
    int primaryLayers() const noexcept
    {
      return buckets.empty() ? 0 : buckets[0].layers;
    }
 
    // Access or lazily create bucket 0 with an owned (format, size).
    // Kept for init-time fallback allocation only — production code
    // goes through findOrCreateBucket() which selects the right bucket
    // for the texture's actual (format, size).
    Bucket& ensurePrimary(QRhiTexture::Format fmt, QSize sz,
                          QRhiTexture::Flags flags = {})
    {
      if(buckets.empty())
        buckets.emplace_back();
      auto& b = buckets[0];
      b.format = fmt;
      b.pixelSize = sz;
      b.flags = flags;
      return b;
    }
 
    // Find a bucket matching (fmt, sz, flags), creating one if none matches and
    // kMaxBuckets is not reached. Returns {bucket_index, pointer}, or
    // {-1, nullptr} on overflow. Bucket identity is the exact tuple, no
    // rounding.
    std::pair<int, Bucket*>
    findOrCreateBucket(QRhiTexture::Format fmt, QSize sz, QRhiTexture::Flags flags)
    {
      for(std::size_t i = 0; i < buckets.size(); ++i)
      {
        if(buckets[i].format == fmt && buckets[i].pixelSize == sz
           && buckets[i].flags == flags)
          return {(int)i, &buckets[i]};
      }
      if((int)buckets.size() >= kMaxBuckets)
        return {-1, nullptr};
      buckets.emplace_back();
      auto& b = buckets.back();
      b.format = fmt;
      b.pixelSize = sz;
      b.flags = flags;
      return {(int)buckets.size() - 1, &b};
    }
  };
 
  struct PrimitiveCloudUploads
  {
    uint64_t rows{0};
    uint64_t meta{0};
  } primitiveCloudUploads;
 
  TextureChannelState& textureChannel(TextureChannel) noexcept
  {
    return m_texturePool;
  }
  const TextureChannelState& textureChannel(TextureChannel) const noexcept
  {
    return m_texturePool;
  }
 
  TextureChannelState& texturePool() noexcept { return m_texturePool; }
  const TextureChannelState& texturePool() const noexcept { return m_texturePool; }
 
  static constexpr const char* kSharedArrayName = "materialArray";
  static constexpr const char* kSharedDynName = "materialDyn";
 
  static const char* textureChannelArrayName(TextureChannel ch) noexcept;
 
  static const char* textureChannelDynBaseName(TextureChannel ch) noexcept;
 
  static QRhiTexture::Flags textureChannelFlags(TextureChannel ch) noexcept;
 
  int resolveDynamicSlot(TextureChannel channel, void* native_handle) noexcept;
 
  // ─── Mesh arena manager ───────────────────────────────────────────
  //
  // Per-mesh slab allocator over the attribute streams of the MDI concatenated
  // geometry -- positions, normals, texcoords, tangents, colors, texcoords1,
  // indices -- each a single growth-capped QRhiBuffer.
  //
  // Indirect-draw correctness invariant: one baseVertex is applied to ALL vertex
  // bindings (VkDrawIndexedIndirectCommand::vertexOffset), so per-mesh byte
  // offsets must satisfy pos/nrm/tan = baseVertex * 16 and uv = baseVertex * 8.
  // Independent per-stream allocators cannot guarantee that once alloc/free
  // traffic fragments the streams: they pick free blocks from different size bins
  // and the offsets diverge, so the vertex shader reads attributes from the wrong
  // slab.
  //
  // Hence TWO shared allocators: m_vertexAllocator in vertex units and
  // m_indexAllocator in index units, both capped at 8M slots. Each slab carries
  // one vertex_slot and one index_slot, and per-stream byte offsets derive as
  // vertex_slot.offset * stride and index_slot.offset * 4, so lockstep is
  // structural. Indirect draw takes baseVertex = vertex_slot.offset and
  // firstIndex = index_slot.offset.
  //
  // Cache: a stable_id hit reuses the slab and skips the upload; a miss allocates.
  // The sweep frees slabs unseen for `grace` frames.
  //
  // Backing buffers are pointer-stable for the registry's lifetime, so downstream
  // bindings resolve once: 128 MB each for positions, normals, tangents and
  // colors, 64 MB per texcoord stream, 32 MB indices.
 
  enum class MeshStream : uint8_t
  {
    Positions  = 0,
    Normals    = 1,
    Texcoords  = 2,  // TEXCOORD_0 (primary UV).
    Tangents   = 3,
    Colors     = 4,  // glTF COLOR_0, vec4 (vec3 sources padded with alpha=1).
    Texcoords1 = 5,  // glTF TEXCOORD_1 (lightmap / secondary UV).
    Indices    = 6,
    Count_     = 7
  };
 
  // Bytes per element per stream, matching the MDI output layout the rasterizer
  // presets consume: positions and normals are vec3 padded to vec4 for std430,
  // tangents and colors vec4 (vec3 sources padded with alpha=1), texcoords vec2,
  // indices uint32.
  static constexpr uint32_t kMeshStride[(std::size_t)MeshStream::Count_]
      = {16, 16, 8, 16, 16, 8, 4};
 
  // Capacity reserved per stream at init: 128 MB for positions, normals, tangents
  // and colors (16 B stride, 8M verts), 64 MB for each texcoord stream (8 B),
  // 32 MB for indices (4 B). A scene exceeding these gets a sentinel allocation
  // from allocate() and the caller skips the mesh.
  static constexpr uint32_t kMeshCapBytes[(std::size_t)MeshStream::Count_]
      = {
          128u * 1024u * 1024u,
          128u * 1024u * 1024u,
           64u * 1024u * 1024u,
          128u * 1024u * 1024u,
          128u * 1024u * 1024u,  // colors
           64u * 1024u * 1024u,  // texcoords1
           32u * 1024u * 1024u,
  };
 
  struct MeshSlab
  {
    uint64_t stable_id{};
    OffsetAllocator::Allocation vertex_slot{};  // offset/size in vertex units
    OffsetAllocator::Allocation index_slot{};   // offset/size in index  units
    uint32_t vertex_count{};
    uint32_t index_count{};
    uint32_t last_seen_frame{};
    bool freshly_allocated{};  // true on the frame the slab was created
  };
 
  MeshSlab* acquireMeshSlab(
      uint64_t stable_id,
      uint32_t vertex_count,
      uint32_t index_count,
      uint32_t current_frame) noexcept;
 
  void markMeshSlabSeen(uint64_t stable_id, uint32_t current_frame) noexcept;
 
  void sweepMeshSlabs(uint32_t current_frame, uint32_t grace = 2) noexcept;
 
  void sweepStaleDynamicTextureSlots() noexcept;
 
  void forgetDynamicTexture(const QRhiTexture* tex) noexcept;
 
  void drainExpiredPendingReleases(
      uint32_t current_frame, uint32_t grace = 2) noexcept;
 
  void releaseMeshSlab(uint64_t stable_id, uint32_t current_frame) noexcept;
 
  uint32_t meshSlabOffsetBytes(
      const MeshSlab& slab, MeshStream stream) const noexcept;
 
  QRhiBuffer* meshStreamBuffer(MeshStream s) const noexcept;
 
  void uploadMeshStream(
      QRhiResourceUpdateBatch& res, const MeshSlab& slab,
      MeshStream s, const void* data, uint32_t size) noexcept;
 
  class MeshUploadBatch
  {
  public:
    explicit MeshUploadBatch(GpuResourceRegistry* r) noexcept
        : m_registry{r}
    {
      if(m_registry)
        ++m_registry->m_meshUploadBatchDepth;
    }
    ~MeshUploadBatch()
    {
      if(m_registry && --m_registry->m_meshUploadBatchDepth == 0)
        m_registry->flushMeshUploads();
    }
    MeshUploadBatch(const MeshUploadBatch&) = delete;
    MeshUploadBatch& operator=(const MeshUploadBatch&) = delete;
 
  private:
    GpuResourceRegistry* m_registry{};
  };
 
  uint32_t meshStreamUsedBytes(MeshStream s) const noexcept;
  uint32_t meshStreamFreeBytes(MeshStream s) const noexcept;
 
private:
  struct ArenaState
  {
    QRhiBuffer* buffer{};
    uint32_t slot_stride{0};   // bytes per slot (arena layout is a packed
                               // std430-compatible array of this stride)
    uint32_t slot_count{0};    // total slots (capacity_bytes = stride × count)
    QRhiBuffer::UsageFlags usage{};
    QRhiBuffer::Type type{QRhiBuffer::Dynamic};
 
    // LIFO stack of free slot indices. Push on free, pop on allocate.
    // O(1) alloc / free, no fragmentation (every slot is the same size).
    std::vector<uint32_t> free_slots;
 
    // Per-slot generation, indexed by slot_index. Sized to slot_count
    // at init() and bumped on every allocate()/free() to that slot.
    // Consumers check the stamped generation in their gpu_slot_ref via
    // isLive().
    std::vector<uint32_t> slot_generations;
  };
 
  std::array<ArenaState, (std::size_t)Arena::Count_> m_arenas{};
 
  TextureChannelState m_texturePool{};
 
  // Per-stream backing buffers, one QRhiBuffer per attribute. Allocation is not
  // per-stream: a single m_vertexAllocator hands out vertex-unit slots that every
  // vertex stream interprets through its own stride, and m_indexAllocator
  // handles indices, keeping byte offsets in lockstep for baseVertex.
  struct MeshStreamState
  {
    QRhiBuffer* buffer{};
    uint32_t capacity_bytes{};
    QRhiBuffer::UsageFlags usage{};
  };
  std::array<MeshStreamState, (std::size_t)MeshStream::Count_> m_meshStreams{};
 
  // Mesh uploads held back by a MeshUploadBatch: one contiguous run per
  // stream, and the batch it goes to.
  struct PendingMeshUpload
  {
    QRhiResourceUpdateBatch* batch{};
    uint32_t offset{};
    std::vector<std::byte> bytes;
  };
  std::array<PendingMeshUpload, (std::size_t)MeshStream::Count_> m_pendingMeshUploads{};
  int m_meshUploadBatchDepth{};
  void flushMeshUpload(MeshStream s) noexcept;
  void flushMeshUploads() noexcept;
 
  // Shared vertex / index allocators (slot units, not bytes).
  // capacity_slots = min(stream_capacity_bytes / stream_stride) across
  // the vertex streams = 8M for the default sizes; index pool
  // capacity = 8M slots.
  std::unique_ptr<OffsetAllocator::Allocator> m_vertexAllocator;
  std::unique_ptr<OffsetAllocator::Allocator> m_indexAllocator;
  uint32_t m_vertexSlotsCapacity{};
  uint32_t m_indexSlotsCapacity{};
  uint32_t m_vertexSlotsUsed{};
  uint32_t m_indexSlotsUsed{};
 
  ossia::hash_map<uint64_t, MeshSlab> m_meshSlabs;
 
  // Slabs whose `released_frame` is set are waiting out the grace
  // period before their OffsetAllocator allocations return to the
  // free list. Prevents use-after-free when an in-flight draw still
  // references the old offset.
  struct PendingRelease
  {
    uint64_t stable_id{};
    uint32_t released_frame{};
    OffsetAllocator::Allocation vertex_slot{};
    OffsetAllocator::Allocation index_slot{};
  };
  std::vector<PendingRelease> m_pendingReleases;
 
  QRhi* m_rhi{};
 
  // Set by seedDefaults() after writing the default-MaterialGPU bytes
  // into Material arena slot 0. Idempotent guard so repeated calls are
  // free.
  bool m_defaults_seeded{false};
};
 
} // namespace score::gfx