RJIT: Store type information in Context

6 files changed, 620 insertions, 134 deletions

diff --git a/lib/ruby_vm/rjit/compiler.rb b/lib/ruby_vm/rjit/compiler.rb
index 39b3c77f3e..64a212adba 100644
--- a/lib/ruby_vm/rjit/compiler.rb
+++ b/lib/ruby_vm/rjit/compiler.rb
@@ -9,6 +9,7 @@ require 'ruby_vm/rjit/insn_compiler'
 require 'ruby_vm/rjit/instruction'
 require 'ruby_vm/rjit/invariants'
 require 'ruby_vm/rjit/jit_state'
+require 'ruby_vm/rjit/type'
 
 module RubyVM::RJIT
   # Compilation status
diff --git a/lib/ruby_vm/rjit/context.rb b/lib/ruby_vm/rjit/context.rb
index 9c8fd183e6..69d7976c38 100644
--- a/lib/ruby_vm/rjit/context.rb
+++ b/lib/ruby_vm/rjit/context.rb
@@ -1,22 +1,46 @@
 module RubyVM::RJIT
+  # Maximum number of temp value types we keep track of
+  MAX_TEMP_TYPES = 8
+  # Maximum number of local variable types we keep track of
+  MAX_LOCAL_TYPES = 8
+
+  # Operand to a YARV bytecode instruction
+  SelfOpnd = :SelfOpnd # The value is self
+  StackOpnd = Data.define(:index) # Temporary stack operand with stack index
+
+  # Potential mapping of a value on the temporary stack to self,
+  # a local variable, or constant so that we can track its type
+  MapToStack = :MapToStack # Normal stack value
+  MapToSelf  = :MapToSelf  # Temp maps to the self operand
+  MapToLocal = Data.define(:local_index) # Temp maps to a local variable with index
+
   class Context < Struct.new(
-    :stack_size,  # @param [Integer] The number of values on the stack
-    :sp_offset,   # @param [Integer] JIT sp offset relative to the interpreter's sp
-    :chain_depth, # @param [Integer] jit_chain_guard depth
+    :stack_size,   # @param [Integer] The number of values on the stack
+    :sp_offset,    # @param [Integer] JIT sp offset relative to the interpreter's sp
+    :chain_depth,  # @param [Integer] jit_chain_guard depth
+    :local_types,  # @param [Array<RubyVM::RJIT::Type>] Local variable types we keep track of
+    :temp_types,   # @param [Array<RubyVM::RJIT::Type>] Temporary variable types we keep track of
+    :self_type,    # @param [RubyVM::RJIT::Type] Type we track for self
+    :temp_mapping, # @param [Array<Symbol>] Mapping of temp stack entries to types we track
   )
-    def initialize(stack_size: 0, sp_offset: 0, chain_depth: 0) = super
-
-    def stack_push(size = 1)
-      self.stack_size += size
-      self.sp_offset += size
-      stack_opnd(0)
-    end
+    def initialize(
+      stack_size:   0,
+      sp_offset:    0,
+      chain_depth:  0,
+      local_types:  [Type::Unknown] * MAX_LOCAL_TYPES,
+      temp_types:   [Type::Unknown] * MAX_TEMP_TYPES,
+      self_type:    Type::Unknown,
+      temp_mapping: [MapToStack] * MAX_TEMP_TYPES
+    ) = super
 
-    def stack_pop(size = 1)
-      opnd = stack_opnd(0)
-      self.stack_size -= size
-      self.sp_offset -= size
-      opnd
+    # Create a new Context instance with a given stack_size and sp_offset adjusted
+    # accordingly. This is useful when you want to virtually rewind a stack_size for
+    # generating a side exit while considering past sp_offset changes on gen_save_sp.
+    def with_stack_size(stack_size)
+      ctx = self.dup
+      ctx.sp_offset -= ctx.stack_size - stack_size
+      ctx.stack_size = stack_size
+      ctx
     end
 
     def stack_opnd(depth_from_top)
@@ -27,14 +51,250 @@ module RubyVM::RJIT
       [SP, (C.VALUE.size * self.sp_offset) + offset_bytes]
     end
 
