From 76f2031884a7857649490f2ef8bcda534bd69c0c Mon Sep 17 00:00:00 2001
From: Takashi Kokubun <takashikkbn@gmail.com>
Date: Tue, 14 Mar 2023 10:26:05 -0700
Subject: [PATCH] YJIT: Allow testing assembler with disasm (#7470)

* YJIT: Allow testing assembler with disasm

* YJIT: Drop new dependencies

* YJIT: Avoid address manipulation

* YJIT: Introduce assert_disasm! macro

* YJIT: Update the comment about assert_disasm
---
 yjit/src/backend/x86_64/mod.rs | 17 +++++++-
 yjit/src/disasm.rs             | 71 ++++++++++++++++++++++++++++++++++
 2 files changed, 86 insertions(+), 2 deletions(-)

diff --git a/yjit/src/backend/x86_64/mod.rs b/yjit/src/backend/x86_64/mod.rs
index 11716e1fca..03ccc10928 100644
--- a/yjit/src/backend/x86_64/mod.rs
+++ b/yjit/src/backend/x86_64/mod.rs
@@ -765,6 +765,10 @@ impl Assembler
 
 #[cfg(test)]
 mod tests {
+    use crate::disasm::{assert_disasm};
+    #[cfg(feature = "disasm")]
+    use crate::disasm::{unindent, disasm_addr_range};
+
     use super::*;
 
     fn setup_asm() -> (Assembler, CodeBlock) {
@@ -938,19 +942,28 @@ 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_disasm!(cb, "488d5b08", {"
+            0x0: lea rbx, [rbx + 8]
+        "});
     }
 
     #[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_disasm!(cb, "488d4308488903", {"
+            0x0: lea rax, [rbx + 8]
+            0x4: mov qword ptr [rbx], rax
+        "});
     }
 
     #[test]
diff --git a/yjit/src/disasm.rs b/yjit/src/disasm.rs
index 54b4db51ed..279faf4ea8 100644
--- a/yjit/src/disasm.rs
+++ b/yjit/src/disasm.rs
@@ -171,6 +171,9 @@ pub fn disasm_addr_range(cb: &CodeBlock, start_addr: usize, end_addr: usize) ->
     // Disassemble the instructions
     let code_size = end_addr - start_addr;
     let code_slice = unsafe { std::slice::from_raw_parts(start_addr as _, code_size) };
+    // Stabilize output for cargo test
+    #[cfg(test)]
+    let start_addr = 0;
     let insns = cs.disasm_all(code_slice, start_addr as u64).unwrap();
 
     // Colorize outlined code in blue
@@ -195,6 +198,74 @@ pub fn disasm_addr_range(cb: &CodeBlock, start_addr: usize, end_addr: usize) ->
     return out;
 }
 
+/// Assert that CodeBlock has the code specified with hex. In addition, if tested with
+/// `cargo test --all-features`, it also checks it generates the specified disasm.
+#[cfg(test)]
+macro_rules! assert_disasm {
+    ($cb:expr, $hex:expr, $disasm:expr) => {
+        assert_eq!(format!("{:x}", $cb), $hex);
+        #[cfg(feature = "disasm")]
+        {
+            let disasm = disasm_addr_range(
+                &$cb,
+                $cb.get_ptr(0).raw_ptr() as usize,
+                $cb.get_write_ptr().raw_ptr() as usize,
+            );
+            assert_eq!(unindent(&disasm, false), unindent(&$disasm, true));
+        }
+    };
+}
+#[cfg(test)]
+pub(crate) use assert_disasm;
+
+/// Remove the minimum indent from every line, skipping the first line if `skip_first`.
+#[cfg(all(feature = "disasm", test))]
+pub fn unindent(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' ' {
+                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]