YJIT: add comments to disassembly

Introduce a new macro `ADD_COMMENT(cb, comment)` that records a comment
for the current write position in the code block.

Co-authored-by: Maxime Chevalier-Boisvert <maximechevalierb@gmail.com>
Co-authored-by: Aaron Patterson <aaron.patterson@shopify.com>

4 files changed, 104 insertions, 28 deletions

diff --git a/yjit.rb b/yjit.rb
index b513e1b0dc..e1d9c16417 100644
--- a/yjit.rb
+++ b/yjit.rb
@@ -1,40 +1,61 @@
 module YJIT
-  def self.disasm(iseq)
-    iseq = RubyVM::InstructionSequence.of(iseq)
+  if defined?(Disasm)
+    def self.disasm(iseq, tty: $stdout && $stdout.tty?)
+      iseq = RubyVM::InstructionSequence.of(iseq)
 
-    blocks = YJIT.blocks_for(iseq)
-    return if blocks.empty?
+      blocks = YJIT.blocks_for(iseq)
+      return if blocks.empty?
 
-    str = ""
-
-    cs = YJIT::Disasm.new
+      str = ""
+      str << iseq.disasm
+      str << "\n"
 
-    str << iseq.disasm
-    str << "\n"
+      # Sort the blocks by increasing addresses
+      sorted_blocks = blocks.sort_by(&:address)
+
+      highlight = ->(str) {
+        if tty
+          "\x1b[1m#{str}\x1b[0m"
+        else
+          str
+        end
+      }
+
+      cs = YJIT::Disasm.new
+      sorted_blocks.each_with_index do |block, i|
+        str << "== BLOCK #{i+1}/#{blocks.length}: #{block.code.length} BYTES, ISEQ RANGE [#{block.iseq_start_index},#{block.iseq_end_index}) ".ljust(80, "=")
+        str << "\n"
+
+        comments = comments_for(block.address, block.address + block.code.length)
+        comment_idx = 0
+        cs.disasm(block.code, block.address).each do |i|
+          while (comment = comments[comment_idx]) && comment.address <= i.address
+            str << "  ;#{highlight.call(comment.comment)}\n"
+            comment_idx += 1
+          end
+
+          str << sprintf(
+            "  %<address>08x:  %<instruction>s\t%<details>s\n",
+            address: i.address,
+            instruction: i.mnemonic,
+            details: i.op_str
+          )
+        end
+      end
 
-    # Sort the blocks by increasing addresses
-    blocks.sort_by(&:address).each_with_index do |block, i|
-      str << "== BLOCK #{i+1}/#{blocks.length}: #{block.code.length} BYTES, ISEQ RANGE [#{block.iseq_start_index},#{block.iseq_end_index}) ".ljust(80, "=")
+      block_sizes = blocks.map { |block| block.code.length }
+      total_bytes = block_sizes.sum
+      str << "\n"
+      str << "Total code size: #{total_bytes} bytes"
       str << "\n"
 
-      cs.disasm(block.code, block.address).each do |i|
-        str << sprintf(
-          "  %<address>08x:  %<instruction>s\t%<details>s\n",
-          address: i.address,
-          instruction: i.mnemonic,
-          details: i.op_str
-        )
-      end
+      str
     end
 
-    block_sizes = blocks.map { |block| block.code.length }
-    total_bytes = block_sizes.reduce(0, :+)
-    str << "\n"
-    str << "Total code size: #{total_bytes} bytes"
-    str << "\n"
-
-    str
-  end if defined?(Disasm)
+    def self.comments_for(start_address, end_address)
+      Primitive.comments_for(start_address, end_address)
+    end
+  end
 
   # Return a hash for statistics generated for the --yjit-stats command line option.
   # Return nil when option is not passed or unavailable.
diff --git a/yjit_codegen.c b/yjit_codegen.c
index fd029b3287..7ff59b94a1 100644
--- a/yjit_codegen.c
+++ b/yjit_codegen.c
@@ -216,9 +216,16 @@ _counted_side_exit(uint8_t *existing_side_exit, int64_t *counter)
     return start;
 }
 
+// Comments for generated machine code
+#define ADD_COMMENT(cb, comment) rb_darray_append(&yjit_code_comments, ((struct yjit_comment){(cb)->write_pos, (comment)}))
+yjit_comment_array_t yjit_code_comments;
+
 #else
+
 #define GEN_COUNTER_INC(cb, counter_name) ((void)0)
 #define COUNTED_EXIT(side_exit, counter_name) side_exit