-    # Create a new Context instance with a given stack_size and sp_offset adjusted
-    # accordingly. This is useful when you want to virtually rewind a stack_size for
-    # generating a side exit while considering past sp_offset changes on gen_save_sp.
-    def with_stack_size(stack_size)
-      ctx = self.dup
-      ctx.sp_offset -= ctx.stack_size - stack_size
-      ctx.stack_size = stack_size
-      ctx
+    # Push one new value on the temp stack with an explicit mapping
+    # Return a pointer to the new stack top
+    def stack_push_mapping(mapping_temp_type)
+      stack_size = self.stack_size
+
+      # Keep track of the type and mapping of the value
+      if stack_size < MAX_TEMP_TYPES
+        mapping, temp_type = mapping_temp_type
+        self.temp_mapping[stack_size] = mapping
+        self.temp_types[stack_size] = temp_type
+
+        case mapping
+        in MapToLocal[idx]
+          assert(idx < MAX_LOCAL_TYPES)
+        else
+        end
+      end
+
+      self.stack_size += 1
+      self.sp_offset += 1
+
+      return self.stack_opnd(0)
+    end
+
+    # Push one new value on the temp stack
+    # Return a pointer to the new stack top
+    def stack_push(val_type)
+      return self.stack_push_mapping([MapToStack, val_type])
+    end
+
+    # Push the self value on the stack
+    def stack_push_self
+      return self.stack_push_mapping([MapToStack, Type::Unknown])
+    end
+
+    # Push a local variable on the stack
+    def stack_push_local(local_idx)
+      if local_idx >= MAX_LOCAL_TYPES
+        return self.stack_push(Type::Unknown)
+      end
+
+      return self.stack_push_mapping([MapToLocal[local_idx], Type::Unknown])
+    end
+
+    # Pop N values off the stack
+    # Return a pointer to the stack top before the pop operation
+    def stack_pop(n = 1)
+      assert(n <= self.stack_size)
+
+      top = self.stack_opnd(0)
+
+      # Clear the types of the popped values
+      n.times do |i|
+        idx = self.stack_size - i - 1
+
+        if idx < MAX_TEMP_TYPES
+          self.temp_types[idx] = Type::Unknown
+          self.temp_mapping[idx] = MapToStack
+        end
+      end
+
+      self.stack_size -= n
+      self.sp_offset -= n
+
+      return top
+    end
+
+    def shift_stack(argc)
+      assert(argc < self.stack_size)
+
+      method_name_index = self.stack_size - argc - 1
+
+      (method_name_index...(self.stack_size - 1)).each do |i|
+        if i + 1 < MAX_TEMP_TYPES
+          self.temp_types[i] = self.temp_types[i + 1]
+          self.temp_mapping[i] = self.temp_mapping[i + 1]
+        end
+      end
+      self.stack_pop(1)
+    end
+
+    # Get the type of an instruction operand
+    def get_opnd_type(opnd)
+      case opnd
+      in SelfOpnd
+        self.self_type
+      in StackOpnd[idx]
+        assert(idx < self.stack_size)
+        stack_idx = self.stack_size - 1 - idx
+
+        # If outside of tracked range, do nothing
+        if stack_idx >= MAX_TEMP_TYPES
+          return Type::Unknown
+        end
+
+        mapping = self.temp_mapping[stack_idx]
+
+        case mapping 
+        in MapToSelf
+          self.self_type
+        in MapToStack
+          self.temp_types[self.stack_size - 1 - idx]
+        in MapToLocal[idx]
+          assert(idx < MAX_LOCAL_TYPES)
+          self.local_types[idx]
+        end
+      end
+    end
+
+    # Get the currently tracked type for a local variable
+    def get_local_type(idx)
+      self.local_types[idx] || Type::Unknown
+    end
+
+    # Upgrade (or "learn") the type of an instruction operand
+    # This value must be compatible and at least as specific as the previously known type.
+    # If this value originated from self, or an lvar, the learned type will be
+    # propagated back to its source.
+    def upgrade_opnd_type(opnd, opnd_type)
+      case opnd
+      in SelfOpnd
+        self.self_type = self.self_type.upgrade(opnd_type)
+      in StackOpnd[idx]
+        assert(idx < self.stack_size)
+        stack_idx = self.stack_size - 1 - idx
+
+        # If outside of tracked range, do nothing
+        if stack_idx >= MAX_TEMP_TYPES
+          return
+        end
+
+        mapping = self.temp_mapping[stack_idx]
+
+        case mapping
+        in MapToSelf
+          self.self_type = self.self_type.upgrade(opnd_type)
+        in MapToStack
+          self.temp_types[stack_idx] = self.temp_types[stack_idx].upgrade(opnd_type)
+        in MapToLocal[idx]
+          assert(idx < MAX_LOCAL_TYPES)
+          self.local_types[idx] = self.local_types[idx].upgrade(opnd_type)
+        end
+      end
+    end
+
+    # Get both the type and mapping (where the value originates) of an operand.
+    # This is can be used with stack_push_mapping or set_opnd_mapping to copy
+    # a stack value's type while maintaining the mapping.
+    def get_opnd_mapping(opnd)
+      opnd_type = self.get_opnd_type(opnd)
+
+      case opnd
+      in SelfOpnd
+        return [MapToSelf, opnd_type]
+      in StackOpnd[idx]
+        assert(idx < self.stack_size)
+        stack_idx = self.stack_size - 1 - idx
+
+        if stack_idx < MAX_TEMP_TYPES
+          return [self.temp_mapping[stack_idx], opnd_type]
+        else
+          # We can't know the source of this stack operand, so we assume it is
+          # a stack-only temporary. type will be UNKNOWN
+          assert(opnd_type == Type::Unknown)
+          return [MapToStack, opnd_type]
+        end
+      end
+    end
+
+    # Overwrite both the type and mapping of a stack operand.
+    def set_opnd_mapping(opnd, mapping_opnd_type)
+      case opnd
+      in SelfOpnd
+        raise 'self always maps to self'
+      in StackOpnd[idx]
+        assert(idx < self.stack_size)
+        stack_idx = self.stack_size - 1 - idx
+
+        # If outside of tracked range, do nothing
+        if stack_idx >= MAX_TEMP_TYPES
+          return
+        end
+
+        mapping, opnd_type = mapping_opnd_type
+        self.temp_mapping[stack_idx] = mapping
+
+        # Only used when mapping == MAP_STACK
+        self.temp_types[stack_idx] = opnd_type
+      end
+    end
+
+    # Set the type of a local variable
+    def set_local_type(local_idx, local_type)
+      if local_idx >= MAX_LOCAL_TYPES
+        return
+      end
+
+      # If any values on the stack map to this local we must detach them
+      MAX_TEMP_TYPES.times do |stack_idx|
+        case self.temp_mapping[stack_idx]
+        in MapToStack
+          # noop
+        in MapToSelf
+          # noop
+        in MapToLocal[local_idx]
+          if stack_idx == local_idx
+            self.temp_types[stack_idx] = self.local_types[local_idx];
+            self.temp_mapping[stack_idx] = MapToStack
+          else
+            # noop
+          end
+        end
+      end
+
+      self.local_types[local_idx] = local_type
+    end
+
+    # Erase local variable type information
+    # eg: because of a call we can't track
+    def clear_local_types
+      # When clearing local types we must detach any stack mappings to those
+      # locals. Even if local values may have changed, stack values will not.
+      MAX_TEMP_TYPES.times do |stack_idx|
+        case self.temp_mapping[stack_idx]
+        in MapToStack
+          # noop
+        in MapToSelf
+          # noop
+        in MapToLocal[local_idx]
+          self.temp_types[stack_idx] = self.local_types[local_idx]
+          self.temp_mapping[stack_idx] = MapToStack
+        end
+      end
+
+      # Clear the local types
+      self.local_types = [Type::Unknown] * MAX_LOCAL_TYPES
+    end
+
+    private
+
+    def assert(cond)
+      unless cond
+        raise "'#{cond.inspect}' was not true"
+      end
     end
   end
 end
diff --git a/lib/ruby_vm/rjit/insn_compiler.rb b/lib/ruby_vm/rjit/insn_compiler.rb
index 6d8961cf20..8c99876c1a 100644
--- a/lib/ruby_vm/rjit/insn_compiler.rb
+++ b/lib/ruby_vm/rjit/insn_compiler.rb
@@ -257,7 +257,7 @@ module RubyVM::RJIT
       asm.write_label(frame_flag_modified)
 
       # Push the proc on the stack
-      stack_ret = ctx.stack_push
+      stack_ret = ctx.stack_push(Type::Unknown)
       ep_reg = :rax
       jit_get_ep(asm, level, reg: ep_reg)
       asm.mov(:rax, [ep_reg, offs])
@@ -328,7 +328,7 @@ module RubyVM::RJIT
         jit_chain_guard(:jnz, jit, starting_context, asm, counted_exit(side_exit, :getblockpp_not_iseq_block))
 
         # Push rb_block_param_proxy. It's a root, so no need to use jit_mov_gc_ptr.
-        top = ctx.stack_push
+        top = ctx.stack_push(Type::BlockParamProxy)
         asm.mov(:rax, C.rb_block_param_proxy)
         asm.mov(top, :rax)
       end
@@ -377,7 +377,7 @@ module RubyVM::RJIT
           asm.call(C.rb_reg_match_last)
         end
 
-        stack_ret = ctx.stack_push
+        stack_ret = ctx.stack_push(Type::Unknown)
         asm.mov(stack_ret, C_RET)
 
         KeepCompiling
@@ -397,7 +397,7 @@ module RubyVM::RJIT
         asm.mov(C_ARGS[1], C_RET) # backref
         asm.call(C.rb_reg_nth_match)
 
-        stack_ret = ctx.stack_push
+        stack_ret = ctx.stack_push(Type::Unknown)
         asm.mov(stack_ret, C_RET)
 
         KeepCompiling
@@ -587,7 +587,7 @@ module RubyVM::RJIT
       asm.mov(C_ARGS[3], jit.operand(1))
       asm.call(C.rb_vm_getclassvariable)
 
-      top = ctx.stack_push
+      top = ctx.stack_push(Type::Unknown)
       asm.mov(top, C_RET)
 
       KeepCompiling
@@ -656,7 +656,7 @@ module RubyVM::RJIT
         asm.mov(:rax, [:rax, C.iseq_inline_constant_cache_entry.offsetof(:value)]) # ic_entry_val
 
         # Push ic->entry->value
