ZJIT: Pessimize locals in the presence of send (with block) (ruby#14374)

We can refine this later by some kind of analysis of the block we're
sending to: maybe it doesn't write to our locals, or at least doesn't
write to all of them.

Author: tekknolagi

zjit/src/hir.rs

@@ -2661,10 +2661,16 @@ fn insn_idx_at_offset(idx: u32, offset: i64) -> u32 {
     ((idx as isize) + (offset as isize)) as u32
 }
 
-fn compute_jump_targets(iseq: *const rb_iseq_t) -> Vec<u32> {
+struct BytecodeInfo {
+    jump_targets: Vec<u32>,
+    has_send: bool,
+}
+
+fn compute_bytecode_info(iseq: *const rb_iseq_t) -> BytecodeInfo {
     let iseq_size = unsafe { get_iseq_encoded_size(iseq) };
     let mut insn_idx = 0;
     let mut jump_targets = HashSet::new();
+    let mut has_send = false;
     while insn_idx < iseq_size {
         // Get the current pc and opcode
         let pc = unsafe { rb_iseq_pc_at_idx(iseq, insn_idx) };
@@ -2688,12 +2694,13 @@ fn compute_jump_targets(iseq: *const rb_iseq_t) -> Vec<u32> {
                     jump_targets.insert(insn_idx);
                 }
             }
+            YARVINSN_send => has_send = true,
             _ => {}
         }
     }
     let mut result = jump_targets.into_iter().collect::<Vec<_>>();
     result.sort();
-    result
+    BytecodeInfo { jump_targets: result, has_send }
 }
 
 #[derive(Debug, PartialEq)]
@@ -2800,7 +2807,7 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
     let mut profiles = ProfileOracle::new(payload);
     let mut fun = Function::new(iseq);
     // Compute a map of PC->Block by finding jump targets
-    let jump_targets = compute_jump_targets(iseq);
+    let BytecodeInfo { jump_targets, has_send } = compute_bytecode_info(iseq);
     let mut insn_idx_to_block = HashMap::new();
     for insn_idx in jump_targets {
         if insn_idx == 0 {
@@ -3120,7 +3127,7 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                 }
                 YARVINSN_getlocal_WC_0 => {
                     let ep_offset = get_arg(pc, 0).as_u32();
-                    if iseq_type == ISEQ_TYPE_EVAL {
+                    if iseq_type == ISEQ_TYPE_EVAL || has_send {
                         // On eval, the locals are always on the heap, so read the local using EP.
                         state.stack_push(fun.push_insn(block, Insn::GetLocal { ep_offset, level: 0 }));
                     } else {
@@ -3138,7 +3145,7 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                     let ep_offset = get_arg(pc, 0).as_u32();
                     let val = state.stack_pop()?;
                     state.setlocal(ep_offset, val);
-                    if iseq_type == ISEQ_TYPE_EVAL {
+                    if iseq_type == ISEQ_TYPE_EVAL || has_send {
                         // On eval, the locals are always on the heap, so write the local using EP.
                         fun.push_insn(block, Insn::SetLocal { val, ep_offset, level: 0 });
                     }
@@ -4674,9 +4681,10 @@ mod tests {
         assert_snapshot!(hir_string("test"), @r"
         fn test@<compiled>:3:
         bb0(v0:BasicObject, v1:BasicObject):
-          v4:BasicObject = Send v1, 0x1000, :each
+          v3:BasicObject = GetLocal l0, EP@3
+          v5:BasicObject = Send v3, 0x1000, :each
           CheckInterrupts
-          Return v4
+          Return v5
         ");
     }
 
@@ -4742,11 +4750,12 @@ mod tests {
         eval("
             def test(a) = foo(&a)
         ");
-        assert_snapshot!(hir_string("test"), @r#"
-            fn test@<compiled>:2:
-            bb0(v0:BasicObject, v1:BasicObject):
-              SideExit UnknownCallType
-        "#);
+        assert_snapshot!(hir_string("test"), @r"
+        fn test@<compiled>:2:
+        bb0(v0:BasicObject, v1:BasicObject):
+          v3:BasicObject = GetLocal l0, EP@3
+          SideExit UnknownCallType
+        ");
     }
 
     #[test]
@@ -7195,6 +7204,30 @@ mod opt_tests {
         ");
     }
 
+    #[test]
+    fn reload_local_across_send() {
+        eval("
+            def foo(&block) = 1
+            def test
+              a = 1
+              foo {|| }
+              a
+            end
+            test
+            test
+        ");
+        assert_snapshot!(hir_string("test"), @r"
+        fn test@<compiled>:4:
+        bb0(v0:BasicObject):
+          v3:Fixnum[1] = Const Value(1)
+          SetLocal l0, EP@3, v3
+          v6:BasicObject = Send v0, 0x1000, :foo
+          v7:BasicObject = GetLocal l0, EP@3
+          CheckInterrupts
+          Return v7
+        ");
+    }
+
     #[test]
     fn dont_specialize_call_to_iseq_with_rest() {
         eval("