Implement --mjit-dump-disasm

6 files changed, 83 insertions, 5 deletions

diff --git a/configure.ac b/configure.ac
index 70946f83d8..64f0a48435 100644
--- a/configure.ac
+++ b/configure.ac
@@ -3855,6 +3855,12 @@ AS_CASE(["${YJIT_SUPPORT}"],
     AC_DEFINE(USE_YJIT, 1)
 ], [AC_DEFINE(USE_YJIT, 0)])
 
+# If YJIT links capstone, libcapstone stops working on the C side.
+# capstone should be linked for MJIT only when YJIT doesn't.
+AS_IF([test x"$MJIT_SUPPORT" = "xyes" -a -z "$CARGO_BUILD_ARGS" ], [
+    AC_CHECK_LIB([capstone], [cs_disasm])
+])
+
 dnl These variables end up in ::RbConfig::CONFIG
 AC_SUBST(YJIT_SUPPORT)dnl what flavor of YJIT the Ruby build includes
 AC_SUBST(RUSTC)dnl Rust compiler command
diff --git a/lib/mjit/x86_64/assembler.rb b/lib/mjit/x86_64/assembler.rb
index 14f414b33b..5458746c9b 100644
--- a/lib/mjit/x86_64/assembler.rb
+++ b/lib/mjit/x86_64/assembler.rb
@@ -6,12 +6,14 @@ class RubyVM::MJIT::Assembler
   end
 
   def compile(compiler)
-    RubyVM::MJIT::C.mjit_mark_writable
-    write_bytes(compiler.write_addr, @bytes)
-    RubyVM::MJIT::C.mjit_mark_executable
+    with_dump_disasm(compiler) do
+      RubyVM::MJIT::C.mjit_mark_writable
+      write_bytes(compiler.write_addr, @bytes)
+      RubyVM::MJIT::C.mjit_mark_executable
 
-    compiler.write_pos += @bytes.size
-    @bytes.clear
+      compiler.write_pos += @bytes.size
+      @bytes.clear
+    end
   end
 
   def mov(_reg, val)
@@ -24,6 +26,17 @@ class RubyVM::MJIT::Assembler
 
   private
 
+  def with_dump_disasm(compiler)
+    from = compiler.write_addr
+    yield
+    to = compiler.write_addr
+    if RubyVM::MJIT::C.mjit_opts.dump_disasm && from < to
+      RubyVM::MJIT::C.dump_disasm(from, to).each do |address, mnemonic, op_str|
+        puts "  0x#{"%p" % address}: #{mnemonic} #{op_str}"
+      end
+    end
+  end
+
   def write_bytes(addr, bytes)
     writer = ByteWriter.new(addr)
     # If you pack bytes containing \x00, Ruby fails to recognize bytes after \x00.
diff --git a/mjit.c b/mjit.c
index 8226404084..b1b7edab62 100644
--- a/mjit.c
+++ b/mjit.c
@@ -304,6 +304,9 @@ mjit_setup_options(const char *s, struct mjit_options *mjit_opt)
     else if (opt_match_noarg(s, l, "pause")) {
         mjit_opt->pause = true;
     }
+    else if (opt_match_noarg(s, l, "dump-disasm")) {
+        mjit_opt->dump_disasm = true;
+    }
     else {
         rb_raise(rb_eRuntimeError,
                  "invalid MJIT option `%s' (--help will show valid MJIT options)", s);
@@ -412,6 +415,10 @@ mjit_init(const struct mjit_options *opts)
     // Normalize options
     if (mjit_opts.call_threshold == 0)
         mjit_opts.call_threshold = DEFAULT_CALL_THRESHOLD;
+#ifndef HAVE_LIBCAPSTONE
+    if (mjit_opts.dump_disasm)
+        verbose(1, "libcapstone has not been linked. Ignoring --mjit-dump-disasm.");
+#endif
 }
 
 void
