RJIT: Use imemo_type_p instead

which seems safer. It seems like imemo_type can cause SEGV.

3 files changed, 6 insertions, 5 deletions

diff --git a/lib/ruby_vm/rjit/compiler.rb b/lib/ruby_vm/rjit/compiler.rb
index 9e0454edaa..697e7f62d6 100644
--- a/lib/ruby_vm/rjit/compiler.rb
+++ b/lib/ruby_vm/rjit/compiler.rb
@@ -139,7 +139,7 @@ module RubyVM::RJIT
     def invalidate_block(block)
       iseq = block.iseq
       # Avoid touching GCed ISEQs. We assume it won't be re-entered.
-      return if C.imemo_type(iseq) != C.imemo_iseq
+      return unless C.imemo_type_p(iseq, C.imemo_iseq)
 
       # Remove this block from the version array
       remove_block(iseq, block)
@@ -296,7 +296,8 @@ module RubyVM::RJIT
 
     def rjit_blocks(iseq)
       # Guard against ISEQ GC at random moments
-      if C.imemo_type(iseq) != C.imemo_iseq
+
+      unless C.imemo_type_p(iseq, C.imemo_iseq)
         return Hash.new { |h, k| h[k] = {} }
       end
 
diff --git a/lib/ruby_vm/rjit/exit_compiler.rb b/lib/ruby_vm/rjit/exit_compiler.rb
index 3e1d35d239..35ebf2ddc4 100644
--- a/lib/ruby_vm/rjit/exit_compiler.rb
+++ b/lib/ruby_vm/rjit/exit_compiler.rb
@@ -83,7 +83,7 @@ module RubyVM::RJIT
     def compile_branch_stub(ctx, asm, branch_stub, target0_p)
       # Call rb_rjit_branch_stub_hit
       iseq = branch_stub.iseq
-      if C.rjit_opts.dump_disasm && C.imemo_type(iseq) == C.imemo_iseq # Guard against ISEQ GC at random moments
+      if C.rjit_opts.dump_disasm && C.imemo_type_p(iseq, C.imemo_iseq) # Guard against ISEQ GC at random moments
         asm.comment("branch stub hit: #{iseq.body.location.label}@#{C.rb_iseq_path(iseq)}:#{iseq_lineno(iseq, target0_p ? branch_stub.target0.pc : branch_stub.target1.pc)}")
       end
       asm.mov(:rdi, to_value(branch_stub))
diff --git a/rjit_c.rb b/rjit_c.rb
index 4eee6bb1dd..746fc03106 100644
--- a/rjit_c.rb
+++ b/rjit_c.rb
@@ -259,9 +259,9 @@ module RubyVM::RJIT # :nodoc: all
       }
     end
 
-    def imemo_type(ptr)
+    def imemo_type_p(ptr, type)
       _ptr = ptr.to_i
-      Primitive.cexpr! 'UINT2NUM(imemo_type((VALUE)NUM2SIZET(_ptr)))'
+      Primitive.cexpr! 'RBOOL(imemo_type_p((VALUE)NUM2SIZET(_ptr), NUM2UINT(type)))'
     end
 
     def rb_iseq_only_optparam_p(iseq)