Let --yjit-dump-disasm=all dump ocb code as well (#6309)

* Let --yjit-dump-disasm=all dump ocb code as well

* Use an enum instead

* Add a None Option to DumpDisasm (#444)

* Add a None Option to DumpDisasm

* Update yjit/src/asm/mod.rs

Co-authored-by: Maxime Chevalier-Boisvert <maximechevalierb@gmail.com>

* Fix a build failure

* Use only a single name

* Only None will be a disabled case

* Fix cargo test

* Fix --yjit-dump-disasm=all to print outlined cb

Co-authored-by: Jimmy Miller <jimmyhmiller@gmail.com>
Co-authored-by: Maxime Chevalier-Boisvert <maximechevalierb@gmail.com>

5 files changed, 37 insertions, 8 deletions

diff --git a/yjit/src/asm/mod.rs b/yjit/src/asm/mod.rs
index 4029e2ca67..22b46e6aca 100644
--- a/yjit/src/asm/mod.rs
+++ b/yjit/src/asm/mod.rs
@@ -275,6 +275,11 @@ impl CodeBlock {
     pub fn mark_all_executable(&mut self) {
         self.mem_block.mark_all_executable();
     }
+
+    #[cfg(feature = "disasm")]
+    pub fn inline(&self) -> bool {
+        !self.outlined
+    }
 }
 
 #[cfg(test)]
diff --git a/yjit/src/backend/ir.rs b/yjit/src/backend/ir.rs
index 33a79a4179..0b96af7f62 100644
--- a/yjit/src/backend/ir.rs
+++ b/yjit/src/backend/ir.rs
@@ -1083,7 +1083,7 @@ impl Assembler
         let gc_offsets = self.compile_with_regs(cb, alloc_regs);
 
         #[cfg(feature = "disasm")]
-        if get_option!(dump_disasm) && !cb.outlined {
+        if get_option!(dump_disasm) == DumpDisasm::All || (get_option!(dump_disasm) == DumpDisasm::Inline && cb.inline()) {
             use crate::disasm::disasm_addr_range;
             let last_ptr = cb.get_write_ptr();
             let disasm = disasm_addr_range(cb, start_addr, last_ptr.raw_ptr() as usize - start_addr as usize);
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index 744495eb29..ac42b4a1e2 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -398,6 +398,7 @@ fn gen_code_for_exit_from_stub(ocb: &mut OutlinedCb) -> CodePtr {
 
     gen_counter_incr!(asm, exit_from_branch_stub);
 
+    asm.comment("exit from branch stub");
     asm.cpop_into(SP);
     asm.cpop_into(EC);
     asm.cpop_into(CFP);
@@ -525,7 +526,7 @@ fn gen_full_cfunc_return(ocb: &mut OutlinedCb) -> CodePtr {
     // This chunk of code expects REG_EC to be filled properly and
     // RAX to contain the return value of the C method.
 
-    // Call full_cfunc_return()
+    asm.comment("full cfunc return");
     asm.ccall(
         rb_full_cfunc_return as *const u8,
         vec![EC, C_RET_OPND]
@@ -562,6 +563,7 @@ fn gen_leave_exit(ocb: &mut OutlinedCb) -> CodePtr {
     // Every exit to the interpreter should be counted
     gen_counter_incr!(asm, leave_interp_return);
 
+    asm.comment("exit from leave");
     asm.cpop_into(SP);
     asm.cpop_into(EC);
     asm.cpop_into(CFP);
@@ -624,7 +626,7 @@ pub fn gen_entry_prologue(cb: &mut CodeBlock, iseq: IseqPtr, insn_idx: u32) -> O
     let code_ptr = cb.get_write_ptr();
 
     let mut asm = Assembler::new();
-    if get_option!(dump_disasm) {
+    if get_option!(dump_disasm).is_enabled() {
         asm.comment(&format!("YJIT entry: {}", iseq_get_location(iseq)));
     } else {
         asm.comment("YJIT entry");
@@ -753,7 +755,7 @@ pub fn gen_single_block(
     let mut asm = Assembler::new();
 
     #[cfg(feature = "disasm")]
-    if get_option!(dump_disasm) {
+    if get_option!(dump_disasm).is_enabled() {
         asm.comment(&format!("Block: {} (ISEQ offset: {})", iseq_get_location(blockid.iseq), blockid.idx));
     }
 
diff --git a/yjit/src/core.rs b/yjit/src/core.rs
index fa82dcc308..687dc21013 100644
--- a/yjit/src/core.rs
+++ b/yjit/src/core.rs
@@ -1779,6 +1779,7 @@ fn get_branch_target(
     let branch_ptr: *const RefCell<Branch> = BranchRef::into_raw(branchref.clone());
 
     let mut asm = Assembler::new();
+    asm.comment("branch stub hit");
 
     // Call branch_stub_hit(branch_ptr, target_idx, ec)
     let jump_addr = asm.ccall(
diff --git a/yjit/src/options.rs b/yjit/src/options.rs
index 2e141445f1..f73dca67de 100644
--- a/yjit/src/options.rs
+++ b/yjit/src/options.rs
@@ -30,8 +30,8 @@ pub struct Options {
     /// Dump compiled and executed instructions for debugging
     pub dump_insns: bool,
 
-    /// Dump all compiled instructions in inlined CodeBlock
-    pub dump_disasm: bool,
+    /// Dump all compiled instructions of target cbs.
+    pub dump_disasm: DumpDisasm,
 
     /// Print when specific ISEQ items are compiled or invalidated
     pub dump_iseq_disasm: Option<String>,
@@ -56,12 +56,28 @@ pub static mut OPTIONS: Options = Options {
     gen_stats: false,
     gen_trace_exits: false,
     dump_insns: false,
-    dump_disasm: false,
+    dump_disasm: DumpDisasm::None,
     verify_ctx: false,
     global_constant_state: false,
     dump_iseq_disasm: None,
 };
 
+#[derive(Copy, Clone, PartialEq, Eq, Debug)]
+pub enum DumpDisasm {
+    // Dump only inline cb
+    Inline,
+    // Dump both inline and outlined cbs
+    All,
+    // Dont dump anything
+    None,
+}
+
+impl DumpDisasm {
+    pub fn is_enabled(&self) -> bool {
+        *self != DumpDisasm::None
+    }
+}
+
 /// Macro to get an option value by name
 macro_rules! get_option {
     // Unsafe is ok here because options are initialized
@@ -123,6 +139,12 @@ pub fn parse_option(str_ptr: *const std::os::raw::c_char) -> Option<()> {
             }
         },
 
+        ("dump-disasm", _) => match opt_val.to_string().as_str() {
+            "all" => unsafe { OPTIONS.dump_disasm = DumpDisasm::All },
+            "" => unsafe { OPTIONS.dump_disasm = DumpDisasm::Inline },
+            _ => return None,
+         },
+
         ("dump-iseq-disasm", _) => unsafe {
             OPTIONS.dump_iseq_disasm = Some(opt_val.to_string());
         },
@@ -132,7 +154,6 @@ pub fn parse_option(str_ptr: *const std::os::raw::c_char) -> Option<()> {
         ("stats", "") => unsafe { OPTIONS.gen_stats = true },
         ("trace-exits", "") => unsafe { OPTIONS.gen_trace_exits = true; OPTIONS.gen_stats = true },
         ("dump-insns", "") => unsafe { OPTIONS.dump_insns = true },
-        ("dump-disasm", "") => unsafe { OPTIONS.dump_disasm = true },
         ("verify-ctx", "") => unsafe { OPTIONS.verify_ctx = true },
         ("global-constant-state", "") => unsafe { OPTIONS.global_constant_state = true },