Add better debugging tools

src/backend.rs

@@ -10,6 +10,7 @@ use std::{
     any::{Any, TypeId},
     collections::HashMap,
     convert::TryFrom,
+    fmt::Display,
     iter::{self, FromIterator},
     mem,
     ops::RangeInclusive,
@@ -518,6 +519,7 @@ pub struct FunctionEnd {
 pub struct CodeGenSession<'module, M> {
     assembler: Assembler,
     pub module_context: &'module M,
+    pub op_offset_map: Vec<(AssemblyOffset, Box<dyn Display + Send + Sync>)>,
     labels: Labels,
     func_starts: Vec<(Option<AssemblyOffset>, DynamicLabel)>,
 }
@@ -531,6 +533,7 @@ impl<'module, M> CodeGenSession<'module, M> {
 
         CodeGenSession {
             assembler,
+            op_offset_map: Default::default(),
             labels: Default::default(),
             func_starts,
             module_context,
@@ -575,6 +578,7 @@ impl<'module, M> CodeGenSession<'module, M> {
         Ok(TranslatedCodeSection {
             exec_buf,
             func_starts,
+            op_offset_map: self.op_offset_map,
             // TODO
             relocatable_accesses: vec![],
         })
@@ -594,11 +598,11 @@ struct RelocateAccess {
     address: RelocateAddress,
 }
 
-#[derive(Debug)]
 pub struct TranslatedCodeSection {
     exec_buf: ExecutableBuffer,
     func_starts: Vec<AssemblyOffset>,
     relocatable_accesses: Vec<RelocateAccess>,
+    op_offset_map: Vec<(AssemblyOffset, Box<dyn Display + Send + Sync>)>,
 }
 
 impl TranslatedCodeSection {
@@ -626,7 +630,7 @@ impl TranslatedCodeSection {
     }
 
     pub fn disassemble(&self) {
-        crate::disassemble::disassemble(&*self.exec_buf).unwrap();
+        crate::disassemble::disassemble(&*self.exec_buf, &self.op_offset_map).unwrap();
     }
 }
 
@@ -651,7 +655,7 @@ type Labels = HashMap<
 >;
 
 pub struct Context<'this, M> {
-    asm: &'this mut Assembler,
+    pub asm: &'this mut Assembler,
     reloc_sink: &'this mut dyn binemit::RelocSink,
     module_context: &'this M,
     current_function: u32,
@@ -2325,10 +2329,22 @@ impl<'this, M: ModuleContext> Context<'this, M> {
             let end_label = self.create_label();
 
             if count > 0 {
+                // TODO: Handle this failing (because we might have all registers used for block
+                //       args we can't rely on this succeeding)
                 let selector_reg = self.into_temp_reg(GPRType::Rq, selector).unwrap();
                 selector = ValueLocation::Reg(selector_reg);
 
-                let tmp = self.take_reg(I64).unwrap();
+                let (tmp, pop_tmp) = if let Some(reg) = self.take_reg(I64) {
+                    (reg, false)
+                } else {
+                    let out_reg = if selector_reg == RAX { RCX } else { RAX };
+
+                    dynasm!(self.asm
+                        ; push Rq(out_reg.rq().unwrap())
+                    );
+
+                    (out_reg, true)
+                };
 
                 self.immediate_to_reg(tmp, (count as u32).into());
                 dynasm!(self.asm
@@ -2339,12 +2355,21 @@ impl<'this, M: ModuleContext> Context<'this, M> {
                         Rq(selector_reg.rq().unwrap()) * 5
                     ]
                     ; add Rq(selector_reg.rq().unwrap()), Rq(tmp.rq().unwrap())
+                );
+
+                if pop_tmp {
+                    dynasm!(self.asm
+                        ; pop Rq(tmp.rq().unwrap())
+                    );
+                }
+
+                self.block_state.regs.release(tmp);
+
+                dynasm!(self.asm
                     ; jmp Rq(selector_reg.rq().unwrap())
                 ; start_label:
                 );
 
-                self.block_state.regs.release(tmp);
-
                 for target in targets {
                     let label = target
                         .map(|target| self.target_to_label(target))
@@ -2743,14 +2768,21 @@ impl<'this, M: ModuleContext> Context<'this, M> {
                         ; mov DWORD [rsp + out_offset], i as i32
                     );
                 } else {
-                    let scratch = self.take_reg(I64).unwrap();
+                    if let Some(scratch) = self.take_reg(I64) {
+                        dynasm!(self.asm
+                            ; mov Rq(scratch.rq().unwrap()), QWORD i
+                            ; mov [rsp + out_offset], Rq(scratch.rq().unwrap())
+                        );
 
-                    dynasm!(self.asm
+                        self.block_state.regs.release(scratch);
-                        ; mov Rq(scratch.rq().unwrap()), QWORD i
+                    } else {
-                        ; mov [rsp + out_offset], Rq(scratch.rq().unwrap())
+                        dynasm!(self.asm
-                    );
+                            ; push rax
-
+                            ; mov rax, QWORD i
-                    self.block_state.regs.release(scratch);
+                            ; mov [rsp + out_offset], rax
+                            ; pop rax
+                        );
+                    }
                 }
             }
             (ValueLocation::Stack(in_offset), CCLoc::Reg(out_reg)) => {
@@ -2978,7 +3010,6 @@ impl<'this, M: ModuleContext> Context<'this, M> {
 
     fn into_temp_loc(&mut self, ty: impl Into<Option<GPRType>>, val: ValueLocation) -> CCLoc {
         match val {
-            ValueLocation::Stack(o) => CCLoc::Stack(o),
             _ => {
                 if let Some(gpr) = self.into_temp_reg(ty, val) {
                     CCLoc::Reg(gpr)

src/disassemble.rs

@@ -1,32 +1,51 @@
-use crate::error::Error;
 use capstone::prelude::*;
-use std::fmt::Write;
+use dynasmrt::AssemblyOffset;
+use std::error::Error;
+use std::fmt::{Display, Write};
 
-pub fn disassemble(mem: &[u8]) -> Result<(), Error> {
+pub fn disassemble(
+    mem: &[u8],
+    mut ops: &[(AssemblyOffset, impl Display)],
+) -> Result<(), Box<dyn Error>> {
     let mut cs = Capstone::new()
         .x86()
         .mode(arch::x86::ArchMode::Mode64)
         .build()?;
 
     println!("{} bytes:", mem.len());
-    let insns = cs.disasm_all(&mem, 0x0).unwrap();
+    let insns = cs.disasm_all(&mem, 0x0)?;
     for i in insns.iter() {
         let mut line = String::new();
 
-        write!(&mut line, "{:4x}:\t", i.address()).unwrap();
+        let address = i.address();
+
+        loop {
+            if let Some((offset, op)) = ops.first() {
+                if offset.0 as u64 <= address {
+                    ops = &ops[1..];
+                    println!("{}", op);
+                } else {
+                    break;
+                }
+            } else {
+                break;
+            }
+        }
+
+        write!(&mut line, "{:4x}:\t", i.address())?;
 
         let mut bytes_str = String::new();
         for b in i.bytes() {
-            write!(&mut bytes_str, "{:02x} ", b).unwrap();
+            write!(&mut bytes_str, "{:02x} ", b)?;
         }
-        write!(&mut line, "{:24}\t", bytes_str).unwrap();
+        write!(&mut line, "{:24}\t", bytes_str)?;
 
         if let Some(s) = i.mnemonic() {
-            write!(&mut line, "{}\t", s).unwrap();
+            write!(&mut line, "{}\t", s)?;
         }
 
         if let Some(s) = i.op_str() {
-            write!(&mut line, "{}", s).unwrap();
+            write!(&mut line, "{}", s)?;
         }
 
         println!("{}", line);

src/function_body.rs

@@ -5,9 +5,10 @@ use crate::error::Error;
 use crate::microwasm::*;
 use crate::module::{ModuleContext, SigType, Signature};
 use cranelift_codegen::binemit;
+use dynasmrt::DynasmApi;
 use either::{Either, Left, Right};
 use multi_mut::HashMapMultiMut;
-use std::{collections::HashMap, hash::Hash};
+use std::{collections::HashMap, fmt, hash::Hash, mem};
 
 #[derive(Debug)]
 struct Block {
@@ -73,7 +74,7 @@ where
     )
 }
 
-pub fn translate<M, I, L>(
+pub fn translate<M, I, L: Send + Sync + 'static>(
     session: &mut CodeGenSession<M>,
     reloc_sink: &mut dyn binemit::RelocSink,
     func_idx: u32,
@@ -107,7 +108,12 @@ where
     let mut body = body.into_iter().peekable();
 
     let module_context = &*session.module_context;
+    let mut op_offset_map = mem::replace(&mut session.op_offset_map, vec![]);
     let ctx = &mut session.new_context(func_idx, reloc_sink);
+    op_offset_map.push((
+        ctx.asm.offset(),
+        Box::new(format!("Function {}:", func_idx)),
+    ));
 
     let params = func_type
         .params()
@@ -126,8 +132,6 @@ where
         Block {
             label: BrTarget::Return,
             params: num_returns as u32,
-            // TODO: This only works for integers
-            //
             calling_convention: Some(Left(BlockCallingConvention::function_start(ret_locs(
                 func_type.returns().iter().map(|t| t.to_microwasm_type()),
             )))),
@@ -146,6 +150,26 @@ where
             block.is_next = true;
         }
 
+        struct DisassemblyOpFormatter<Label>(Operator<Label>);
+
+        impl<Label> fmt::Display for DisassemblyOpFormatter<Label>
+        where
+            Operator<Label>: fmt::Display,
+        {
+            fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+                match self.0 {
+                    Operator::Label(_) => write!(f, "{}", self.0),
+                    Operator::Block { .. } => write!(f, "{:5}\t{}", "", self.0),
+                    _ => write!(f, "{:5}\t  {}", "", self.0),
+                }
+            }
+        }
+
+        op_offset_map.push((
+            ctx.asm.offset(),
+            Box::new(DisassemblyOpFormatter(op.clone())),
+        ));
+
         match op {
             Operator::Unreachable => {
                 ctx.trap();
@@ -764,5 +788,7 @@ where
 
     ctx.epilogue();
 
+    mem::replace(&mut session.op_offset_map, op_offset_map);
+
     Ok(())
 }

src/microwasm.rs

@@ -305,7 +305,7 @@ pub enum NameTag {
 
 pub type WasmLabel = (u32, NameTag);
 
-type OperatorFromWasm = Operator<WasmLabel>;
+pub type OperatorFromWasm = Operator<WasmLabel>;
 
 #[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)]
 pub enum BrTarget<L> {

src/tests.rs

@@ -55,7 +55,7 @@ mod op32 {
                     }
 
                     fn lit_reg(a: i32, b: i32) -> bool {
-                                                let translated = translate_wat(&format!("
+                        let translated = translate_wat(&format!("
                             (module (func (param i32) (result i32)
                                 (i32.{op} (i32.const {left}) (get_local 0))))
                         ", op = OP, left = a));
@@ -66,7 +66,7 @@ mod op32 {
                     }
 
                     fn reg_lit(a: i32, b: i32) -> bool {
-                                                let translated = translate_wat(&format!("
+                        let translated = translate_wat(&format!("
                             (module (func (param i32) (result i32)
                                 (i32.{op} (get_local 0) (i32.const {right}))))
                         ", op = OP, right = b));
@@ -101,7 +101,7 @@ mod op32 {
                     }
 
                     fn lit(a: u32) -> bool {
-                                                let translated = translate_wat(&format!(concat!("
+                        let translated = translate_wat(&format!(concat!("
                             (module (func (result i32)
                                 (i32.",stringify!($name)," (i32.const {val}))))
                         "), val = a));