diff --git a/mjit.h b/mjit.h
index cea1677e9b..824d91d076 100644
--- a/mjit.h
+++ b/mjit.h
@@ -63,6 +63,8 @@ struct mjit_options {
     bool pause;
     // [experimental] Call custom RubyVM::MJIT.compile instead of MJIT.
     bool custom;
+    // Enable disasm of all JIT code
+    bool dump_disasm;
 };
 
 // State of optimization switches
diff --git a/mjit_c.c b/mjit_c.c
index 9ba023e84b..0af2e4b239 100644
--- a/mjit_c.c
+++ b/mjit_c.c
@@ -38,6 +38,49 @@
 #define SIZEOF(type) RB_SIZE2NUM(sizeof(type))
 #define SIGNED_TYPE_P(type) RBOOL((type)(-1) < (type)(1))
 
+// macOS: brew install capstone
+// Ubuntu/Debian: apt-get install libcapstone-dev
+// Fedora: dnf -y install capstone-devel
+//#ifdef HAVE_LIBCAPSTONE
+#if 1
+#include <capstone/capstone.h>
+
+#define CODE "\x55\x48\x8b\x05\xb8\x13\x00\x00"
+
+// Return an array of [address, mnemonic, op_str]
+static VALUE
+dump_disasm(rb_execution_context_t *ec, VALUE self, VALUE from, VALUE to)
+{
+    // Prepare for calling cs_disasm
+    static csh handle;
+    if (cs_open(CS_ARCH_X86, CS_MODE_64, &handle) != CS_ERR_OK) {
+        rb_raise(rb_eRuntimeError, "failed to make Capstone handle");
+    }
+    size_t from_addr = NUM2SIZET(from);
+    size_t to_addr = NUM2SIZET(to);
+
+    // Call cs_disasm and convert results to a Ruby array
+    cs_insn *insns;
+    size_t count = cs_disasm(handle, (const uint8_t *)from_addr, to_addr - from_addr, from_addr, 0, &insns);
+    VALUE result = rb_ary_new_capa(count);
+    for (size_t i = 0; i < count; i++) {
+        VALUE vals = rb_ary_new_from_args(3, LONG2NUM(insns[i].address), rb_str_new2(insns[i].mnemonic), rb_str_new2(insns[i].op_str));
+        rb_ary_push(result, vals);
+    }
+
+    // Free memory used by capstone
+    cs_free(insns, count);
+    cs_close(&handle);
+    return result;
+}
+#else
+static VALUE
+mjit_disasm(VALUE self, VALUE from, VALUE to)
+{
+    return rb_ary_new();
+}
+#endif
+
 #include "mjit_c.rbinc"
 
 #endif // USE_MJIT
diff --git a/mjit_c.rb b/mjit_c.rb
index 7e352c61d7..506e15f76e 100644
--- a/mjit_c.rb
+++ b/mjit_c.rb
@@ -25,6 +25,12 @@ module RubyVM::MJIT # :nodoc: all
       }
     end
 
+    # @param from [Integer] - From address
+    # @param to [Integer]   - To address
+    def dump_disasm(from, to)
+      Primitive.dump_disasm(from, to)
+    end
+
     #========================================================================================
     #
     # Old stuff
@@ -372,6 +378,7 @@ module RubyVM::MJIT # :nodoc: all
       max_cache_size: [CType::Immediate.parse("int"), Primitive.cexpr!("OFFSETOF((*((struct mjit_options *)NULL)), max_cache_size)")],
       pause: [self._Bool, Primitive.cexpr!("OFFSETOF((*((struct mjit_options *)NULL)), pause)")],
       custom: [self._Bool, Primitive.cexpr!("OFFSETOF((*((struct mjit_options *)NULL)), custom)")],
+      dump_disasm: [self._Bool, Primitive.cexpr!("OFFSETOF((*((struct mjit_options *)NULL)), dump_disasm)")],
     )
   end