Render sampleBuffer to MTKView with coordination convert support.

Render.swift

extension CGAffineTransform {
    func convertToSIMD() -> simd_float3x3 {
        return simd_float3x3(
            SIMD3(Float(self.a),  Float(self.b),  0),
            SIMD3(Float(self.c),  Float(self.d),  0),
            SIMD3(Float(self.tx), Float(self.ty), 1)
        )
    }
}

// View <- Drawable <- SampleBuffer
class Renderer: NSObject, MTKViewDelegate {
    enum ContentMode : Int, Sendable {
        case scaleAspectFit = 1
        case scaleAspectFill = 2
    }

    let device: MTLDevice?
    let commandQueue: MTLCommandQueue?

    var computePipelineState: MTLComputePipelineState?

    var sampleBufferTextureCache: CVMetalTextureCache?
    var sampleBuffer: CMSampleBuffer?

    // transfer the drawable's texture gid to buffer's texture coord.
    @Published var textureTransform: CGAffineTransform = .identity

    /// only alpha channel. range 0 - 1.
    /// 0: the sample buffer color
    /// 1: the other color
    var maskTexture: (any MTLTexture)?

    // rect in view bound coordinate.
    // maskRectInViewBounds -> drawable size -> maskTexture.
    var maskRectInViewBounds: CGRect? {
        didSet {
            updateMask()
        }
    }

    // this value will clamp between 0 - maskRectInViewBounds / 2
    var maskCornerRadiusInViewBounds: CGFloat = .zero {
        didSet {
            updateMask()
        }
    }

    var maskImage: UIImage? = nil {
        didSet {
            updateMask()
        }
    }

    var drawableScale: CGFloat = .zero

    var drawableSize: CGSize = .zero {
        didSet {
            if oldValue != drawableSize {
                updateTransform()
                updateMask()
            }
        }
    }

    var sampleBufferAngle: CGFloat = 0 {
        didSet { if oldValue != sampleBufferAngle { updateTransform() }}
    }

    var sampleBufferSize: CGSize = .zero {
        didSet { if oldValue != sampleBufferSize { updateTransform() }}
    }

    var contentMode: ContentMode = .scaleAspectFill {
        didSet { if oldValue != contentMode { updateTransform() }}
    }

    init(view: MTKView) {
        self.device = view.device
        self.drawableScale = view.contentScaleFactor
        self.commandQueue = view.device?.makeCommandQueue()
        self.drawableSize = view.drawableSize
        super.init()

//        view.autoResizeDrawable = false
        view.framebufferOnly = false
        view.delegate = self

        CVMetalTextureCacheCreate(kCFAllocatorDefault, nil, device!, nil, &sampleBufferTextureCache)
        setupComputePipeline()
    }

    func setupComputePipeline() {
        guard let device = device else { return }

        let library = device.makeDefaultLibrary()
        let kernelFunction = library?.makeFunction(name: "gaussianBlur")

        do {
            computePipelineState = try device.makeComputePipelineState(function: kernelFunction!)
        } catch {
            print("Failed to create compute pipeline state: \(error)")
        }
    }

    func mtkView(_ view: MTKView, drawableSizeWillChange size: CGSize) {
        drawableSize = size
    }

