Merge pull request ruby#198 from Shopify/noah_asm_comments

noahgibbs · web-flow · commit cd6aa4b6717b · 2022-03-01T23:23:39.000Z
Assembly comments on CodeBlocks
diff --git a/yjit/src/asm/mod.rs b/yjit/src/asm/mod.rs
@@ -87,6 +87,9 @@ pub struct CodeBlock
     // References to labels
     label_refs: Vec<LabelRef>,
 
+    // Comments for assembly instructions, if that feature is enabled
+    asm_comments: Vec<(usize, String)>,
+
     // Keep track of the current aligned write position.
     // Used for changing protection when writing to the JIT buffer
     current_aligned_write_pos: usize,
@@ -116,6 +119,7 @@ impl CodeBlock
             label_addrs: Vec::new(),
             label_names: Vec::new(),
             label_refs: Vec::new(),
+            asm_comments: Vec::new(),
             current_aligned_write_pos: ALIGNED_WRITE_POSITION_NONE,
             page_size: 4096,
             dropped_bytes: false
@@ -131,6 +135,7 @@ impl CodeBlock
             label_addrs: Vec::new(),
             label_names: Vec::new(),
             label_refs: Vec::new(),
+            asm_comments: Vec::new(),
             current_aligned_write_pos: ALIGNED_WRITE_POSITION_NONE,
             page_size,
             dropped_bytes: false
@@ -142,6 +147,34 @@ impl CodeBlock
         self.write_pos + num_bytes < self.mem_size
     }
 
+    /// Add an assembly comment if the feature is on.
+    /// If not, this becomes an inline no-op.
+    #[inline]
+    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)));
+            }
+        }
+    }
+
+    /// 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)));
+    }
+
+    /// 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 get_write_pos(&self) -> usize {
         self.write_pos
     }
diff --git a/yjit/src/asm/x86_64/tests.rs b/yjit/src/asm/x86_64/tests.rs
@@ -420,3 +420,20 @@ fn basic_capstone_usage() -> std::result::Result<(), capstone::Error> {
         )),
     }
 }
+
+#[test]
+#[cfg(feature = "asm_comments")]
+fn block_comments() {
+    let mut cb = super::CodeBlock::new_dummy(4096);
+
+    cb.add_comment("Beginning");
+    xor(&mut cb, EAX, EAX); // 2 bytes long
+    cb.add_comment("Two bytes in");
+    cb.add_comment("Two bytes in"); // Duplicate, should be ignored
+    test(&mut cb, mem_opnd(64, RSI, 64), imm_opnd(!0x08)); // 8 bytes long
+    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);
+}
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
@@ -239,23 +239,7 @@ fn add_comment(cb: &mut CodeBlock, comment_str: &str)
 {
     #[cfg(feature = "asm_comments")]
     {
-        /*
-        // We can't add comments to the outlined code block
-        if (cb == ocb)
-            return;
-
-        // Avoid adding duplicate comment strings (can happen due to deferred codegen)
-        size_t num_comments = rb_darray_size(yjit_code_comments);
-        if (num_comments > 0) {
-            struct yjit_comment last_comment = rb_darray_get(yjit_code_comments, num_comments - 1);
-            if (last_comment.offset == cb->write_pos && strcmp(last_comment.comment, comment_str) == 0) {
-                return;
-            }
-        }
-
-        struct yjit_comment new_comment = (struct yjit_comment){ cb->write_pos, comment_str };
-        rb_darray_append(&yjit_code_comments, new_comment);
-        */
+        cb.add_comment(comment_str);
     }
 }
 

Original file line number	Diff line number	Diff line change
`@@ -239,23 +239,7 @@ fn add_comment(cb: &mut CodeBlock, comment_str: &str)`
`239`	`239`	`{`
`240`	`240`	`#[cfg(feature = "asm_comments")]`
`241`	`241`	`{`
`242`		`- /*`
`243`		`- // We can't add comments to the outlined code block`
`244`		`- if (cb == ocb)`
`245`		`- return;`
`246`		`-`
`247`		`- // Avoid adding duplicate comment strings (can happen due to deferred codegen)`
`248`		`- size_t num_comments = rb_darray_size(yjit_code_comments);`
`249`		`- if (num_comments > 0) {`
`250`		`- struct yjit_comment last_comment = rb_darray_get(yjit_code_comments, num_comments - 1);`
`251`		`- if (last_comment.offset == cb->write_pos && strcmp(last_comment.comment, comment_str) == 0) {`
`252`		`- return;`
`253`		`- }`
`254`		`- }`
`255`		`-`
`256`		`- struct yjit_comment new_comment = (struct yjit_comment){ cb->write_pos, comment_str };`
`257`		`- rb_darray_append(&yjit_code_comments, new_comment);`
`258`		`- */`
	`242`	`+ cb.add_comment(comment_str);`
`259`	`243`	`}`
`260`	`244`	`}`
`261`	`245`