-        stack_top = ctx.stack_push
+        stack_top = ctx.stack_push(Type::Unknown)
         asm.mov(stack_top, :rax)
       else # without cref
         # TODO: implement this
@@ -694,7 +694,7 @@ module RubyVM::RJIT
       asm.mov(C_ARGS[3], allow_nil_opnd)
       asm.call(C.rb_vm_get_ev_const)
 
-      top = ctx.stack_push
+      top = ctx.stack_push(Type::Unknown)
       asm.mov(top, C_RET)
 
       KeepCompiling
@@ -714,7 +714,7 @@ module RubyVM::RJIT
       asm.mov(C_ARGS[0], gid)
       asm.call(C.rb_gvar_get)
 
-      top = ctx.stack_push
+      top = ctx.stack_push(Type::Unknown)
       asm.mov(top, C_RET)
 
       KeepCompiling
@@ -733,7 +733,7 @@ module RubyVM::RJIT
     # @param ctx [RubyVM::RJIT::Context]
     # @param asm [RubyVM::RJIT::Assembler]
     def putself(jit, ctx, asm)
-      stack_top = ctx.stack_push
+      stack_top = ctx.stack_push_self
       asm.mov(:rax, [CFP, C.rb_control_frame_t.offsetof(:self)])
       asm.mov(stack_top, :rax)
       KeepCompiling
@@ -744,14 +744,14 @@ module RubyVM::RJIT
     # @param asm [RubyVM::RJIT::Assembler]
     def putobject(jit, ctx, asm, val: jit.operand(0))
       # Push it to the stack
-      stack_top = ctx.stack_push
+      val_type = Type.from(C.to_ruby(val))
+      stack_top = ctx.stack_push(val_type)
       if asm.imm32?(val)
         asm.mov(stack_top, val)
       else # 64-bit immediates can't be directly written to memory
         asm.mov(:rax, val)
         asm.mov(stack_top, :rax)
       end
-      # TODO: GC offsets?
 
       KeepCompiling
     end
@@ -762,7 +762,7 @@ module RubyVM::RJIT
     def putspecialobject(jit, ctx, asm)
       object_type = jit.operand(0)
       if object_type == C::VM_SPECIAL_OBJECT_VMCORE
-        stack_top = ctx.stack_push
+        stack_top = ctx.stack_push(Type::UnknownHeap)
         asm.mov(:rax, C.rb_mRubyVMFrozenCore)
         asm.mov(stack_top, :rax)
         KeepCompiling
@@ -786,7 +786,7 @@ module RubyVM::RJIT
       asm.mov(C_ARGS[1], to_value(put_val))
       asm.call(C.rb_ec_str_resurrect)
 
-      stack_top = ctx.stack_push
+      stack_top = ctx.stack_push(Type::CString)
       asm.mov(stack_top, C_RET)
 
       KeepCompiling
@@ -809,7 +809,7 @@ module RubyVM::RJIT
       asm.call(C.rb_str_concat_literals)
 
       ctx.stack_pop(n)
-      stack_ret = ctx.stack_push
+      stack_ret = ctx.stack_push(Type::CString)
       asm.mov(stack_ret, C_RET)
 
       KeepCompiling
@@ -830,7 +830,7 @@ module RubyVM::RJIT
       asm.call(C.rb_obj_as_string_result)
 
       # Push the return value
-      stack_ret = ctx.stack_push
+      stack_ret = ctx.stack_push(Type::TString)
       asm.mov(stack_ret, C_RET)
 
       KeepCompiling
@@ -869,7 +869,7 @@ module RubyVM::RJIT
       asm.pop(:rcx) # ary
 
       # The value we want to push on the stack is in RAX right now
-      stack_ret = ctx.stack_push
+      stack_ret = ctx.stack_push(Type::UnknownHeap)
       asm.mov(stack_ret, C_RET)
 
       # Clear the temp array.
@@ -891,7 +891,7 @@ module RubyVM::RJIT
       asm.call(C.rb_str_intern)
 
       # Push the return value
-      stack_ret = ctx.stack_push
+      stack_ret = ctx.stack_push(Type::Unknown)
       asm.mov(stack_ret, C_RET)
 
       KeepCompiling
@@ -924,7 +924,7 @@ module RubyVM::RJIT
       asm.call(C.rb_ec_ary_new_from_values)
 
       ctx.stack_pop(n)
-      stack_ret = ctx.stack_push
+      stack_ret = ctx.stack_push(Type::CArray)
       asm.mov(stack_ret, C_RET)
 
       KeepCompiling
@@ -946,7 +946,7 @@ module RubyVM::RJIT
       asm.mov(C_ARGS[0], ary)
       asm.call(C.rb_ary_resurrect)
 
-      stack_ret = ctx.stack_push
+      stack_ret = ctx.stack_push(Type::CArray)
       asm.mov(stack_ret, C_RET)
 
       KeepCompiling
@@ -976,16 +976,26 @@ module RubyVM::RJIT
 
       side_exit = side_exit(jit, ctx)
 
-      array_opnd = ctx.stack_pop(1)
+      array_opnd = ctx.stack_opnd(0)
 
       # num is the number of requested values. If there aren't enough in the
       # array then we're going to push on nils.
-      # TODO: implement this
+      if ctx.get_opnd_type(StackOpnd[0]) == Type::Nil
+        ctx.stack_pop(1) # pop after using the type info
+        # special case for a, b = nil pattern
+        # push N nils onto the stack
+        num.times do
+          push_opnd = ctx.stack_push(Type::Nil)
+          asm.mov(push_opnd, Qnil)
+        end
+        return KeepCompiling
+      end
 
       # Move the array from the stack and check that it's an array.
       asm.mov(:rax, array_opnd)
       guard_object_is_heap(asm, :rax, counted_exit(side_exit, :expandarray_not_array))
       guard_object_is_array(asm, :rax, :rcx, counted_exit(side_exit, :expandarray_not_array))
+      ctx.stack_pop(1) # pop after using the type info
 
       # If we don't actually want any values, then just return.
       if num == 0
@@ -1014,7 +1024,7 @@ module RubyVM::RJIT
 
       # Loop backward through the array and push each element onto the stack.
       (num - 1).downto(0).each do |i|
-        top = ctx.stack_push
+        top = ctx.stack_push(Type::Unknown)
         asm.mov(:rax, [:rcx, i * C.VALUE.size])
         asm.mov(top, :rax)
       end
@@ -1039,7 +1049,7 @@ module RubyVM::RJIT
       asm.mov(C_ARGS[1], ary2st_opnd)
       asm.call(C.rb_vm_concat_array)
 
-      stack_ret = ctx.stack_push
+      stack_ret = ctx.stack_push(Type::TArray)
       asm.mov(stack_ret, C_RET)
 
       KeepCompiling
@@ -1063,7 +1073,7 @@ module RubyVM::RJIT
       asm.mov(C_ARGS[1], ary_opnd)
       asm.call(C.rb_vm_splat_array)
 
-      stack_ret = ctx.stack_push
+      stack_ret = ctx.stack_push(Type::TArray)
       asm.mov(stack_ret, C_RET)
 
       KeepCompiling
@@ -1100,12 +1110,12 @@ module RubyVM::RJIT
         asm.pop(:rax)
 
         ctx.stack_pop(num)
