Add dis_fn testing function (#155)

* Add dis_fn testing function

* Update spirv-builder/src/test/mod.rs

Co-authored-by: XAMPPRocky <4464295+XAMPPRocky@users.noreply.github.com>

Co-authored-by: XAMPPRocky <4464295+XAMPPRocky@users.noreply.github.com>

Cargo.lock

@@ -2168,6 +2168,7 @@ version = "0.1.0"
 dependencies = [
  "lazy_static",
  "memchr",
+ "pretty_assertions",
  "raw-string",
  "rustc_codegen_spirv",
  "serde",

rustc_codegen_spirv/src/lib.rs

@@ -75,6 +75,7 @@ mod symbols;
 
 use builder::Builder;
 use codegen_cx::CodegenCx;
+pub use rspirv;
 use rspirv::binary::Assemble;
 use rustc_ast::expand::allocator::AllocatorKind;
 use rustc_codegen_ssa::back::lto::{LtoModuleCodegen, SerializedModule, ThinModule};

spirv-builder/Cargo.toml

@@ -17,3 +17,4 @@ rustc_codegen_spirv = { path = "../rustc_codegen_spirv" }
 
 [dev-dependencies]
 lazy_static = "1.4"
+pretty_assertions = "0.6"

spirv-builder/src/test/basic.rs

@@ -1,4 +1,4 @@
-use super::val;
+use super::{dis_fn, val};
 
 #[test]
 fn hello_world() {
@@ -9,3 +9,27 @@ pub fn main() {
 }
 "#);
 }
+
+#[test]
+fn add_two_ints() {
+    dis_fn(
+        r#"
+fn add_two_ints(x: u32, y: u32) -> u32 {
+    x + y
+}
+#[allow(unused_attributes)]
+#[spirv(fragment)]
+pub fn main() {
+    add_two_ints(2, 3);
+}
+"#,
+        "add_two_ints",
+        r#"%1 = OpFunction %2 None %3
+%4 = OpFunctionParameter %2
+%5 = OpFunctionParameter %2
+%6 = OpLabel
+%7 = OpIAdd %2 %4 %5
+OpReturnValue %7
+OpFunctionEnd"#,
+    );
+}

spirv-builder/src/test/mod.rs

@@ -1,10 +1,21 @@
 mod basic;
 
 use lazy_static::lazy_static;
+use rustc_codegen_spirv::rspirv;
 use std::error::Error;
 use std::path::{Path, PathBuf};
 use std::sync::Mutex;
 
+// https://github.com/colin-kiegel/rust-pretty-assertions/issues/24
+#[derive(PartialEq, Eq)]
+pub struct PrettyString<'a>(pub &'a str);
+/// Make diff to display string as multi-line string
+impl<'a> std::fmt::Debug for PrettyString<'a> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.write_str(self.0)
+    }
+}
+
 // Tests need to run serially, since they write project files to disk and whatnot. We don't want to
 // create a new temp dir for every test, though, since then every test would need to build libcore.
 // We could require the user to pass --thread-count 1 to cargo test, but that affects other tests.
@@ -37,8 +48,6 @@ use spirv_std::*;
 fn panic(_: &PanicInfo) -> ! {
     loop {}
 }
-#[lang = "eh_personality"]
-extern "C" fn rust_eh_personality() {}
 "#;
 
 fn setup(src: &str) -> Result<PathBuf, Box<dyn Error>> {
@@ -62,8 +71,74 @@ fn build(src: &str) -> PathBuf {
         .expect("Failed to build test")
 }
 
+fn read_module(path: &Path) -> Result<rspirv::dr::Module, Box<dyn Error>> {
+    let bytes = std::fs::read(path)?;
+    let mut loader = rspirv::dr::Loader::new();
+    rspirv::binary::parse_bytes(&bytes, &mut loader)?;
+    Ok(loader.module())
+}
+
 fn val(src: &str) {
     let _lock = GLOBAL_MUTEX.lock().unwrap();
     // spirv-val is included in building
     build(src);
 }
+
+fn assert_str_eq(expected: &str, result: &str) {
+    let expected = expected
+        .split('\n')
+        .map(|l| l.trim())
+        .collect::<Vec<_>>()
+        .join("\n");
+
+    let result = result
+        .split('\n')
+        .map(|l| l.trim().replace("  ", " ")) // rspirv outputs multiple spaces between operands
+        .collect::<Vec<_>>()
+        .join("\n");
+
+    pretty_assertions::assert_eq!(PrettyString(&expected), PrettyString(&result))
+}
+
+fn dis_fn(src: &str, func: &str, expect: &str) {
+    let _lock = GLOBAL_MUTEX.lock().unwrap();
+    let module = read_module(&build(src)).unwrap();
+    let id = module
+        .debugs
+        .iter()
+        .find(|inst| inst.operands[1].unwrap_literal_string() == func)
+        .expect("No function with that name found")
+        .operands[0]
+        .unwrap_id_ref();
+    let mut func = module
+        .functions
+        .into_iter()
+        .find(|f| f.def_id().unwrap() == id)
+        .unwrap();
+    // Compact to make IDs more stable
+    compact_ids(&mut func);
+    use rspirv::binary::Disassemble;
+    assert_str_eq(expect, &func.disassemble())
+}
+
+fn compact_ids(module: &mut rspirv::dr::Function) -> u32 {
+    let mut remap = std::collections::HashMap::new();
+    let mut insert = |current_id: &mut u32| {
+        let len = remap.len();
+        *current_id = *remap.entry(*current_id).or_insert_with(|| len as u32 + 1)
+    };
+    module.all_inst_iter_mut().for_each(|inst| {
+        if let Some(ref mut result_id) = &mut inst.result_id {
+            insert(result_id)
+        }
+        if let Some(ref mut result_type) = &mut inst.result_type {
+            insert(result_type)
+        }
+        inst.operands.iter_mut().for_each(|op| {
+            if let Some(w) = op.id_ref_any_mut() {
+                insert(w)
+            }
+        })
+    });
+    remap.len() as u32 + 1
+}