Introduce the wasmtime-explorer crate (#5975)

This implements Godbolt Compiler Explorer-like functionality for Wasmtime and
Cranelift. Given a Wasm module, it compiles the module to native code and then
writes a standalone HTML file that gives a split pane view between the WAT and
ASM disassemblies.

Cargo.lock

@@ -3533,6 +3533,7 @@ dependencies = [
  "wasmtime-component-util",
  "wasmtime-cranelift",
  "wasmtime-environ",
+ "wasmtime-explorer",
  "wasmtime-runtime",
  "wasmtime-wasi",
  "wasmtime-wasi-crypto",
@@ -3647,6 +3648,19 @@ dependencies = [
  "wat",
 ]
 
+[[package]]
+name = "wasmtime-explorer"
+version = "8.0.0"
+dependencies = [
+ "anyhow",
+ "capstone",
+ "serde",
+ "serde_json",
+ "target-lexicon",
+ "wasmprinter",
+ "wasmtime",
+]
+
 [[package]]
 name = "wasmtime-fiber"
 version = "8.0.0"

Cargo.toml

@@ -27,6 +27,7 @@ wasmtime-cache = { workspace = true }
 wasmtime-cli-flags = { workspace = true }
 wasmtime-cranelift = { workspace = true }
 wasmtime-environ = { workspace = true }
+wasmtime-explorer = { workspace = true }
 wasmtime-wast = { workspace = true }
 wasmtime-wasi = { workspace = true, features = ["exit"] }
 wasmtime-wasi-crypto = { workspace = true, optional = true }
@@ -39,8 +40,8 @@ humantime = "2.0.0"
 once_cell = { workspace = true }
 listenfd = "1.0.0"
 wat = { workspace = true }
-serde = "1.0.94"
-serde_json = "1.0.26"
+serde = { workspace = true }
+serde_json = { workspace = true }
 wasmparser = { workspace = true }
 wasm-coredump-builder = { version = "0.1.11" }
 
@@ -70,8 +71,8 @@ component-macro-test = { path = "crates/misc/component-macro-test" }
 component-test-util = { workspace = true }
 bstr = "0.2.17"
 libc = "0.2.60"
-serde = "1.0"
-serde_json = "1.0"
+serde = { workspace = true }
+serde_json = { workspace = true }
 
 [target.'cfg(windows)'.dev-dependencies]
 windows-sys = { workspace = true, features = ["Win32_System_Memory"] }
@@ -120,6 +121,7 @@ wasmtime-cli-flags = { path = "crates/cli-flags", version = "=8.0.0" }
 wasmtime-cranelift = { path = "crates/cranelift", version = "=8.0.0" }
 wasmtime-cranelift-shared = { path = "crates/cranelift-shared", version = "=8.0.0" }
 wasmtime-environ = { path = "crates/environ", version = "=8.0.0" }
+wasmtime-explorer = { path = "crates/explorer", version = "=8.0.0" }
 wasmtime-fiber = { path = "crates/fiber", version = "=8.0.0" }
 wasmtime-types = { path = "crates/types", version = "8.0.0" }
 wasmtime-jit = { path = "crates/jit", version = "=8.0.0" }
@@ -196,6 +198,7 @@ heck = "0.4"
 similar = "2.1.0"
 toml = "0.5.9"
 serde = "1.0.94"
+serde_json = "1.0.80"
 glob = "0.3.0"
 
 [features]

crates/explorer/Cargo.toml

@@ -0,0 +1,18 @@
+[package]
+name = "wasmtime-explorer"
+authors.workspace = true
+description = "Compiler explorer for Wasmtime and Cranelift"
+documentation = "https://docs.rs/wasmtime-explorer/"
+edition.workspace = true
+license = "Apache-2.0 WITH LLVM-exception"
+repository = "https://github.com/bytecodealliance/wasmtime"
+version.workspace = true
+
+[dependencies]
+anyhow = { workspace = true }
+capstone = { workspace = true }
+serde = { workspace = true }
+serde_json = { workspace = true }
+target-lexicon = { workspace = true }
+wasmprinter = { workspace = true }
+wasmtime = { workspace = true, features = ["cranelift"] }

crates/explorer/src/.eslintrc.yml

@@ -0,0 +1,8 @@
+root: true
+
+env:
+  browser: true
+  es2022: true
+
+extends:
+  - "eslint:recommended"

crates/explorer/src/index.css

@@ -0,0 +1,26 @@
+* {
+    margin: 0;
+    padding: 0;
+}
+
+.hbox {
+    display: flex;
+    flex-direction: row;
+}
+
+html, body {
+    width: 100%;
+    height: 100%;
+}
+
+#wat {
+    width: 50%;
+    height: 100%;
+    overflow: scroll;
+}
+
+#asm {
+    width: 50%;
+    height: 100%;
+    overflow: scroll;
+}

