Add --stack-maps to wasmtime objdump (#10439)

* Add `--stack-maps` to `wasmtime objdump`

Follow-up to #10404 and #10405

* Enable traps/stack maps by default in objdump

Author: alexcrichton

src/commands/objdump.rs

@@ -15,7 +15,7 @@ use std::path::{Path, PathBuf};
 use std::str::FromStr;
 use termcolor::{Color, ColorChoice, ColorSpec, StandardStream, WriteColor};
 use wasmtime::Engine;
-use wasmtime_environ::{obj, FilePos, Trap};
+use wasmtime_environ::{obj, FilePos, StackMap, Trap};
 
 /// A helper utility in wasmtime to explore the compiled object file format of
 /// a `*.cwasm` file.
@@ -36,11 +36,11 @@ pub struct ObjdumpCommand {
     address_jumps: bool,
 
     /// What functions should be printed (all|wasm|trampoline|builtin|libcall, default: wasm)
-    #[arg(long, value_parser = Func::from_str)]
+    #[arg(long, value_parser = Func::from_str, value_name = "KIND")]
     funcs: Vec<Func>,
 
     /// String filter to apply to function names to only print some functions.
-    #[arg(long)]
+    #[arg(long, value_name = "STR")]
     filter: Option<String>,
 
     /// Whether or not instruction bytes are disassembled.
@@ -52,19 +52,43 @@ pub struct ObjdumpCommand {
     color: ColorChoice,
 
     /// Whether or not to interleave instructions with address maps.
-    #[arg(long)]
-    addrmap: bool,
+    #[arg(long, require_equals = true, value_name = "true|false")]
+    addrmap: Option<Option<bool>>,
 
     /// Column width of how large an address is rendered as.
-    #[arg(long, default_value = "10")]
+    #[arg(long, default_value = "10", value_name = "N")]
     address_width: usize,
 
     /// Whether or not to show information about what instructions can trap.
-    #[arg(long)]
-    traps: bool,
+    #[arg(long, require_equals = true, value_name = "true|false")]
+    traps: Option<Option<bool>>,
+
+    /// Whether or not to show information about stack maps.
+    #[arg(long, require_equals = true, value_name = "true|false")]
+    stack_maps: Option<Option<bool>>,
+}
+
+fn optional_flag_with_default(flag: Option<Option<bool>>, default: bool) -> bool {
+    match flag {
+        None => default,
+        Some(None) => true,
+        Some(Some(val)) => val,
+    }
 }
 
 impl ObjdumpCommand {
+    fn addrmap(&self) -> bool {
+        optional_flag_with_default(self.addrmap, false)
+    }
+
+    fn traps(&self) -> bool {
+        optional_flag_with_default(self.traps, true)
+    }
+
+    fn stack_maps(&self) -> bool {
+        optional_flag_with_default(self.stack_maps, true)
+    }
+
     /// Executes the command.
     pub fn execute(self) -> Result<()> {
         // Setup stdout handling color options. Also build some variables used
@@ -107,6 +131,11 @@ impl ObjdumpCommand {
                 .and_then(|section| section.data().ok())
                 .and_then(|bytes| wasmtime_environ::iterate_traps(bytes))
                 .map(|i| (Box::new(i) as Box<dyn Iterator<Item = _>>).peekable()),
+            stack_maps: elf
+                .section_by_name(obj::ELF_WASMTIME_STACK_MAP)
+                .and_then(|section| section.data().ok())
+                .and_then(|bytes| StackMap::iter(bytes))
+                .map(|i| (Box::new(i) as Box<dyn Iterator<Item = _>>).peekable()),
             objdump: &self,
         };
 
@@ -485,16 +514,18 @@ struct Decorator<'a> {
     objdump: &'a ObjdumpCommand,
     addrmap: Option<Peekable<Box<dyn Iterator<Item = (u32, FilePos)> + 'a>>>,
     traps: Option<Peekable<Box<dyn Iterator<Item = (u32, Trap)> + 'a>>>,
+    stack_maps: Option<Peekable<Box<dyn Iterator<Item = (u32, StackMap<'a>)> + 'a>>>,
 }
 
 impl Decorator<'_> {
     fn decorate(&mut self, address: u64, list: &mut Vec<String>) {
         self.addrmap(address, list);
         self.traps(address, list);
+        self.stack_maps(address, list);
     }
 
     fn addrmap(&mut self, address: u64, list: &mut Vec<String>) {
-        if !self.objdump.addrmap {
+        if !self.objdump.addrmap() {
             return;
         }
         let Some(addrmap) = &mut self.addrmap else {
@@ -511,7 +542,7 @@ impl Decorator<'_> {
     }
 
     fn traps(&mut self, address: u64, list: &mut Vec<String>) {
-        if !self.objdump.traps {
+        if !self.objdump.traps() {
             return;
         }
         let Some(traps) = &mut self.traps else {
@@ -524,4 +555,25 @@ impl Decorator<'_> {
             list.push(format!("trap: {trap:?}"));
         }
     }
+
+    fn stack_maps(&mut self, address: u64, list: &mut Vec<String>) {
+        if !self.objdump.stack_maps() {
+            return;
+        }
+        let Some(stack_maps) = &mut self.stack_maps else {
+            return;
+        };
+        while let Some((addr, stack_map)) =
+            stack_maps.next_if(|(addr, _pos)| u64::from(*addr) <= address)
+        {
+            if u64::from(addr) != address {
+                continue;
+            }
+            list.push(format!(
+                "stack_map: frame_size={}, frame_offsets={:?}",
+                stack_map.frame_size(),
+                stack_map.offsets().collect::<Vec<_>>()
+            ));
+        }
+    }
 }

tests/disas.rs

@@ -268,8 +268,13 @@ fn assert_output(test: &Test, output: CompileOutput) -> Result<()> {
                 .stdin(Stdio::piped())
                 .stdout(Stdio::piped())
                 .stderr(Stdio::piped());
-            if let Some(args) = &test.config.objdump {
-                cmd.args(args.to_vec());
+            match &test.config.objdump {
+                Some(args) => {
+                    cmd.args(args.to_vec());
+                }
+                None => {
+                    cmd.arg("--traps=false");
+                }
             }
 
             let mut child = cmd.spawn().context("failed to run wasmtime")?;

tests/disas/gc/struct-new-stack-map.wat

@@ -0,0 +1,76 @@
+;;! target = "x86_64"
+;;! flags = "-W function-references,gc"
+;;! test = "compile"
+
+(module
+  (type $ty (struct (field (mut f32))
+                    (field (mut i8))
+                    (field (mut anyref))))
+
+  (func (param f32 i32 anyref) (result (ref $ty))
+    (struct.new $ty (local.get 0) (local.get 1) (local.get 2))
+  )
+)
+;; wasm[0]::function[0]:
+;;       pushq   %rbp
+;;       movq    %rsp, %rbp
+;;       movq    8(%rdi), %r10
+;;       movq    0x10(%r10), %r10
+;;       addq    $0x50, %r10
+;;       cmpq    %rsp, %r10
+;;       ja      0xe4
+;;   19: subq    $0x40, %rsp
+;;       movq    %r13, 0x20(%rsp)
+;;       movq    %r14, 0x28(%rsp)
+;;       movq    %r15, 0x30(%rsp)
+;;       movq    %rdx, %r15
+;;       movdqu  %xmm0, 8(%rsp)
+;;       leaq    (%rsp), %r14
+;;       movl    %ecx, (%r14)
+;;       movl    $0xb0000001, %esi
+;;       xorl    %edx, %edx
+;;       movl    $0x28, %ecx
+;;       movl    $8, %r8d
+;;       movq    %rdi, %r13
+;;       callq   0x195
+;;       movq    0x28(%r13), %r9
+;;       ╰─╼ stack_map: frame_size=64, frame_offsets=[0]
+;;       movq    %rax, %r8
+;;       movl    %r8d, %r10d
+;;       movdqu  8(%rsp), %xmm0
+;;       movss   %xmm0, 0x1c(%r9, %r10)
+;;       movq    %r15, %rdx
+;;       movb    %dl, 0x20(%r9, %r10)
+;;       movl    (%r14), %r11d
+;;       movq    %r11, %rdx
+;;       andl    $1, %edx
+;;       testl   %r11d, %r11d
+;;       sete    %sil
+;;       movzbl  %sil, %esi
+;;       orl     %esi, %edx
+;;       testl   %edx, %edx
+;;       jne     0xc1
+;;   93: movl    %r11d, %edi
+;;       addq    $8, %rdi
+;;       jb      0xe6
+;;   a0: movq    %rdi, %rcx
+;;       addq    $8, %rcx
+;;       jb      0xe8
+;;   ad: cmpq    0x30(%r13), %rcx
+;;       ja      0xea
+;;   b7: movl    $1, %r11d
+;;       addq    %r11, (%r9, %rdi)
+;;       movl    (%r14), %r11d
+;;       movl    %r11d, 0x18(%r9, %r10)
+;;       movq    %r8, %rax
+;;       movq    0x20(%rsp), %r13
+;;       movq    0x28(%rsp), %r14
+;;       movq    0x30(%rsp), %r15
+;;       addq    $0x40, %rsp
+;;       movq    %rbp, %rsp
+;;       popq    %rbp
+;;       retq
+;;   e4: ud2
+;;   e6: ud2
+;;   e8: ud2
+;;   ea: ud2