From a48588b11b1d3bfd57e4dfbc539ab74ec630201f Mon Sep 17 00:00:00 2001
From: Eduard-Mihai Burtescu <eddyb@lyken.rs>
Date: Mon, 24 Jul 2023 18:52:24 +0300
Subject: [PATCH] Update `spirt` for better `OpExtInst` support
 (https://github.com/EmbarkStudios/spirt/pull/45).

---
 CHANGELOG.md                                  |  3 +
 Cargo.lock                                    |  2 +-
 .../rustc_codegen_spirv/src/custom_insts.rs   | 56 +++++++++++++++++--
 crates/rustc_codegen_spirv/src/linker/mod.rs  |  1 +
 4 files changed, 55 insertions(+), 7 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index f5ed85dc6b..7aa18accd4 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -46,6 +46,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   (see its documentation for more details about each available panic handling strategy)
 
 ### Changed 🛠
+- [PR#1083](https://github.com/EmbarkStudios/rust-gpu/pull/1083) updated SPIR-T to get pretty-printer
+  improvements (especially for `OpExtInst`, including Rust-GPU's custom ones), and started more
+  aggressively deduplicating custom debuginfo instructions (to make SPIR-T dumps more readable)
 - [PR#1079](https://github.com/EmbarkStudios/rust-gpu/pull/1079) revised `spirv-builder`'s `README.md`,
   and added a way for `docs.rs` to be able to build it (via `cargo +stable doc --no-default-features`)
 - [PR#1070](https://github.com/EmbarkStudios/rust-gpu/pull/1070) made panics (via the `abort` intrinsic)
diff --git a/Cargo.lock b/Cargo.lock
index 2f917f4d93..065dfe5ca2 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2280,7 +2280,7 @@ dependencies = [
 [[package]]
 name = "spirt"
 version = "0.2.0"
-source = "git+https://github.com/EmbarkStudios/spirt?branch=main#cd21d968127035ad3707829849d0466368649299"
+source = "git+https://github.com/EmbarkStudios/spirt?branch=main#82daf2516710504986cdc35e0e27455326aeef90"
 dependencies = [
  "arrayvec",
  "bytemuck",
diff --git a/crates/rustc_codegen_spirv/src/custom_insts.rs b/crates/rustc_codegen_spirv/src/custom_insts.rs
index 8b7c9db1bf..d5721c2485 100644
--- a/crates/rustc_codegen_spirv/src/custom_insts.rs
+++ b/crates/rustc_codegen_spirv/src/custom_insts.rs
@@ -15,6 +15,18 @@ use smallvec::SmallVec;
 /// See `CUSTOM_EXT_INST_SET`'s docs for further constraints on the full name.
 pub const CUSTOM_EXT_INST_SET_PREFIX: &str = concat!("Rust.", env!("CARGO_PKG_NAME"), ".");
 
+macro_rules! join_cargo_pkg_version_major_minor_patch {
+    ($sep:literal) => {
+        concat!(
+            env!("CARGO_PKG_VERSION_MAJOR"),
+            $sep,
+            env!("CARGO_PKG_VERSION_MINOR"),
+            $sep,
+            env!("CARGO_PKG_VERSION_PATCH"),
+        )
+    };
+}
+
 lazy_static! {
     /// `OpExtInstImport` "instruction set" name for all Rust-GPU instructions.
     ///
@@ -30,10 +42,6 @@ lazy_static! {
     ///   if the definitions of the custom instructions have changed - this is
     ///   achieved by hashing the `SCHEMA` constant from `def_custom_insts!` below
     pub static ref CUSTOM_EXT_INST_SET: String = {
-        const VER_MAJOR: &str = env!("CARGO_PKG_VERSION_MAJOR");
-        const VER_MINOR: &str = env!("CARGO_PKG_VERSION_MINOR");
-        const VER_PATCH: &str = env!("CARGO_PKG_VERSION_PATCH");
-
         let schema_hash = {
             use rustc_data_structures::stable_hasher::StableHasher;
             use std::hash::Hash;
@@ -43,11 +51,47 @@ lazy_static! {
             let (lo, hi) = hasher.finalize();
             (lo as u128) | ((hi as u128) << 64)
         };
-
-        format!("{CUSTOM_EXT_INST_SET_PREFIX}{VER_MAJOR}_{VER_MINOR}_{VER_PATCH}.{schema_hash:x}")
+        let version = join_cargo_pkg_version_major_minor_patch!("_");
+        format!("{CUSTOM_EXT_INST_SET_PREFIX}{version}.{schema_hash:x}")
     };
 }
 
+pub fn register_to_spirt_context(cx: &spirt::Context) {
+    use spirt::spv::spec::{ExtInstSetDesc, ExtInstSetInstructionDesc};
+    cx.register_custom_ext_inst_set(
+        &CUSTOM_EXT_INST_SET,
+        ExtInstSetDesc {
+            // HACK(eddyb) this is the most compact form I've found, that isn't
+            // outright lossy by omitting "Rust vs Rust-GPU" or the version.
+            short_alias: Some(
+                concat!("Rust-GPU ", join_cargo_pkg_version_major_minor_patch!(".")).into(),
+            ),
+            instructions: SCHEMA
+                .iter()
+                .map(|&(i, name, operand_names)| {
+                    (
+                        i,
+                        ExtInstSetInstructionDesc {
+                            name: name.into(),
+                            operand_names: operand_names
+                                .iter()
+                                .map(|name| {
+                                    name.strip_prefix("..")
+                                        .unwrap_or(name)
+                                        .replace('_', " ")
+                                        .into()
+                                })
+                                .collect(),
+                            is_debuginfo: name.contains("Debug")
+                                || name.contains("InlinedCallFrame"),
+                        },
+                    )
+                })
+                .collect(),
+        },
+    );
+}
+
 macro_rules! def_custom_insts {
     ($($num:literal => $name:ident $({ $($field:ident),+ $(, ..$variadic_field:ident)? $(,)? })?),+ $(,)?) => {
         const SCHEMA: &[(u32, &str, &[&str])] = &[
diff --git a/crates/rustc_codegen_spirv/src/linker/mod.rs b/crates/rustc_codegen_spirv/src/linker/mod.rs
index 19a488043c..3606779ee9 100644
--- a/crates/rustc_codegen_spirv/src/linker/mod.rs
+++ b/crates/rustc_codegen_spirv/src/linker/mod.rs
@@ -409,6 +409,7 @@ pub fn link(
             spirv_tools::binary::from_binary(&spv_words).to_vec()
         };
         let cx = std::rc::Rc::new(spirt::Context::new());
+        crate::custom_insts::register_to_spirt_context(&cx);
         let mut module = {
             let _timer = sess.timer("spirt::Module::lower_from_spv_file");
             match spirt::Module::lower_from_spv_bytes(cx.clone(), spv_bytes) {