MetalCVPixelBuffer.swift

MetalCVPixelBuffer.swift

import CoreVideo
import Alloy
import Foundation
import Accelerate

enum MetalCVPixelBuffer {
    static func transform(texture input: MTLTexture,
                          transformer: @escaping (CVPixelBuffer) throws -> CVPixelBuffer,
                          cache: CVMetalTextureCache) throws -> MTLTexture {
        guard let inputPixelBuffer = input.pixelBuffer else {
            throw StyleError.pixelBufferCreationFailed
        }

        let outputPixelBuffer = try transformer(inputPixelBuffer)

        let outputPixelFormat = CVPixelBufferGetPixelFormatType(outputPixelBuffer)

        guard let metalCompatibleBuffer = outputPixelBuffer.copyToMetalCompatible(pixelFormat: outputPixelFormat) else {
            throw StyleError.pixelBufferCreationFailed
        }

        guard let texturePixelFormat = MTLPixelFormat(cvPixelFormat: outputPixelFormat) else {
            throw StyleError.unsupportedPixelFormat
        }

        guard let output = metalCompatibleBuffer.metalTexture(using: cache, pixelFormat: texturePixelFormat) else {
            throw StyleError.textureCreationFailed
        }

        return output
    }
}

private extension MTLPixelFormat {
    init?(cvPixelFormat: OSType) {
        switch cvPixelFormat {
        case kCVPixelFormatType_32BGRA:
            self = .bgra8Unorm
        case kCVPixelFormatType_420YpCbCr8BiPlanarVideoRange,
             kCVPixelFormatType_420YpCbCr8BiPlanarFullRange:
            self = .bgra8Unorm
        case kCVPixelFormatType_32RGBA:
            self = .rgba8Unorm
        case kCVPixelFormatType_DisparityFloat16,
             kCVPixelFormatType_DepthFloat16,
             kCVPixelFormatType_OneComponent16Half:
            self = .r16Float
        case kCVPixelFormatType_DisparityFloat32,
             kCVPixelFormatType_DepthFloat32,
             kCVPixelFormatType_OneComponent32Float:
            self = .r32Float
        case kCVPixelFormatType_OneComponent8:
            self = .r8Unorm
        default:
            return nil
        }
    }
}

private func metalCompatiblityAttributes() -> [String: Any] {
    let attributes: [String: Any] = [
        String(kCVPixelBufferMetalCompatibilityKey): true,
        String(kCVPixelBufferOpenGLCompatibilityKey): true,
        String(kCVPixelBufferIOSurfacePropertiesKey): [
            String(kCVPixelBufferIOSurfaceOpenGLESTextureCompatibilityKey): true,
            String(kCVPixelBufferIOSurfaceOpenGLESFBOCompatibilityKey): true,
            String(kCVPixelBufferIOSurfaceCoreAnimationCompatibilityKey): true
        ]
    ]
    return attributes
}

private extension CVPixelBuffer {
    func copyToMetalCompatible(pixelFormat: OSType) -> CVPixelBuffer? {
        return deepCopy(withAttributes: metalCompatiblityAttributes(), pixelFormat: pixelFormat)
    }

    func deepCopy(withAttributes attributes: [String: Any] = [:], pixelFormat: OSType) -> CVPixelBuffer? {
        let srcPixelBuffer = self
        let srcFlags: CVPixelBufferLockFlags = .readOnly
        guard kCVReturnSuccess == CVPixelBufferLockBaseAddress(srcPixelBuffer, srcFlags) else {
            return nil
        }
        defer { CVPixelBufferUnlockBaseAddress(srcPixelBuffer, srcFlags) }

        var combinedAttributes: [String: Any] = [:]

        // Copy attachment attributes.
        if let attachments = CVBufferGetAttachments(srcPixelBuffer, .shouldPropagate) as? [String: Any] {
            for (key, value) in attachments {
                combinedAttributes[key] = value
            }
        }

        // Add user attributes.
        combinedAttributes = combinedAttributes.merging(attributes) { $1 }

        var maybePixelBuffer: CVPixelBuffer?
        let status = CVPixelBufferCreate(kCFAllocatorDefault,
                                         CVPixelBufferGetWidth(srcPixelBuffer),
                                         CVPixelBufferGetHeight(srcPixelBuffer),
                                         pixelFormat,
                                         combinedAttributes as CFDictionary,
                                         &maybePixelBuffer)

        guard status == kCVReturnSuccess, let dstPixelBuffer = maybePixelBuffer else {
            return nil
        }

        let dstFlags = CVPixelBufferLockFlags(rawValue: 0)
        guard kCVReturnSuccess == CVPixelBufferLockBaseAddress(dstPixelBuffer, dstFlags) else {
            return nil
        }
        defer { CVPixelBufferUnlockBaseAddress(dstPixelBuffer, dstFlags) }

        for plane in 0...max(0, CVPixelBufferGetPlaneCount(srcPixelBuffer) - 1) {
            if let srcAddr = CVPixelBufferGetBaseAddressOfPlane(srcPixelBuffer, plane),
               let dstAddr = CVPixelBufferGetBaseAddressOfPlane(dstPixelBuffer, plane) {
                let srcBytesPerRow = CVPixelBufferGetBytesPerRowOfPlane(srcPixelBuffer, plane)
                let dstBytesPerRow = CVPixelBufferGetBytesPerRowOfPlane(dstPixelBuffer, plane)

                for h in 0..<CVPixelBufferGetHeightOfPlane(srcPixelBuffer, plane) {
                    let srcPtr = srcAddr.advanced(by: h*srcBytesPerRow)
                    let dstPtr = dstAddr.advanced(by: h*dstBytesPerRow)
                    dstPtr.copyMemory(from: srcPtr, byteCount: srcBytesPerRow)
                }
            }
        }
        return dstPixelBuffer
    }
}