YJIT: Split the block on optimized getlocal/setlocal (#13331)

k0kubun · web-flow · commit d8e7ab99a30f · 2025-05-13T20:10:15.000-07:00
diff --git a/test/-ext-/debug/test_debug.rb b/test/-ext-/debug/test_debug.rb
@@ -75,19 +75,58 @@ def test_lazy_block
     end
     assert_equal true, x, '[Bug #15105]'
   end
+end
+
+# This is a YJIT test, but we can't test this without a C extension that calls
+# rb_debug_inspector_open(), so we're testing it using "-test-/debug" here.
+class TestDebugWithYJIT < Test::Unit::TestCase
+  class LocalSetArray
+    def to_a
+      Bug::Debug.inspector.each do |_, binding,|
+        binding.local_variable_set(:local, :ok) if binding
+      end
+      [:ok]
+    end
+  end
+
+  class DebugArray
+    def to_a
+      Bug::Debug.inspector
+      [:ok]
+    end
+  end
+
+  def test_yjit_invalidates_getlocal_after_splatarray
+    val = getlocal_after_splatarray(LocalSetArray.new)
+    assert_equal [:ok, :ok], val
+  end
 
-  # This is a YJIT test, but we can't test this without a C extension that calls
-  # rb_debug_inspector_open(), so we're testing it using "-test-/debug" here.
-  def test_yjit_invalidates_setlocal_after_inspector_call
+  def test_yjit_invalidates_setlocal_after_splatarray
+    val = setlocal_after_splatarray(DebugArray.new)
+    assert_equal [:ok], val
+  end
+
+  def test_yjit_invalidates_setlocal_after_proc_call
     val = setlocal_after_proc_call(proc { Bug::Debug.inspector; :ok })
     assert_equal :ok, val
-  end if defined?(RubyVM::YJIT) && RubyVM::YJIT.enabled?
+  end
 
   private
 
+  def getlocal_after_splatarray(array)
+    local = 1
+    [*array, local]
+  end
+
+  def setlocal_after_splatarray(array)
+    local = *array # setlocal followed by splatarray
+    itself # split a block using a C call
+    local # getlocal
+  end
+
   def setlocal_after_proc_call(block)
     local = block.call # setlocal followed by OPTIMIZED_METHOD_TYPE_CALL
     itself # split a block using a C call
     local # getlocal
   end
-end
+end if defined?(RubyVM::YJIT) && RubyVM::YJIT.enabled?
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
@@ -2445,6 +2445,11 @@ fn gen_getlocal_generic(
     ep_offset: u32,
     level: u32,
 ) -> Option<CodegenStatus> {
+    // Split the block if we need to invalidate this instruction when EP escapes
+    if level == 0 && !jit.escapes_ep() && !jit.at_compile_target() {
+        return jit.defer_compilation(asm);
+    }
+
     let local_opnd = if level == 0 && jit.assume_no_ep_escape(asm) {
         // Load the local using SP register
         asm.local_opnd(ep_offset)
@@ -2535,6 +2540,11 @@ fn gen_setlocal_generic(
         return Some(KeepCompiling);
     }
 
+    // Split the block if we need to invalidate this instruction when EP escapes
+    if level == 0 && !jit.escapes_ep() && !jit.at_compile_target() {
+        return jit.defer_compilation(asm);
+    }
+
     let (flags_opnd, local_opnd) = if level == 0 && jit.assume_no_ep_escape(asm) {
         // Load flags and the local using SP register
         let flags_opnd = asm.ctx.ep_opnd(VM_ENV_DATA_INDEX_FLAGS as i32);
