YJIT: Skip dumping code for the other cb on --yjit-dump-disasm (#6592)

YJIT: Skip dumping code for the other cb

on --yjit-dump-disasm

3 files changed, 17 insertions, 14 deletions

diff --git a/yjit/src/asm/mod.rs b/yjit/src/asm/mod.rs
index 271f11defc..61c1523aab 100644
--- a/yjit/src/asm/mod.rs
+++ b/yjit/src/asm/mod.rs
@@ -213,9 +213,8 @@ impl CodeBlock {
         self.page_end_reserve = old_page_end_reserve;
     }
 
-    #[cfg(target_arch = "aarch64")]
-    #[cfg(not(test))]
     /// Return the address ranges of a given address range that this CodeBlock can write.
+    #[cfg(any(feature = "disasm", target_arch = "aarch64"))]
     pub fn writable_addrs(&self, start_ptr: CodePtr, end_ptr: CodePtr) -> Vec<(usize, usize)> {
         let mut addrs = vec![];
         let mut start = start_ptr.raw_ptr() as usize;
diff --git a/yjit/src/backend/ir.rs b/yjit/src/backend/ir.rs
index df0600b2ff..64e5805d9e 100644
--- a/yjit/src/backend/ir.rs
+++ b/yjit/src/backend/ir.rs
@@ -1094,7 +1094,7 @@ impl Assembler
     pub fn compile(self, cb: &mut CodeBlock) -> Vec<u32>
     {
         #[cfg(feature = "disasm")]
-        let start_addr = cb.get_write_ptr().raw_ptr();
+        let start_addr = cb.get_write_ptr();
 
         let alloc_regs = Self::get_alloc_regs();
         let gc_offsets = self.compile_with_regs(cb, alloc_regs);
@@ -1102,7 +1102,7 @@ impl Assembler
         #[cfg(feature = "disasm")]
         if let Some(dump_disasm) = get_option_ref!(dump_disasm) {
             use crate::disasm::dump_disasm_addr_range;
-            let end_addr = cb.get_write_ptr().raw_ptr();
+            let end_addr = cb.get_write_ptr();
             dump_disasm_addr_range(cb, start_addr, end_addr, dump_disasm)
         }
         gc_offsets
diff --git a/yjit/src/disasm.rs b/yjit/src/disasm.rs
index 101afb681b..7795372696 100644
--- a/yjit/src/disasm.rs
+++ b/yjit/src/disasm.rs
@@ -4,6 +4,8 @@ use crate::yjit::yjit_enabled_p;
 #[cfg(feature = "disasm")]
 use crate::asm::CodeBlock;
 #[cfg(feature = "disasm")]
+use crate::codegen::CodePtr;
+#[cfg(feature = "disasm")]
 use crate::options::DumpDisasm;
 
 #[cfg(feature = "disasm")]
@@ -120,19 +122,21 @@ pub fn disasm_iseq_insn_range(iseq: IseqPtr, start_idx: u32, end_idx: u32) -> St
 }
 
 #[cfg(feature = "disasm")]
-pub fn dump_disasm_addr_range(cb: &CodeBlock, start_addr: *const u8, end_addr: *const u8, dump_disasm: &DumpDisasm) {
+pub fn dump_disasm_addr_range(cb: &CodeBlock, start_addr: CodePtr, end_addr: CodePtr, dump_disasm: &DumpDisasm) {
     use std::fs::File;
     use std::io::Write;
 
-    let disasm = disasm_addr_range(cb, start_addr, end_addr);
-    if disasm.len() > 0 {
-        match dump_disasm {
-            DumpDisasm::Stdout => println!("{disasm}"),
-            DumpDisasm::File(path) => {
-                let mut f = File::options().append(true).create(true).open(path).unwrap();
-                f.write_all(disasm.as_bytes()).unwrap();
-            }
-        };
+    for (start_addr, end_addr) in cb.writable_addrs(start_addr, end_addr) {
+        let disasm = disasm_addr_range(cb, start_addr as *const u8, end_addr as *const u8);
+        if disasm.len() > 0 {
+            match dump_disasm {
+                DumpDisasm::Stdout => println!("{disasm}"),
+                DumpDisasm::File(path) => {
+                    let mut f = File::options().create(true).append(true).open(path).unwrap();
+                    f.write_all(disasm.as_bytes()).unwrap();
+                }
+            };
+        }
     }
 }