+#define ADD_COMMENT(cb, comment) ((void)0)
+
 #endif // if RUBY_DEBUG
 
 /*
@@ -323,6 +330,9 @@ yjit_gen_block(ctx_t* ctx, block_t* block, rb_execution_context_t* ec)
         // Note that the increment happens even when the output takes side exit.
         GEN_COUNTER_INC(cb, exec_instruction);
 
+        // Add a comment for the name of the YARV instruction
+        ADD_COMMENT(cb, insn_name(opcode));
+
         // Call the code generation function
         bool continue_generating = p_desc->gen_fn(&jit, ctx);
 
@@ -751,6 +761,7 @@ gen_getinstancevariable(jitstate_t* jit, ctx_t* ctx)
 
             // Guard that self is embedded
             // TODO: BT and JC is shorter
+            ADD_COMMENT(cb, "guard embedded getivar");
             x86opnd_t flags_opnd = member_opnd(REG0, struct RBasic, flags);
             test(cb, flags_opnd, imm_opnd(ROBJECT_EMBED));
             jit_chain_guard(JCC_JZ, jit, ctx, GETIVAR_MAX_DEPTH, side_exit);
@@ -772,6 +783,7 @@ gen_getinstancevariable(jitstate_t* jit, ctx_t* ctx)
 
             // Guard that self is *not* embedded
             // See ROBJECT_IVPTR() from include/ruby/internal/core/robject.h
+            ADD_COMMENT(cb, "guard extended getivar");
             x86opnd_t flags_opnd = member_opnd(REG0, struct RBasic, flags);
             test(cb, flags_opnd, imm_opnd(ROBJECT_EMBED));
             jit_chain_guard(JCC_JNZ, jit, ctx, GETIVAR_MAX_DEPTH, side_exit);
diff --git a/yjit_iface.c b/yjit_iface.c
index 9e708b5bb8..3d1a617757 100644
--- a/yjit_iface.c
+++ b/yjit_iface.c
@@ -30,6 +30,7 @@ static int64_t vm_insns_count = 0;
 static int64_t exit_op_count[VM_INSTRUCTION_SIZE] = { 0 };
 int64_t rb_compiled_iseq_count = 0;
 struct rb_yjit_runtime_counters yjit_runtime_counters = { 0 };
+static VALUE cYjitCodeComment;
 #endif
 
 // Machine code blocks (executable memory)
@@ -521,6 +522,7 @@ yjit_blocks_for(VALUE mod, VALUE rb_iseq)
         rb_darray_for(versions, block_idx) {
             block_t *block = rb_darray_get(versions, block_idx);
 
+            // FIXME: The object craeted here can outlive the block itself
             VALUE rb_block = TypedData_Wrap_Struct(cYjitBlock, &yjit_block_type, block);
             rb_ary_push(all_versions, rb_block);
         }
@@ -655,6 +657,37 @@ at_exit_print_stats(RB_BLOCK_CALL_FUNC_ARGLIST(yieldarg, data))
     return Qnil;
 }
 
+// Primitive called in yjit.rb. Export all machine code comments as a Ruby array.
+static VALUE
+comments_for(rb_execution_context_t *ec, VALUE self, VALUE start_address, VALUE end_address)
+{
+    VALUE comment_array = rb_ary_new();
+#if RUBY_DEBUG
+    uint8_t *start = (void *)NUM2ULL(start_address);
+    uint8_t *end = (void *)NUM2ULL(end_address);
+
+    rb_darray_for(yjit_code_comments, i) {
+        struct yjit_comment comment = rb_darray_get(yjit_code_comments, i);
+        uint8_t *comment_pos = cb_get_ptr(cb, comment.offset);
+
+        if (comment_pos >= end) {
+            break;
+        }
+        if (comment_pos >= start) {
+            VALUE vals = rb_ary_new_from_args(
+                2,
+                LL2NUM((long long) comment_pos),
+                rb_str_new_cstr(comment.comment)
+            );
+            rb_ary_push(comment_array, rb_struct_alloc(cYjitCodeComment, vals));
+        }
+    }
+
+#endif // if RUBY_DEBUG
+
+    return comment_array;
+}
+
 // Primitive called in yjit.rb. Export all runtime counters as a Ruby hash.
 static VALUE
 get_stat_counters(rb_execution_context_t *ec, VALUE self)
@@ -962,6 +995,7 @@ rb_yjit_init(struct rb_yjit_options *options)
     rb_define_alloc_func(cYjitDisasm, yjit_disasm_init);
     rb_define_method(cYjitDisasm, "disasm", yjit_disasm, 2);
     cYjitDisasmInsn = rb_struct_define_under(cYjitDisasm, "Insn", "address", "mnemonic", "op_str", NULL);
+    cYjitCodeComment = rb_struct_define_under(cYjitDisasm, "Comment", "address", "comment");
 #endif
 
     if (RUBY_DEBUG && rb_yjit_opts.gen_stats) {
diff --git a/yjit_iface.h b/yjit_iface.h
index 6428906b28..63871d24c2 100644
--- a/yjit_iface.h
+++ b/yjit_iface.h
@@ -67,6 +67,15 @@ YJIT_DECLARE_COUNTERS(
 
 #undef YJIT_DECLARE_COUNTERS
 
+struct yjit_comment {
+    int32_t offset;
+    const char *comment;
+};
+
+typedef rb_darray(struct yjit_comment) yjit_comment_array_t;
+
+extern yjit_comment_array_t yjit_code_comments;
+
 #endif // if RUBY_DEBUG
 
 RUBY_EXTERN struct rb_yjit_options rb_yjit_opts;