rlx-vulkan 0.2.10

//! Compile the GLSL compute shaders under `shaders/*.comp` to SPIR-V at
//! build time using `naga` (pure Rust). The resulting `.spv` blobs are
//! written to `OUT_DIR` and a generated `shaders_generated.rs` embeds them
//! via `include_bytes!`. No external toolchain (glslang / shaderc / Vulkan
//! SDK) is required, and there is no runtime shader compilation.

use std::fmt::Write as _;
use std::fs;
use std::path::Path;

fn main() {
    let manifest_dir = std::env::var("CARGO_MANIFEST_DIR").unwrap();
    let shader_dir = Path::new(&manifest_dir).join("shaders");
    let out_dir = std::env::var("OUT_DIR").unwrap();

    println!("cargo:rerun-if-changed=shaders");
    println!("cargo:rerun-if-changed=build.rs");

    let mut entries: Vec<String> = Vec::new(); // shader names (sorted)

    let mut files: Vec<_> = fs::read_dir(&shader_dir)
        .unwrap_or_else(|e| panic!("rlx-vulkan: cannot read {}: {e}", shader_dir.display()))
        .filter_map(Result::ok)
        .map(|e| e.path())
        .filter(|p| p.extension().map(|x| x == "comp").unwrap_or(false))
        .collect();
    files.sort();

    for path in &files {
        println!("cargo:rerun-if-changed={}", path.display());
        let name = path.file_stem().unwrap().to_string_lossy().to_string();
        let src = fs::read_to_string(path)
            .unwrap_or_else(|e| panic!("rlx-vulkan: read {}: {e}", path.display()));

        let words = compile_glsl_to_spirv(&name, &src);

        // Write SPIR-V words as little-endian bytes.
        let mut bytes = Vec::with_capacity(words.len() * 4);
        for w in &words {
            bytes.extend_from_slice(&w.to_le_bytes());
        }
        let spv_path = Path::new(&out_dir).join(format!("{name}.spv"));
        fs::write(&spv_path, &bytes)
            .unwrap_or_else(|e| panic!("rlx-vulkan: write {}: {e}", spv_path.display()));
        entries.push(name);
    }

    // Pre-compiled SPIR-V kernels under `shaders/precompiled/*.spv` — kernels
    // naga can't build from GLSL (cooperative matrix / f16). The `.comp` source
    // lives beside each `.spv` for reference, but the committed `.spv` is the
    // build input (no glslang/SDK needed here). Copy into OUT_DIR so they embed
    // with the same `include_bytes!` pattern as the naga-compiled kernels.
    let precompiled_dir = shader_dir.join("precompiled");
    if precompiled_dir.is_dir() {
        let mut spvs: Vec<_> = fs::read_dir(&precompiled_dir)
            .unwrap_or_else(|e| {
                panic!("rlx-vulkan: cannot read {}: {e}", precompiled_dir.display())
            })
            .filter_map(Result::ok)
            .map(|e| e.path())
            .filter(|p| p.extension().map(|x| x == "spv").unwrap_or(false))
            .collect();
        spvs.sort();
        for path in &spvs {
            println!("cargo:rerun-if-changed={}", path.display());
            let name = path.file_stem().unwrap().to_string_lossy().to_string();
            let bytes = fs::read(path)
                .unwrap_or_else(|e| panic!("rlx-vulkan: read {}: {e}", path.display()));
            let spv_path = Path::new(&out_dir).join(format!("{name}.spv"));
            fs::write(&spv_path, &bytes)
                .unwrap_or_else(|e| panic!("rlx-vulkan: write {}: {e}", spv_path.display()));
            entries.push(name);
        }
    }

    // Emit the registry source. Reference the `.spv` blobs RELATIVE to OUT_DIR
    // (resolved by `env!("OUT_DIR")` at the crate's compile time) rather than
    // baking absolute paths — so the embed survives a moved/relocated target
    // dir (e.g. a Docker volume mounted at a different path).
    let mut out_src = String::new();
    out_src.push_str("// @generated by build.rs — GLSL→SPIR-V compute kernels.\n");
    out_src.push_str("/// (kernel name, SPIR-V byte blob) for every shader under `shaders/`.\n");
    out_src.push_str("pub static SHADER_BLOBS: &[(&str, &[u8])] = &[\n");
    for name in &entries {
        writeln!(
            out_src,
            "    ({name:?}, include_bytes!(concat!(env!(\"OUT_DIR\"), \"/{name}.spv\"))),"
        )
        .unwrap();
    }
    out_src.push_str("];\n");
    let gen_path = Path::new(&out_dir).join("shaders_generated.rs");
    fs::write(&gen_path, out_src).unwrap();
}

fn compile_glsl_to_spirv(name: &str, src: &str) -> Vec<u32> {
    use naga::ShaderStage;
    use naga::back::spv;
    use naga::front::glsl::{Frontend, Options};
    use naga::valid::{Capabilities, ValidationFlags, Validator};

    let options = Options::from(ShaderStage::Compute);
    let module = Frontend::default()
        .parse(&options, src)
        .unwrap_or_else(|e| panic!("rlx-vulkan: GLSL parse error in {name}.comp: {e:?}"));

    let info = Validator::new(ValidationFlags::all(), Capabilities::all())
        .validate(&module)
        .unwrap_or_else(|e| panic!("rlx-vulkan: validation error in {name}.comp: {e:?}"));

    let spv_opts = spv::Options::default();
    let pipe_opts = spv::PipelineOptions {
        shader_stage: ShaderStage::Compute,
        entry_point: "main".to_string(),
    };
    spv::write_vec(&module, &info, &spv_opts, Some(&pipe_opts))
        .unwrap_or_else(|e| panic!("rlx-vulkan: SPIR-V emit error in {name}.comp: {e:?}"))
}