    func draw(in view: MTKView) {
        guard
            let device,
            let sampleBuffer,
            let drawable = view.currentDrawable,
            let commandQueue,
            let sampleBufferTextureCache,
            let pipelineState = computePipelineState
        else {
            return
        }
        view.contentScaleFactor = self.drawableScale // contentScaleFactor may change unexpectly.

        assert(drawableSize == view.drawableSize, "Drawable Size Has Changed. Current \(view.drawableSize), but Render still \(drawableSize).")
        assert(view.bounds.width * self.drawableScale == self.drawableSize.width, "Bounds(\(view.bounds)), Scale(\(drawableScale)), Drawable Size(\(drawableSize)) are not match.")

        // 1. convert input sample buffer to input texture
        var inputTexture: CVMetalTexture?
        do {
            let imageBuffer = CMSampleBufferGetImageBuffer(sampleBuffer)!

            let imageWidth = CVPixelBufferGetWidth(imageBuffer)
            let imageHeight = CVPixelBufferGetHeight(imageBuffer)

            CVMetalTextureCacheCreateTextureFromImage(kCFAllocatorDefault, sampleBufferTextureCache,
                                                      imageBuffer, nil, .bgra8Unorm,
                                                      imageWidth, imageHeight, 0, &inputTexture)
        }

        guard let inputMTLTexture = CVMetalTextureGetTexture(inputTexture!) else { return }

        // 2. create current metal render command
        let commandBuffer = commandQueue.makeCommandBuffer()!

        // 3. Create Gaussian Blur Texture by using MPSImageGaussianBlur
        let blurTextureDesc = MTLTextureDescriptor.texture2DDescriptor(pixelFormat: inputMTLTexture.pixelFormat,
                                                                       width: inputMTLTexture.width,
                                                                       height: inputMTLTexture.height,
                                                                       mipmapped: false)
        blurTextureDesc.usage = [.shaderRead, .shaderWrite]
        guard let blurTexture = device.makeTexture(descriptor: blurTextureDesc) else { return }

        let blurKernel = MPSImageGaussianBlur(device: device, sigma: 20.0)
        blurKernel.edgeMode = .clamp
        blurKernel.encode(commandBuffer: commandBuffer,
                          sourceTexture: inputMTLTexture,
                          destinationTexture: blurTexture)

        // 4. create texture transform to simd format so that we can pass into metal function
        var transform = textureTransform.convertToSIMD()

        // 5. setup metal function
        let computeEncoder = commandBuffer.makeComputeCommandEncoder()!
        computeEncoder.setComputePipelineState(pipelineState)
        computeEncoder.setTexture(drawable.texture, index: 0) // render texture into drawable
        computeEncoder.setTexture(inputMTLTexture, index: 1)
        computeEncoder.setTexture(blurTexture, index: 2)
        computeEncoder.setTexture(maskTexture, index: 3)
        computeEncoder.setBytes(&transform, length: MemoryLayout<simd_float3x3>.stride, index: 0)

        // https://developer.apple.com/documentation/metal/calculating-threadgroup-and-grid-sizes#Calculate-Threads-per-Threadgroup
        let tWidth = pipelineState.threadExecutionWidth
        let tHeight = pipelineState.maxTotalThreadsPerThreadgroup / tWidth
        let threadsPerGroup = MTLSizeMake(tWidth, tHeight, 1)
        let dWidth = Int(view.drawableSize.width)
        let dHeight = Int(view.drawableSize.height)
        let threadsPerGrid = MTLSizeMake(dWidth, dHeight, 1)

        if device.supportsFamily(.apple4) {
            computeEncoder.dispatchThreads(threadsPerGrid, threadsPerThreadgroup: threadsPerGroup)
        } else {
            let threadgroupsPerGrid = MTLSize(width:  (threadsPerGrid.width  + threadsPerGroup.width  - 1) / threadsPerGroup.width,
                                              height: (threadsPerGrid.height + threadsPerGroup.height - 1) / threadsPerGroup.height,
                                              depth:  1)
            computeEncoder.dispatchThreadgroups(threadgroupsPerGrid, threadsPerThreadgroup: threadsPerGroup)
        }
        computeEncoder.endEncoding()

        // 6. draw the result
        commandBuffer.present(drawable)
        commandBuffer.commit()
    }
}


extension Renderer {
    var sampleBufferSizeWithAngleTransfermed: CGSize {
        let rect = CGRect(origin: .zero, size: self.sampleBufferSize)
        let rotatedSize = rect.applying(.identity.rotated(by: self.sampleBufferAngle / 180 * .pi)).size
        return CGSize(width: Int(round(rotatedSize.width)), height: Int(round(rotatedSize.height)))
    }


    func updateTransform() {
        guard sampleBufferSize != .zero, drawableSize != .zero else {
            textureTransform = .identity
            return
        }

        let sampleBufferSize = sampleBufferSizeWithAngleTransfermed

        let widthScale  = sampleBufferSize.width / drawableSize.width
        let heightScale = sampleBufferSize.height / drawableSize.height

        let scale = switch contentMode {
        case .scaleAspectFit:
            max(widthScale, heightScale)
        case .scaleAspectFill:
            min(widthScale, heightScale)
        }

        let deltaX = (sampleBufferSize.width - drawableSize.width * scale) / 2
        let deltaY = (sampleBufferSize.height - drawableSize.height * scale) / 2

        textureTransform = CGAffineTransform.identity
            .scaledBy(x: scale, y: scale)
            .translatedBy(x: deltaX, y: deltaY)
    }

    /// Convert the point from the view's coordinate to the sample buffer's normalized coordinate.
    ///
    /// This function will take drawableScale, drawableSize, sample angle, sample size into consideration.
    func convert(fromViewPoint point: CGPoint) -> CGPoint {
        let transformed = point
            .applying(.identity.scaledBy(x: drawableScale,
                                         y: drawableScale))
            .applying(textureTransform)

        return .init(x: transformed.x / sampleBufferSizeWithAngleTransfermed.width,
                     y: transformed.y / sampleBufferSizeWithAngleTransfermed.height)
    }

