Rotate a CVPixelBuffer with ARGB pixels by 90 degrees using the Accelerate framework

CVPixelBuffer+Rotate.swift

//
//  CVPixelBuffer+Rotate.swift
//
//  Abstract:
//  This CVPixelBuffer extension provides a method that will rotate an
//  ARGB pixel buffer by a multiple of 90 degrees using the high performance
//  Accelerate framework provided on Apple platforms. This method is
//  designed to be compatible with Core Video APIs on iOS, which may require
//  the pixel buffer to be memory aligned and backed by an IOSurface.
//
//  Created by Richard Henry on 11/1/23.
//  Copyright (c) 2023 Telepath Inc.
//
//  This is free and unencumbered software released into the public domain.
//
//  Anyone is free to copy, modify, publish, use, compile, sell, or
//  distribute this software, either in source code form or as a compiled
//  binary, for any purpose, commercial or non-commercial, and by any
//  means.
//
//  In jurisdictions that recognize copyright laws, the author or authors
//  of this software dedicate any and all copyright interest in the
//  software to the public domain. We make this dedication for the benefit
//  of the public at large and to the detriment of our heirs and
//  successors. We intend this dedication to be an overt act of
//  relinquishment in perpetuity of all present and future rights to this
//  software under copyright law.
//
//  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
//  EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
//  MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
//  IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR
//  OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
//  ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
//  OTHER DEALINGS IN THE SOFTWARE.
//
//  For more information, please refer to <https://unlicense.org>
//

import Accelerate

/// The `CVPixelBuffer` rotation to apply.
public enum PixelBufferRotation {
    /// Rotate the pixel buffer by 90 degrees clockwise.
    case clockwise90Degrees
    /// Rotate the pixel buffer by 180 degrees clockwise.
    case clockwise180Degrees
    /// Rotate the pixel buffer by 270 degrees clockwise.
    case clockwise270Degrees

    /// The associated vImage rotation constant.
    var vImageRotationConstant: UInt8 {
        switch self {
        case .clockwise90Degrees:
            return UInt8(kRotate90DegreesClockwise)
        case .clockwise180Degrees:
            return UInt8(kRotate180DegreesClockwise)
        case .clockwise270Degrees:
            return UInt8(kRotate270DegreesClockwise)
        }
    }

    var isSizeRotated: Bool {
        switch self {
        case .clockwise90Degrees:
            return true
        case .clockwise180Degrees:
            return false
        case .clockwise270Degrees:
            return true
        }
    }
}

/// A `CVPixelBuffer` rotation extension error.
public enum PixelBufferRotationError: Error {
    case getSrcBaseAddressFailed
    case getDestBaseAddressFailed
    case mallocFailed
    case vImageRotate90Error(code: Int, planeIndex: Int?)
    case cvPixelBufferCreateError(Int32)
}

extension CVPixelBuffer {

    /**
     Rotate a `CVPixelBuffer` with an ARGB pixel format (e.g. 32ARGB or 32BGRA) by a multiple of 90 degrees.

     - Parameters:
        - rotation: The pixel buffer rotation to apply.
     - Returns: A new `CVPixelBuffer` with the rotation applied.
     - Throws: A `PixelBufferRotationError` if an error occurred.
     */
    func rotateARGBWithAccelerate(_ rotation: PixelBufferRotation) throws -> CVPixelBuffer {
        let format = CVPixelBufferGetPixelFormatType(self)

        CVPixelBufferLockBaseAddress(self, .readOnly)
        defer { CVPixelBufferUnlockBaseAddress(self, .readOnly) }

        guard let src = CVPixelBufferGetBaseAddress(self) else {
            throw PixelBufferRotationError.getSrcBaseAddressFailed
        }

        let srcWidth = CVPixelBufferGetWidth(self)
        let srcHeight = CVPixelBufferGetHeight(self)
        let srcRowBytes = CVPixelBufferGetBytesPerRow(self)

        var srcBuffer = vImage_Buffer(
            data: src,
            height: vImagePixelCount(srcHeight),
            width: vImagePixelCount(srcWidth),
            rowBytes: srcRowBytes
        )

        let destWidth = rotation.isSizeRotated ? srcHeight : srcWidth
        let destHeight = rotation.isSizeRotated ? srcWidth : srcHeight
        let unalignedDestRowBytes = destWidth * 4  // ARGB is 4 bytes per pixel

        guard let temp = malloc(destHeight * unalignedDestRowBytes) else {
            throw PixelBufferRotationError.mallocFailed
        }
        defer { free(temp) }

        var destBuffer = vImage_Buffer(
            data: temp,
            height: vImagePixelCount(destHeight),
            width: vImagePixelCount(destWidth),
            rowBytes: unalignedDestRowBytes
        )

        let noFlags = vImage_Flags(kvImageNoFlags)

        var backColor: UInt8 = 0
        let rotateStatus = vImageRotate90_ARGB8888(&srcBuffer, &destBuffer, rotation.vImageRotationConstant, &backColor, noFlags)
        guard rotateStatus == kvImageNoError else {
            throw PixelBufferRotationError.vImageRotate90Error(code: rotateStatus, planeIndex: nil)
        }
        
        var attributes = CVPixelBufferCopyCreationAttributes(self) as! [CFString : Any]
        attributes[kCVPixelBufferIOSurfacePropertiesKey] = [:]  // enable IOSurface backing

        var destPixelBuffer: CVPixelBuffer?
        let createStatus = CVPixelBufferCreate(nil, destWidth, destHeight, format, attributes as CFDictionary, &destPixelBuffer)
        guard createStatus == kCVReturnSuccess, let destPixelBuffer = destPixelBuffer else {
            throw PixelBufferRotationError.cvPixelBufferCreateError(createStatus)
        }

        CVPixelBufferLockBaseAddress(destPixelBuffer, [])
        defer { CVPixelBufferUnlockBaseAddress(destPixelBuffer, []) }

        guard let dest = CVPixelBufferGetBaseAddress(destPixelBuffer) else {
            throw PixelBufferRotationError.getDestBaseAddressFailed
        }

        // get the bytes per row in the destination pixel buffer, which includes memory alignment
        let alignedBytesPerRow = CVPixelBufferGetBytesPerRow(destPixelBuffer)
        for row in 0..<destHeight {
            // copy each row, including the alignment padding at the end of the row
            let srcRowAddress = temp.advanced(by: row * unalignedDestRowBytes)
            let dstRowAddress = dest.advanced(by: row * alignedBytesPerRow)
            memcpy(dstRowAddress, srcRowAddress, Int(destWidth * 4))
        }

        // copy attachments, including (importantly) the color profile
        CVBufferPropagateAttachments(self, destPixelBuffer)

        return destPixelBuffer
    }

}