-        stack_ret = ctx.stack_push
+        stack_ret = ctx.stack_push(Type::Hash)
         asm.mov(stack_ret, :rax)
       else
         # val = rb_hash_new();
         asm.call(C.rb_hash_new)
-        stack_ret = ctx.stack_push
+        stack_ret = ctx.stack_push(Type::Hash)
         asm.mov(stack_ret, C_RET)
       end
 
@@ -1128,7 +1138,7 @@ module RubyVM::RJIT
       asm.call(C.rb_range_new)
 
       ctx.stack_pop(2)
-      stack_ret = ctx.stack_push
+      stack_ret = ctx.stack_push(Type::UnknownHeap)
       asm.mov(stack_ret, C_RET)
 
       KeepCompiling
@@ -1146,10 +1156,13 @@ module RubyVM::RJIT
     # @param ctx [RubyVM::RJIT::Context]
     # @param asm [RubyVM::RJIT::Assembler]
     def dup(jit, ctx, asm)
-      val1 = ctx.stack_opnd(0)
-      val2 = ctx.stack_push
-      asm.mov(:rax, val1)
-      asm.mov(val2, :rax)
+      dup_val = ctx.stack_opnd(0)
+      mapping, tmp_type = ctx.get_opnd_mapping(StackOpnd[0])
+
+      loc0 = ctx.stack_push_mapping([mapping, tmp_type])
+      asm.mov(:rax, dup_val)
+      asm.mov(loc0, :rax)
+
       KeepCompiling
     end
 
@@ -1167,11 +1180,14 @@ module RubyVM::RJIT
       opnd1 = ctx.stack_opnd(1)
       opnd0 = ctx.stack_opnd(0)
 
-      dst1 = ctx.stack_push
+      mapping1 = ctx.get_opnd_mapping(StackOpnd[1])
+      mapping0 = ctx.get_opnd_mapping(StackOpnd[0])
+
+      dst1 = ctx.stack_push_mapping(mapping1)
       asm.mov(:rax, opnd1)
       asm.mov(dst1, :rax)
 
-      dst0 = ctx.stack_push
+      dst0 = ctx.stack_push_mapping(mapping0)
       asm.mov(:rax, opnd0)
       asm.mov(dst0, :rax)
 
@@ -1195,7 +1211,8 @@ module RubyVM::RJIT
       n = jit.operand(0)
 
       top_n_val = ctx.stack_opnd(n)
-      loc0 = ctx.stack_push
+      mapping = ctx.get_opnd_mapping(StackOpnd[n])
+      loc0 = ctx.stack_push_mapping(mapping)
       asm.mov(:rax, top_n_val)
       asm.mov(loc0, :rax)
 
@@ -1254,7 +1271,12 @@ module RubyVM::RJIT
       asm.cmovnz(:rax, :rcx)
 
       # Push the return value onto the stack
-      stack_ret = ctx.stack_push
+      out_type = if C::SPECIAL_CONST_P(pushval)
+        Type::UnknownImm
+      else
+        Type::Unknown
+      end
+      stack_ret = ctx.stack_push(out_type)
       asm.mov(stack_ret, :rax)
 
       KeepCompiling
@@ -1303,7 +1325,8 @@ module RubyVM::RJIT
         asm.cmovnz(:rax, :rcx)
 
         # Push the return value onto the stack
-        stack_ret = ctx.stack_push
+        out_type = C::SPECIAL_CONST_P(pushval) ? Type::UnknownImm : Type::Unknown
+        stack_ret = ctx.stack_push(out_type)
         asm.mov(stack_ret, :rax)
 
         return KeepCompiling
@@ -1367,7 +1390,7 @@ module RubyVM::RJIT
       asm.mov(:rcx, Qtrue)
       asm.cmovz(:rax, :rcx)
 
-      stack_ret = ctx.stack_push
+      stack_ret = ctx.stack_push(Type::UnknownImm)
       asm.mov(stack_ret, :rax)
 
       KeepCompiling
@@ -1460,7 +1483,7 @@ module RubyVM::RJIT
       str = jit.operand(0, ruby: true)
 
       # Push the return value onto the stack
-      stack_ret = ctx.stack_push
+      stack_ret = ctx.stack_push(Type::CString)
       asm.mov(:rax, to_value(str))
       asm.mov(stack_ret, :rax)
 
@@ -1496,7 +1519,7 @@ module RubyVM::RJIT
       asm.call(C.rb_vm_opt_newarray_min)
 
       ctx.stack_pop(num)
-      stack_ret = ctx.stack_push
+      stack_ret = ctx.stack_push(Type::Unknown)
       asm.mov(stack_ret, C_RET)
 
       KeepCompiling
@@ -1611,6 +1634,9 @@ module RubyVM::RJIT
       Invariants.assume_method_lookup_stable(jit, me)
       Invariants.assume_method_lookup_stable(jit, cme)
 
+      # Method calls may corrupt types
+      ctx.clear_local_types
+
       calling = build_calling(ci:, block_handler: block)
       case cme_def_type
       in C::VM_METHOD_TYPE_ISEQ
@@ -1706,9 +1732,12 @@ module RubyVM::RJIT
         asm.call(C.rb_vm_yield_with_cfunc)
 
         ctx.stack_pop(calling.argc)
-        stack_ret = ctx.stack_push
+        stack_ret = ctx.stack_push(Type::Unknown)
         asm.mov(stack_ret, C_RET)
 
+        # cfunc calls may corrupt types
+        ctx.clear_local_types
+
         # Share the successor with other chains
         jump_to_next_insn(jit, ctx, asm)
         EndBlock
@@ -2040,7 +2069,7 @@ module RubyVM::RJIT
         asm.add(:rax, :rcx)
         asm.jo(side_exit)
 
-        dst_opnd = ctx.stack_push
+        dst_opnd = ctx.stack_push(Type::Fixnum)
         asm.mov(dst_opnd, :rax)
 
         KeepCompiling
@@ -2086,7 +2115,7 @@ module RubyVM::RJIT
         asm.jo(side_exit)
         asm.add(:rax, 1) # re-tag
 
-        dst_opnd = ctx.stack_push
+        dst_opnd = ctx.stack_push(Type::Fixnum)
         asm.mov(dst_opnd, :rax)
 
         KeepCompiling
@@ -2144,7 +2173,7 @@ module RubyVM::RJIT
         asm.call(C.rb_fix_mod_fix)
 
         # Push the return value onto the stack
-        stack_ret = ctx.stack_push
+        stack_ret = ctx.stack_push(Type::Fixnum)
         asm.mov(stack_ret, C_RET)
 
         KeepCompiling
@@ -2245,7 +2274,7 @@ module RubyVM::RJIT
         asm.and(:rax, arg1)
 
         # Push the return value onto the stack
-        dst = ctx.stack_push
+        dst = ctx.stack_push(Type::Fixnum)
         asm.mov(dst, :rax)
 
         KeepCompiling
@@ -2285,7 +2314,7 @@ module RubyVM::RJIT
         asm.or(:rax, arg1)
 
         # Push the return value onto the stack
-        dst = ctx.stack_push
+        dst = ctx.stack_push(Type::Fixnum)
         asm.mov(dst, :rax)
 
         KeepCompiling
