concinnity_toolchain/

metal_shaders.rs

// Build-time precompilation of the engine's static Metal shaders.
//
// The device crate's build script hands this module the shader directory, the
// names that must stay source-only (runtime-assembled templates), and the
// shared fragments spliced into shaders at a marker. Every other `.metal` file
// is assembled into OUT_DIR and compiled to a `.metallib` with the same
// `xcrun metal` + `xcrun metallib` pair the cook's shader toolchain uses, and a
// lookup function mapping shader name to embedded bytes is generated there too.
//
// When the Metal compiler is not installed (Command Line Tools without a full
// Xcode) the generated lookup returns `None` for every name and the renderer
// falls back to compiling its embedded source at startup, exactly as before.
// A cargo warning flags the slower path so it is never a silent regression.

use std::path::{Path, PathBuf};
use std::process::Command;

use concinnity_slang as slang;

/// One single-source shader library to precompile to a metallib: the `.slang`
/// file under the slang shader directory, the entry points linked into the
/// library, and the `#define`s that select its variant. `name` is the key the
/// renderer's lookup uses (distinct from `file` when one source yields several
/// variant libraries).
pub struct SlangLibSpec {
    /// Lookup key the renderer resolves this library by.
    pub name: &'static str,
    /// The `.slang` file, relative to the shader directory.
    pub file: &'static str,
    /// Entry points linked into the library.
    pub entries: &'static [&'static str],
    /// `#define`s selecting this variant.
    pub defines: &'static [(&'static str, &'static str)],
}

/// The single-source half of the precompile: where the `.slang` files live, the
/// shared declarations spliced into the ones carrying a marker, and one spec per
/// metallib variant. Grouped because they always travel together, and because
/// the splice table has to match the renderer's `slang_source::assemble` exactly
/// -- the two produce the same text or the content-addressed cache serves one
/// path's bytes to the other.
pub struct SlangShaders<'a> {
    /// Directory holding the `.slang` sources.
    pub dir: &'a Path,
    /// Shared declarations spliced into sources carrying a marker.
    pub fragments: &'a [(&'a str, &'a str)],
    /// One spec per metallib variant.
    pub specs: &'a [SlangLibSpec],
}

/// Precompile every eligible `.metal` under `shaders_dir`, plus every `.slang`
/// spec in `slang`, into OUT_DIR and generate `engine_metallibs.rs` there.
/// `fragments` pairs a source marker with the file under `shaders_dir` that
/// replaces it, matching the substitution the renderer applies when it compiles
/// the same shader from source. Panics if the Metal toolchain is present but a
/// shader fails to compile: a broken shader must fail the build, not surface at
/// renderer init.
pub fn precompile_metal_shaders(
    shaders_dir: &Path,
    source_only: &[&str],
    fragments: &[(&str, &str)],
    slang: &SlangShaders,
) {
    println!("cargo:rerun-if-changed={}", shaders_dir.display());
    println!("cargo:rerun-if-changed={}", slang.dir.display());
    let out_dir = PathBuf::from(std::env::var("OUT_DIR").expect("OUT_DIR"));
    let generated = out_dir.join("engine_metallibs.rs");

    let shaders = eligible_shaders(shaders_dir, source_only);
    if !metal_toolchain_present() {
        println!(
            "cargo:warning=Metal compiler not found (full Xcode required); engine shaders \
             will compile from source at startup"
        );
        std::fs::write(&generated, stub_lookup_source()).expect("write engine_metallibs.rs");
        return;
    }

    let fragments = read_fragments(shaders_dir, fragments);
    let slang_fragments = read_fragments(slang.dir, slang.fragments);

    let lib_dir = out_dir.join("engine_shaders");
    std::fs::create_dir_all(&lib_dir).expect("create engine_shaders dir");
    let mut entries = Vec::with_capacity(shaders.len() + slang.specs.len());
    for path in &shaders {
        let name = path
            .file_name()
            .and_then(|n| n.to_str())
            .expect("utf8 shader filename")
            .to_string();
        let source = std::fs::read_to_string(path)
            .unwrap_or_else(|e| panic!("read {}: {e}", path.display()));
        let assembled = lib_dir.join(&name);
        std::fs::write(&assembled, assemble(&source, &fragments))
            .unwrap_or_else(|e| panic!("write {}: {e}", assembled.display()));
        let lib_path = lib_dir.join(&name).with_extension("metallib");
        compile_metallib(&assembled, &lib_path);
        entries.push((name, lib_path));
    }
    entries.extend(precompile_slang_libs(
        slang.dir,
        slang.specs,
        &slang_fragments,
        &lib_dir,
    ));
    std::fs::write(&generated, metallib_lookup_source(&entries))
        .expect("write engine_metallibs.rs");
}