    /// Convert the point from sample buffer's normalized coordinate to the view's coordinate.
    ///
    /// This function will take drawableScale, drawableSize, sample angle, sample size into consideration.
    func convert(fromSampleBufferPoint point: CGPoint) -> CGPoint {
        let pixelPoint: CGPoint = .init(x: point.x * sampleBufferSizeWithAngleTransfermed.width,
                                        y: point.y * sampleBufferSizeWithAngleTransfermed.height)

        return pixelPoint
            .applying(textureTransform.inverted())
            .applying(.identity.scaledBy(x: 1 / drawableScale,
                                         y: 1 / drawableScale))
    }

    func convert(fromViewRect rect: CGRect) -> CGRect {
        let min = convert(fromViewPoint: rect.origin)
        let max = convert(fromViewPoint: CGPoint(x: rect.maxX, y: rect.maxY))

        return CGRect(origin: min, size: CGSize(width: max.x - min.x, height: max.y - min.y))
    }
}


extension Renderer {
    func updateMask() {
        guard drawableSize != .zero, let device else {
            return
        }

        let width = Int(drawableSize.width)
        let height = Int(drawableSize.height)

        let textureDescriptor = MTLTextureDescriptor.texture2DDescriptor(
            pixelFormat: .r8Unorm,
            width: width,
            height: height,
            mipmapped: false
        )

        guard let texture = device.makeTexture(descriptor: textureDescriptor) else {
            return
        }

        let buffer:[UInt8]? = if let maskImage {
            buildImageMask(image: maskImage)
        } else if let rect = maskRectInViewBounds {
            buildRoundedRectangleMask(roundedRect: rect.applying(.identity.scaledBy(x: drawableScale, y: drawableScale)),
                                      cornerRadius: maskCornerRadiusInViewBounds * drawableScale)
        } else {
            nil
        }

        guard let buffer else {
            self.maskTexture = nil
            return
        }

        let region = MTLRegionMake2D(0, 0, width, height)
        texture.replace(region: region, mipmapLevel: 0, withBytes: buffer, bytesPerRow: width)

        self.maskTexture = texture
    }
}

extension Renderer {
    /// using MTLPixelFormat.r8Unorm
    static func buildMaskBufferInGraySpace(size: CGSize, drawOnContext: (CGContext, CGSize) -> Void) -> [UInt8]? {
        let start = DispatchTime.now().uptimeNanoseconds

        let width = Int(size.width)
        let height = Int(size.height)

        var buffer = [UInt8](repeating: 0, count: width * height)

        let colorSpace = CGColorSpaceCreateDeviceGray()
        guard let context = CGContext(
            data: &buffer,
            width: width,
            height: height,
            bitsPerComponent: 8,
            bytesPerRow: width,
            space: colorSpace,
            bitmapInfo: CGImageAlphaInfo.alphaOnly.rawValue
        ) else {
            return nil
        }

        context.translateBy(x: 0, y: CGFloat(height))
        context.scaleBy(x: 1, y: -1)

        drawOnContext(context, size)
        let end = DispatchTime.now().uptimeNanoseconds
        print("[*]: \((end - start) / 1_000_000)ms")
        return buffer
    }
}

extension Renderer {
    // about 5 ms
    func buildRoundedRectangleMask(roundedRect: CGRect, cornerRadius: CGFloat) -> [UInt8]? {
        let buffer = Self.buildMaskBufferInGraySpace(size: drawableSize) { context, size in
            let path = UIBezierPath(roundedRect: roundedRect, cornerRadius: cornerRadius)

            context.addPath(path.cgPath)
            context.setFillColor(gray: 1, alpha: 1) // change gray from 0 - 1
            context.fillPath()
        }
        return buffer
    }

    /// Create a mask from the image's alpha channel.
    ///
    /// This func spend about 70ms.
    ///
    /// Exmaple:
    /// ```Swift
    /// let image = UIGraphicsImageRenderer(size: size).image { rendererCtx in
    ///     let ctx = rendererCtx.cgContext
    ///
    ///     let scaledRect = rect.applying(.identity.scaledBy(x: scale, y: scale))
    ///     let scaledCornerRadius = cornerRadis * scale
    ///
    ///     let path = UIBezierPath(roundedRect: scaledRect, cornerRadius: scaledCornerRadius)
    ///     ctx.addPath(path.cgPath)
    ///
    ///     ctx.setFillColor(gray: 1, alpha: 0.5) // change gray from 0 - 1
    ///     ctx.fillPath()
    /// }
    /// ```
    func buildImageMask(image: UIImage) -> [UInt8]? {
        guard let cgImage = image.cgImage else {
            return nil
        }

        let buffer = Self.buildMaskBufferInGraySpace(size: drawableSize) { context, size in
            let rect = CGRect(origin: .zero, size: size)
            context.draw(cgImage, in: rect) // TODO: Scale the image
        }

        return buffer
    }
}