luma/Sources/LumaShaderCompiler/main.swift

import Foundation

#if canImport(CShaderTranslate)
import CShaderTranslate
#endif

LumaShaderCompiler.main()

/// Translates each authored effect in `Shaders/` into the forms the frontends
/// consume: a Swift catalog of the GLSL, which the GTK host hands to its
/// `luma_shader_effect` widget, and a Metal fragment function for the SwiftUI
/// host. Authors write the effect once, body-only, against the uniform
/// vocabulary the preamble below declares.
struct LumaShaderCompiler {
    static let preamble = """
        #version 450
        layout(location = 0) in vec2 v_uv;
        layout(location = 0) out vec4 frag_color;
        layout(binding = 0) uniform ShaderEffectUniforms {
            vec2 u_resolution;
            float u_time;
            float u_scheme;
            float u_activity;
            float u_pulse;
            float u_data_count;
            // Values the caller fed the canvas, up to 64 of them, packed four
            // to a vec4. Read them through dataAt.
            vec4 u_data[16];
            // Where the author's vertices land: orthographic for a flat
            // drawing, perspective for one with depth. Identity until set.
            mat4 u_mvp;
        };
        float dataAt(int i) { return u_data[i >> 2][i & 3]; }

        """

    static func main() {
        let options = Options(arguments: CommandLine.arguments)
        let effects = effects(in: options.shaderDir)
        // A checkout on Windows carries the shaders with carriage returns, and
        // one inside a multi-line literal is a line the compiler cannot indent.
        let bodies = effects.map {
            try! String(contentsOf: $0.source, encoding: .utf8)
                .replacingOccurrences(of: "\r\n", with: "\n")
        }

        if let swiftOut = options.swiftOut {
            write(swiftCatalog(of: effects, bodies: bodies), to: swiftOut)
        }
        if let metalDir = options.metalDir {
            for (effect, body) in zip(effects, bodies) {
                let fragment = metalFragment(named: effect.metalFunction, body: body)
                write(fragment, to: metalDir.appendingPathComponent("\(effect.metalFileName).metal"))
            }
        }
    }

    private static func effects(in directory: URL) -> [Effect] {
        let names = try! FileManager.default.contentsOfDirectory(atPath: directory.path)
        return names
            .filter { $0.hasSuffix(Effect.suffix) }
            .sorted()
            .map { Effect(fileName: $0, directory: directory) }
    }

    private static func swiftCatalog(of effects: [Effect], bodies: [String]) -> String {
        let entries = zip(effects, bodies).map { effect, body in
            let escaped = body
                .replacingOccurrences(of: "\\", with: "\\\\")
                .trimmingCharacters(in: .newlines)
                .split(separator: "\n", omittingEmptySubsequences: false)
                .map { $0.isEmpty ? "" : "        \($0)" }
                .joined(separator: "\n")
            return """
                    public static let \(effect.swiftName) = \"\"\"
                \(escaped)
                        \"\"\"
                """
        }
        let names = effects.map { "        \"\($0.swiftName)\": \($0.swiftName)," }
        return """
            // This file is auto-generated by LumaShaderCompiler.

            /// The authored GLSL of each effect in `Shaders/`, body-only: the host
            /// widget supplies the version directive and the uniform declarations.
            public enum ShaderEffects {
            \(entries.joined(separator: "\n\n"))

                /// Every effect by name, for hosts that choose one at runtime.
                public static let all: [String: String] = [
            \(names.joined(separator: "\n"))
                ]

                /// The Metal library function the same effect compiles to.
                public static func metalFunction(named name: String) -> String {
                    name + "Fragment"
                }
            }

            """
    }

    /// The same libraries that translate a shader written at run time, rather
    /// than the command-line tools built beside them: one way of doing it, and
    /// nothing to install.
    private static func metalFragment(named name: String, body: String) -> String {
        #if canImport(CShaderTranslate)
        var message: UnsafeMutablePointer<CChar>?
        let translated = luma_shader_translate_to_msl(preamble + body, 1, name, &message)
        guard let translated else {
            let detail = message.map { String(cString: $0) } ?? "no detail"
            FileHandle.standardError.write(Data("LumaShaderCompiler: \(name): \(detail)\n".utf8))
            exit(1)
        }
        defer { free(translated) }
        return String(cString: translated)
        #else
        FileHandle.standardError.write(
            Data("LumaShaderCompiler: --metal-dir needs the shader toolchain\n".utf8))
        exit(1)
        #endif
    }

    private static func write(_ contents: String, to destination: URL) {
        try! FileManager.default.createDirectory(
            at: destination.deletingLastPathComponent(),
            withIntermediateDirectories: true
        )
        guard (try? String(contentsOf: destination, encoding: .utf8)) != contents else { return }
        try! contents.write(to: destination, atomically: true, encoding: .utf8)
    }

    struct Effect {
        static let suffix = ".frag.glsl"

        let swiftName: String
        let metalFunction: String
        let metalFileName: String
        let source: URL

        init(fileName: String, directory: URL) {
            let words = fileName.dropLast(Effect.suffix.count).split(separator: "-")
            swiftName = words.first! + words.dropFirst().map { $0.capitalized }.joined()
            metalFunction = swiftName + "Fragment"
            metalFileName = words.map { $0.capitalized }.joined() + "Fragment"
            source = directory.appendingPathComponent(fileName)
        }
    }

    /// `--metal-dir` is what pulls in glslang and spirv-cross, so hosts with no
    /// Metal to feed leave it off and need neither tool installed.

    private struct Options {
        let shaderDir: URL
        let swiftOut: URL?
        let metalDir: URL?

        init(arguments: [String]) {
            var values: [String: String] = [:]
            for index in stride(from: 1, to: arguments.count - 1, by: 2) {
                values[arguments[index]] = arguments[index + 1]
            }
            shaderDir = URL(fileURLWithPath: values["--shader-dir"]!)
            swiftOut = values["--swift-out"].map { URL(fileURLWithPath: $0) }
            metalDir = values["--metal-dir"].map { URL(fileURLWithPath: $0) }
        }
    }
}