@@ -2343,7 +2372,7 @@ module RubyVM::RJIT
         ctx.stack_pop(2)
 
         # Push the return value onto the stack
-        stack_ret = ctx.stack_push
+        stack_ret = ctx.stack_push(Type::Unknown)
         asm.mov(stack_ret, C_RET)
 
         # Let guard chains share the same successor
@@ -2373,7 +2402,7 @@ module RubyVM::RJIT
         # Pop the key and the receiver
         ctx.stack_pop(2)
 
-        stack_ret = ctx.stack_push
+        stack_ret = ctx.stack_push(Type::Unknown)
         asm.mov(stack_ret, C_RET)
 
         # Let guard chains share the same successor
@@ -2431,7 +2460,7 @@ module RubyVM::RJIT
 
         # Push the return value onto the stack
         ctx.stack_pop(3)
-        stack_ret = ctx.stack_push
+        stack_ret = ctx.stack_push(Type::Unknown)
         asm.mov(:rax, val)
         asm.mov(stack_ret, :rax)
 
@@ -2457,7 +2486,7 @@ module RubyVM::RJIT
 
         # Push the return value onto the stack
         ctx.stack_pop(3)
-        stack_ret = ctx.stack_push
+        stack_ret = ctx.stack_push(Type::Unknown)
         asm.mov(stack_ret, C_RET)
 
         jump_to_next_insn(jit, ctx, asm)
@@ -2514,6 +2543,9 @@ module RubyVM::RJIT
 
     # invokebuiltin
 
+    # @param jit [RubyVM::RJIT::JITState]
+    # @param ctx [RubyVM::RJIT::Context]
+    # @param asm [RubyVM::RJIT::Assembler]
     def opt_invokebuiltin_delegate(jit, ctx, asm)
       bf = C.rb_builtin_function.new(jit.operand(0))
       bf_argc = bf.argc
@@ -2546,7 +2578,7 @@ module RubyVM::RJIT
       asm.call(bf.func_ptr)
 
       # Push the return value
-      stack_ret = ctx.stack_push
+      stack_ret = ctx.stack_push(Type::Unknown)
       asm.mov(stack_ret, C_RET)
 
       KeepCompiling
@@ -2583,9 +2615,9 @@ module RubyVM::RJIT
     # @param asm [RubyVM::RJIT::Assembler]
     def jit_rb_true(jit, ctx, asm, argc, _known_recv_class)
       return false if argc != 0
-      asm.comment('nil? == true');
+      asm.comment('nil? == true')
       ctx.stack_pop(1)
-      stack_ret = ctx.stack_push
+      stack_ret = ctx.stack_push(Type::True)
       asm.mov(stack_ret, Qtrue)
       true
     end
@@ -2595,9 +2627,9 @@ module RubyVM::RJIT
     # @param asm [RubyVM::RJIT::Assembler]
     def jit_rb_false(jit, ctx, asm, argc, _known_recv_class)
       return false if argc != 0
-      asm.comment('nil? == false');
+      asm.comment('nil? == false')
       ctx.stack_pop(1)
-      stack_ret = ctx.stack_push
+      stack_ret = ctx.stack_push(Type::False)
       asm.mov(stack_ret, Qfalse)
       true
     end
@@ -2640,10 +2672,11 @@ module RubyVM::RJIT
 
       ctx.stack_pop(2)
 
-      stack_ret = ctx.stack_push
       if sample_is_a
+        stack_ret = ctx.stack_push(Type::True)
         asm.mov(stack_ret, Qtrue)
       else
+        stack_ret = ctx.stack_push(Type::False)
         asm.mov(stack_ret, Qfalse)
       end
       return true
@@ -2693,10 +2726,11 @@ module RubyVM::RJIT
 
       ctx.stack_pop(2)
 
-      stack_ret = ctx.stack_push
       if sample_instance_of
+        stack_ret = ctx.stack_push(Type::True)
         asm.mov(stack_ret, Qtrue)
       else
+        stack_ret = ctx.stack_push(Type::False)
         asm.mov(stack_ret, Qfalse)
       end
       return true;
@@ -2707,17 +2741,33 @@ module RubyVM::RJIT
     # @param asm [RubyVM::RJIT::Assembler]
     def jit_rb_obj_not(jit, ctx, asm, argc, _known_recv_class)
       return false if argc != 0
-      asm.comment('rb_obj_not')
+      recv_type = ctx.get_opnd_type(StackOpnd[0])
 
-      recv = ctx.stack_pop
-      # This `test` sets ZF only for Qnil and Qfalse, which let cmovz set.
-      asm.test(recv, ~Qnil)
-      asm.mov(:rax, Qfalse)
-      asm.mov(:rcx, Qtrue)
-      asm.cmovz(:rax, :rcx)
+      case recv_type.known_truthy
+      in false
+        asm.comment('rb_obj_not(nil_or_false)')
+        ctx.stack_pop(1)
+        out_opnd = ctx.stack_push(Type::True)
+        asm.mov(out_opnd, Qtrue)
+      in true
+        # Note: recv_type != Type::Nil && recv_type != Type::False.
+        asm.comment('rb_obj_not(truthy)')
+        ctx.stack_pop(1)
+        out_opnd = ctx.stack_push(Type::False)
+        asm.mov(out_opnd, Qfalse)
+      in nil
+        asm.comment('rb_obj_not')
+
+        recv = ctx.stack_pop
+        # This `test` sets ZF only for Qnil and Qfalse, which let cmovz set.
+        asm.test(recv, ~Qnil)
+        asm.mov(:rax, Qfalse)
+        asm.mov(:rcx, Qtrue)
+        asm.cmovz(:rax, :rcx)
 
-      stack_ret = ctx.stack_push
-      asm.mov(stack_ret, :rax)
+        stack_ret = ctx.stack_push(Type::UnknownImm)
+        asm.mov(stack_ret, :rax)
+      end
       true
     end
 
@@ -2737,7 +2787,7 @@ module RubyVM::RJIT
       asm.mov(:rcx, Qtrue)
       asm.cmove(:rax, :rcx)
 
-      stack_ret = ctx.stack_push
+      stack_ret = ctx.stack_push(Type::UnknownImm)
       asm.mov(stack_ret, :rax)
       true
     end
@@ -2771,7 +2821,7 @@ module RubyVM::RJIT
 
       # Return the result
       ctx.stack_pop(2)
-      stack_ret = ctx.stack_push
+      stack_ret = ctx.stack_push(Type::UnknownImm)
       asm.mov(stack_ret, C_RET)
 
       return true
@@ -2797,7 +2847,7 @@ module RubyVM::RJIT
       asm.mov(:rcx, Qtrue)
       asm.cmove(:rax, :rcx)
 
-      stack_ret = ctx.stack_push
+      stack_ret = ctx.stack_push(Type::UnknownImm)
       asm.mov(stack_ret, :rax)
       true
     end
@@ -2819,7 +2869,7 @@ module RubyVM::RJIT
       asm.mov(C_ARGS[1], y_opnd)
       asm.call(C.rb_fix_mul_fix)
 
-      ret_opnd = ctx.stack_push
+      ret_opnd = ctx.stack_push(Type::Unknown)
       asm.mov(ret_opnd, C_RET)
       true
     end
@@ -2842,7 +2892,7 @@ module RubyVM::RJIT
       asm.mov(C_ARGS[1], :rax)
       asm.call(C.rb_fix_div_fix)
 
