spirv-builder: add RUSTGPU_RUSTFLAGS/RUSTGPU_CODEGEN_ARGS env vars.

CHANGELOG.md

@@ -29,6 +29,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added ⭐
+
+- Added two `spirv-builder` environment variables to customize *only* the `rustc` invocations for shader crates and their dependencies:
+    - `RUSTGPU_RUSTFLAGS="..."` for shader `RUSTFLAGS="..."`
+    - `RUSTGPU_CODEGEN_ARGS="..."` for shader `RUSTFLAGS=-Cllvm-args="..."`  
+      (run with `RUSTGPU_CODEGEN_ARGS=--help` to see a list of options)
+
 ### Changed 🛠️
 
 - Updated toolchain to `nightly-2022-10-29`

crates/rustc_codegen_spirv/src/codegen_cx/mod.rs

@@ -31,6 +31,7 @@ use rustc_target::abi::call::{FnAbi, PassMode};
 use rustc_target::abi::{HasDataLayout, TargetDataLayout};
 use rustc_target::spec::{HasTargetSpec, Target};
 use std::cell::{Cell, RefCell};
+use std::collections::BTreeSet;
 use std::iter::once;
 use std::path::Path;
 use std::rc::Rc;
@@ -308,12 +309,46 @@ impl CodegenArgs {
 
         let matches = opts.parse(args)?;
 
-        if matches.opt_present("h") || matches.opt_present("help") {
-            // FIXME(eddyb) make this somehow nicer wrt end-users
-            // (mentioning `-Cllvm-args` is pretty bad, we'd want to hide that).
+        let help_flag_positions: BTreeSet<_> = ["h", "help"]
+            .iter()
+            .flat_map(|&name| matches.opt_positions(name))
+            .collect();
+        if !help_flag_positions.is_empty() {
+            // HACK(eddyb) this tries to be a bit nicer to end-users, when they
+            // use `spirv-builder` (and so the `RUSTGPU_CODEGEN_ARGS` env var,
+            // to set codegen args), as mentioning `-Cllvm-args` is suboptimal.
+            let spirv_builder_env_var = "RUSTGPU_CODEGEN_ARGS";
+            let help_flag_comes_from_spirv_builder_env_var = std::env::var(spirv_builder_env_var)
+                .ok()
+                .and_then(|args_from_env| {
+                    let args_from_env: Vec<_> = args_from_env.split_whitespace().collect();
+                    if args_from_env.is_empty() {
+                        return None;
+                    }
+
+                    // HACK(eddyb) this may be a bit inefficient but we want to
+                    // make sure that *at least one* of the `-h`/`--help` flags
+                    // came from the `spirv-builder`-supported env var *and*
+                    // that the env var's contents are fully contained in the
+                    // `-C llvm-args` this `rustc` invocation is seeing.
+                    args.windows(args_from_env.len())
+                        .enumerate()
+                        .filter(|&(_, w)| w == args_from_env)
+                        .map(|(w_start, w)| w_start..w_start + w.len())
+                        .flat_map(|w_range| help_flag_positions.range(w_range))
+                        .next()
+                })
+                .is_some();
+            let codegen_args_lhs = if help_flag_comes_from_spirv_builder_env_var {
+                spirv_builder_env_var
+            } else {
+                "rustc -Cllvm-args"
+            };
             println!(
                 "{}",
-                opts.usage("Usage: rustc -Cllvm-args=\"[OPTIONS]\" with:")
+                opts.usage(&format!(
+                    "Usage: {codegen_args_lhs}=\"...\" with `...` from:"
+                ))
             );
             // HACK(eddyb) this avoids `Cargo` continuing after the message is printed.
             std::process::exit(1);

crates/spirv-builder/src/lib.rs

@@ -428,36 +428,51 @@ fn invoke_rustc(builder: &SpirvBuilder) -> Result<PathBuf, SpirvBuilderError> {
         "-Zcrate-attr=register_tool(rust_gpu)".to_string(),
     ];
 
+    // Wrapper for `env::var` that appropriately informs Cargo of the dependency.
+    let tracked_env_var_get = |name| {
+        if let MetadataPrintout::Full | MetadataPrintout::DependencyOnly = builder.print_metadata {
+            println!("cargo:rerun-if-env-changed={name}");
+        }
+        env::var(name)
+    };
+
     let mut llvm_args = vec![];
     if builder.multimodule {
-        llvm_args.push("--module-output=multiple");
+        llvm_args.push("--module-output=multiple".to_string());
     }
     match builder.spirv_metadata {
         SpirvMetadata::None => (),
-        SpirvMetadata::NameVariables => llvm_args.push("--spirv-metadata=name-variables"),
-        SpirvMetadata::Full => llvm_args.push("--spirv-metadata=full"),
+        SpirvMetadata::NameVariables => {
+            llvm_args.push("--spirv-metadata=name-variables".to_string());
+        }
+        SpirvMetadata::Full => llvm_args.push("--spirv-metadata=full".to_string()),
     }
     if builder.relax_struct_store {
-        llvm_args.push("--relax-struct-store");
+        llvm_args.push("--relax-struct-store".to_string());
     }
     if builder.relax_logical_pointer {
-        llvm_args.push("--relax-logical-pointer");
+        llvm_args.push("--relax-logical-pointer".to_string());
     }
     if builder.relax_block_layout {
-        llvm_args.push("--relax-block-layout");
+        llvm_args.push("--relax-block-layout".to_string());
     }
     if builder.uniform_buffer_standard_layout {
-        llvm_args.push("--uniform-buffer-standard-layout");
+        llvm_args.push("--uniform-buffer-standard-layout".to_string());
     }
     if builder.scalar_block_layout {
-        llvm_args.push("--scalar-block-layout");
+        llvm_args.push("--scalar-block-layout".to_string());
     }
     if builder.skip_block_layout {
-        llvm_args.push("--skip-block-layout");
+        llvm_args.push("--skip-block-layout".to_string());
     }
     if builder.preserve_bindings {
-        llvm_args.push("--preserve-bindings");
+        llvm_args.push("--preserve-bindings".to_string());
     }
+
+    if let Ok(extra_codegen_args) = tracked_env_var_get("RUSTGPU_CODEGEN_ARGS") {
+        llvm_args.extend(extra_codegen_args.split_whitespace().map(|s| s.to_string()));
+    }
+
     let llvm_args = join_checking_for_separators(llvm_args, " ");
     if !llvm_args.is_empty() {
         rustflags.push(["-Cllvm-args=", &llvm_args].concat());
@@ -475,6 +490,10 @@ fn invoke_rustc(builder: &SpirvBuilder) -> Result<PathBuf, SpirvBuilderError> {
         rustflags.push("-Dwarnings".to_string());
     }
 
+    if let Ok(extra_rustflags) = tracked_env_var_get("RUSTGPU_RUSTFLAGS") {
+        rustflags.extend(extra_rustflags.split_whitespace().map(|s| s.to_string()));
+    }
+
     let mut cargo = Command::new("cargo");
     cargo.args([
         "build",