Provide a set of MTLTextures from a heap

MetalTextureHeap.swift

//
// Created by Neil Clayton on 19/10/23.
// Copyright (c) 2023 Neil Clayton. All rights reserved.
//

import Foundation
import MetalKit

public class MetalTextureHeap {
    public private(set) var currentAllocation = 0
    public private(set) var maxAllocation = 0
    private var textureHeap: MTLHeap!
    private var eachTextureHasSize: Int = 0

    public enum Errors: Error {
        case noDefaultMetalDevice
        case failedToCreateHeap
        case failedToReCreateHeap
        case failedToGetTextureAfterRecreatingHeap
    }

    private var textureDescriptor: MTLTextureDescriptor

    init?(size: CGSize, maxTextures: Int) throws {
        assert(maxTextures > 2, "maxTextures must be > 2")
        maxAllocation = maxTextures
        textureDescriptor = MTLTextureDescriptor.texture2DDescriptor(pixelFormat: .rgba8Unorm, width: Int(size.width), height: Int(size.height), mipmapped: false)
        // .private means the texture is GPU only. No sync blit will be done during rendering.
        textureDescriptor.storageMode = .private
        textureDescriptor.usage = [.shaderRead, .renderTarget]

        textureHeap = try allocateTextureHeap(numberOfTextures: 2)
    }

    public func allocateTextureHeap(numberOfTextures: Int) throws -> MTLHeap {
        let heapDescriptor = MTLHeapDescriptor()
        heapDescriptor.storageMode = .private

        if let defaultDevice = MTLCreateSystemDefaultDevice() {
            var sizeAndAlign = defaultDevice.heapTextureSizeAndAlign(descriptor: textureDescriptor)

            sizeAndAlign.size = Self.alignUp(sizeAndAlign.size, sizeAndAlign.align)
            heapDescriptor.size += sizeAndAlign.size * numberOfTextures
            eachTextureHasSize = sizeAndAlign.size

            guard let newHeap = defaultDevice.makeHeap(descriptor: heapDescriptor) else {
                throw Errors.failedToCreateHeap
            }
            currentAllocation = numberOfTextures
            return newHeap
        } else {
            throw Errors.noDefaultMetalDevice
        }
    }

    public var usedSize: Int {
        textureHeap.usedSize / eachTextureHasSize
    }

    static func alignUp(_ inSize: Int, _ align: Int) -> Int {
        // Assert if align is not a power of 2
        assert((align - 1) & align == 0)

        let alignmentMask = align - 1

        return (inSize + alignmentMask) & ~alignmentMask
    }

    func newTexture(descriptor: MTLTextureDescriptor) throws -> MTLTexture? {
        let texture = textureHeap.makeTexture(descriptor: descriptor)
        if texture == nil, currentAllocation < maxAllocation {
            let newAllocation = currentAllocation + 1
            // try allocating more, if we can
            do {
                textureHeap = try allocateTextureHeap(numberOfTextures: newAllocation)
                let texture = textureHeap.makeTexture(descriptor: descriptor)
                if texture == nil {
                    throw Errors.failedToGetTextureAfterRecreatingHeap
                }
            } catch {
                throw Errors.failedToReCreateHeap
            }
        }
        return texture
    }

    func returnTexture(texture: MTLTexture) {
        // Do nothing. You just need to deallocate it to have it returned
//        texture.makeAliasable()
    }
}