motionlab/native/Sources/MotionLab/Renderer.swift

import MetalKit

/// One effect file in the chain.
final class Layer {
    let url: URL
    let name: String
    var enabled = false
    var order = 50
    var params: [EffectParam] = []
    var error: String?
    fileprivate var pipeline: MTLRenderPipelineState?
    fileprivate var modified: Date?
    fileprivate var usesParticles = false
    fileprivate var targets: [MTLTexture] = []  // [this frame, previous frame], swapped after every draw

    init(url: URL) {
        self.url = url
        self.name = url.deletingPathExtension().lastPathComponent
    }
}

private let passthrough = "float4 effect(float2 uv, Inputs in) { return in.image.sample(in.smp, uv); }"

/// Runs every enabled effect in order on each new camera frame, shows the result in the window,
/// and hands it to whatever outputs are active (recording, virtual camera, screenshots).
final class Renderer: NSObject, MTKViewDelegate {
    private let device: MTLDevice
    private let queue: MTLCommandQueue
    private let camera: Camera
    private let effectsDirectory: URL
    private let copyPipeline: MTLRenderPipelineState
    private var analysis: Analysis?
    private var plain: MTLTexture?   // holds the camera picture when no effect is enabled
    private var result: MTLTexture?  // the finished frame
    private var lastSerial = -1
    private let startTime = CACurrentMediaTime()
    private var frameCount = 0
    private var pixels: [UInt8] = []

    private(set) var layers: [Layer] = []
    /// Called when layers were added, removed, reloaded or toggled.
    var onLayersChanged: (() -> Void)?
    /// Called when recording or the virtual camera starts or stops.
    var onStatusChanged: (() -> Void)?

    /// When set, the app exits after this many camera frames, first saving that frame as a PNG if a url is given
    /// (used to check effects without watching).
    var snapshot: (frames: Int, url: URL?)?
    /// When set, recording into this file starts with the first frame.
    var recordURL: URL?
    /// When true, the OBS virtual camera is opened on the first frame, at the size of the output.
    var wantsVirtualCamera = false
    private var virtualCamera: VirtualCamera? { didSet { onStatusChanged?() } }
    private var recorder: Recorder? { didSet { onStatusChanged?() } }
    private var screenshotRequested = false

    var status: String {
        [recorder != nil ? "REC" : nil, virtualCamera != nil ? "CAM" : nil].compactMap { $0 }.joined(separator: "  ")
    }

    init(device: MTLDevice, camera: Camera, effectsDirectory: URL) throws {
        self.device = device
        self.camera = camera
        self.effectsDirectory = effectsDirectory
        self.queue = device.makeCommandQueue()!
        self.copyPipeline = try makePipeline(device, source: effectShader(body: passthrough, params: []),
                                             fragment: "effectFragment", format: .bgra8Unorm)
        super.init()
        refreshFiles()
    }

    // MARK: effect files

    /// Picks up new, deleted and edited effect files. On a compile error the last working version keeps running.
    private func refreshFiles() {
        let files = ((try? FileManager.default.contentsOfDirectory(at: effectsDirectory, includingPropertiesForKeys: nil)) ?? [])
            .filter { $0.pathExtension == "metal" }
        var changed = false

        let before = layers.count
        layers.removeAll { layer in !files.contains(layer.url) }
        changed = layers.count != before
        for url in files where !layers.contains(where: { $0.url == url }) {
            layers.append(Layer(url: url))
            changed = true
        }
        for layer in layers where reload(layer) { changed = true }

        if changed {
            layers.sort { ($0.order, $0.name) < ($1.order, $1.name) }
            onLayersChanged?()
        }
    }