// Load each `(marker, file)` pair's replacement text from `dir`.
fn read_fragments<'a>(dir: &Path, fragments: &'a [(&'a str, &'a str)]) -> Vec<(&'a str, String)> {
    fragments
        .iter()
        .map(|(marker, file)| {
            let path = dir.join(file);
            let text = std::fs::read_to_string(&path)
                .unwrap_or_else(|e| panic!("read fragment {}: {e}", path.display()));
            (*marker, text)
        })
        .collect()
}

// Compile each single-source spec to a metallib via slangc. With slangc absent
// the specs are skipped behind a cargo warning: their names miss the generated
// lookup and the renderer compiles them at startup instead (which needs slangc
// at runtime and errors clearly when it is missing there too).
fn precompile_slang_libs(
    slang_dir: &Path,
    specs: &[SlangLibSpec],
    fragments: &[(&str, String)],
    lib_dir: &Path,
) -> Vec<(String, PathBuf)> {
    if specs.is_empty() {
        return Vec::new();
    }
    if slang::slangc_path().is_none() {
        println!(
            "cargo:warning=slangc not found; single-source engine shaders will compile at \
             startup (install the Vulkan SDK or a standalone Slang release)"
        );
        return Vec::new();
    }
    specs
        .iter()
        .map(|spec| {
            let path = slang_dir.join(spec.file);
            let source = std::fs::read_to_string(&path)
                .unwrap_or_else(|e| panic!("read {}: {e}", path.display()));
            let source = assemble(&source, fragments);
            let job = slang::SlangJob {
                source: &concinnity_render::slang_source::inject_defines(&source, spec.defines),
                file_name: spec.name,
                entries: spec.entries,
                target: slang::SlangTarget::Metallib,
            };
            let bytes = slang::compile(&job, lib_dir)
                .unwrap_or_else(|e| panic!("slang metallib precompile failed: {e}"));
            let lib_path = lib_dir.join(spec.name).with_extension("metallib");
            std::fs::write(&lib_path, bytes)
                .unwrap_or_else(|e| panic!("write {}: {e}", lib_path.display()));
            (spec.name.to_string(), lib_path)
        })
        .collect()
}

// Replace every fragment marker in `source`. Kept pure for unit testing.
fn assemble(source: &str, fragments: &[(&str, String)]) -> String {
    let mut out = source.to_string();
    for (marker, text) in fragments {
        out = out.replace(marker, text);
    }
    out
}

// Every `.metal` in the directory except the runtime-assembled template
// fragments, sorted for a deterministic generated file.
fn eligible_shaders(shaders_dir: &Path, source_only: &[&str]) -> Vec<PathBuf> {
    let mut shaders: Vec<PathBuf> = std::fs::read_dir(shaders_dir)
        .unwrap_or_else(|e| panic!("read {}: {e}", shaders_dir.display()))
        .filter_map(|entry| Some(entry.ok()?.path()))
        .filter(|path| {
            path.extension().is_some_and(|ext| ext == "metal")
                && path
                    .file_name()
                    .and_then(|n| n.to_str())
                    .is_some_and(|name| !source_only.contains(&name))
        })
        .collect();
    shaders.sort();
    shaders
}

fn metal_toolchain_present() -> bool {
    Command::new("xcrun")
        .args(["--sdk", "macosx", "-f", "metal"])
        .output()
        .is_ok_and(|out| out.status.success())
}

// Same two-step pipeline the cook's Metal toolchain runs: source to AIR, AIR
// linked into a single-file metallib.
fn compile_metallib(source: &Path, lib_path: &Path) {
    let air_path = lib_path.with_extension("air");
    run_step(
        Command::new("xcrun")
            .args(["--sdk", "macosx", "metal", "-c"])
            .arg(source)
            .arg("-o")
            .arg(&air_path),
        source,
        "xcrun metal",
    );
    run_step(
        Command::new("xcrun")
            .args(["--sdk", "macosx", "metallib"])
            .arg(&air_path)
            .arg("-o")
            .arg(lib_path),
        source,
        "xcrun metallib",
    );
}

