Add comments to YJIT disassembly (ruby#231)

Noah Gibbs · web-flow · commit 3f5b30164bc3 · 2022-03-14T13:22:52.000-04:00
* Add comments to YJIT disassembly. Also add checks to avoid panics when
disassembling without YJIT active.

* Fix test to use new asm comments interface

* Use self.enabled? instead of primitive expr

* Comments on appending to Vec&lt;String&gt; in add_comment

* Directly use the global inline codeblock in disasm, don't abstract it

* Use usize instead of u64 for addresses
diff --git a/yjit.rb b/yjit.rb
@@ -34,9 +34,13 @@ def self.disasm(iseq)
     # If a method or proc is passed in, get its iseq
     iseq = RubyVM::InstructionSequence.of(iseq)
 
-    # Produce the disassembly string
-    # Include the YARV iseq disasm in the string for additional context
-    iseq.disasm + "\n" + Primitive.rb_yjit_disasm_iseq(iseq)
+    if self.enabled?
+      # Produce the disassembly string
+      # Include the YARV iseq disasm in the string for additional context
+      iseq.disasm + "\n" + Primitive.rb_yjit_disasm_iseq(iseq)
+    else
+      iseq.disasm
+    end
   end
 
   #def self.simulate_oom!
diff --git a/yjit/src/asm/mod.rs b/yjit/src/asm/mod.rs
@@ -1,4 +1,5 @@
 use std::mem;
+use std::collections::{BTreeMap};
 
 pub mod x86_64;
 
@@ -19,6 +20,11 @@ impl CodePtr {
         let CodePtr(ptr) = self;
         *ptr as i64
     }
+
+    fn into_usize(&self) -> usize {
+        let CodePtr(ptr) = self;
+        *ptr as usize
+    }
 }
 
 impl From<*mut u8> for CodePtr {
@@ -88,7 +94,7 @@ pub struct CodeBlock
     label_refs: Vec<LabelRef>,
 
     // Comments for assembly instructions, if that feature is enabled
-    asm_comments: Vec<(usize, String)>,
+    asm_comments: BTreeMap<usize, Vec<String>>,
 
     // Keep track of the current aligned write position.
     // Used for changing protection when writing to the JIT buffer
@@ -119,7 +125,7 @@ impl CodeBlock
             label_addrs: Vec::new(),
             label_names: Vec::new(),
             label_refs: Vec::new(),
-            asm_comments: Vec::new(),
+            asm_comments: BTreeMap::new(),
             current_aligned_write_pos: ALIGNED_WRITE_POSITION_NONE,
             page_size: 4096,
             dropped_bytes: false
@@ -135,7 +141,7 @@ impl CodeBlock
             label_addrs: Vec::new(),
             label_names: Vec::new(),
             label_refs: Vec::new(),
-            asm_comments: Vec::new(),
+            asm_comments: BTreeMap::new(),
             current_aligned_write_pos: ALIGNED_WRITE_POSITION_NONE,
             page_size,
             dropped_bytes: false
@@ -153,26 +159,26 @@ impl CodeBlock
     pub fn add_comment(&mut self, comment: &str) {
         #[cfg(feature="asm_comments")]
         {
-            let is_dup = if let Some((last_pos, last_comment)) = self.asm_comments.last() {
-                *last_pos == self.write_pos && *last_comment == comment
-            } else { false };
-            if !is_dup {
-                self.asm_comments.push((self.write_pos, String::from(comment)));
+            let cur_ptr = self.get_write_ptr().into_usize();
+            let this_line_comments = self.asm_comments.get(&cur_ptr);
+
+            // If there's no current list of comments for this line number, add one.
+            if this_line_comments.is_none() {
+                let new_comments = Vec::new();
+                self.asm_comments.insert(cur_ptr, new_comments);
             }
-        }
-    }
+            let this_line_comments = self.asm_comments.get_mut(&cur_ptr).unwrap();
 
-    /// Add an assembly comment at a specific byte offset if the feature is on.
-    /// If not, this becomes an inline no-op.
-    #[inline]
-    pub fn add_comment_at(&mut self, pos:usize, comment: &str) {
-        #[cfg(feature="asm_comments")]
-        self.asm_comments.push((pos, String::from(comment)));
+            // Unless this comment is the same as the last one at this same line, add it.
+            let string_comment = String::from(comment);
+            if this_line_comments.last() != Some(&string_comment) {
+                this_line_comments.push(string_comment);
+            }
+        }
     }
 
-    /// Get a slice (readonly ref) of assembly comments - if the feature is off, this will be empty.
-    pub fn get_comments(&self) -> &[(usize, String)] {
-        return self.asm_comments.as_slice();
+    pub fn comments_at(&self, pos: usize) -> Option<&Vec<String>> {
+        self.asm_comments.get(&pos)
     }
 
     pub fn get_write_pos(&self) -> usize {
diff --git a/yjit/src/asm/x86_64/tests.rs b/yjit/src/asm/x86_64/tests.rs
@@ -430,14 +430,18 @@ fn basic_capstone_usage() -> std::result::Result<(), capstone::Error> {
 fn block_comments() {
     let mut cb = super::CodeBlock::new_dummy(4096);
 
+    let first_write_ptr = cb.get_write_ptr().into_usize();
     cb.add_comment("Beginning");
     xor(&mut cb, EAX, EAX); // 2 bytes long
+    let second_write_ptr = cb.get_write_ptr().into_usize();
     cb.add_comment("Two bytes in");
-    cb.add_comment("Two bytes in"); // Duplicate, should be ignored
+    cb.add_comment("Still two bytes in");
+    cb.add_comment("Still two bytes in"); // Duplicate, should be ignored
     test(&mut cb, mem_opnd(64, RSI, 64), imm_opnd(!0x08)); // 8 bytes long
+    let third_write_ptr = cb.get_write_ptr().into_usize();
     cb.add_comment("Ten bytes in");
 
-    let comments = cb.get_comments();
-    let expected:Vec<(usize,String)> = vec!( (0, "Beginning".to_string()), (2, "Two bytes in".to_string()), (10, "Ten bytes in".to_string()) );
-    assert_eq!(expected, comments);
+    assert_eq!(&vec!( "Beginning".to_string() ), cb.comments_at(first_write_ptr).unwrap());
+    assert_eq!(&vec!( "Two bytes in".to_string(), "Still two bytes in".to_string() ), cb.comments_at(second_write_ptr).unwrap());
+    assert_eq!(&vec!( "Ten bytes in".to_string() ), cb.comments_at(third_write_ptr).unwrap());
 }
diff --git a/yjit/src/disasm.rs b/yjit/src/disasm.rs
@@ -1,6 +1,9 @@
 use std::fmt::Write;
 use crate::cruby::*;
+use crate::codegen::*;
+use crate::asm::*;
 use crate::core::*;
+use crate::yjit::yjit_enabled_p;
 
 /// Primitive called in yjit.rb
 /// Produce a string representing the disassembly for an ISEQ
@@ -19,6 +22,10 @@ pub extern "C" fn rb_yjit_disasm_iseq(ec: EcPtr, ruby_self: VALUE, iseqw: VALUE)
         //    return Qnil;
         //}
 
+        if !yjit_enabled_p() {
+            return Qnil;
+        }
+
         // Get the iseq pointer from the wrapper
         let iseq = unsafe { rb_iseqw_to_iseq(iseqw) };
 
@@ -35,6 +42,9 @@ fn disasm_iseq(iseq: IseqPtr) -> String {
     // Get a list of block versions generated for this iseq
     let mut block_list = get_iseq_block_list(iseq);
 
+    // Get a list of codeblocks relevant to this iseq
+    let global_cb = CodegenGlobals::get_inline_cb();
+
     // Sort the blocks by increasing start addresses
     block_list.sort_by(|a, b| {
         use std::cmp::Ordering;
@@ -57,7 +67,7 @@ fn disasm_iseq(iseq: IseqPtr) -> String {
     // Compute total code size in bytes for all blocks in the function
     let mut total_code_size = 0;
     for blockref in &block_list {
-        total_code_size = blockref.borrow().code_size();
+        total_code_size += blockref.borrow().code_size();
     }
 
     // Initialize capstone
@@ -99,6 +109,12 @@ fn disasm_iseq(iseq: IseqPtr) -> String {
 
         // For each instruction in this block
         for insn in insns.as_ref() {
+            // Comments for this block
+            if let Some(comment_list) = global_cb.comments_at(insn.address() as usize) {
+                for comment in comment_list {
+                    out.push_str(&format!("  // {}\n", comment));
+                }
+            }
             out.push_str(&format!("  {}\n", insn));
         }
 