    private func reload(_ layer: Layer) -> Bool {
        let modified = (try? layer.url.resourceValues(forKeys: [.contentModificationDateKey]))?.contentModificationDate
        guard let modified, modified != layer.modified else { return false }
        let firstLoad = layer.modified == nil
        layer.modified = modified
        do {
            let source = try String(contentsOf: layer.url, encoding: .utf8)
            var info = parseEffect(source)
            layer.pipeline = try makePipeline(device, source: effectShader(body: source, params: info.params),
                                              fragment: "effectFragment", format: .bgra8Unorm)
            for i in info.params.indices {  // keep slider positions across reloads
                if let old = layer.params.first(where: { $0.name == info.params[i].name }) {
                    info.params[i].value = min(max(old.value, info.params[i].min), info.params[i].max)
                }
            }
            layer.params = info.params
            layer.order = info.order
            layer.usesParticles = source.contains("in.particles")
            layer.error = nil
            if firstLoad { layer.enabled = info.startsOn }
            print("[\(layer.name)] loaded")
        } catch {
            layer.error = error.localizedDescription
            print("[\(layer.name)] shader error:\n\(error.localizedDescription)")
        }
        return true
    }

    var failedLayers: [String] { layers.filter { $0.error != nil }.map(\.name) }

    /// Enables exactly the named layers. Returns the names that don't exist.
    func enableOnly(_ names: [String]) -> [String] {
        for layer in layers { layer.enabled = names.contains(layer.name) }
        return names.filter { name in !layers.contains { $0.name == name } }
    }

    func setEnabled(_ index: Int, _ enabled: Bool) {
        guard layers.indices.contains(index) else { return }
        layers[index].enabled = enabled
        if !enabled { layers[index].targets = [] }  // free its memory, and start clean next time
        onLayersChanged?()
    }

    func setParam(layer: Int, param: Int, value: Float) {
        guard layers.indices.contains(layer), layers[layer].params.indices.contains(param) else { return }
        let p = layers[layer].params[param]
        layers[layer].params[param].value = min(max(value, p.min), p.max)
    }

    // MARK: outputs

    func requestScreenshot() { screenshotRequested = true }

    func toggleRecording() {
        if let recorder {
            self.recorder = nil
            recorder.finish()
            print("saved \(recorder.url.path)")
        } else if let result {
            do {
                recorder = try Recorder(url: capturesURL(extension: "mp4"), width: result.width, height: result.height)
                print("recording \(recorder!.url.path)")
            } catch {
                print("could not start recording: \(error.localizedDescription)")
            }
        }
    }

    func finish() {
        if recorder != nil { toggleRecording() }
    }

    private func capturesURL(extension ext: String) -> URL {
        let directory = URL(fileURLWithPath: "captures")
        try? FileManager.default.createDirectory(at: directory, withIntermediateDirectories: true)
        let stamp = DateFormatter()
        stamp.dateFormat = "yyyyMMdd-HHmmss"
        return directory.appendingPathComponent("motionlab-\(stamp.string(from: Date())).\(ext)")
    }

    /// Runs once per camera frame, after the GPU work was submitted.
    private func deliver(_ commands: MTLCommandBuffer) {
        frameCount += 1
        let snapshotDue = snapshot.map { frameCount >= $0.frames } ?? false
        guard let result else { return }
        let width = result.width, height = result.height
        if let url = recordURL {
            recordURL = nil
            do { recorder = try Recorder(url: url, width: width, height: height) } catch { print("could not start recording: \(error.localizedDescription)") }
        }
        if wantsVirtualCamera {
            wantsVirtualCamera = false
            do {
                virtualCamera = try VirtualCamera(width: width, height: height)
                print("virtual camera on: pick 'OBS Virtual Camera' as the camera in your call or streaming app")
            } catch {
                print("virtual camera off: \(error)")
            }
        }
        guard recorder != nil || virtualCamera != nil || screenshotRequested || snapshotDue else { return }

        commands.waitUntilCompleted()
        if pixels.count != width * height * 4 { pixels = [UInt8](repeating: 0, count: width * height * 4) }
        result.getBytes(&pixels, bytesPerRow: width * 4, from: MTLRegionMake2D(0, 0, width, height), mipmapLevel: 0)

        recorder?.append(bgra: pixels, at: CACurrentMediaTime())
        virtualCamera?.send(bgra: pixels, width: width, height: height)
        if screenshotRequested {
            screenshotRequested = false
            let url = capturesURL(extension: "png")
            writePNG(bgra: &pixels, width: width, height: height, to: url)
            print("saved \(url.path)")
        }
        if snapshotDue {
            if let url = snapshot?.url {
                writePNG(bgra: &pixels, width: width, height: height, to: url)
                print("\(frameCount) frames -> \(url.path)")
            }
            finish()
            exit(0)
        }
    }

