ZJIT: Scan rescue/ensure ISEQs in locals_written_in_block

locals_written_in_block was only scanning blockiseqs reachable through
send instructions, missing rescue/ensure handler ISEQs referenced from
the catch table. This caused locals written inside rescue blocks to not
be reloaded after send-with-block, leading to stale values.

Author: tekknolagi

jit.c

@@ -366,6 +366,19 @@ rb_get_iseq_body_param_opt_table(const rb_iseq_t *iseq)
     return iseq->body->param.opt_table;
 }
 
+unsigned int
+rb_get_iseq_body_catch_table_size(const rb_iseq_t *iseq)
+{
+    const struct iseq_catch_table *ct = iseq->body->catch_table;
+    return ct ? ct->size : 0;
+}
+
+const rb_iseq_t *
+rb_get_iseq_catch_table_entry_iseq(const rb_iseq_t *iseq, unsigned int idx)
+{
+    return iseq->body->catch_table->entries[idx].iseq;
+}
+
 struct rb_control_frame_struct *
 rb_get_ec_cfp(const rb_execution_context_t *ec)
 {

zjit/bindgen/src/main.rs

@@ -404,6 +404,8 @@ fn main() {
         .allowlist_function("rb_get_iseq_body_param_lead_num")
         .allowlist_function("rb_get_iseq_body_param_opt_num")
         .allowlist_function("rb_get_iseq_body_param_opt_table")
+        .allowlist_function("rb_get_iseq_body_catch_table_size")
+        .allowlist_function("rb_get_iseq_catch_table_entry_iseq")
         .allowlist_function("rb_get_cikw_keyword_len")
         .allowlist_function("rb_get_cikw_keywords_idx")
         .allowlist_function("rb_get_call_data_ci")

zjit/src/cruby_bindings.inc.rs

@@ -6280,6 +6280,17 @@ fn locals_written_in_block(parent_iseq: IseqPtr, blockiseq: IseqPtr, target_dept
       }
       insn_idx += insn_len(opcode as usize);
   }
+
+  // Also scan rescue/ensure handler ISEQs referenced from the catch table.
+  // These are child ISEQs (one level deeper than blockiseq) that can write
+  // to the parent's locals via setlocal with higher level values.
+  let catch_size = unsafe { rb_get_iseq_body_catch_table_size(blockiseq) };
+  for i in 0..catch_size {
+      let handler_iseq = unsafe { rb_get_iseq_catch_table_entry_iseq(blockiseq, i) };
+      if !handler_iseq.is_null() {
+          locals_written_in_block(parent_iseq, handler_iseq, target_depth + 1, written);
+      }
+  }
 }
 
 /// Reload locals that may have been modified by the blockiseq.

zjit/src/hir.rs

@@ -4257,6 +4257,89 @@ pub mod hir_build_tests {
           Return v35
         ");
     }
+
+    #[test]
+    fn reload_local_written_in_rescue() {
+        eval(r#"
+            def test
+              a = 1
+              b = 2
+              tap do
+                begin
+                  raise
+                rescue
+                  b = 3
+                end
+              end
+              a + b
+            end
+        "#);
+        assert_contains_opcode("test", YARVINSN_send);
+        assert_snapshot!(hir_string("test"), @r"
+        fn test@<compiled>:3:
+        bb0():
+          EntryPoint interpreter
+          v1:BasicObject = LoadSelf
+          v2:NilClass = Const Value(nil)
+          v3:NilClass = Const Value(nil)
+          Jump bb2(v1, v2, v3)
+        bb1(v6:BasicObject):
+          EntryPoint JIT(0)
+          v7:NilClass = Const Value(nil)
+          v8:NilClass = Const Value(nil)
+          Jump bb2(v6, v7, v8)
+        bb2(v10:BasicObject, v11:NilClass, v12:NilClass):
+          v16:Fixnum[1] = Const Value(1)
+          v20:Fixnum[2] = Const Value(2)
+          v25:BasicObject = Send v10, 0x1000, :tap # SendFallbackReason: Uncategorized(send)
+          v26:BasicObject = GetLocal :b, l0, EP@3
+          PatchPoint NoEPEscape(test)
+          v35:BasicObject = SendWithoutBlock v16, :+, v26 # SendFallbackReason: Uncategorized(opt_plus)
+          CheckInterrupts
+          Return v35
+        ");
+    }
+
+    #[test]
+    fn dont_reload_local_not_written_in_rescue() {
+        eval(r#"
+            def test
+              a = 1
+              b = 2
+              tap do
+                begin
+                  raise
+                rescue
+                  a
+                end
+              end
+              a + b
+            end
+        "#);
+        assert_contains_opcode("test", YARVINSN_send);
+        assert_snapshot!(hir_string("test"), @r"
+        fn test@<compiled>:3:
+        bb0():
+          EntryPoint interpreter
+          v1:BasicObject = LoadSelf
+          v2:NilClass = Const Value(nil)
+          v3:NilClass = Const Value(nil)
+          Jump bb2(v1, v2, v3)
+        bb1(v6:BasicObject):
+          EntryPoint JIT(0)
+          v7:NilClass = Const Value(nil)
+          v8:NilClass = Const Value(nil)
+          Jump bb2(v6, v7, v8)
+        bb2(v10:BasicObject, v11:NilClass, v12:NilClass):
+          v16:Fixnum[1] = Const Value(1)
+          v20:Fixnum[2] = Const Value(2)
+          v25:BasicObject = Send v10, 0x1000, :tap # SendFallbackReason: Uncategorized(send)
+          PatchPoint NoEPEscape(test)
+          v34:BasicObject = SendWithoutBlock v16, :+, v20 # SendFallbackReason: Uncategorized(opt_plus)
+          CheckInterrupts
+          Return v34
+        ");
+    }
  }
 
  /// Test successor and predecessor set computations.