motionlab/native/Sources/MotionLab/Shaders.swift

import Metal

/// A slider an effect file declares with a comment line:  `// slider <name> <min> <max> <default>`
struct EffectParam {
    let name: String
    let min: Float
    let max: Float
    var value: Float
}

/// Everything an effect file says about itself in `//` comment lines:
///
///     // slider <name> <min> <max> <default>    a slider, read in the shader as in.p.<name>
///     // order <number>                         position in the chain, low runs first (default 50)
///     // on                                     enabled when the app starts
struct EffectInfo {
    var params: [EffectParam] = []
    var order = 50
    var startsOn = false
}

let maxParams = 16

func parseEffect(_ source: String) -> EffectInfo {
    var info = EffectInfo()
    for line in source.split(separator: "\n") {
        let parts = line.split(separator: " ")
        guard parts.count >= 2, parts[0] == "//" else { continue }
        switch parts[1] {
        case "slider" where parts.count == 6:
            if let min = Float(parts[3]), let max = Float(parts[4]), let value = Float(parts[5]) {
                info.params.append(EffectParam(name: String(parts[2]), min: min, max: max, value: value))
            }
        case "order" where parts.count == 3:
            info.order = Int(parts[2]) ?? info.order
        case "on" where parts.count == 2:
            info.startsOn = true
        default:
            break
        }
    }
    info.params = Array(info.params.prefix(maxParams))
    return info
}

/// Shared by every shader in the app: a triangle that covers the whole target.
let shaderPrelude = """
#include <metal_stdlib>
using namespace metal;

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

vertex VertexOut fullscreenVertex(uint id [[vertex_id]]) {
    float2 p = float2((id << 1) & 2, id & 2);
    VertexOut out;
    out.position = float4(p.x * 2.0 - 1.0, 1.0 - p.y * 2.0, 0.0, 1.0);
    out.uv = p;
    return out;
}

"""

/// Wraps an effect file into a complete Metal shader. An effect file only defines:
///
///     float4 effect(float2 uv, Inputs in) { ... }
///
/// `uv` runs from (0,0) top-left to (1,1) bottom-right.
func effectShader(body: String, params: [EffectParam]) -> String {
    let fields = params.isEmpty ? "float unused;" : params.map { "float \($0.name);" }.joined(separator: " ")
    let values = params.isEmpty ? "0.0" : params.indices.map { "u.params[\($0)]" }.joined(separator: ", ")
    return shaderPrelude + """
    struct Uniforms {
        float time;
        float2 resolution;
        float params[\(maxParams)];
    };

    struct Params { \(fields) };

    struct Inputs {
        texture2d<float> image;      // the picture so far: output of the effect above, or the camera for the first one
        texture2d<float> camera;     // the untouched camera frame
        texture2d<float> last;       // the camera frame before this one (compare with camera to find motion)
        texture2d<float> prev;       // this effect's own output from the previous frame (for trails)
        texture2d<float> flow;       // .xy = which way and how far things moved since the last frame, in uv units
        texture2d<float> particles;  // .r = brightness of particles that drift along with the motion
        sampler smp;
        float time;                  // seconds since start
        float2 resolution;           // output size in pixels
        Params p;                    // this effect's sliders
    };

    #line 1
    \(body)

    fragment float4 effectFragment(VertexOut v [[stage_in]],
                                   texture2d<float> image [[texture(0)]],
                                   texture2d<float> camera [[texture(1)]],
                                   texture2d<float> last [[texture(2)]],
                                   texture2d<float> prev [[texture(3)]],
                                   texture2d<float> flow [[texture(4)]],
                                   texture2d<float> particles [[texture(5)]],
                                   constant Uniforms &u [[buffer(0)]]) {
        constexpr sampler smp(filter::linear, address::clamp_to_edge);
        Params p = { \(values) };
        Inputs in = { image, camera, last, prev, flow, particles, smp, u.time, u.resolution, p };
        return effect(v.uv, in);
    }
    """
}

func makePipeline(_ device: MTLDevice, source: String, vertex: String = "fullscreenVertex", fragment: String,
                  format: MTLPixelFormat, additive: Bool = false) throws -> MTLRenderPipelineState {
    let library = try device.makeLibrary(source: source, options: nil)
    let descriptor = MTLRenderPipelineDescriptor()
    descriptor.vertexFunction = library.makeFunction(name: vertex)
    descriptor.fragmentFunction = library.makeFunction(name: fragment)
    let attachment = descriptor.colorAttachments[0]!
    attachment.pixelFormat = format
    if additive {
        attachment.isBlendingEnabled = true
        attachment.sourceRGBBlendFactor = .one
        attachment.destinationRGBBlendFactor = .one
        attachment.sourceAlphaBlendFactor = .one
        attachment.destinationAlphaBlendFactor = .one
    }
    return try device.makeRenderPipelineState(descriptor: descriptor)
}

func makeTarget(_ device: MTLDevice, width: Int, height: Int, format: MTLPixelFormat = .bgra8Unorm) -> MTLTexture {
    let descriptor = MTLTextureDescriptor.texture2DDescriptor(pixelFormat: format, width: width, height: height, mipmapped: false)
    descriptor.usage = [.renderTarget, .shaderRead]
    return device.makeTexture(descriptor: descriptor)!
}

extension MTLCommandBuffer {
    /// One render pass into `target`. With `clear`, the target starts black; otherwise its old contents are ignored.
    func pass(into target: MTLTexture, clear: Bool = false, _ body: (MTLRenderCommandEncoder) -> Void = { _ in }) {
        let descriptor = MTLRenderPassDescriptor()
        descriptor.colorAttachments[0].texture = target
        descriptor.colorAttachments[0].loadAction = clear ? .clear : .dontCare
        descriptor.colorAttachments[0].storeAction = .store
        guard let encoder = makeRenderCommandEncoder(descriptor: descriptor) else { return }
        body(encoder)
        encoder.endEncoding()
    }
}

extension MTLRenderCommandEncoder {
    func fullscreen(_ pipeline: MTLRenderPipelineState, textures: [MTLTexture], uniforms: [Float] = []) {
        setRenderPipelineState(pipeline)
        setFragmentTextures(textures, range: 0..<textures.count)
        if !uniforms.isEmpty {
            setFragmentBytes(uniforms, length: uniforms.count * MemoryLayout<Float>.stride, index: 0)
        }
        drawPrimitives(type: .triangle, vertexStart: 0, vertexCount: 3)
    }
}