fn run_step(cmd: &mut Command, source: &Path, what: &str) {
    let output = cmd
        .output()
        .unwrap_or_else(|e| panic!("{what} failed to launch for {}: {e}", source.display()));
    if !output.status.success() {
        panic!(
            "{what} failed for {}:\n{}\n{}",
            source.display(),
            String::from_utf8_lossy(&output.stdout),
            String::from_utf8_lossy(&output.stderr),
        );
    }
}

// Generated lookup mapping a registered shader name to its precompiled
// metallib bytes. Kept pure (paths in, source out) for unit testing.
fn metallib_lookup_source(entries: &[(String, PathBuf)]) -> String {
    let mut src = String::from(
        "// @generated by concinnity-toolchain::precompile_metal_shaders\n\
         pub(crate) fn embedded_metallib(name: &str) -> Option<&'static [u8]> {\n\
         \x20   match name {\n",
    );
    for (name, lib_path) in entries {
        src.push_str(&format!(
            "        {name:?} => Some(include_bytes!({:?})),\n",
            lib_path.display().to_string()
        ));
    }
    src.push_str("        _ => None,\n    }\n}\n");
    src
}

// Fallback when the Metal toolchain is unavailable at build time.
fn stub_lookup_source() -> String {
    "// @generated by concinnity-toolchain::precompile_metal_shaders (no Metal toolchain)\n\
     pub(crate) fn embedded_metallib(_name: &str) -> Option<&'static [u8]> {\n\
     \x20   None\n\
     }\n"
    .to_string()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn lookup_source_maps_each_entry_and_falls_through() {
        let entries = vec![
            (
                "post.metal".to_string(),
                PathBuf::from("/out/post.metallib"),
            ),
            ("taa.metal".to_string(), PathBuf::from("/out/taa.metallib")),
        ];
        let src = metallib_lookup_source(&entries);
        assert!(src.contains("\"post.metal\" => Some(include_bytes!(\"/out/post.metallib\"))"));
        assert!(src.contains("\"taa.metal\" => Some(include_bytes!(\"/out/taa.metallib\"))"));
        assert!(src.contains("_ => None"));
        assert!(src.contains("fn embedded_metallib"));
    }

    #[test]
    fn stub_lookup_source_returns_none_for_everything() {
        let src = stub_lookup_source();
        assert!(src.contains("fn embedded_metallib"));
        assert!(src.contains("None"));
        assert!(!src.contains("include_bytes!"));
    }

    #[test]
    fn assemble_substitutes_every_marker() {
        let fragments = vec![
            (
                "{OBJECT_DATA}",
                "struct GpuObjectData { float4x4 model; };".to_string(),
            ),
            ("{OTHER}", "// other".to_string()),
        ];
        let out = assemble("head\n{OBJECT_DATA}\nmid\n{OTHER}\ntail\n", &fragments);
        assert!(out.contains("struct GpuObjectData"));
        assert!(out.contains("// other"));
        assert!(!out.contains("{OBJECT_DATA}"));
        assert!(!out.contains("{OTHER}"));
        assert!(out.starts_with("head\n") && out.ends_with("tail\n"));
    }

    #[test]
    fn assemble_replaces_every_occurrence_and_leaves_markerless_source_alone() {
        let fragments = vec![("{OBJECT_DATA}", "RECORD".to_string())];
        assert_eq!(
            assemble("{OBJECT_DATA} a {OBJECT_DATA}", &fragments),
            "RECORD a RECORD"
        );
        assert_eq!(assemble("no markers", &fragments), "no markers");
    }

    #[test]
    fn eligible_shaders_excludes_source_only_and_non_metal() {
        let dir =
            std::env::temp_dir().join(format!("cn_metal_shaders_test_{}", std::process::id()));
        std::fs::create_dir_all(&dir).unwrap();
        for name in ["a.metal", "b.metal", "template.metal", "notes.txt"] {
            std::fs::write(dir.join(name), "").unwrap();
        }
        let shaders = eligible_shaders(&dir, &["template.metal"]);
        let names: Vec<_> = shaders
            .iter()
            .map(|p| p.file_name().unwrap().to_str().unwrap().to_string())
            .collect();
        assert_eq!(names, ["a.metal", "b.metal"]);
        std::fs::remove_dir_all(&dir).unwrap();
    }
}