RJIT: Upgrade type to Fixnum after guard

2 files changed, 60 insertions, 8 deletions

diff --git a/lib/ruby_vm/rjit/insn_compiler.rb b/lib/ruby_vm/rjit/insn_compiler.rb
index 2c8086345d..f11801c007 100644
--- a/lib/ruby_vm/rjit/insn_compiler.rb
+++ b/lib/ruby_vm/rjit/insn_compiler.rb
@@ -3667,15 +3667,48 @@ module RubyVM::RJIT
       arg1 = ctx.stack_opnd(0)
       arg0 = ctx.stack_opnd(1)
 
-      asm.comment('guard arg0 fixnum')
-      asm.test(arg0, C::RUBY_FIXNUM_FLAG)
-      jit_chain_guard(:jz, jit, ctx, asm, side_exit)
-      # TODO: upgrade type, and skip the check when possible
+      # Get the stack operand types
+      arg1_type = ctx.get_opnd_type(StackOpnd[0])
+      arg0_type = ctx.get_opnd_type(StackOpnd[1])
 
-      asm.comment('guard arg1 fixnum')
-      asm.test(arg1, C::RUBY_FIXNUM_FLAG)
-      jit_chain_guard(:jz, jit, ctx, asm, side_exit)
-      # TODO: upgrade type, and skip the check when possible
+      if arg0_type.heap? || arg1_type.heap?
+        asm.comment('arg is heap object')
+        asm.jmp(side_exit)
+        return
+      end
+
+      if arg0_type != Type::Fixnum && arg0_type.specific?
+        asm.comment('arg0 not fixnum')
+        asm.jmp(side_exit)
+        return
+      end
+
+      if arg1_type != Type::Fixnum && arg1_type.specific?
+        asm.comment('arg1 not fixnum')
+        asm.jmp(side_exit)
+        return
+      end
+
+      assert(!arg0_type.heap?)
+      assert(!arg1_type.heap?)
+      assert(arg0_type == Type::Fixnum || arg0_type.unknown?)
+      assert(arg1_type == Type::Fixnum || arg1_type.unknown?)
+
+      # If not fixnums at run-time, fall back
+      if arg0_type != Type::Fixnum
+        asm.comment('guard arg0 fixnum')
+        asm.test(arg0, C::RUBY_FIXNUM_FLAG)
+        jit_chain_guard(:jz, jit, ctx, asm, side_exit)
+      end
+      if arg1_type != Type::Fixnum
+        asm.comment('guard arg1 fixnum')
+        asm.test(arg1, C::RUBY_FIXNUM_FLAG)
+        jit_chain_guard(:jz, jit, ctx, asm, side_exit)
+      end
+
+      # Set stack types in context
+      ctx.upgrade_opnd_type(StackOpnd[0], Type::Fixnum)
+      ctx.upgrade_opnd_type(StackOpnd[1], Type::Fixnum)
     end
 
     # @param jit [RubyVM::RJIT::JITState]
@@ -5642,6 +5675,10 @@ module RubyVM::RJIT
       asm.cmovnz(ary_opnd, array_reg)
     end
 
+    def assert(cond)
+      assert_equal(cond, true)
+    end
+
     def assert_equal(left, right)
       if left != right
         raise "'#{left.inspect}' was not '#{right.inspect}'"
diff --git a/lib/ruby_vm/rjit/type.rb b/lib/ruby_vm/rjit/type.rb
index b17d94cc8c..5ae70c5779 100644
--- a/lib/ruby_vm/rjit/type.rb
+++ b/lib/ruby_vm/rjit/type.rb
@@ -15,6 +15,21 @@ module RubyVM::RJIT
       end
     end
 
+    # Returns true when the type is not specific.
+    def unknown?
+      case self
+      in Type::Unknown | Type::UnknownImm | Type::UnknownHeap then true
+      else false
+      end
+    end
+
+    # Returns true when we know the VALUE is a specific handle type,
+    # such as a static symbol ([Type::ImmSymbol], i.e. true from RB_STATIC_SYM_P()).
+    # Opposite of [Self::is_unknown].
+    def specific?
+      !self.unknown?
+    end
+
     # Check if the type is a heap object
     def heap?
       case self