-      ret_opnd = ctx.stack_push
+      ret_opnd = ctx.stack_push(Type::Unknown)
       asm.mov(ret_opnd, C_RET)
       true
     end
@@ -2865,7 +2915,7 @@ module RubyVM::RJIT
       asm.mov(C_ARGS[1], y_opnd)
       asm.call(C.rb_fix_aref)
 
-      ret_opnd = ctx.stack_push
+      ret_opnd = ctx.stack_push(Type::UnknownImm)
       asm.mov(ret_opnd, C_RET)
       true
     end
@@ -2879,7 +2929,7 @@ module RubyVM::RJIT
       # `C.RString.offsetof(:as, :embed, :len)` doesn't work because of USE_RVARGC=0 CI
 
       recv_opnd = ctx.stack_pop(1)
-      out_opnd = ctx.stack_push
+      out_opnd = ctx.stack_push(Type::UnknownImm)
 
       asm.comment('get string length')
       asm.mov(:rax, recv_opnd)
@@ -2918,7 +2968,7 @@ module RubyVM::RJIT
       asm.mov(C_ARGS[0], recv)
       asm.call(C.rb_str_bytesize)
 
-      out_opnd = ctx.stack_push
+      out_opnd = ctx.stack_push(Type::Fixnum)
       asm.mov(out_opnd, C_RET)
 
       true
@@ -2945,7 +2995,7 @@ module RubyVM::RJIT
       guard_object_is_string(asm, :rax, :rcx, side_exit)
 
       # Guard buffers from GC since rb_str_buf_append may allocate. During the VM lock on GC,
-      # other Ractors may trigger global invalidation, so we need record_boundary_patch_point.
+      # other Ractors may trigger global invalidation, so we need ctx.clear_local_types.
       # PC is used on errors like Encoding::CompatibilityError raised by rb_str_buf_append.
       jit_prepare_routine_call(jit, ctx, asm)
 
@@ -2969,7 +3019,7 @@ module RubyVM::RJIT
       asm.test(recv_reg, C::RUBY_ENCODING_MASK)
 
       # Push once, use the resulting operand in both branches below.
-      stack_ret = ctx.stack_push
+      stack_ret = ctx.stack_push(Type::CString)
 
       enc_mismatch = asm.new_label('enc_mismatch')
       asm.jnz(enc_mismatch)
@@ -3014,7 +3064,7 @@ module RubyVM::RJIT
       ret_label = asm.new_label('stack_ret')
 
       # String#+@ can only exist on T_STRING
-      stack_ret = ctx.stack_push
+      stack_ret = ctx.stack_push(Type::TString)
 
       # If the string isn't frozen, we just return it.
       asm.mov(stack_ret, :rax) # recv_opnd
@@ -3043,7 +3093,7 @@ module RubyVM::RJIT
       asm.mov(C_ARGS[1], index_opnd)
       asm.call(C.rb_str_getbyte)
 
-      ret_opnd = ctx.stack_push
+      ret_opnd = ctx.stack_push(Type::Fixnum)
       asm.mov(ret_opnd, C_RET)
       true
     end
@@ -3061,7 +3111,7 @@ module RubyVM::RJIT
       asm.mov(:rcx, Qtrue)
       asm.cmovz(:rax, :rcx)
 
-      out_opnd = ctx.stack_push
+      out_opnd = ctx.stack_push(Type::UnknownImm)
       asm.mov(out_opnd, :rax)
 
       return true
@@ -3082,7 +3132,7 @@ module RubyVM::RJIT
       asm.mov(C_ARGS[1], item_opnd)
       asm.call(C.rb_ary_push)
 
-      ret_opnd = ctx.stack_push
+      ret_opnd = ctx.stack_push(Type::TArray)
       asm.mov(ret_opnd, C_RET)
       true
     end
@@ -3187,7 +3237,7 @@ module RubyVM::RJIT
       asm.mov(:rax, [:rax, C.VALUE.size * C::VM_ENV_DATA_INDEX_SPECVAL]) # block_handler
 
       ctx.stack_pop(1)
-      out_opnd = ctx.stack_push
+      out_opnd = ctx.stack_push(Type::UnknownImm)
 
       # Return `block_handler != VM_BLOCK_HANDLER_NONE`
       asm.cmp(:rax, C::VM_BLOCK_HANDLER_NONE)
@@ -3213,7 +3263,7 @@ module RubyVM::RJIT
       # thread->self
       asm.mov(:rax, [:rax, C.rb_thread_struct.offsetof(:self)])
 
-      stack_ret = ctx.stack_push
+      stack_ret = ctx.stack_push(Type::UnknownHeap)
       asm.mov(stack_ret, :rax)
       true
     end
@@ -3304,7 +3354,12 @@ module RubyVM::RJIT
       asm.mov(:rax, [ep_reg, -idx * C.VALUE.size])
 
       # Write the local at SP
-      stack_top = ctx.stack_push
+      stack_top = if level == 0
+        local_idx = ep_offset_to_local_idx(jit.iseq, idx)
+        ctx.stack_push_local(local_idx)
+      else
+        ctx.stack_push(Type::Unknown)
+      end
 
       asm.mov(stack_top, :rax)
       KeepCompiling
@@ -3337,6 +3392,30 @@ module RubyVM::RJIT
     end
 
     # Compute the index of a local variable from its slot index
+    def ep_offset_to_local_idx(iseq, ep_offset)
+      # Layout illustration
+      # This is an array of VALUE
+      #                                           | VM_ENV_DATA_SIZE |
+      #                                           v                  v
+      # low addr <+-------+-------+-------+-------+------------------+
+      #           |local 0|local 1|  ...  |local n|       ....       |
+      #           +-------+-------+-------+-------+------------------+
+      #           ^       ^                       ^                  ^
+      #           +-------+---local_table_size----+         cfp->ep--+
+      #                   |                                          |
+      #                   +------------------ep_offset---------------+
+      #
+      # See usages of local_var_name() from iseq.c for similar calculation.
+
+      # Equivalent of iseq->body->local_table_size
+      local_table_size = iseq.body.local_table_size
+      op = ep_offset - C::VM_ENV_DATA_SIZE
+      local_idx = local_table_size - op - 1
+      assert_equal(true, local_idx >= 0 && local_idx < local_table_size)
+      local_idx
+    end
+
+    # Compute the index of a local variable from its slot index
     def slot_to_local_idx(iseq, slot_idx)
       # Layout illustration
       # This is an array of VALUE
@@ -3582,7 +3661,7 @@ module RubyVM::RJIT
         asm.mov(:rcx, Qtrue)
         asm.public_send(opcode, :rax, :rcx)
 
-        dst_opnd = ctx.stack_push
+        dst_opnd = ctx.stack_push(Type::UnknownImm)
         asm.mov(dst_opnd, :rax)
 
         KeepCompiling
@@ -3621,7 +3700,7 @@ module RubyVM::RJIT
 
         # Push the output on the stack
         ctx.stack_pop(2)
-        dst = ctx.stack_push
+        dst = ctx.stack_push(Type::UnknownImm)
         asm.mov(dst, :rax)
 
         true
@@ -3655,7 +3734,7 @@ module RubyVM::RJIT
 
         # Push the output on the stack
         ctx.stack_pop(2)
