YJIT: dump-disasm: Print comments and bytes in release builds

This change implements a fallback mode for the `--yjit-dump-disasm`
development command-line option to make it usable in release builds.
Previously, using the option with release builds of YJIT yielded only
a warning asking the user to build with `--enable-yjit=dev`.

While builds that use the `disasm` feature still give the best output,
just having the comments is useful enough for many kinds of debugging.
Having it usable in release builds is nice for new hackers, too, since
this allows for tinkering without having to learn how to build YJIT in
development mode.

Sample output on A64:

```
  # regenerate_branch
  # Insn: 0001 opt_send_without_block (stack_size: 1)
  # guard known object with singleton class
  0x11f7e0034: 4b 00 00 58 03 00 00 14 08 ce 9c 04 01 00 00
  0x11f7e0043: 00 3f 00 0b eb 81 06 01 54 1f 20 03 d5
  # RUBY_VM_CHECK_INTS(ec)
  0x11f7e0050: 8b 02 42 b8 cb 07 01 35
  # stack overflow check
  0x11f7e0058: ab 62 02 91 7f 02 0b eb 69 07 01 54
  # save PC to CFP
  0x11f7e0064: 0b 3b 9a d2 2b 2f a0 f2 0b 00 cc f2 6b 02 00
  0x11f7e0073: f8 ab 82 00 91
```

To ensure this feature doesn't incur too much cost when running without
the `--yjit-dump-disasm` option, I checked that there is no significant
impact to compile time and memory usage with the `compile_time_ns` and
`yjit_alloc_size` entry in `RubyVM::YJIT.runtime_stats`. For each
sample, I ran 3 iterations of the `lobsters` YJIT benchmark. The
statistics summary and done with the `summary` function in R.

Compile time, sample size of 60, lower is better:

```
       Before              After
 Min.   :2.054e+09   Min.   :2.028e+09
 1st Qu.:2.069e+09   1st Qu.:2.044e+09
 Median :2.081e+09   Median :2.060e+09
 Mean   :2.089e+09   Mean   :2.066e+09
 3rd Qu.:2.109e+09   3rd Qu.:2.085e+09
 Max.   :2.146e+09   Max.   :2.144e+09
```

Allocation size, sample size of 20, lower is better:

```
       Before             After
 Min.   :21804742   Min.   :21794082
 1st Qu.:21826682   1st Qu.:21816282
 Median :21844042   Median :21826814
 Mean   :21960664   Mean   :22026291
 3rd Qu.:21861228   3rd Qu.:22040439
 Max.   :22587426   Max.   :22930614
```

The `yjit_alloc_size` samples are noisy, but since the average increased
by only 0.3%, and the median is lower, I feel safe saying that there is
no significant change.

Author: XrXr

yjit/src/asm/mod.rs

@@ -2,16 +2,14 @@ use std::cell::RefCell;
 use std::fmt;
 use std::mem;
 use std::rc::Rc;
+use std::collections::BTreeMap;
+
 use crate::core::IseqPayload;
 use crate::core::for_each_off_stack_iseq_payload;
 use crate::core::for_each_on_stack_iseq_payload;
 use crate::invariants::rb_yjit_tracing_invalidate_all;
 use crate::stats::incr_counter;
 use crate::virtualmem::WriteError;
-
-#[cfg(feature = "disasm")]
-use std::collections::BTreeMap;
-
 use crate::codegen::CodegenGlobals;
 use crate::virtualmem::{VirtualMem, CodePtr};
 