    // MARK: drawing

    private func process(_ frames: (current: MTLTexture, last: MTLTexture, serial: Int), _ commands: MTLCommandBuffer) throws {
        let width = frames.current.width, height = frames.current.height
        if analysis == nil { analysis = try Analysis(device: device, queue: queue, width: width, height: height) }
        guard let analysis else { return }
        analysis.update(commands, camera: frames.current, particles: layers.contains { $0.enabled && $0.usesParticles })

        let time = Float(CACurrentMediaTime() - startTime)
        var image = frames.current
        for layer in layers where layer.enabled {
            guard let pipeline = layer.pipeline else { continue }
            if layer.targets.first?.width != width || layer.targets.first?.height != height {
                layer.targets = [makeTarget(device, width: width, height: height), makeTarget(device, width: width, height: height)]
                for target in layer.targets { commands.pass(into: target, clear: true) }
            }
            // must match `struct Uniforms` in the shader: time, padding, resolution.xy, params[16]
            var uniforms: [Float] = [time, 0, Float(width), Float(height)]
            uniforms += layer.params.map(\.value) + [Float](repeating: 0, count: maxParams - layer.params.count)
            let textures = [image, frames.current, frames.last, layer.targets[1], analysis.flow, analysis.particles]
            commands.pass(into: layer.targets[0]) { $0.fullscreen(pipeline, textures: textures, uniforms: uniforms) }
            image = layer.targets[0]
            layer.targets.swapAt(0, 1)
        }

        if image === frames.current {
            if plain?.width != width || plain?.height != height { plain = makeTarget(device, width: width, height: height) }
            commands.pass(into: plain!) { $0.fullscreen(copyPipeline, textures: [image], uniforms: [0, 0, 0, 0]) }
            image = plain!
        }
        result = image
    }

    /// The largest rectangle with the output's aspect ratio that fits the window, centred.
    private func fitted(_ output: MTLTexture, in size: CGSize) -> MTLViewport {
        let scale = min(size.width / Double(output.width), size.height / Double(output.height))
        let width = Double(output.width) * scale, height = Double(output.height) * scale
        return MTLViewport(originX: (size.width - width) / 2, originY: (size.height - height) / 2,
                           width: width, height: height, znear: 0, zfar: 1)
    }

    func mtkView(_ view: MTKView, drawableSizeWillChange size: CGSize) {}

    func draw(in view: MTKView) {
        guard let frames = camera.latest, let commands = queue.makeCommandBuffer() else { return }

        // effects advance once per camera frame; in between, the window just shows the last result again
        let isNewFrame = frames.serial != lastSerial
        if isNewFrame {
            lastSerial = frames.serial
            if frameCount % 15 == 0 { refreshFiles() }
            do {
                try process(frames, commands)
            } catch {
                print("render setup failed: \(error.localizedDescription)")
                exit(1)
            }
        }

        if let result, let drawable = view.currentDrawable, let windowPass = view.currentRenderPassDescriptor,
           let encoder = commands.makeRenderCommandEncoder(descriptor: windowPass) {
            encoder.setViewport(fitted(result, in: view.drawableSize))
            encoder.fullscreen(copyPipeline, textures: [result], uniforms: [0, 0, 0, 0])
            encoder.endEncoding()
            commands.present(drawable)
        }
        commands.commit()
        if isNewFrame { deliver(commands) }
    }
}