-        dst = ctx.stack_push
+        dst = ctx.stack_push(Type::UnknownImm)
         asm.mov(dst, C_RET)
         asm.jmp(ret_label)
 
@@ -3678,6 +3757,10 @@ module RubyVM::RJIT
       jit.record_boundary_patch_point = true
       jit_save_pc(jit, asm)
       jit_save_sp(ctx, asm)
+
+      # In case the routine calls Ruby methods, it can set local variables
+      # through Kernel#binding and other means.
+      ctx.clear_local_types
     end
 
     # NOTE: This clobbers :rax
@@ -3797,7 +3880,7 @@ module RubyVM::RJIT
         end
 
         # Push the ivar on the stack
-        out_opnd = ctx.stack_push
+        out_opnd = ctx.stack_push(Type::Unknown)
         asm.mov(out_opnd, C_RET)
 
         # Jump to next instruction. This allows guard chains to share the same successor.
@@ -3818,8 +3901,18 @@ module RubyVM::RJIT
       asm.cmp(DwordPtr[:rax, C.rb_shape_id_offset], shape_id)
       jit_chain_guard(:jne, jit, starting_ctx, asm, counted_exit(side_exit, :getivar_megamorphic))
 
+      if obj_opnd
+        ctx.stack_pop # pop receiver for attr_reader
+      end
+
       index = C.rb_shape_get_iv_index(shape_id, ivar_id)
-      if index
+      # If there is no IVAR index, then the ivar was undefined
+      # when we entered the compiler.  That means we can just return
+      # nil for this shape + iv name
+      if index.nil?
+        stack_opnd = ctx.stack_push(Type::Nil)
+        val_opnd = Qnil
+      else
         asm.comment('ROBJECT_IVPTR')
         if C::FL_TEST_RAW(comptime_obj, C::ROBJECT_EMBED)
           # Access embedded array
@@ -3830,15 +3923,9 @@ module RubyVM::RJIT
           # Read the table
           asm.mov(:rax, [:rax, index * C.VALUE.size])
         end
+        stack_opnd = ctx.stack_push(Type::Unknown)
         val_opnd = :rax
-      else
-        val_opnd = Qnil
       end
-
-      if obj_opnd
-        ctx.stack_pop # pop receiver for attr_reader
-      end
-      stack_opnd = ctx.stack_push
       asm.mov(stack_opnd, val_opnd)
 
       # Let guard chains share the same successor
@@ -4325,7 +4412,7 @@ module RubyVM::RJIT
 
         remaining_opt.times do
           # We need to push nil for the optional arguments
-          stack_ret = ctx.stack_push
+          stack_ret = ctx.stack_push(Type::Unknown)
           asm.mov(stack_ret, Qnil)
         end
       end
@@ -4367,7 +4454,7 @@ module RubyVM::RJIT
             asm.call(C.rb_ary_unshift_m)
             ctx.stack_pop(diff)
 
-            stack_ret = ctx.stack_push
+            stack_ret = ctx.stack_push(Type::TArray)
             asm.mov(stack_ret, C_RET)
             # We now should have the required arguments
             # and an array of all the rest arguments
@@ -4383,7 +4470,7 @@ module RubyVM::RJIT
             asm.mov(C_ARGS[0], array)
             asm.mov(C_ARGS[1], diff)
             asm.call(C.rjit_rb_ary_subseq_length)
-            stack_ret = ctx.stack_push
+            stack_ret = ctx.stack_push(Type::TArray)
             asm.mov(stack_ret, C_RET)
 
             # We now should have the required arguments
@@ -4392,7 +4479,7 @@ module RubyVM::RJIT
           else
             # The arguments are equal so we can just push to the stack
             assert_equal(non_rest_arg_count, required_num)
-            stack_ret = ctx.stack_push
+            stack_ret = ctx.stack_push(Type::TArray)
             asm.mov(stack_ret, array)
           end
         else
@@ -4416,7 +4503,7 @@ module RubyVM::RJIT
           asm.call(C.rb_ec_ary_new_from_values)
 
           ctx.stack_pop(n)
-          stack_ret = ctx.stack_push
+          stack_ret = ctx.stack_push(Type::CArray)
           asm.mov(stack_ret, C_RET)
         end
       end
@@ -4481,7 +4568,7 @@ module RubyVM::RJIT
             # filling in (which is done in the next loop). Also increments
             # argc so that the callee's SP is recorded correctly.
             argc += 1
-            default_arg = ctx.stack_push
+            default_arg = ctx.stack_push(Type::Unknown)
 
             # callee_idx - keyword->required_num is used in a couple of places below.
             req_num = keyword.required_num
@@ -4618,7 +4705,7 @@ module RubyVM::RJIT
       asm.call(builtin_func.func_ptr)
 
       # Push the return value
-      stack_ret = ctx.stack_push
+      stack_ret = ctx.stack_push(Type::Unknown)
       asm.mov(stack_ret, C_RET)
       return true
     end
@@ -4830,7 +4917,7 @@ module RubyVM::RJIT
       Invariants.record_global_inval_patch(asm, full_cfunc_return)
 
       # Push the return value on the Ruby stack
-      stack_ret = ctx.stack_push
+      stack_ret = ctx.stack_push(Type::Unknown)
       asm.mov(stack_ret, C_RET)
 
       # Pop the stack frame (ec->cfp++)
@@ -4838,7 +4925,10 @@ module RubyVM::RJIT
       # register
       asm.mov([EC, C.rb_execution_context_t.offsetof(:cfp)], CFP)
 
-      # Note: the return block of gen_send_iseq() has ctx->sp_offset == 1
+      # cfunc calls may corrupt types
+      ctx.clear_local_types
+
+      # Note: the return block of jit_call_iseq has ctx->sp_offset == 1
       # which allows for sharing the same successor.
 
       # Jump (fall through) to the call continuation block
@@ -4898,7 +4988,7 @@ module RubyVM::RJIT
       asm.mov(C_ARGS[2], val_opnd)
       asm.call(C.rb_vm_set_ivar_id)
 
-      out_opnd = ctx.stack_push
+      out_opnd = ctx.stack_push(Type::Unknown)
       asm.mov(out_opnd, C_RET)
 
       KeepCompiling
@@ -5093,7 +5183,7 @@ module RubyVM::RJIT
 
       ctx.stack_pop(argc + 1)
 
-      stack_ret = ctx.stack_push
+      stack_ret = ctx.stack_push(Type::Unknown)
       asm.mov(stack_ret, C_RET)
       return KeepCompiling
     end
@@ -5135,7 +5225,7 @@ module RubyVM::RJIT
         asm.mov(:rax, [:rax, C.VALUE.size * off])
       end
 
-      ret = ctx.stack_push
+      ret = ctx.stack_push(Type::Unknown)
       asm.mov(ret, :rax)
 
       jump_to_next_insn(jit, ctx, asm)
@@ -5157,7 +5247,7 @@ module RubyVM::RJIT
         asm.mov(opnd2, :rax)
       end
 
-      ctx.stack_pop(1)
+      ctx.shift_stack(1)
     end
 
     # vm_call_symbol
@@ -5295,10 +5385,13 @@ module RubyVM::RJIT
 
       # cfp->jit_return is used only for ISEQs
       if iseq