@@ -77,8 +75,10 @@ pub struct CodeBlock {
     // References to labels
     label_refs: Vec<LabelRef>,
 
+    // A switch for keeping comments. They take up memory.
+    allow_comments: bool,
+
     // Comments for assembly instructions, if that feature is enabled
-    #[cfg(feature = "disasm")]
     asm_comments: BTreeMap<usize, Vec<String>>,
 
     // True for OutlinedCb
@@ -107,7 +107,7 @@ impl CodeBlock {
     const PREFERRED_CODE_PAGE_SIZE: usize = 16 * 1024;
 
     /// Make a new CodeBlock
-    pub fn new(mem_block: Rc<RefCell<VirtualMem>>, outlined: bool, freed_pages: Rc<Option<Vec<usize>>>) -> Self {
+    pub fn new(mem_block: Rc<RefCell<VirtualMem>>, outlined: bool, freed_pages: Rc<Option<Vec<usize>>>, allow_comments: bool) -> Self {
         // Pick the code page size
         let system_page_size = mem_block.borrow().system_page_size();
         let page_size = if 0 == Self::PREFERRED_CODE_PAGE_SIZE % system_page_size {
@@ -128,7 +128,7 @@ impl CodeBlock {
             label_addrs: Vec::new(),
             label_names: Vec::new(),
             label_refs: Vec::new(),
-            #[cfg(feature = "disasm")]
+            allow_comments,
             asm_comments: BTreeMap::new(),
             outlined,
             dropped_bytes: false,
@@ -366,9 +366,11 @@ impl CodeBlock {
     }
 
     /// Add an assembly comment if the feature is on.
-    /// If not, this becomes an inline no-op.
-    #[cfg(feature = "disasm")]
     pub fn add_comment(&mut self, comment: &str) {
+        if !self.allow_comments {
+            return;
+        }
+
         let cur_ptr = self.get_write_ptr().raw_addr(self);
 
         // If there's no current list of comments for this line number, add one.
@@ -379,28 +381,21 @@ impl CodeBlock {
             this_line_comments.push(comment.to_string());
         }
     }
-    #[cfg(not(feature = "disasm"))]
-    #[inline]
-    pub fn add_comment(&mut self, _: &str) {}
 
-    #[cfg(feature = "disasm")]
     pub fn comments_at(&self, pos: usize) -> Option<&Vec<String>> {
         self.asm_comments.get(&pos)
     }
 
-    #[allow(unused_variables)]
-    #[cfg(feature = "disasm")]
     pub fn remove_comments(&mut self, start_addr: CodePtr, end_addr: CodePtr) {
+        if self.asm_comments.is_empty() {
+            return;
+        }
         for addr in start_addr.raw_addr(self)..end_addr.raw_addr(self) {
             self.asm_comments.remove(&addr);
         }
     }
-    #[cfg(not(feature = "disasm"))]
-    #[inline]
-    pub fn remove_comments(&mut self, _: CodePtr, _: CodePtr) {}
 
     pub fn clear_comments(&mut self) {
-        #[cfg(feature = "disasm")]
         self.asm_comments.clear();
     }
 
@@ -693,7 +688,7 @@ impl CodeBlock {
         let mem_start: *const u8 = alloc.mem_start();
         let virt_mem = VirtualMem::new(alloc, 1, NonNull::new(mem_start as *mut u8).unwrap(), mem_size);
 
-        Self::new(Rc::new(RefCell::new(virt_mem)), false, Rc::new(None))
+        Self::new(Rc::new(RefCell::new(virt_mem)), false, Rc::new(None), true)
     }
 
     /// Stubbed CodeBlock for testing conditions that can arise due to code GC. Can't execute generated code.
@@ -711,7 +706,7 @@ impl CodeBlock {
         let mem_start: *const u8 = alloc.mem_start();
         let virt_mem = VirtualMem::new(alloc, 1, NonNull::new(mem_start as *mut u8).unwrap(), mem_size);
 
-        Self::new(Rc::new(RefCell::new(virt_mem)), false, Rc::new(Some(freed_pages)))
+        Self::new(Rc::new(RefCell::new(virt_mem)), false, Rc::new(Some(freed_pages)), true)
     }
 }
 

yjit/src/backend/arm64/mod.rs

@@ -924,9 +924,7 @@ impl Assembler
 
             match insn {
                 Insn::Comment(text) => {
-                    if cfg!(feature = "disasm") {
-                        cb.add_comment(text);
-                    }
+                    cb.add_comment(text);
                 },
                 Insn::Label(target) => {
                     cb.write_label(target.unwrap_label_idx());

yjit/src/backend/ir.rs

@@ -1565,13 +1565,10 @@ impl Assembler
     #[must_use]
     pub fn compile(self, cb: &mut CodeBlock, ocb: Option<&mut OutlinedCb>) -> Option<(CodePtr, Vec<u32>)>
     {
-        #[cfg(feature = "disasm")]
         let start_addr = cb.get_write_ptr();
-
         let alloc_regs = Self::get_alloc_regs();
         let ret = self.compile_with_regs(cb, ocb, alloc_regs);
 
-        #[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();
@@ -2057,10 +2054,10 @@ impl Assembler {
 }
 
 /// Macro to use format! for Insn::Comment, which skips a format! call
-/// when disasm is not supported.
+/// when not dumping disassembly.
 macro_rules! asm_comment {
     ($asm:expr, $($fmt:tt)*) => {
-        if cfg!(feature = "disasm") {
+        if $crate::options::get_option_ref!(dump_disasm).is_some() {
             $asm.push_insn(Insn::Comment(format!($($fmt)*)));
         }
     };

yjit/src/backend/x86_64/mod.rs

@@ -492,9 +492,7 @@ impl Assembler
 
             match insn {
                 Insn::Comment(text) => {
-                    if cfg!(feature = "disasm") {
-                        cb.add_comment(text);
-                    }
+                    cb.add_comment(text);
                 },
 
                 // Write the label at the current position

yjit/src/codegen.rs

@@ -10322,8 +10322,10 @@ impl CodegenGlobals {
             let mem_block = Rc::new(RefCell::new(mem_block));
 
             let freed_pages = Rc::new(None);
-            let cb = CodeBlock::new(mem_block.clone(), false, freed_pages.clone());
-            let ocb = OutlinedCb::wrap(CodeBlock::new(mem_block, true, freed_pages));
+
+            let asm_comments = get_option_ref!(dump_disasm).is_some();
+            let cb = CodeBlock::new(mem_block.clone(), false, freed_pages.clone(), asm_comments);
+            let ocb = OutlinedCb::wrap(CodeBlock::new(mem_block, true, freed_pages, asm_comments));
 
             (cb, ocb)
         };

yjit/src/disasm.rs

@@ -1,14 +1,10 @@
 use crate::core::*;
 use crate::cruby::*;
 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")]
 use std::fmt::Write;
 
 /// Primitive called in yjit.rb
@@ -116,7 +112,6 @@ pub fn disasm_iseq_insn_range(iseq: IseqPtr, start_idx: u16, end_idx: u16) -> St
 }
 
 /// Dump dissassembly for a range in a [CodeBlock]. VM lock required.
-#[cfg(feature = "disasm")]
 pub fn dump_disasm_addr_range(cb: &CodeBlock, start_addr: CodePtr, end_addr: CodePtr, dump_disasm: &DumpDisasm) {
     for (start_addr, end_addr) in cb.writable_addrs(start_addr, end_addr) {
         let disasm = disasm_addr_range(cb, start_addr, end_addr);
@@ -192,6 +187,44 @@ pub fn disasm_addr_range(cb: &CodeBlock, start_addr: usize, end_addr: usize) ->
     return out;
 }
 
+#[cfg(not(feature = "disasm"))]
+pub fn disasm_addr_range(cb: &CodeBlock, start_addr: usize, end_addr: usize) -> String {
+    let mut out = String::new();
+    let mut line_byte_idx = 0;
+    const MAX_BYTES_PER_LINE: usize = 16;
+
+    for addr in start_addr..end_addr {
+        if let Some(comment_list) = cb.comments_at(addr) {
+            // Start a new line if we're in the middle of one
+            if line_byte_idx != 0 {
+                writeln!(&mut out).unwrap();
+                line_byte_idx = 0;
+            }
+            for comment in comment_list {
+                writeln!(&mut out, "  \x1b[1m# {comment}\x1b[22m").unwrap(); // Make comments bold
+            }
+        }
+        if line_byte_idx == 0 {
+            write!(&mut out, "  0x{addr:x}: ").unwrap();
+        } else {
+            write!(&mut out, " ").unwrap();
+        }
+        let byte = unsafe { (addr as *const u8).read() };
+        write!(&mut out, "{byte:02x}").unwrap();
+        line_byte_idx += 1;
+        if line_byte_idx == MAX_BYTES_PER_LINE - 1 {
+            writeln!(&mut out).unwrap();
+            line_byte_idx = 0;
+        }
+    }
+
+    if !out.is_empty() {
+        writeln!(&mut out).unwrap();
+    }
+
+    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)]

yjit/src/options.rs

@@ -250,7 +250,7 @@ pub fn parse_option(str_ptr: *const std::os::raw::c_char) -> Option<()> {
 
         ("dump-disasm", _) => {
             if !cfg!(feature = "disasm") {
-                eprintln!("WARNING: the {} option is only available when YJIT is built in dev mode, i.e. ./configure --enable-yjit=dev", opt_name);
+                eprintln!("WARNING: the {} option works best when YJIT is built in dev mode, i.e. ./configure --enable-yjit=dev", opt_name);
             }
 
             match opt_val {