crates/explorer/src/index.js

@@ -0,0 +1,238 @@
+/*** State *********************************************************************/
+
+class State {
+  constructor(wat, asm) {
+    this.wat = wat;
+    this.asm = asm;
+  }
+}
+
+const state = window.STATE = new State(window.WAT, window.ASM);
+
+/*** Hues for Offsets **********************************************************/
+
+const hues = [
+  80,
+  160,
+  240,
+  320,
+  40,
+  120,
+  200,
+  280,
+  20,
+  100,
+  180,
+  260,
+  340,
+  60,
+  140,
+  220,
+  300,
+];
+
+const nextHue = (function () {
+  let i = 0;
+  return () => {
+    return hues[++i % hues.length];
+  };
+}());
+
+// NB: don't just assign hues based on something simple like `hues[offset %
+// hues.length]` since that can suffer from bias due to certain alignments
+// happening more or less frequently.
+const offsetToHue = new Map();
+
+// Get the hue for the given offset, or assign it a new one if it doesn't have
+// one already.
+const hueForOffset = offset => {
+  if (offsetToHue.has(offset)) {
+    return offsetToHue.get(offset);
+  } else {
+    let hue = nextHue();
+    offsetToHue.set(offset, hue);
+    return hue;
+  }
+};
+
+// Get the hue for the given offset, only if the offset has already been
+// assigned a hue.
+const existingHueForOffset = offset => {
+  return offsetToHue.get(offset);
+};
+
+// Get WAT chunk elements by Wasm offset.
+const watByOffset = new Map();
+
+// Get asm instruction elements by Wasm offset.
+const asmByOffset = new Map();
+
+// Get all (WAT chunk or asm instruction) elements by offset.
+const anyByOffset = new Map();
+
+const addWatElem = (offset, elem) => {
+  if (!watByOffset.has(offset)) {
+    watByOffset.set(offset, []);
+  }
+  watByOffset.get(offset).push(elem);
+
+  if (!anyByOffset.has(offset)) {
+    anyByOffset.set(offset, []);
+  }
+  anyByOffset.get(offset).push(elem);
+};
+
+const addAsmElem = (offset, elem) => {
+  if (!asmByOffset.has(offset)) {
+    asmByOffset.set(offset, []);
+  }
+  asmByOffset.get(offset).push(elem);
+
+  if (!anyByOffset.has(offset)) {
+    anyByOffset.set(offset, []);
+  }
+  anyByOffset.get(offset).push(elem);
+};
+
+/*** Event Handlers ************************************************************/
+
+const watElem = document.getElementById("wat");
+watElem.addEventListener("click", event => {
+  if (event.target.dataset.wasmOffset == null) {
+    return;
+  }
+
+  const offset = parseInt(event.target.dataset.wasmOffset);
+  if (!asmByOffset.get(offset)) {
+    return;
+  }
+
+  const firstAsmElem = asmByOffset.get(offset)[0];
+  firstAsmElem.scrollIntoView({
+    behavior: "smooth",
+    block: "center",
+    inline: "nearest",
+  });
+}, { passive: true });
+
+const asmElem = document.getElementById("asm");
+asmElem.addEventListener("click", event => {
+  if (event.target.dataset.wasmOffset == null) {
+    return;
+  }
+
+  const offset = parseInt(event.target.dataset.wasmOffset);
+  if (!watByOffset.get(offset)) {
+    return;
+  }
+
+  const firstWatElem = watByOffset.get(offset)[0];
+  firstWatElem.scrollIntoView({
+    behavior: "smooth",
+    block: "center",
+    inline: "nearest",
+  });
+}, { passive: true });
+
+const onMouseEnter = event => {
+  if (event.target.dataset.wasmOffset == null) {
+    return;
+  }
+
+  const offset = parseInt(event.target.dataset.wasmOffset);
+  const hue = hueForOffset(offset);
+  for (const elem of anyByOffset.get(offset)) {
+    elem.style.backgroundColor = `hsl(${hue} 75% 80%)`;
+  }
+};
+
+const onMouseLeave = event => {
+  if (event.target.dataset.wasmOffset == null) {
+    return;
+  }
+
+  const offset = parseInt(event.target.dataset.wasmOffset);
+  const hue = hueForOffset(offset);
+  for (const elem of anyByOffset.get(offset)) {
+    elem.style.backgroundColor = `hsl(${hue} 50% 95%)`;
+  }
+};
+
+/*** Rendering *****************************************************************/
+
+const repeat = (s, n) => {
+  return s.repeat(n >= 0 ? n : 0);
+};
+
+const renderAddress = addr => {
+  let hex = addr.toString(16);
+  return repeat("0", 8 - hex.length) + hex;
+};
+
+const renderBytes = bytes => {
+  let s = "";
+  for (let i = 0; i < bytes.length; i++) {
+    if (i != 0) {
+      s += " ";
+    }
+    const hexByte = bytes[i].toString(16);
+    s += hexByte.length == 2 ? hexByte : "0" + hexByte;
+  }
+  return s + repeat(" ", 30 - s.length);
+};
+
+const renderInst = (mnemonic, operands) => {
+  if (operands.length == 0) {
+    return mnemonic;
+  } else {
+    return mnemonic + " " + operands;
+  }
+};
+
+// Render the ASM.
+
+let nthFunc = 0;
+for (const func of state.asm.functions) {
+  const funcElem = document.createElement("div");
+
+  const funcHeader = document.createElement("h3");
+  funcHeader.textContent = `Defined Function ${nthFunc}`;
+  funcElem.appendChild(funcHeader);
+
+  const bodyElem = document.createElement("pre");
+  for (const inst of func.instructions) {
+    const instElem = document.createElement("span");
+    instElem.textContent = `${renderAddress(inst.address)}    ${renderBytes(inst.bytes)}    ${renderInst(inst.mnemonic, inst.operands)}\n`;
+    if (inst.wasm_offset != null) {
+      instElem.setAttribute("data-wasm-offset", inst.wasm_offset);
+      const hue = hueForOffset(inst.wasm_offset);
+      instElem.style.backgroundColor = `hsl(${hue} 50% 90%)`;
+      instElem.addEventListener("mouseenter", onMouseEnter);
+      instElem.addEventListener("mouseleave", onMouseLeave);
+      addAsmElem(inst.wasm_offset, instElem);
+    }
+    bodyElem.appendChild(instElem);
+  }
+  funcElem.appendChild(bodyElem);
+
+  asmElem.appendChild(funcElem);
+  nthFunc++;
+}
+
+// Render the WAT.
+
+for (const chunk of state.wat.chunks) {
+  const chunkElem = document.createElement("span");
+  if (chunk.wasm_offset != null) {
+    chunkElem.dataset.wasmOffset = chunk.wasm_offset;
+    const hue = existingHueForOffset(chunk.wasm_offset);
+    if (hue) {
+      chunkElem.style.backgroundColor = `hsl(${hue} 50% 95%)`;
+      chunkElem.addEventListener("mouseenter", onMouseEnter);
+      chunkElem.addEventListener("mouseleave", onMouseLeave);
+      addWatElem(chunk.wasm_offset, chunkElem);
+    }
+  }
+  chunkElem.textContent = chunk.wat;
+  watElem.appendChild(chunkElem);
+}

