From 2884f53519c4b86072d5fc3b41a71cee697af8ba Mon Sep 17 00:00:00 2001
From: Alan Wu <XrXr@users.noreply.github.com>
Date: Fri, 12 Dec 2025 13:42:00 -0500
Subject: [PATCH] YJIT: Add missing local variable type update for fallback
 setlocal blocks

Previously, the chain_depth>0 version of setlocal blocks did not
update the type of the local variable in the context. This can leave
the context with stale type information and trigger panics like in
[Bug #21772] or lead to miscompilation.

To trigger the issue, YJIT needs to see the same ISEQ before and after
environment escape and have tracked type info before the escape. To
trigger in ISEQs that do not send with a block, it probably requires
Kernel#binding or the use of include/ruby/debug.h APIs.
---
 bootstraptest/test_yjit.rb | 19 +++++++++++++++++++
 yjit/src/codegen.rs        |  7 +++++++
 2 files changed, 26 insertions(+)

diff --git a/bootstraptest/test_yjit.rb b/bootstraptest/test_yjit.rb
index bd44f3ac3e..fc592520e8 100644
--- a/bootstraptest/test_yjit.rb
+++ b/bootstraptest/test_yjit.rb
@@ -5416,3 +5416,22 @@ assert_equal 'foo', %{
 
   10.times.map { foo.dup }.last
 }
+
+# regression test for [Bug #21772]
+# local variable type tracking desync
+assert_normal_exit %q{
+  def some_method = 0
+
+  def test_body(key)
+    some_method
+    key = key.to_s # setting of local relevant
+
+    key == "symbol"
+  end
+
+  def jit_caller = test_body("session_id")
+
+  jit_caller # first iteration, non-escaped environment
+  alias some_method binding # induce environment escape
+  test_body(:symbol)
+}
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index 865f80ca4f..4a53e8834c 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -2518,6 +2518,7 @@ fn gen_setlocal_generic(
     ep_offset: u32,
     level: u32,
 ) -> Option<CodegenStatus> {
+    // Post condition: The type of of the set local is updated in the Context.
     let value_type = asm.ctx.get_opnd_type(StackOpnd(0));
 
     // Fallback because of write barrier
@@ -2539,6 +2540,11 @@ fn gen_setlocal_generic(
         );
         asm.stack_pop(1);
 
+        // Set local type in the context
+        if level == 0 {
+            let local_idx = ep_offset_to_local_idx(jit.get_iseq(), ep_offset).as_usize();
+            asm.ctx.set_local_type(local_idx, value_type);
+        }
         return Some(KeepCompiling);
     }
 
@@ -2591,6 +2597,7 @@ fn gen_setlocal_generic(
         );
     }
 
+    // Set local type in the context
     if level == 0 {
         let local_idx = ep_offset_to_local_idx(jit.get_iseq(), ep_offset).as_usize();
         asm.ctx.set_local_type(local_idx, value_type);