Always emit json metadata, and emit entry point names (#622)

* Always emit json metadata

* Codegen shader entry point names

* Fix tests

crates/rustc_codegen_spirv/src/compile_result.rs

@@ -0,0 +1,169 @@
+use rustc_data_structures::fx::FxHashMap;
+use serde::{Deserialize, Serialize};
+use std::fmt::Write;
+use std::path::{Path, PathBuf};
+
+#[derive(Debug, Serialize, Deserialize)]
+#[serde(untagged)]
+pub enum ModuleResult {
+    SingleModule(PathBuf),
+    MultiModule(FxHashMap<String, PathBuf>),
+}
+
+impl ModuleResult {
+    pub fn unwrap_single(&self) -> &Path {
+        match self {
+            ModuleResult::SingleModule(result) => result,
+            ModuleResult::MultiModule(_) => {
+                panic!("called `ModuleResult::unwrap_single()` on a `MultiModule` result")
+            }
+        }
+    }
+
+    pub fn unwrap_multi(&self) -> &FxHashMap<String, PathBuf> {
+        match self {
+            ModuleResult::MultiModule(result) => result,
+            ModuleResult::SingleModule(_) => {
+                panic!("called `ModuleResult::unwrap_multi()` on a `SingleModule` result")
+            }
+        }
+    }
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct CompileResult {
+    pub module: ModuleResult,
+    pub entry_points: Vec<String>,
+}
+
+impl CompileResult {
+    pub fn codegen_entry_point_strings(&self) -> String {
+        let trie = Trie::create_from(self.entry_points.iter().map(|x| x as &str));
+        let mut builder = String::new();
+        trie.emit(&mut builder, String::new(), 0);
+        builder
+    }
+}
+
+#[derive(Default)]
+struct Trie {
+    present: bool,
+    children: FxHashMap<String, Box<Trie>>,
+}
+
+impl Trie {
+    fn create_from<'a>(entry_points: impl IntoIterator<Item = &'a str>) -> Self {
+        let mut result = Trie::default();
+        for entry in entry_points {
+            result.insert(entry.split("::").map(|x| x.to_owned()));
+        }
+        result
+    }
+
+    fn insert(&mut self, mut sequence: impl Iterator<Item = String>) {
+        match sequence.next() {
+            None => self.present = true,
+            Some(next) => self
+                .children
+                .entry(next)
+                .or_insert_with(Default::default)
+                .insert(sequence),
+        }
+    }
+
+    fn emit(&self, builder: &mut String, full_name: String, indent: usize) {
+        let mut children = self.children.iter().collect::<Vec<_>>();
+        children.sort_unstable_by(|(k1, _), (k2, _)| k1.cmp(k2));
+        for (child_name, child) in children {
+            let full_child_name = if full_name.is_empty() {
+                child_name.to_string()
+            } else {
+                format!("{}::{}", full_name, child_name)
+            };
+            if child.present {
+                assert!(child.children.is_empty());
+                writeln!(
+                    builder,
+                    "{:indent$}#[allow(non_upper_case_globals)]",
+                    "",
+                    indent = indent * 4
+                )
+                .unwrap();
+                writeln!(
+                    builder,
+                    "{:indent$}pub const {}: &str = \"{}\";",
+                    "",
+                    child_name,
+                    full_child_name,
+                    indent = indent * 4
+                )
+                .unwrap();
+            } else {
+                writeln!(
+                    builder,
+                    "{:indent$}pub mod {} {{",
+                    "",
+                    child_name,
+                    indent = indent * 4
+                )
+                .unwrap();
+                child.emit(builder, full_child_name, indent + 1);
+                writeln!(builder, "{:indent$}}}", "", indent = indent * 4).unwrap();
+            }
+        }
+    }
+}
+
+#[allow(non_upper_case_globals)]
+pub const a: &str = "x::a";
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use std::array::IntoIter;
+
+    fn test<const N: usize>(arr: [&str; N], expected: &str) {
+        let trie = Trie::create_from(IntoIter::new(arr));
+        let mut builder = String::new();
+        trie.emit(&mut builder, String::new(), 0);
+        assert_eq!(builder, expected);
+    }
+
+    #[test]
+    fn basic() {
+        test(
+            ["a", "b"],
+            r#"#[allow(non_upper_case_globals)]
+pub const a: &str = "a";
+#[allow(non_upper_case_globals)]
+pub const b: &str = "b";
+"#,
+        );
+    }
+
+    #[test]
+    fn modules() {
+        test(
+            ["a", "x::a", "x::b", "x::y::a", "y::z::a"],
+            r#"#[allow(non_upper_case_globals)]
+pub const a: &str = "a";
+pub mod x {
+    #[allow(non_upper_case_globals)]
+    pub const a: &str = "x::a";
+    #[allow(non_upper_case_globals)]
+    pub const b: &str = "x::b";
+    pub mod y {
+        #[allow(non_upper_case_globals)]
+        pub const a: &str = "x::y::a";
+    }
+}
+pub mod y {
+    pub mod z {
+        #[allow(non_upper_case_globals)]
+        pub const a: &str = "y::z::a";
+    }
+}
+"#,
+        );
+    }
+}

