RJIT: Propagate self's type information

1 files changed, 10 insertions, 3 deletions

diff --git a/lib/ruby_vm/rjit/insn_compiler.rb b/lib/ruby_vm/rjit/insn_compiler.rb
index 9aa643037a..1a0514e75a 100644
--- a/lib/ruby_vm/rjit/insn_compiler.rb
+++ b/lib/ruby_vm/rjit/insn_compiler.rb
@@ -4799,6 +4799,16 @@ module RubyVM::RJIT
         return CantCompile
       end
 
+      # Create a context for the callee
+      callee_ctx = Context.new
+
+      recv_type = if calling.block_handler == :captured
+        Type::Unknown # we don't track the type information of captured->self for now
+      else
+        ctx.get_opnd_type(StackOpnd[argc])
+      end
+      callee_ctx.upgrade_opnd_type(SelfOpnd, recv_type)
+
       # Setup the new frame
       frame_type ||= C::VM_FRAME_MAGIC_METHOD | C::VM_ENV_FLAG_LOCAL
       jit_push_frame(
@@ -4810,9 +4820,6 @@ module RubyVM::RJIT
         doing_kw_call:,
       )
 
-      # Create a context for the callee
-      callee_ctx = Context.new
-
       # Directly jump to the entry point of the callee
       pc = (iseq.body.iseq_encoded + start_pc_offset).to_i
       jit_direct_jump(iseq, pc, callee_ctx, asm)