+        # The callee might change locals through Kernel#binding and other means.
+        ctx.clear_local_types
+
         # Stub cfp->jit_return
         return_ctx = ctx.dup
-        return_ctx.stack_size -= argc + ((block_handler == :captured) ? 0 : 1) # Pop args and receiver. blockarg has been popped
-        return_ctx.stack_size += 1 # push callee's return value
+        return_ctx.stack_pop(argc + ((block_handler == :captured) ? 0 : 1)) # Pop args and receiver. blockarg has been popped
+        return_ctx.stack_push(Type::Unknown) # push callee's return value
         return_ctx.sp_offset = 1 # SP is in the position after popping a receiver and arguments
         return_ctx.chain_depth = 0
         branch_stub = BranchStub.new(
@@ -5388,7 +5481,7 @@ module RubyVM::RJIT
       asm.cmovnz(ary_opnd, :rcx)
 
       num_args.times do |i|
-        top = ctx.stack_push
+        top = ctx.stack_push(Type::Unknown)
         asm.mov(:rcx, [ary_opnd, i * C.VALUE.size])
         asm.mov(top, :rcx)
       end
@@ -5452,7 +5545,7 @@ module RubyVM::RJIT
         ary_opnd = :rax
 
         (0...required_args).each do |i|
-          top = ctx.stack_push
+          top = ctx.stack_push(Type::Unknown)
           asm.mov(:rcx, [ary_opnd, i * C.VALUE.size])
           asm.mov(top, :rcx)
         end
diff --git a/lib/ruby_vm/rjit/type.rb b/lib/ruby_vm/rjit/type.rb
new file mode 100644
index 0000000000..80157fc9c2
--- /dev/null
+++ b/lib/ruby_vm/rjit/type.rb
@@ -0,0 +1,122 @@
+module RubyVM::RJIT
+  # Represent the type of a value (local/stack/self) in RJIT
+  Type = Data.define(:type) do
+    # Returns a boolean representing whether the value is truthy if known, otherwise nil
+    def known_truthy
+      case self
+      in Type::Nil | Type::False
+        false
+      in Type::UnknownHeap
+        true
+      in Type::Unknown | Type::UnknownImm
+        nil
+      else
+        true
+      end
+    end
+
+    def diff(dst)
+      if self == dst
+        return TypeDiff::Compatible[0]
+      end
+
+      return TypeDiff::Incompatible
+    end
+
+    def upgrade(new_type)
+      assert(new_type.diff(self) != TypeDiff::Incompatible)
+      new_type
+    end
+
+    private
+
+    def assert(cond)
+      unless cond
+        raise "'#{cond.inspect}' was not true"
+      end
+    end
+  end
+
+  # This returns an appropriate Type based on a known value
+  class << Type
+    def from(val)
+      if C::SPECIAL_CONST_P(val)
+        if fixnum?(val)
+          Type::Fixnum
+        elsif val.nil?
+          Type::Nil
+        elsif val == true
+          Type::True
+        elsif val == false
+          Type::False
+        elsif static_symbol?(val)
+          Type::ImmSymbol
+        elsif flonum?(val)
+          Type::Flonum
+        else
+          raise "Illegal value: #{val.inspect}"
+        end
+      else
+        val_class = C.to_value(C.rb_class_of(val))
+        if val_class == C.rb_cString
+          return Type::CString
+        end
+        if val_class == C.rb_cArray
+          return Type::CArray
+        end
+        if C.to_value(val) == C.rb_block_param_proxy
+          return Type::BlockParamProxy
+        end
+        case C::BUILTIN_TYPE(val)
+        in C::RUBY_T_ARRAY
+          Type::TArray
+        in C::RUBY_T_HASH
+          Type::Hash
+        in C::RUBY_T_STRING
+          Type::TString
+        else
+          Type::UnknownHeap
+        end
+      end
+    end
+
+    private
+
+    def fixnum?(obj)
+      (C.to_value(obj) & C::RUBY_FIXNUM_FLAG) == C::RUBY_FIXNUM_FLAG
+    end
+
+    def flonum?(obj)
+      (C.to_value(obj) & C::RUBY_FLONUM_MASK) == C::RUBY_FLONUM_FLAG
+    end
+
+    def static_symbol?(obj)
+      (C.to_value(obj) & 0xff) == C::RUBY_SYMBOL_FLAG
+    end
+  end
+
+  # List of types
+  Type::Unknown     = Type[:Unknown]
+  Type::UnknownImm  = Type[:UnknownImm]
+  Type::UnknownHeap = Type[:UnknownHeap]
+  Type::Nil         = Type[:Nil]
+  Type::True        = Type[:True]
+  Type::False       = Type[:False]
+  Type::Fixnum      = Type[:Fixnum]
+  Type::Flonum      = Type[:Flonum]
+  Type::Hash        = Type[:Hash]
+  Type::ImmSymbol   = Type[:ImmSymbol]
+  Type::HeapSymbol  = Type[:HeapSymbol]
+
+  Type::TString = Type[:TString] # An object with the T_STRING flag set, possibly an rb_cString
+  Type::CString = Type[:CString] # An un-subclassed string of type rb_cString (can have instance vars in some cases)
+  Type::TArray  = Type[:TArray]  # An object with the T_ARRAY flag set, possibly an rb_cArray
+  Type::CArray  = Type[:CArray]  # An un-subclassed string of type rb_cArray (can have instance vars in some cases)
+
+  Type::BlockParamProxy = Type[:BlockParamProxy] # A special sentinel value indicating the block parameter should be read from
+
+  module TypeDiff
+    Compatible = Data.define(:diversion) # The smaller, the more compatible.
+    Incompatible = :Incompatible
+  end
+end
diff --git a/rjit_c.rb b/rjit_c.rb
index 9bbc2bbcda..fa4ea2dddc 100644
--- a/rjit_c.rb
+++ b/rjit_c.rb
@@ -480,6 +480,10 @@ module RubyVM::RJIT # :nodoc: all
     Primitive.cexpr! %q{ SIZET2NUM(rb_block_param_proxy) }
   end
 
+  def C.rb_cArray
+    Primitive.cexpr! %q{ SIZET2NUM(rb_cArray) }
+  end
+
   def C.rb_cFalseClass
     Primitive.cexpr! %q{ SIZET2NUM(rb_cFalseClass) }
   end
@@ -496,6 +500,10 @@ module RubyVM::RJIT # :nodoc: all
     Primitive.cexpr! %q{ SIZET2NUM(rb_cNilClass) }
   end
 
+  def C.rb_cString
+    Primitive.cexpr! %q{ SIZET2NUM(rb_cString) }
+  end
+
   def C.rb_cSymbol
     Primitive.cexpr! %q{ SIZET2NUM(rb_cSymbol) }
   end
diff --git a/tool/rjit/bindgen.rb b/tool/rjit/bindgen.rb
index 13b8ec3c07..a1b6c87249 100755
--- a/tool/rjit/bindgen.rb
+++ b/tool/rjit/bindgen.rb
@@ -494,10 +494,12 @@ generator = BindingGenerator.new(
       imemo_iseq
       imemo_callinfo
       rb_block_param_proxy
+      rb_cArray
       rb_cFalseClass
       rb_cFloat
       rb_cInteger
       rb_cNilClass
+      rb_cString
       rb_cSymbol
       rb_cTrueClass
       rb_rjit_global_events