How do I use MetalKit texture loader with Metal heaps?

I have a set of Metal textures that are stored in an Xcode Assets Catalog as Texture Sets. I'm loading these using MTKTextureLoader.newTexture(name:scaleFactor:bundle:options).

I then use a MTLArgumentEncoder to encode all of the textures into a Metal 2 argument buffer.

This works great. However, the Introducing Metal 2 WWDC 2017 session recommends combining argument buffers with resource heaps for even better performance, and I'm quite keen to try this. According to the Argument Buffer documentation, instead of having to call MTLRenderCommandEncoder.useResource on each texture in the argument buffer, you just call useHeap on the heap that the textures were allocated from.

However, I haven't found a straightforward way to use MTKTextureLoader together with MTLHeap. It doesn't seem to have a loading option to allocate the texture from a heap.

I'm guessing that the approach would be:

• load the textures with MTKTextureLoader
• reverse-engineer a set of MTLTextureDescriptor objects for each texture
• use the texture descriptors to create an appropriately sized MTLHeap
• assign a new set of textures from the MTLHeap
• use some method to copy the textures across, perhaps replaceBytes or maybe even a MTLBlitCommandEncoder
• deallocate the original textures loaded with the MTKTextureLoader

It seems like a fairly long-winded approach, and i've not seen any examples of this, so I thought I'd ask here first in case I'm missing something obvious.

Should I abandon MTKTextureLoader, and search out some pre-MetalKit art on loading textures from asset catalogs?

I'm using Swift, but happy to accept Objective-C answers.

Well, the method I outlined above seems to work. As predicted, it's pretty long-winded. I'd be very interested to know if anyone has anything more elegant.

enum MetalError: Error {
    case anErrorOccured
}

extension MTLTexture {
    var descriptor: MTLTextureDescriptor {
        let descriptor = MTLTextureDescriptor()
        descriptor.width = width
        descriptor.height = height
        descriptor.depth = depth
        descriptor.textureType = textureType
        descriptor.cpuCacheMode = cpuCacheMode
        descriptor.storageMode = storageMode
        descriptor.pixelFormat = pixelFormat
        descriptor.arrayLength = arrayLength
        descriptor.mipmapLevelCount = mipmapLevelCount
        descriptor.sampleCount = sampleCount
        descriptor.usage = usage
        return descriptor
    }

    var size: MTLSize {
        return MTLSize(width: width, height: height, depth: depth)
    }
}

extension MTKTextureLoader {
    func newHeap(withTexturesNamed names: [String], queue: MTLCommandQueue, scaleFactor: CGFloat, bundle: Bundle?, options: [MTKTextureLoader.Option : Any]?, onCompletion: (([MTLTexture]) -> Void)?) throws -> MTLHeap {
        let device = queue.device
        let sourceTextures = try names.map { name in
            return try newTexture(name: name, scaleFactor: scaleFactor, bundle: bundle, options: options)
        }
        let storageMode: MTLStorageMode = .private
        let descriptors: [MTLTextureDescriptor] = sourceTextures.map { source in
            let desc = source.descriptor
            desc.storageMode = storageMode
            return desc
        }
        let sizeAligns = descriptors.map { device.heapTextureSizeAndAlign(descriptor: $0) }
        let heapDescriptor = MTLHeapDescriptor()
        heapDescriptor.size = sizeAligns.reduce(0) { $0 + $1.size }
        heapDescriptor.cpuCacheMode = descriptors[0].cpuCacheMode
        heapDescriptor.storageMode = storageMode
        guard let heap = device.makeHeap(descriptor: heapDescriptor),
            let buffer = queue.makeCommandBuffer(),
            let blit = buffer.makeBlitCommandEncoder()
            else {
            throw MetalError.anErrorOccured
        }
        let destTextures = descriptors.map { descriptor in
            return heap.makeTexture(descriptor: descriptor)
        }
        let origin = MTLOrigin()
        zip(sourceTextures, destTextures).forEach {(source, dest) in
            blit.copy(from: source, sourceSlice: 0, sourceLevel: 0, sourceOrigin: origin, sourceSize: source.size, to: dest, destinationSlice: 0, destinationLevel: 0, destinationOrigin: origin)
            blit.generateMipmaps(for: dest)
        }
        blit.endEncoding()
        buffer.addCompletedHandler { _ in
            onCompletion?(destTextures)
        }
        buffer.commit()
        return heap
    }
}