Use capstone to validate precise-output tests (#5780)

Use the capstone library to disassemble precise-output tests, in addition to pretty-printing their vcode.

cranelift/codegen/src/isa/aarch64/mod.rs

@@ -90,7 +90,7 @@ impl TargetIsa for AArch64Backend {
         Ok(CompiledCodeStencil {
             buffer,
             frame_size,
-            disasm: emit_result.disasm,
+            vcode: emit_result.disasm,
             value_labels_ranges,
             sized_stackslot_offsets,
             dynamic_stackslot_offsets,
@@ -195,6 +195,22 @@ impl TargetIsa for AArch64Backend {
         // 4-byte alignment.
         32
     }
+
+    #[cfg(feature = "disas")]
+    fn to_capstone(&self) -> Result<capstone::Capstone, capstone::Error> {
+        use capstone::prelude::*;
+        let mut cs = Capstone::new()
+            .arm64()
+            .mode(arch::arm64::ArchMode::Arm)
+            .build()?;
+        // AArch64 uses inline constants rather than a separate constant pool right now.
+        // Without this option, Capstone will stop disassembling as soon as it sees
+        // an inline constant that is not also a valid instruction. With this option,
+        // Capstone will print a `.byte` directive with the bytes of the inline constant
+        // and continue to the next instruction.
+        cs.set_skipdata(true)?;
+        Ok(cs)
+    }
 }
 
 impl fmt::Display for AArch64Backend {

cranelift/codegen/src/isa/mod.rs

@@ -307,6 +307,12 @@ pub trait TargetIsa: fmt::Display + Send + Sync {
     {
         Arc::new(self)
     }
+
+    /// Generate a `Capstone` context for disassembling bytecode for this architecture.
+    #[cfg(feature = "disas")]
+    fn to_capstone(&self) -> Result<capstone::Capstone, capstone::Error> {
+        Err(capstone::Error::UnsupportedArch)
+    }
 }
 
 /// Methods implemented for free for target ISA!

cranelift/codegen/src/isa/riscv64/mod.rs

@@ -91,7 +91,7 @@ impl TargetIsa for Riscv64Backend {
         Ok(CompiledCodeStencil {
             buffer,
             frame_size,
-            disasm: emit_result.disasm,
+            vcode: emit_result.disasm,
             value_labels_ranges,
             sized_stackslot_offsets,
             dynamic_stackslot_offsets,
@@ -169,6 +169,20 @@ impl TargetIsa for Riscv64Backend {
     fn function_alignment(&self) -> u32 {
         4
     }
+
+    #[cfg(feature = "disas")]
+    fn to_capstone(&self) -> Result<capstone::Capstone, capstone::Error> {
+        use capstone::prelude::*;
+        let mut cs = Capstone::new()
+            .riscv()
+            .mode(arch::riscv::ArchMode::RiscV64)
+            .build()?;
+        // Similar to AArch64, RISC-V uses inline constants rather than a separate
+        // constant pool. We want to skip dissasembly over inline constants instead
+        // of stopping on invalid bytes.
+        cs.set_skipdata(true)?;
+        Ok(cs)
+    }
 }
 
 impl fmt::Display for Riscv64Backend {

cranelift/codegen/src/isa/s390x/mod.rs

@@ -87,7 +87,7 @@ impl TargetIsa for S390xBackend {
         Ok(CompiledCodeStencil {
             buffer,
             frame_size,
-            disasm: emit_result.disasm,
+            vcode: emit_result.disasm,
             value_labels_ranges,
             sized_stackslot_offsets,
             dynamic_stackslot_offsets,
@@ -170,6 +170,19 @@ impl TargetIsa for S390xBackend {
     fn function_alignment(&self) -> u32 {
         4
     }
+
+    #[cfg(feature = "disas")]
+    fn to_capstone(&self) -> Result<capstone::Capstone, capstone::Error> {
+        use capstone::prelude::*;
+        let mut cs = Capstone::new()
+            .sysz()
+            .mode(arch::sysz::ArchMode::Default)
+            .build()?;
+
+        cs.set_skipdata(true)?;
+
+        Ok(cs)
+    }
 }
 
 impl fmt::Display for S390xBackend {

cranelift/codegen/src/isa/x64/mod.rs

@@ -84,7 +84,7 @@ impl TargetIsa for X64Backend {
         Ok(CompiledCodeStencil {
             buffer,
             frame_size,
-            disasm: emit_result.disasm,
+            vcode: emit_result.disasm,
             value_labels_ranges,
             sized_stackslot_offsets,
             dynamic_stackslot_offsets,
@@ -171,6 +171,16 @@ impl TargetIsa for X64Backend {
     fn function_alignment(&self) -> u32 {
         16
     }
+
+    #[cfg(feature = "disas")]
+    fn to_capstone(&self) -> Result<capstone::Capstone, capstone::Error> {
+        use capstone::prelude::*;
+        Capstone::new()
+            .x86()
+            .mode(arch::x86::ArchMode::Mode64)
+            .syntax(arch::x86::ArchSyntax::Att)
+            .build()
+    }
 }
 
 impl fmt::Display for X64Backend {

cranelift/codegen/src/machinst/mod.rs

@@ -287,7 +287,7 @@ pub struct CompiledCodeBase<T: CompilePhase> {
     /// Size of stack frame, in bytes.
     pub frame_size: u32,
     /// Disassembly, if requested.
-    pub disasm: Option<String>,
+    pub vcode: Option<String>,
     /// Debug info: value labels to registers/stackslots at code offsets.
     pub value_labels_ranges: ValueLabelsRanges,
     /// Debug info: stackslots to stack pointer offsets.
@@ -317,7 +317,7 @@ impl CompiledCodeStencil {
         CompiledCode {
             buffer: self.buffer.apply_base_srcloc(params.base_srcloc()),
             frame_size: self.frame_size,
-            disasm: self.disasm,
+            vcode: self.vcode,
             value_labels_ranges: self.value_labels_ranges,
             sized_stackslot_offsets: self.sized_stackslot_offsets,
             dynamic_stackslot_offsets: self.dynamic_stackslot_offsets,
@@ -340,6 +340,71 @@ impl<T: CompilePhase> CompiledCodeBase<T> {
     pub fn code_buffer(&self) -> &[u8] {
         self.buffer.data()
     }
+
+    /// Get the disassembly of the buffer, using the given capstone context.
+    #[cfg(feature = "disas")]
+    pub fn disassemble(
+        &self,
+        params: Option<&crate::ir::function::FunctionParameters>,
+        cs: &capstone::Capstone,
+    ) -> Result<String, anyhow::Error> {
+        use std::fmt::Write;
+
+        let mut buf = String::new();
+
+        let relocs = self.buffer.relocs();
+        let traps = self.buffer.traps();
+        let labels = self.bb_starts.as_slice();
+
+        let insns = cs.disasm_all(self.buffer.data(), 0x0).map_err(map_caperr)?;
+        for i in insns.iter() {
+            if let Some((n, off)) = labels
+                .iter()
+                .copied()
+                .enumerate()
+                .find(|(_, val)| *val == i.address() as u32)
+            {
+                writeln!(buf, "block{}: ; offset 0x{:x}", n, off)?;
+            }
+
+            write!(buf, "  ")?;
+
+            let op_str = i.op_str().unwrap_or("");
+            if let Some(s) = i.mnemonic() {
+                write!(buf, "{}", s)?;
+                if !op_str.is_empty() {
+                    write!(buf, " ")?;
+                }
+            }
+
+            write!(buf, "{}", op_str)?;
+
+            let end = i.address() + i.bytes().len() as u64;
+            let contains = |off| i.address() <= off && off < end;
+
+            if let Some(reloc) = relocs.iter().find(|reloc| contains(reloc.offset as u64)) {
+                write!(
+                    buf,
+                    " ; reloc_external {} {} {}",
+                    reloc.kind,
+                    reloc.name.display(params),
+                    reloc.addend,
+                )?;
+            }
+
+            if let Some(trap) = traps.iter().find(|trap| contains(trap.offset as u64)) {
+                write!(buf, " ; trap: {}", trap.code)?;
+            }
+
+            writeln!(buf)?;
+        }
+
+        return Ok(buf);
+
+        fn map_caperr(err: capstone::Error) -> anyhow::Error {
+            anyhow::format_err!("{}", err)
+        }
+    }
 }
 
 /// Result of compiling a `FunctionStencil`, before applying `FunctionParameters` onto it.