metal_image.swift

///
/// This is a 'single' function way to get a UIImage from a Metal shader
/// You still need to set up an MTLibrary in your app, but you should be
/// able to fairly readily drop this into a standard XCode app and as long
/// as you have the vertex and fragment shaders compiled into the library
/// you should be good to go.
///

/** 
    In your shader file:

    struct Vertex
    {
        float4 position [[ position ]];
        float2 uv;
    };

    vertex
    Vertex
    vertex_func (
        const device Vertex * vertices [[ buffer(0) ]],
        uint id [[ vertex_id ]] )
    {
        return vertices[id];
    }

    fragment
    float4
    test_func (Vertex v [[ stage_in ]])
    {
        // Re-map 0,0 to the centre and rescale from -1,1 vertically
        float w,h = 240.
        float x = v.position.x/w;
        float y = v.position.y/h;
        float ratio = w/h;
        x = ((2.0*x)-1.0)*ratio;
        y = (2.0*y)-1.0;

        // Colour-code distance to the centre
        return {abs(x),abs(y),0,1};
    }
**/

import MetalKit

func generate_image (_ drawing_function:String, width:Int, height:Int) -> UIImage
{
    /// Very rough timing of the function
    let start = Date(); defer { print("\(#function) took \(-start.timeIntervalSinceNow)") }
    
    
    /// Simple geometry to get fragments to work with
    /// (Don't strictly need the uv variable, but hey)
    struct Vertex
    {
        var position : simd_float4 = [0,0,0,0]
        var uv       : simd_float2 = [0,0]
    }
    
    let view_corner_verts : [Vertex] = [
        Vertex(position: simd_float4(-1.0,-1.0, 0.0, 1.0), uv: simd_float2(0,1)),
        Vertex(position: simd_float4(-1.0, 1.0, 0.0, 1.0), uv: simd_float2(0,0)),
        Vertex(position: simd_float4( 1.0,-1.0, 0.0, 1.0), uv: simd_float2(1,1)),
        Vertex(position: simd_float4( 1.0, 1.0, 0.0, 1.0), uv: simd_float2(1,0)),
    ]
    
    
    /// Set up the stuff Metal needs
    let device = MTLCreateSystemDefaultDevice()!
    let library = device.makeDefaultLibrary()!
    let command_queue = device.makeCommandQueue()!
    let command_buffer = command_queue.makeCommandBuffer()!

    
    /// Create an MTLTexture to render into
    let texture_descriptor = MTLTextureDescriptor()
    texture_descriptor.width = width
    texture_descriptor.height = height
    texture_descriptor.pixelFormat = .bgra8Unorm
    texture_descriptor.usage = [.renderTarget]
    let texture_output = device.makeTexture(descriptor: texture_descriptor)!
    
    
    /// Create a render pipeline to do the rendering
    let pass_descriptor = MTLRenderPassDescriptor()
    pass_descriptor.colorAttachments[0].texture = texture_output
    
    let pipeline_descriptor = MTLRenderPipelineDescriptor()
    pipeline_descriptor.colorAttachments[0].pixelFormat = texture_output.pixelFormat
    pipeline_descriptor.vertexFunction   = library.makeFunction(name: "vertex_func")!
    pipeline_descriptor.fragmentFunction = library.makeFunction(name: drawing_function)!
    let pass_pipeline = try! device.makeRenderPipelineState(descriptor: pipeline_descriptor)
    
    /// Do the rendering
    let picture_command_encoder = command_buffer.makeRenderCommandEncoder(descriptor: pass_descriptor)!
    picture_command_encoder.setRenderPipelineState(pass_pipeline)
    picture_command_encoder.setVertexBytes(view_corner_verts, length: MemoryLayout<Vertex>.stride * view_corner_verts.count, index: 0)
    picture_command_encoder.drawPrimitives(type: .triangleStrip, vertexStart: 0, vertexCount: view_corner_verts.count)
    picture_command_encoder.endEncoding()
    
    
    /// Wait until the GPU is done
    command_buffer.commit()
    command_buffer.waitUntilCompleted()
    
    
    /// Convert the MTLTexture into a UIImage and return it
    let cii   = CIImage(mtlTexture: texture_output, options: nil)!
    let image = UIImage(ciImage: cii)
    return image
}