crates/explorer/src/lib.rs

@@ -0,0 +1,175 @@
+use anyhow::Result;
+use capstone::arch::BuildsCapstone;
+use serde::Serialize;
+use std::{io::Write, str::FromStr};
+
+pub fn generate(
+    config: &wasmtime::Config,
+    target: Option<&str>,
+    wasm: &[u8],
+    dest: &mut dyn Write,
+) -> Result<()> {
+    let target = match target {
+        None => target_lexicon::Triple::host(),
+        Some(target) => target_lexicon::Triple::from_str(target)?,
+    };
+
+    let wat = annotate_wat(wasm)?;
+    let wat_json = serde_json::to_string(&wat)?;
+    let asm = annotate_asm(config, &target, wasm)?;
+    let asm_json = serde_json::to_string(&asm)?;
+
+    let index_css = include_str!("./index.css");
+    let index_js = include_str!("./index.js");
+
+    write!(
+        dest,
+        r#"
+<!DOCTYPE html>
+<html>
+  <head>
+    <title>Wasmtime Compiler Explorer</title>
+    <style>
+      {index_css}
+    </style>
+  </head>
+  <body class="hbox">
+    <pre id="wat"></pre>
+    <div id="asm"></div>
+    <script>
+      window.WAT = {wat_json};
+      window.ASM = {asm_json};
+    </script>
+    <script>
+      {index_js}
+    </script>
+  </body>
+</html>
+        "#
+    )?;
+    Ok(())
+}
+
+#[derive(Serialize, Clone, Copy, Debug)]
+struct WasmOffset(u32);
+
+#[derive(Serialize, Debug)]
+struct AnnotatedWat {
+    chunks: Vec<AnnotatedWatChunk>,
+}
+
+#[derive(Serialize, Debug)]
+struct AnnotatedWatChunk {
+    wasm_offset: Option<WasmOffset>,
+    wat: String,
+}
+
+fn annotate_wat(wasm: &[u8]) -> Result<AnnotatedWat> {
+    let mut printer = wasmprinter::Printer::new();
+    let chunks = printer
+        .offsets_and_lines(wasm)?
+        .map(|(offset, wat)| AnnotatedWatChunk {
+            wasm_offset: offset.map(|o| WasmOffset(u32::try_from(o).unwrap())),
+            wat: wat.to_string(),
+        })
+        .collect();
+    Ok(AnnotatedWat { chunks })
+}
+
+#[derive(Serialize, Debug)]
+struct AnnotatedAsm {
+    functions: Vec<AnnotatedFunction>,
+}
+
+#[derive(Serialize, Debug)]
+struct AnnotatedFunction {
+    instructions: Vec<AnnotatedInstruction>,
+}
+
+#[derive(Serialize, Debug)]
+struct AnnotatedInstruction {
+    wasm_offset: Option<WasmOffset>,
+    address: u32,
+    bytes: Vec<u8>,
+    mnemonic: Option<String>,
+    operands: Option<String>,
+}
+
+fn annotate_asm(
+    config: &wasmtime::Config,
+    target: &target_lexicon::Triple,
+    wasm: &[u8],
+) -> Result<AnnotatedAsm> {
+    let engine = wasmtime::Engine::new(config)?;
+    let module = wasmtime::Module::new(&engine, wasm)?;
+
+    let text = module.text();
+    let address_map: Vec<_> = module
+        .address_map()
+        .ok_or_else(|| anyhow::anyhow!("address maps must be enabled in the config"))?
+        .collect();
+
+    let mut address_map_iter = address_map.into_iter().peekable();
+    let mut current_entry = address_map_iter.next();
+    let mut wasm_offset_for_address = |address: u32| -> Option<WasmOffset> {
+        while address_map_iter.peek().map_or(false, |next_entry| {
+            u32::try_from(next_entry.0).unwrap() < address
+        }) {
+            current_entry = address_map_iter.next();
+        }
+        current_entry.and_then(|entry| entry.1.map(WasmOffset))
+    };
+
+    let functions = module
+        .function_locations()
+        .into_iter()
+        .map(|(start, len)| {
+            let body = &text[start..][..len];
+
+            let cs = match target.architecture {
+                target_lexicon::Architecture::Aarch64(_) => capstone::Capstone::new()
+                    .arm64()
+                    .mode(capstone::arch::arm64::ArchMode::Arm)
+                    .build()
+                    .map_err(|e| anyhow::anyhow!("{e}"))?,
+                target_lexicon::Architecture::Riscv64(_) => capstone::Capstone::new()
+                    .riscv()
+                    .mode(capstone::arch::riscv::ArchMode::RiscV64)
+                    .build()
+                    .map_err(|e| anyhow::anyhow!("{e}"))?,
+                target_lexicon::Architecture::S390x => capstone::Capstone::new()
+                    .sysz()
+                    .mode(capstone::arch::sysz::ArchMode::Default)
+                    .build()
+                    .map_err(|e| anyhow::anyhow!("{e}"))?,
+                target_lexicon::Architecture::X86_64 => capstone::Capstone::new()
+                    .x86()
+                    .mode(capstone::arch::x86::ArchMode::Mode64)
+                    .build()
+                    .map_err(|e| anyhow::anyhow!("{e}"))?,
+                _ => anyhow::bail!("Unsupported target: {target}"),
+            };
+
+            let instructions = cs
+                .disasm_all(body, start as u64)
+                .map_err(|e| anyhow::anyhow!("{e}"))?;
+            let instructions = instructions
+                .iter()
+                .map(|inst| {
+                    let address = u32::try_from(inst.address()).unwrap();
+                    let wasm_offset = wasm_offset_for_address(address);
+                    Ok(AnnotatedInstruction {
+                        wasm_offset,
+                        address,
+                        bytes: inst.bytes().to_vec(),
+                        mnemonic: inst.mnemonic().map(ToString::to_string),
+                        operands: inst.op_str().map(ToString::to_string),
+                    })
+                })
+                .collect::<Result<Vec<_>>>()?;
+            Ok(AnnotatedFunction { instructions })
+        })
+        .collect::<Result<Vec<_>>>()?;
+
+    Ok(AnnotatedAsm { functions })
+}

