YJIT: Allow testing assembler with disasm

yjit/Cargo.toml

@@ -18,6 +18,9 @@ crate-type = ["staticlib"]
 capstone = { version = "0.10.0", optional = true }
 stats_alloc = { version = "0.1.10", optional = true }
 
+[dev-dependencies]
+capstone = "0.10.0"
+
 [features]
 # NOTE: Development builds select a set of these via configure.ac
 # For debugging, `make V=1` shows exact cargo invocation.

yjit/src/asm/mod.rs

@@ -12,7 +12,7 @@ use crate::core::for_each_on_stack_iseq_payload;
 use crate::invariants::rb_yjit_tracing_invalidate_all;
 use crate::virtualmem::WriteError;
 
-#[cfg(feature = "disasm")]
+#[cfg(any(feature = "disasm", test))]
 use std::collections::BTreeMap;
 
 use crate::codegen::CodegenGlobals;
@@ -75,7 +75,7 @@ pub struct CodeBlock {
     label_refs: Vec<LabelRef>,
 
     // Comments for assembly instructions, if that feature is enabled
-    #[cfg(feature = "disasm")]
+    #[cfg(any(feature = "disasm", test))]
     asm_comments: BTreeMap<usize, Vec<String>>,
 
     // True for OutlinedCb
@@ -123,7 +123,7 @@ impl CodeBlock {
             label_addrs: Vec::new(),
             label_names: Vec::new(),
             label_refs: Vec::new(),
-            #[cfg(feature = "disasm")]
+            #[cfg(any(feature = "disasm", test))]
             asm_comments: BTreeMap::new(),
             outlined,
             dropped_bytes: false,
@@ -367,7 +367,7 @@ impl CodeBlock {
     #[inline]
     pub fn add_comment(&mut self, _: &str) {}
 
-    #[cfg(feature = "disasm")]
+    #[cfg(any(feature = "disasm", test))]
     pub fn comments_at(&self, pos: usize) -> Option<&Vec<String>> {
         self.asm_comments.get(&pos)
     }

yjit/src/backend/x86_64/mod.rs

@@ -756,6 +756,8 @@ impl Assembler
 
 #[cfg(test)]
 mod tests {
+    use crate::disasm::{disasm_code_block, unindent, unindent_string};
+
     use super::*;
 
     fn setup_asm() -> (Assembler, CodeBlock) {
@@ -929,18 +931,25 @@ mod tests {
     #[test]
     fn test_merge_lea_reg() {
         let (mut asm, mut cb) = setup_asm();
+
         let sp = asm.lea(Opnd::mem(64, SP, 8));
         asm.mov(SP, sp); // should be merged to lea
         asm.compile_with_num_regs(&mut cb, 1);
-        assert_eq!(format!("{:x}", cb), "488d5b08");
+
+        assert_eq!(disasm_code_block(&cb), "0x0: lea rbx, [rbx + 8]\n");
     }
 
     #[test]
     fn test_merge_lea_mem() {
         let (mut asm, mut cb) = setup_asm();
+
         let sp = asm.lea(Opnd::mem(64, SP, 8));
         asm.mov(Opnd::mem(64, SP, 0), sp); // should NOT be merged to lea
         asm.compile_with_num_regs(&mut cb, 1);
-        assert_eq!(format!("{:x}", cb), "488d4308488903");
+
+        assert_eq!(disasm_code_block(&cb), unindent! {"
+            0x0: lea rax, [rbx + 8]
+            0x4: mov qword ptr [rbx], rax
+        "});
     }
 }

yjit/src/disasm.rs

@@ -1,14 +1,14 @@
 use crate::core::*;
 use crate::cruby::*;
 use crate::yjit::yjit_enabled_p;
-#[cfg(feature = "disasm")]
+#[cfg(any(feature = "disasm", test))]
 use crate::asm::CodeBlock;
 #[cfg(feature = "disasm")]
 use crate::codegen::CodePtr;
 #[cfg(feature = "disasm")]
 use crate::options::DumpDisasm;
 
-#[cfg(feature = "disasm")]
+#[cfg(any(feature = "disasm", test))]
 use std::fmt::Write;
 
 /// Primitive called in yjit.rb
@@ -144,7 +144,7 @@ pub fn dump_disasm_addr_range(cb: &CodeBlock, start_addr: CodePtr, end_addr: Cod
     }
 }
 
-#[cfg(feature = "disasm")]
+#[cfg(any(feature = "disasm", test))]
 pub fn disasm_addr_range(cb: &CodeBlock, start_addr: usize, end_addr: usize) -> String {
     let mut out = String::from("");
 
@@ -195,6 +195,83 @@ pub fn disasm_addr_range(cb: &CodeBlock, start_addr: usize, end_addr: usize) ->
     return out;
 }
 
+/// Disassemble the entire code block for testing
+#[cfg(test)]
+pub fn disasm_code_block(cb: &CodeBlock) -> String {
+    let start_addr = cb.get_ptr(0).raw_ptr() as usize;
+    let end_addr = cb.get_write_ptr().raw_ptr() as usize;
+
+    let mut disasm = disasm_addr_range(cb, start_addr, end_addr);
+    disasm = unindent_string(&disasm, false);
+
+    // Code addresses will be different every time you test it. So this loop replaces
+    // code addresses that appear in disasm with `start_addr`-origin values.
+    for addr in start_addr..end_addr {
+        let from = format!("{:#x}", addr);
+        let to = format!("{:#x}", addr - start_addr);
+        disasm = disasm.replace(&from, &to);
+    }
+    disasm
+}
+
+/// Macro useful for raw string literals
+#[cfg(test)]
+macro_rules! unindent {
+    ($str:expr) => {
+        unindent_string($str, true)
+    };
+}
+#[cfg(test)]
+pub(crate) use unindent;
+
+/// Remove the minimum indent from every line, skipping the first line if `skip_first`.
+#[cfg(test)]
+pub fn unindent_string(string: &str, trim_lines: bool) -> String {
+    fn split_lines(string: &str) -> Vec<String> {
+        let mut result: Vec<String> = vec![];
+        let mut buf: Vec<u8> = vec![];
+        for byte in string.as_bytes().iter() {
+            buf.push(*byte);
+            if *byte == b'\n' {
+                result.push(String::from_utf8(buf).unwrap());
+                buf = vec![];
+            }
+        }
+        if !buf.is_empty() {
+            result.push(String::from_utf8(buf).unwrap());
+        }
+        result
+    }
+
+    // Break up a string into multiple lines
+    let mut lines = split_lines(string);
+    if trim_lines { // raw string literals come with extra lines
+        lines.remove(0);
+        lines.remove(lines.len() - 1);
+    }
+
+    // Count the minimum number of spaces
+    let spaces = lines.iter().filter_map(|line| {
+        for (i, ch) in line.as_bytes().iter().enumerate() {
+            if *ch != b' ' && *ch != b'\t' {
+                return Some(i);
+            }
+        }
+        None
+    }).min().unwrap_or(0);
+
+    // Join lines, removing spaces
+    let mut unindented: Vec<u8> = vec![];
+    for line in lines.iter() {
+        if line.len() > spaces {
+            unindented.extend_from_slice(&line.as_bytes()[spaces..]);
+        } else {
+            unindented.extend_from_slice(&line.as_bytes());
+        }
+    }
+    String::from_utf8(unindented).unwrap()
+}
+
 /// Primitive called in yjit.rb
 /// Produce a list of instructions compiled for an isew
 #[no_mangle]