From 76808b1ee45db247ad2aad9cc950a3a3a6e888eb Mon Sep 17 00:00:00 2001
From: Takashi Kokubun <takashikkbn@gmail.com>
Date: Fri, 10 Mar 2023 13:19:05 -0800
Subject: [PATCH] RJIT: Start testing Assembler

---
 lib/ruby_vm/rjit/code_block.rb   | 15 ++++----
 rjit_c.rb                        | 10 +++++
 test/ruby/rjit/test_assembler.rb | 66 ++++++++++++++++++++++++++++++++
 3 files changed, 84 insertions(+), 7 deletions(-)
 create mode 100644 test/ruby/rjit/test_assembler.rb

diff --git a/lib/ruby_vm/rjit/code_block.rb b/lib/ruby_vm/rjit/code_block.rb
index 33b96334b6..196e42d6a8 100644
--- a/lib/ruby_vm/rjit/code_block.rb
+++ b/lib/ruby_vm/rjit/code_block.rb
@@ -58,19 +58,20 @@ module RubyVM::RJIT
       (@mem_block...(@mem_block + @mem_size)).include?(addr)
     end
 
-    private
-
-    def dump_disasm(from, to)
+    def dump_disasm(from, to, io: STDOUT, color: true)
       C.dump_disasm(from, to).each do |address, mnemonic, op_str|
         @comments.fetch(address, []).each do |comment|
-          puts colorize("  # #{comment}", bold: true)
+          io.puts colorize("  # #{comment}", bold: true, color:)
         end
-        puts colorize("  0x#{format("%x", address)}: #{mnemonic} #{op_str}")
+        io.puts colorize("  0x#{format("%x", address)}: #{mnemonic} #{op_str}", color:)
       end
-      puts
+      io.puts
     end
 
-    def colorize(text, bold: false)
+    private
+
+    def colorize(text, bold: false, color:)
+      return text unless color
       buf = +''
       buf << "\e[1m" if bold
       buf << "\e[34m" if @outlined
diff --git a/rjit_c.rb b/rjit_c.rb
index 6f5017a3ea..7b69d043bd 100644
--- a/rjit_c.rb
+++ b/rjit_c.rb
@@ -474,6 +474,16 @@ module RubyVM::RJIT # :nodoc: all
       Primitive.cexpr! '(VALUE)NUM2PTR(value)'
     end
 
+    def HAVE_LIBCAPSTONE
+      Primitive.cstmt! %{
+        #ifdef HAVE_LIBCAPSTONE
+          return Qtrue;
+        #else
+          return Qfalse;
+        #endif
+      }
+    end
+
     #
     # Utilities: Not used by RJIT, but useful for debugging
     #
diff --git a/test/ruby/rjit/test_assembler.rb b/test/ruby/rjit/test_assembler.rb
new file mode 100644
index 0000000000..12d73b2410
--- /dev/null
+++ b/test/ruby/rjit/test_assembler.rb
@@ -0,0 +1,66 @@
+require 'test/unit'
+require_relative '../../lib/jit_support'
+
+return unless JITSupport.rjit_supported?
+return unless RubyVM::RJIT.enabled?
+return unless RubyVM::RJIT::C.HAVE_LIBCAPSTONE
+
+require 'stringio'
+require 'ruby_vm/rjit/assembler'
+
+module RubyVM::RJIT
+  class TestAssembler < Test::Unit::TestCase
+    MEM_SIZE = 16 * 1024
+
+    def setup
+      @mem_block ||= C.mmap(MEM_SIZE)
+      @cb = CodeBlock.new(mem_block: @mem_block, mem_size: MEM_SIZE)
+    end
+
+    def test_add
+      asm = Assembler.new
+      asm.add(:rax, 255)
+      assert_compile(asm, '0x0: add rax, 0xff')
+    end
+
+    def test_jmp
+      asm = Assembler.new
+      label = asm.new_label('label')
+      asm.jmp(label)
+      asm.write_label(label)
+      asm.jmp(label)
+      assert_compile(asm, <<~EOS)
+        0x0: jmp 0x2
+        0x2: jmp 0x2
+      EOS
+    end
+
+    private
+
+    def assert_compile(asm, expected)
+      actual = compile(asm)
+      assert_equal expected, actual, "---\n#{actual}---"
+    end
+
+    def compile(asm)
+      start_addr = @cb.write_addr
+      @cb.write(asm)
+      end_addr = @cb.write_addr
+
+      io = StringIO.new
+      @cb.dump_disasm(start_addr, end_addr, io:, color: false)
+      io.seek(0)
+      disasm = io.read
+
+      disasm.gsub!(/^  /, '')
+      disasm.sub!(/\n\z/, '')
+      if disasm.lines.size == 1
+        disasm.rstrip!
+      end
+      (start_addr...end_addr).each do |addr|
+        disasm.gsub!("0x#{addr.to_s(16)}", "0x#{(addr - start_addr).to_s(16)}")
+      end
+      disasm
+    end
+  end
+end