[iOS] Extension for creating UIImage from Metal texture

UIImageMetalExtensions.swift

// http://blog.human-friendly.com/drawing-images-from-pixel-data-in-swift
// https://github.com/FlexMonkey/MetalReactionDiffusion/blob/1ea9aa4a841d20e0b247505fdf716cd5fe1a01fd/MetalReactionDiffusion/ViewController.swift

public struct PixelData {
    var a:UInt8 = 255
    var r:UInt8
    var g:UInt8
    var b:UInt8
}

private let rgbColorSpace = CGColorSpaceCreateDeviceRGB()
private let bitmapInfo:CGBitmapInfo = CGBitmapInfo(CGImageAlphaInfo.PremultipliedFirst.toRaw())

public func imageFromARGB32Bitmap(pixels:[PixelData], width:UInt, height:UInt)->UIImage {
    let bitsPerComponent:UInt = 8
    let bitsPerPixel:UInt = 32
    
    assert(pixels.count == Int(width * height))
    
    var data = pixels // Copy to mutable []
    let providerRef = CGDataProviderCreateWithCFData(
            NSData(bytes: &data, length: data.count * sizeof(PixelData))
        )

    let cgim = CGImageCreate(
            width,
            height,
            bitsPerComponent,
            bitsPerPixel,
            width * UInt(sizeof(PixelData)),
            rgbColorSpace,
            bitmapInfo,
            providerRef,
            nil,
            true,
            kCGRenderingIntentDefault
        )
    return UIImage(CGImage: cgim)
}

extension UIImage {
  init(texture: MTLTexture) {
    let bitsPerComponent = 8
    let bitsPerPixel = 32
    let bytesPerRow: UInt = texture.width * 4
    let rgbColorSpace = CGColorSpaceCreateDeviceRGB()
    let bitmapInfo:CGBitmapInfo = [.ByteOrder32Big, CGBitmapInfo(rawValue: CGImageAlphaInfo.PremultipliedFirst.rawValue)]
    
    let cgim = CGImageCreate(
            texture.width,
            texture.height,
            bitsPerComponent,
            bitsPerPixel,
            bytesPerRow,
            rgbColorSpace,
            bitmapInfo,
            dataProviderRefFrom(texture),
            nil,
            false,
            kCGRenderingIntentDefault
        )
        init(cgim)
  }
  
  func dataProviderRefFrom(texture: MTLTexture) -> CGDataProviderRef {
    let region = MTLRegionMake2D(0, 0, Int(texture.width), Int(texture.height))
    let pixelCount: Int = texture.width * texture.height
    var imageBytes = [UInt8](count: pixelCount * 4, repeatedValue: 0)
    texture.getBytes(&imageBytes, bytesPerRow: 4 * texture.width, from: region, mipmapLevel: 0)
    let providerRef = CGDataProviderCreateWithCFData(NSData(bytes: &imageBytes, length: pixelCount * 4 * sizeof(UInt8)))
    return providerRef
  }
}