From 5d01603390f4c1f03dff84d2ef8ea9d4f536cd5a Mon Sep 17 00:00:00 2001
From: Benjamin Bouvier <public@benj.me>
Date: Mon, 8 Jun 2020 17:17:11 +0200
Subject: [PATCH] mach backend: allow snapshotting IR graphs with the
 SNAPSHOT_REGALLOC env variable;

This also requires the serde feature, which isn't enabled by default,
thus it must be passed as a command-line argument to cargo.
---
 Cargo.lock                                |  2 ++
 cranelift/codegen/Cargo.toml              |  5 +++++
 cranelift/codegen/src/machinst/compile.rs | 13 +++++++++++++
 3 files changed, 20 insertions(+)

diff --git a/Cargo.lock b/Cargo.lock
index ce9b1894c6..2fbb6a80c6 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -368,6 +368,7 @@ dependencies = [
 name = "cranelift-codegen"
 version = "0.64.0"
 dependencies = [
+ "bincode",
  "byteorder",
  "cranelift-bforest",
  "cranelift-codegen-meta",
@@ -1669,6 +1670,7 @@ checksum = "7c03092d79e0fd610932d89ed53895a38c0dd3bcd317a0046e69940de32f1d95"
 dependencies = [
  "log",
  "rustc-hash",
+ "serde",
  "smallvec",
 ]
 
diff --git a/cranelift/codegen/Cargo.toml b/cranelift/codegen/Cargo.toml
index 63c9e5ba53..e2cce0f85f 100644
--- a/cranelift/codegen/Cargo.toml
+++ b/cranelift/codegen/Cargo.toml
@@ -20,6 +20,7 @@ hashbrown = { version = "0.7", optional = true }
 target-lexicon = "0.10"
 log = { version = "0.4.6", default-features = false }
 serde = { version = "1.0.94", features = ["derive"], optional = true }
+bincode = { version = "1.2.1", optional = true }
 gimli = { version = "0.21.0", default-features = false, features = ["write"], optional = true }
 smallvec = { version = "1.0.0" }
 thiserror = "1.0.4"
@@ -74,6 +75,10 @@ all-arch = [
 # For dependent crates that want to serialize some parts of cranelift
 enable-serde = ["serde"]
 
+# Allow snapshotting regalloc test cases. Useful only to report bad register
+# allocation failures, or for regalloc.rs developers.
+regalloc-snapshot = ["bincode", "regalloc/enable-serde"]
+
 # Recompile our optimizations that are written in the `peepmatic` DSL into a
 # compact finite-state transducer automaton.
 rebuild-peephole-optimizers = ["peepmatic"]
diff --git a/cranelift/codegen/src/machinst/compile.rs b/cranelift/codegen/src/machinst/compile.rs
index 508e242cd7..8254fff77b 100644
--- a/cranelift/codegen/src/machinst/compile.rs
+++ b/cranelift/codegen/src/machinst/compile.rs
@@ -44,6 +44,19 @@ where
         }
     };
 
+    #[cfg(feature = "regalloc-snapshot")]
+    {
+        use std::fs;
+        use std::path::Path;
+        if let Some(path) = std::env::var("SERIALIZE_REGALLOC").ok() {
+            let snapshot = regalloc::IRSnapshot::from_function(&vcode, b.reg_universe());
+            let serialized = bincode::serialize(&snapshot).expect("couldn't serialize snapshot");
+
+            let file_path = Path::new(&path).join(Path::new(&format!("ir{}.bin", f.name)));
+            fs::write(file_path, &serialized).expect("couldn't write IR snapshot file");
+        }
+    }
+
     let result = {
         let _tt = timing::regalloc();
         allocate_registers_with_opts(