scripts/publish.rs

@@ -70,6 +70,7 @@ const CRATES_TO_PUBLISH: &[&str] = &[
     "wasmtime-wasi-threads",
     "wasmtime-wast",
     "wasmtime-cli-flags",
+    "wasmtime-explorer",
     "wasmtime-cli",
 ];
 

src/bin/wasmtime.rs

@@ -6,7 +6,7 @@
 use anyhow::Result;
 use clap::{ErrorKind, Parser};
 use wasmtime_cli::commands::{
-    CompileCommand, ConfigCommand, RunCommand, SettingsCommand, WastCommand,
+    CompileCommand, ConfigCommand, ExploreCommand, RunCommand, SettingsCommand, WastCommand,
 };
 
 /// Wasmtime WebAssembly Runtime
@@ -35,6 +35,8 @@ enum Wasmtime {
     Config(ConfigCommand),
     /// Compiles a WebAssembly module.
     Compile(CompileCommand),
+    /// Explore the compilation of a WebAssembly module to native code.
+    Explore(ExploreCommand),
     /// Runs a WebAssembly module
     Run(RunCommand),
     /// Displays available Cranelift settings for a target.
@@ -49,6 +51,7 @@ impl Wasmtime {
         match self {
             Self::Config(c) => c.execute(),
             Self::Compile(c) => c.execute(),
+            Self::Explore(c) => c.execute(),
             Self::Run(c) => c.execute(),
             Self::Settings(c) => c.execute(),
             Self::Wast(c) => c.execute(),

src/commands.rs

@@ -2,8 +2,9 @@
 
 mod compile;
 mod config;
+mod explore;
 mod run;
 mod settings;
 mod wast;
 
-pub use self::{compile::*, config::*, run::*, settings::*, wast::*};
+pub use self::{compile::*, config::*, explore::*, run::*, settings::*, wast::*};

src/commands/explore.rs

@@ -0,0 +1,51 @@
+//! The module that implements the `wasmtime explore` command.
+
+use anyhow::{Context, Result};
+use clap::Parser;
+use std::path::PathBuf;
+use wasmtime_cli_flags::CommonOptions;
+
+/// Explore the compilation of a WebAssembly module to native code.
+#[derive(Parser)]
+#[clap(name = "explore")]
+pub struct ExploreCommand {
+    #[clap(flatten)]
+    common: CommonOptions,
+
+    /// The target triple; default is the host triple
+    #[clap(long, value_name = "TARGET")]
+    target: Option<String>,
+
+    /// The path of the WebAssembly module to compile
+    #[clap(required = true, value_name = "MODULE")]
+    module: PathBuf,
+
+    /// The path of the explorer output (derived from the MODULE name if none
+    /// provided)
+    #[clap(short, long)]
+    output: Option<PathBuf>,
+}
+
+impl ExploreCommand {
+    /// Executes the command.
+    pub fn execute(&self) -> Result<()> {
+        self.common.init_logging();
+
+        let config = self.common.config(self.target.as_deref())?;
+
+        let wasm = std::fs::read(&self.module)
+            .with_context(|| format!("failed to read Wasm module: {}", self.module.display()))?;
+
+        let output = self
+            .output
+            .clone()
+            .unwrap_or_else(|| self.module.with_extension("explore.html"));
+        let output_file = std::fs::File::create(&output)
+            .with_context(|| format!("failed to create file: {}", output.display()))?;
+        let mut output_file = std::io::BufWriter::new(output_file);
+
+        wasmtime_explorer::generate(&config, self.target.as_deref(), &wasm, &mut output_file)?;
+        println!("Exploration written to {}", output.display());
+        Ok(())
+    }
+}