crates/rustc_codegen_spirv/src/lib.rs

@@ -121,6 +121,7 @@ mod attr;
 mod builder;
 mod builder_spirv;
 mod codegen_cx;
+mod compile_result;
 mod decorations;
 mod link;
 mod linker;
@@ -132,6 +133,7 @@ mod target_feature;
 
 use builder::Builder;
 use codegen_cx::{CodegenArgs, CodegenCx, ModuleOutputType};
+pub use compile_result::*;
 pub use rspirv;
 use rspirv::binary::Assemble;
 use rustc_ast::expand::allocator::AllocatorKind;

crates/rustc_codegen_spirv/src/link.rs

@@ -1,4 +1,6 @@
-use crate::{linker, SpirvCodegenBackend, SpirvModuleBuffer, SpirvThinBuffer};
+use crate::{
+    linker, CompileResult, ModuleResult, SpirvCodegenBackend, SpirvModuleBuffer, SpirvThinBuffer,
+};
 use rustc_codegen_ssa::back::lto::{LtoModuleCodegen, SerializedModule, ThinModule, ThinShared};
 use rustc_codegen_ssa::back::write::CodegenContext;
 use rustc_codegen_ssa::{CodegenResults, NativeLib};
@@ -139,30 +141,54 @@ fn link_exe(
 
     let cg_args = crate::codegen_cx::CodegenArgs::from_session(sess);
 
-    use rspirv::binary::Assemble;
-    match spv_binary {
-        linker::LinkResult::SingleModule(spv_binary) => {
-            post_link_single_module(sess, spv_binary.assemble(), out_filename);
-            cg_args.do_disassemble(&spv_binary);
-        }
-        linker::LinkResult::MultipleModules(map) => {
     let mut root_file_name = out_filename.file_name().unwrap().to_owned();
     root_file_name.push(".dir");
     let out_dir = out_filename.with_file_name(root_file_name);
     if !out_dir.is_dir() {
         std::fs::create_dir_all(&out_dir).unwrap();
     }
+
+    use rspirv::binary::Assemble;
+    let compile_result = match spv_binary {
+        linker::LinkResult::SingleModule(spv_binary) => {
+            let mut module_filename = out_dir;
+            module_filename.push("module");
+            post_link_single_module(sess, spv_binary.assemble(), &module_filename);
+            cg_args.do_disassemble(&spv_binary);
+            let module_result = ModuleResult::SingleModule(module_filename);
+            CompileResult {
+                module: module_result,
+                entry_points: entry_points(&spv_binary),
+            }
+        }
+        linker::LinkResult::MultipleModules(map) => {
             let mut hashmap = FxHashMap::default();
+            let entry_points = map.keys().cloned().collect();
             for (name, spv_binary) in map {
                 let mut module_filename = out_dir.clone();
                 module_filename.push(sanitize_filename::sanitize(&name));
                 post_link_single_module(sess, spv_binary.assemble(), &module_filename);
                 hashmap.insert(name, module_filename);
             }
+            let module_result = ModuleResult::MultiModule(hashmap);
+            CompileResult {
+                module: module_result,
+                entry_points,
+            }
+        }
+    };
+
     let file = File::create(out_filename).unwrap();
-            serde_json::to_writer(BufWriter::new(file), &hashmap).unwrap();
-        }
+    serde_json::to_writer(BufWriter::new(file), &compile_result).unwrap();
 }
+
+fn entry_points(module: &rspirv::dr::Module) -> Vec<String> {
+    module
+        .entry_points
+        .iter()
+        .filter(|inst| inst.class.opcode == rspirv::spirv::Op::EntryPoint)
+        .map(|inst| inst.operands[2].unwrap_literal_string().to_string())
+        .collect()
 }
 
 fn post_link_single_module(sess: &Session, spv_binary: Vec<u32>, out_filename: &Path) {

crates/spirv-builder/src/lib.rs

@@ -66,8 +66,11 @@ use std::io::BufReader;
 use std::path::{Path, PathBuf};
 use std::process::{Command, Stdio};
 
+pub use rustc_codegen_spirv::{CompileResult, ModuleResult};
+
 #[derive(Debug)]
 pub enum SpirvBuilderError {
+    CratePathDoesntExist(PathBuf),
     BuildFailed,
     MultiModuleWithPrintMetadata,
     MetadataFileMissing(std::io::Error),
@@ -77,6 +80,9 @@ pub enum SpirvBuilderError {
 impl fmt::Display for SpirvBuilderError {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         match self {
+            SpirvBuilderError::CratePathDoesntExist(path) => {
+                write!(f, "Crate path {} does not exist", path.display())
+            }
             SpirvBuilderError::BuildFailed => f.write_str("Build failed"),
             SpirvBuilderError::MultiModuleWithPrintMetadata => {
                 f.write_str("Multi-module build cannot be used with print_metadata = true")
@@ -105,6 +111,7 @@ pub struct SpirvBuilder {
     release: bool,
     target: String,
     bindless: bool,
+    multimodule: bool,
 }
 
 impl SpirvBuilder {
@@ -113,8 +120,9 @@ impl SpirvBuilder {
             path_to_crate: path_to_crate.as_ref().to_owned(),
             print_metadata: true,
             release: true,
-            bindless: false,
             target: target.into(),
+            bindless: false,
+            multimodule: false,
         }
     }
 
@@ -137,27 +145,40 @@ impl SpirvBuilder {
         self
     }
 
-    /// Builds the module. Returns the path to the built spir-v file. If `print_metadata` is true,
-    /// you usually don't have to inspect the path, as the environment variable will already be
-    /// set.
-    pub fn build(self) -> Result<PathBuf, SpirvBuilderError> {
-        let spirv_module = invoke_rustc(&self, false)?;
-        let env_var = spirv_module.file_name().unwrap().to_str().unwrap();
+    /// Splits the resulting SPIR-V file into one module per entry point. This is useful in cases
+    /// where ecosystem tooling has bugs around multiple entry points per module - having all entry
+    /// points bundled into a single file is the preferred system.
+    pub fn multimodule(mut self, v: bool) -> Self {
+        self.multimodule = v;
+        self
+    }
+
+    /// Builds the module. If `print_metadata` is true, you usually don't have to inspect the path
+    /// in the result, as the environment variable for the path to the module will already be set.
+    pub fn build(self) -> Result<CompileResult, SpirvBuilderError> {
+        if self.print_metadata && self.multimodule {
+            return Err(SpirvBuilderError::MultiModuleWithPrintMetadata);
+        }
+        if !self.path_to_crate.is_dir() {
+            return Err(SpirvBuilderError::CratePathDoesntExist(self.path_to_crate));
+        }
+        let metadata_file = invoke_rustc(&self)?;
+        let metadata_contents =
+            File::open(&metadata_file).map_err(SpirvBuilderError::MetadataFileMissing)?;
+        let metadata: CompileResult = serde_json::from_reader(BufReader::new(metadata_contents))
+            .map_err(SpirvBuilderError::MetadataFileMalformed)?;
+        match &metadata.module {
+            ModuleResult::SingleModule(spirv_module) => {
+                assert!(!self.multimodule);
+                let env_var = metadata_file.file_name().unwrap().to_str().unwrap();
                 if self.print_metadata {
                     println!("cargo:rustc-env={}={}", env_var, spirv_module.display());
                 }
-        Ok(spirv_module)
             }
-
-    pub fn build_multimodule(self) -> Result<HashMap<String, PathBuf>, SpirvBuilderError> {
-        if self.print_metadata {
-            return Err(SpirvBuilderError::MultiModuleWithPrintMetadata);
+            ModuleResult::MultiModule(_) => {
+                assert!(self.multimodule);
+            }
         }
-        let metadata_file = invoke_rustc(&self, true)?;
-        let metadata_contents =
-            File::open(metadata_file).map_err(SpirvBuilderError::MetadataFileMissing)?;
-        let metadata = serde_json::from_reader(BufReader::new(metadata_contents))
-            .map_err(SpirvBuilderError::MetadataFileMalformed)?;
         Ok(metadata)
     }
 }
@@ -194,8 +215,8 @@ fn find_rustc_codegen_spirv() -> PathBuf {
     panic!("Could not find {} in library path", filename);
 }
 
-// Note: in case of multimodule, returns path to the metadata json
-fn invoke_rustc(builder: &SpirvBuilder, multimodule: bool) -> Result<PathBuf, SpirvBuilderError> {
+// Returns path to the metadata json.
+fn invoke_rustc(builder: &SpirvBuilder) -> Result<PathBuf, SpirvBuilderError> {
     // Okay, this is a little bonkers: in a normal world, we'd have the user clone
     // rustc_codegen_spirv and pass in the path to it, and then we'd invoke cargo to build it, grab
     // the resulting .so, and pass it into -Z codegen-backend. But that's really gross: the user
@@ -205,7 +226,8 @@ fn invoke_rustc(builder: &SpirvBuilder, multimodule: bool) -> Result<PathBuf, Sp
     // rustc expects a full path, instead of a filename looked up via LD_LIBRARY_PATH, so we need
     // to copy cargo's understanding of library lookup and find the library and its full path.
     let rustc_codegen_spirv = find_rustc_codegen_spirv();
-    let llvm_args = multimodule
+    let llvm_args = builder
+        .multimodule
         .then(|| " -C llvm-args=--module-output=multiple")
         .unwrap_or_default();
 

examples/multibuilder/src/main.rs

@@ -3,7 +3,8 @@ use spirv_builder::SpirvBuilder;
 fn main() {
     let result = SpirvBuilder::new("../shaders/sky-shader", "spirv-unknown-spv1.3")
         .print_metadata(false)
-        .build_multimodule()
+        .multimodule(true)
+        .build()
         .unwrap();
     println!("{:#?}", result);
 }

examples/runners/ash/src/main.rs

@@ -176,7 +176,10 @@ pub fn compile_shaders() -> Vec<SpvFile> {
             SpirvBuilder::new("examples/shaders/sky-shader", "spirv-unknown-vulkan1.1")
                 .print_metadata(false)
                 .build()
-                .unwrap(),
+                .unwrap()
+                .module
+                .unwrap_single()
+                .to_path_buf(),
         ];
     let mut spv_files = Vec::<SpvFile>::with_capacity(spv_paths.len());
     for path in spv_paths.iter() {

examples/runners/wgpu/builder/src/main.rs

@@ -1,15 +1,23 @@
 use spirv_builder::SpirvBuilder;
+use std::env;
 use std::error::Error;
+use std::fs;
+use std::path::Path;
 
-fn build_shader(path_to_create: &str) -> Result<(), Box<dyn Error>> {
-    SpirvBuilder::new(path_to_create, "spirv-unknown-vulkan1.0").build()?;
+fn build_shader(path_to_create: &str, codegen_names: bool) -> Result<(), Box<dyn Error>> {
+    let result = SpirvBuilder::new(path_to_create, "spirv-unknown-vulkan1.0").build()?;
+    if codegen_names {
+        let out_dir = env::var_os("OUT_DIR").unwrap();
+        let dest_path = Path::new(&out_dir).join("entry_points.rs");
+        fs::write(&dest_path, result.codegen_entry_point_strings()).unwrap();
+    }
     Ok(())
 }
 
 fn main() -> Result<(), Box<dyn Error>> {
-    build_shader("../../shaders/sky-shader")?;
-    build_shader("../../shaders/simplest-shader")?;
-    build_shader("../../shaders/compute-shader")?;
-    build_shader("../../shaders/mouse-shader")?;
+    build_shader("../../shaders/sky-shader", true)?;
+    build_shader("../../shaders/simplest-shader", false)?;
+    build_shader("../../shaders/compute-shader", false)?;
+    build_shader("../../shaders/mouse-shader", false)?;
     Ok(())
 }

examples/runners/wgpu/src/graphics.rs

@@ -6,6 +6,21 @@ use winit::{
     window::Window,
 };
 
+#[cfg(not(any(target_os = "android", target_arch = "wasm32")))]
+mod shaders {
+    // The usual usecase of code generation is always building in build.rs, and so the codegen
+    // always happens. However, we want to both test code generation (on android) and runtime
+    // compilation (on desktop), so manually fill in what would have been codegenned for desktop.
+    #[allow(non_upper_case_globals)]
+    pub const main_fs: &str = "main_fs";
+    #[allow(non_upper_case_globals)]
+    pub const main_vs: &str = "main_vs";
+}
+#[cfg(any(target_os = "android", target_arch = "wasm32"))]
+mod shaders {
+    include!(concat!(env!("OUT_DIR"), "/entry_points.rs"));
+}
+
 unsafe fn any_as_u8_slice<T: Sized>(p: &T) -> &[u8] {
     ::std::slice::from_raw_parts((p as *const T) as *const u8, ::std::mem::size_of::<T>())
 }
@@ -81,7 +96,7 @@ async fn run(
         layout: Some(&pipeline_layout),
         vertex: wgpu::VertexState {
             module: &module,
-            entry_point: "main_vs",
+            entry_point: shaders::main_vs,
             buffers: &[],
         },
         primitive: wgpu::PrimitiveState {
@@ -99,7 +114,7 @@ async fn run(
         },
         fragment: Some(wgpu::FragmentState {
             module: &module,
-            entry_point: "main_fs",
+            entry_point: shaders::main_fs,
             targets: &[wgpu::ColorTargetState {
                 format: swapchain_format,
                 alpha_blend: wgpu::BlendState::REPLACE,

examples/runners/wgpu/src/lib.rs

@@ -87,11 +87,12 @@ fn shader_module(shader: RustGPUShader) -> wgpu::ShaderModuleDescriptor<'static>
         .iter()
         .copied()
         .collect::<PathBuf>();
-        let result = SpirvBuilder::new(crate_path, "spirv-unknown-vulkan1.0")
+        let compile_result = SpirvBuilder::new(crate_path, "spirv-unknown-vulkan1.0")
             .print_metadata(false)
             .build()
             .unwrap();
-        let data = std::fs::read(result).unwrap();
+        let module_path = compile_result.module.unwrap_single();
+        let data = std::fs::read(module_path).unwrap();
         let spirv = wgpu::util::make_spirv(&data);
         let spirv = match spirv {
             wgpu::ShaderSource::Wgsl(cow) => wgpu::ShaderSource::Wgsl(Cow::Owned(cow.into_owned())),

tests/ui/lang/consts/nested-ref-in-composite.stderr

@@ -21,7 +21,7 @@ error: constant arrays/structs cannot contain pointers to other constants
 error: error:0:0 - No OpEntryPoint instruction was found. This is only allowed if the Linkage capability is being used.
   |
   = note: spirv-val failed
-  = note: module `$TEST_BUILD_DIR/lang/consts/nested-ref-in-composite.stage-id.spv`
+  = note: module `$TEST_BUILD_DIR/lang/consts/nested-ref-in-composite.stage-id.spv.dir/module`
 
 error: aborting due to 3 previous errors
 

tests/ui/lang/core/ptr/allocate_const_scalar.stderr

@@ -7,7 +7,7 @@ error: pointer has non-null integer address
 error: error:0:0 - No OpEntryPoint instruction was found. This is only allowed if the Linkage capability is being used.
   |
   = note: spirv-val failed
-  = note: module `$TEST_BUILD_DIR/lang/core/ptr/allocate_const_scalar.stage-id.spv`
+  = note: module `$TEST_BUILD_DIR/lang/core/ptr/allocate_const_scalar.stage-id.spv.dir/module`
 
 error: aborting due to 2 previous errors