From 28591d8132d2dc692af9a281ae27a03fa68d84e5 Mon Sep 17 00:00:00 2001
From: Alex Crichton <alex@alexcrichton.com>
Date: Wed, 31 Jan 2024 13:45:30 -0800
Subject: [PATCH] Add a debugging utility for Rust `generate!`

Lift the approach I added in bytecodealliance/wasmtime#7794
---
 crates/rust-macro/build.rs   |  5 +++++
 crates/rust-macro/src/lib.rs | 27 ++++++++++++++++++++++++++-
 2 files changed, 31 insertions(+), 1 deletion(-)
 create mode 100644 crates/rust-macro/build.rs

diff --git a/crates/rust-macro/build.rs b/crates/rust-macro/build.rs
new file mode 100644
index 000000000..be0ad5bc9
--- /dev/null
+++ b/crates/rust-macro/build.rs
@@ -0,0 +1,5 @@
+fn main() {
+    println!("cargo:rerun-if-changed=build.rs");
+    let out_dir = std::env::var("OUT_DIR").unwrap();
+    println!("cargo:rustc-env=DEBUG_OUTPUT_DIR={out_dir}");
+}
diff --git a/crates/rust-macro/src/lib.rs b/crates/rust-macro/src/lib.rs
index a2dfe4e42..0832f4970 100644
--- a/crates/rust-macro/src/lib.rs
+++ b/crates/rust-macro/src/lib.rs
@@ -2,6 +2,7 @@ use proc_macro2::{Span, TokenStream};
 use quote::ToTokens;
 use std::collections::HashMap;
 use std::path::{Path, PathBuf};
+use std::sync::atomic::{AtomicUsize, Ordering::Relaxed};
 use syn::parse::{Error, Parse, ParseStream, Result};
 use syn::punctuated::Punctuated;
 use syn::{braced, token, Token};
@@ -154,7 +155,31 @@ impl Config {
             .generate(&self.resolve, self.world, &mut files)
             .map_err(|e| Error::new(Span::call_site(), e))?;
         let (_, src) = files.iter().next().unwrap();
-        let src = std::str::from_utf8(src).unwrap();
+        let mut src = std::str::from_utf8(src).unwrap().to_string();
+
+        // If a magical `WIT_BINDGEN_DEBUG` environment variable is set then
+        // place a formatted version of the expanded code into a file. This file
+        // will then show up in rustc error messages for any codegen issues and can
+        // be inspected manually.
+        if std::env::var("WIT_BINDGEN_DEBUG").is_ok() {
+            static INVOCATION: AtomicUsize = AtomicUsize::new(0);
+            let root = Path::new(env!("DEBUG_OUTPUT_DIR"));
+            let world_name = &self.resolve.worlds[self.world].name;
+            let n = INVOCATION.fetch_add(1, Relaxed);
+            let path = root.join(format!("{world_name}{n}.rs"));
+
+            std::fs::write(&path, &src).unwrap();
+
+            // optimistically format the code but don't require success
+            drop(
+                std::process::Command::new("rustfmt")
+                    .arg(&path)
+                    .arg("--edition=2021")
+                    .output(),
+            );
+
+            src = format!("include!({path:?});");
+        }
         let mut contents = src.parse::<TokenStream>().unwrap();
 
         // Include a dummy `include_str!` for any files we read so rustc knows that