Pretty-print OpExtInst using official extinst.*.grammar.json and/or custom Context-registered ones.

CHANGELOG.md

@@ -33,6 +33,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased] - ReleaseDate
 
+### Added ⭐
+- [PR#45](https://github.com/EmbarkStudios/spirt/pull/45) added the abilit to
+  pretty-print `OpExtInst`s (SPIR-V "extended instructions") using official
+  `extinst.*.grammar.json` descriptions and/or custom ones (registered via `Context`)
+
 ### Changed 🛠
 - [PR#43](https://github.com/EmbarkStudios/spirt/pull/43) tweaked several pretty-printing
   details to improve visual cohesion ("named arguments" in `module.{dialect,debuginfo}`)

build.rs

@@ -0,0 +1,52 @@
+fn main() {
+    // HACK(eddyb) disable the default of re-running the build script on *any*
+    // change to *the entire source tree* (i.e. the default is roughly `./`).
+    println!("cargo:rerun-if-changed=build.rs");
+
+    let khr_spv_include_dir =
+        std::env::current_dir().unwrap().join("khronos-spec/SPIRV-Headers/include/spirv/unified1");
+    println!("cargo:rerun-if-changed={}", khr_spv_include_dir.display());
+
+    if !std::fs::metadata(&khr_spv_include_dir).map_or(false, |m| m.is_dir()) {
+        eprintln!(" error: {} is not a directory", khr_spv_include_dir.display());
+        eprintln!("  help: git submodules are required to build from a git checkout");
+        eprintln!("  help: run `git submodule update --init`");
+        eprintln!("  note: if the error persists, please open an issue");
+        std::process::exit(1);
+    }
+
+    let core_grammar = "spirv.core.grammar.json";
+    let mut extinsts_names_and_grammars: Vec<_> = std::fs::read_dir(&khr_spv_include_dir)
+        .unwrap()
+        .filter_map(|e| {
+            let file_name = e.unwrap().file_name();
+            Some((
+                file_name
+                    .to_str()?
+                    .strip_prefix("extinst.")?
+                    .strip_suffix(".grammar.json")?
+                    .to_string(),
+                file_name,
+            ))
+        })
+        .collect();
+    extinsts_names_and_grammars.sort();
+
+    let all_jsons = format!(
+        "pub(super) const SPIRV_CORE_GRAMMAR: &str = include_str!({:?});\n\
+         pub(super) const EXTINST_NAMES_AND_GRAMMARS: &[(&str, &str)] = &[\n{}];",
+        khr_spv_include_dir.join(core_grammar),
+        extinsts_names_and_grammars
+            .into_iter()
+            .map(|(name, grammar)| {
+                format!("({:?}, include_str!({:?})),\n", name, khr_spv_include_dir.join(grammar),)
+            })
+            .collect::<String>()
+    );
+    std::fs::write(
+        std::path::PathBuf::from(std::env::var_os("OUT_DIR").unwrap())
+            .join("khr_spv_grammar_jsons.rs"),
+        all_jsons,
+    )
+    .unwrap();
+}

src/context.rs

@@ -1,5 +1,6 @@
 //! [`Context`](struct.Context.html) and related types/traits.
 
+use crate::spv::spec::ExtInstSetDesc;
 use rustc_hash::FxHashMap;
 use std::hash::Hash;
 use std::mem;
@@ -15,10 +16,45 @@ use std::ops::{Deref, DerefMut};
 ///   * the *definition* of an entity isn't kept in the [`Context`], but rather in
 ///     some [`EntityDefs`] collection somewhere in a [`Module`](crate::Module) (or further nested),
 ///     with only the entity *indices* being allocated by the [`Context`]
+/// * custom SPIR-V "extended instruction set" descriptions, which can be
+///   dynamically registered, to account for any such "extended instruction set"
+///   not covered by [`spv::spec::Spec`](crate::spv::spec::Spec)'s built-in list
+///   (e.g. non-standard tool-specific sets), and only used for pretty-printing
 #[derive(Default)]
 pub struct Context {
     interners: Interners,
     entity_allocs: EntityAllocs,
+
+    custom_spv_ext_inst_set_descs: elsa::FrozenBTreeMap<String, Box<ExtInstSetDesc>>,
+}
+
+impl Context {
+    /// Register a custom [`ExtInstSetDesc`] with name `ext_inst_set_name`,
+    /// to be used by pretty-printing when using this [`Context`].
+    pub fn register_custom_ext_inst_set(
+        &self,
+        ext_inst_set_name: &str,
+        ext_inst_set_desc: ExtInstSetDesc,
+    ) {
+        let lowercase_ext_inst_set_name = ext_inst_set_name.to_ascii_lowercase();
+        assert!(
+            self.get_custom_ext_inst_set_by_lowercase_name(&lowercase_ext_inst_set_name).is_none(),
+            "already registered {lowercase_ext_inst_set_name:?} \
+             (name before lowercasing: {ext_inst_set_name})"
+        );
+        self.custom_spv_ext_inst_set_descs
+            .insert(lowercase_ext_inst_set_name, Box::new(ext_inst_set_desc));
+    }
+
+    /// Return a custom [`ExtInstSetDesc`], if one was registered on this [`Context`],
+    /// for this `OpExtInstImport` name (required to be lowercase, due to Khronos'
+    /// choice of case insensitivity, but **not checked by this function**).
+    pub fn get_custom_ext_inst_set_by_lowercase_name(
+        &self,
+        lowercase_ext_inst_set_name: &str,
+    ) -> Option<&ExtInstSetDesc> {
+        self.custom_spv_ext_inst_set_descs.get(lowercase_ext_inst_set_name)
+    }
 }
 
 /// Private module containing traits (and related types) used in public APIs,