mtl_nsimage.m

#import "NSImage+MTLTexture.h"
#import <Foundation/Foundation.h>
#import <Metal/MTLDevice.h>
#import <Accelerate/Accelerate.h>

@implementation NSImage(MTLTexture)

// ref: https://gist.github.com/cbirkhold/52bfdcfd2db7777b64b8
// The original gist bore the MIT license

-(id<MTLTexture>) createMTLTextureWithDevice:(id<MTLDevice>)device flip:(BOOL)flip
{
    if (!device)
        return nil;
    
    NSSize sz = [self size];
    const int bitsPerComponent = 8;
    const int bytesPerRow = (bitsPerComponent * 4 / 8) * sz.width;
    const CGBitmapInfo bitmapInfo = (CGBitmapInfo) kCGImageAlphaPremultipliedLast;
    const CGColorSpaceRef colorSpace = CGColorSpaceCreateDeviceRGB();
    const CGContextRef context = CGBitmapContextCreate(NULL, sz.width, sz.height,
                                                       bitsPerComponent, bytesPerRow, colorSpace, bitmapInfo);
    CGColorSpaceRelease(colorSpace);

    if (!context)
        return nil;
    
    if (flip)
    {
        CGContextTranslateCTM(context, sz.width, sz.height);
        CGContextScaleCTM(context, -1.0, -1.0);
    }
    
    const CGRect viewport = CGRectMake(0.0f, 0.0f, sz.width, sz.height);
    NSRect viewportRect = NSRectFromCGRect(viewport);
    CGContextClearRect(context, viewport);
    CGImageRef img = [self CGImageForProposedRect:&viewportRect context:nil hints:nil];
    CGContextDrawImage(context, viewport, img);
    
    const void *const data = CGBitmapContextGetData(context);

    MTLTextureDescriptor *textureDescriptor = [[MTLTextureDescriptor alloc] init];
    
    // Indicate that each pixel has a r, g, b, and alpha channel, where each channel is
    // an 8-bit unsigned normalized value (i.e. 0 maps to 0.0 and 255 maps to 1.0)
    textureDescriptor.pixelFormat = MTLPixelFormatRGBA8Unorm;
    
    // Set the pixel dimensions of the texture
    textureDescriptor.width = sz.width;
    textureDescriptor.height = sz.height;
    
    // Create the texture from the device by using the descriptor
    id<MTLTexture> texture = [device newTextureWithDescriptor:textureDescriptor];
    
    [texture replaceRegion:MTLRegionMake2D(0, 0, sz.width, sz.height)
               mipmapLevel:0 withBytes:data bytesPerRow:bytesPerRow];
    
    CGContextRelease(context);
    
    return texture;
}

// ref: https://gist.github.com/codelynx/4e56758fb89e94d0d1a58b40ddaade45

+ (NSImage*) NSImageFromMTLTexture:(id<MTLTexture>) texture
{
    BOOL swizzle = texture.pixelFormat == MTLPixelFormatBGRA8Unorm ||
                    texture.pixelFormat == MTLPixelFormatBGRA8Unorm_sRGB;
    
    // read texture as byte array
    size_t rowBytes = texture.width * 4;
    size_t length = rowBytes * texture.height;
    UInt8* bgraBytes = (UInt8*) malloc(length);
    MTLRegion region = MTLRegionMake2D(0, 0, texture.width, texture.height);
    [texture getBytes:bgraBytes bytesPerRow:rowBytes fromRegion:region mipmapLevel:0];

    UInt8* rgbaBytes;
    UInt8* freeRgbaBytes = NULL;
    if (swizzle)
    {
        // use Accelerate framework to convert from BGRA to RGBA
        vImage_Buffer bgraBuffer = {bgraBytes, texture.height, texture.width, rowBytes};
        rgbaBytes = (UInt8*) malloc(length);
        freeRgbaBytes = rgbaBytes;
        
        vImage_Buffer rgbaBuffer = {rgbaBytes, texture.height, texture.width, rowBytes};
        
        UInt8 map[4] = { 2, 1, 0, 3 };
        vImagePermuteChannels_ARGB8888(&bgraBuffer, &rgbaBuffer, map, 0);
    }
    else
        rgbaBytes = bgraBytes;
    
    // create CGImage with RGBA
    CGColorSpaceRef sRgb = CGColorSpaceCreateDeviceRGB();
    CGBitmapInfo bitmapInfo = (CGBitmapInfo) kCGImageAlphaPremultipliedLast;
    CFDataRef data = CFDataCreate(nil, rgbaBytes, length);
    if (!data)
    {
        CFRelease(sRgb);
        return nil;
    }
    
    CGDataProviderRef dataProvider = CGDataProviderCreateWithCFData(data);
    if (!dataProvider)
    {
        CFRelease(sRgb);
        CFRelease(data);
        return nil;
    }
    
    CGImageRef cgImage = CGImageCreate(texture.width, texture.height,
                                       8 /*bits per component*/, 32 /*bpp*/, rowBytes,
                                       sRgb, bitmapInfo, dataProvider, nil,
                                       true, kCGRenderingIntentDefault);
    
    NSSize sz = { CGImageGetWidth(cgImage), CGImageGetHeight(cgImage) };
    NSImage* result = [[NSImage alloc] initWithCGImage:cgImage size:sz];
    
    CFRelease(dataProvider);
    CFRelease(data);
    CFRelease(sRgb);
    CFRelease(cgImage);
    
    free(bgraBytes);
    if (freeRgbaBytes)
        free(rgbaBytes);
    
    return result;
}


@end