1 files changed, 185 insertions, 19 deletions
diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
index fca8c237fd..7a9bb132aa 100644
--- a/zjit/src/hir.rs
+++ b/zjit/src/hir.rs
@@ -496,6 +496,9 @@ pub enum Insn {
     FixnumGt   { left: InsnId, right: InsnId },
     FixnumGe   { left: InsnId, right: InsnId },
 
+    // Distinct from `SendWithoutBlock` with `mid:to_s` because does not have a patch point for String to_s being redefined
+    ObjToString { val: InsnId, call_info: CallInfo, cd: *const rb_call_data, state: InsnId },
+
     /// Side-exit if val doesn't have the expected type.
     GuardType { val: InsnId, guard_type: Type, state: InsnId },
     /// Side-exit if val is not the expected VALUE.
@@ -695,6 +698,7 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> {
             Insn::ToNewArray { val, .. } => write!(f, "ToNewArray {val}"),
             Insn::ArrayExtend { left, right, .. } => write!(f, "ArrayExtend {left}, {right}"),
             Insn::ArrayPush { array, val, .. } => write!(f, "ArrayPush {array}, {val}"),
+            Insn::ObjToString { val, .. } => { write!(f, "ObjToString {val}") },
             Insn::SideExit { .. } => write!(f, "SideExit"),
             Insn::PutSpecialObject { value_type } => {
                 write!(f, "PutSpecialObject {}", value_type)
@@ -1013,6 +1017,12 @@ impl Function {
             FixnumLt { left, right } => FixnumLt { left: find!(*left), right: find!(*right) },
             FixnumLe { left, right } => FixnumLe { left: find!(*left), right: find!(*right) },
             PutSpecialObject { value_type } => PutSpecialObject { value_type: *value_type },
+            ObjToString { val, call_info, cd, state } => ObjToString {
+                val: find!(*val),
+                call_info: call_info.clone(),
+                cd: *cd,
+                state: *state,
+            },
             SendWithoutBlock { self_val, call_info, cd, args, state } => SendWithoutBlock {
                 self_val: find!(*self_val),
                 call_info: call_info.clone(),
@@ -1143,6 +1153,7 @@ impl Function {
             Insn::GetIvar { .. } => types::BasicObject,
             Insn::ToNewArray { .. } => types::ArrayExact,
             Insn::ToArray { .. } => types::ArrayExact,
+            Insn::ObjToString { .. } => types::BasicObject,
         }
     }
 
@@ -1386,6 +1397,15 @@ impl Function {
                         let replacement = self.push_insn(block, Insn::Const { val: Const::Value(unsafe { (*ice).value }) });
                         self.make_equal_to(insn_id, replacement);
                     }
+                    Insn::ObjToString { val, call_info, cd, state, .. } => {
+                        if self.is_a(val, types::StringExact) {
+                            // behaves differently from `SendWithoutBlock` with `mid:to_s` because ObjToString should not have a patch point for String to_s being redefined
+                            self.make_equal_to(insn_id, val);
+                        } else {
+                            let replacement = self.push_insn(block, Insn::SendWithoutBlock { self_val: val, call_info, cd, args: vec![], state });
+                            self.make_equal_to(insn_id, replacement)
+                        }
+                    }
                     _ => { self.push_insn_id(block, insn_id); }
                 }
             }
@@ -1758,6 +1778,10 @@ impl Function {
                     worklist.push_back(val);
                     worklist.push_back(state);
                 }
+                Insn::ObjToString { val, state, .. } => {
+                    worklist.push_back(val);
+                    worklist.push_back(state);
+                }
                 Insn::GetGlobal { state, .. } |
                 Insn::SideExit { state } => worklist.push_back(state),
             }
@@ -1768,6 +1792,67 @@ impl Function {
         }
     }
 
+    fn absorb_dst_block(&mut self, num_in_edges: &Vec<u32>, block: BlockId) -> bool {
+        let Some(terminator_id) = self.blocks[block.0].insns.last()
+            else { return false };
+        let Insn::Jump(BranchEdge { target, args }) = self.find(*terminator_id)
+            else { return false };
+        if target == block {
+            // Can't absorb self
+            return false;
+        }
+        if num_in_edges[target.0] != 1 {
+            // Can't absorb block if it's the target of more than one branch
+            return false;
+        }
+        // Link up params with block args
+        let params = std::mem::take(&mut self.blocks[target.0].params);
+        assert_eq!(args.len(), params.len());
+        for (arg, param) in args.iter().zip(params) {
+            self.make_equal_to(param, *arg);
+        }
+        // Remove branch instruction
+        self.blocks[block.0].insns.pop();
+        // Move target instructions into block
+        let target_insns = std::mem::take(&mut self.blocks[target.0].insns);
+        self.blocks[block.0].insns.extend(target_insns);
+        true
+    }
+
+    /// Clean up linked lists of blocks A -> B -> C into A (with B's and C's instructions).
+    fn clean_cfg(&mut self) {
+        // num_in_edges is invariant throughout cleaning the CFG:
+        // * we don't allocate new blocks
+        // * blocks that get absorbed are not in RPO anymore
+        // * blocks pointed to by blocks that get absorbed retain the same number of in-edges
+        let mut num_in_edges = vec![0; self.blocks.len()];
+        for block in self.rpo() {
+            for &insn in &self.blocks[block.0].insns {
+                if let Insn::IfTrue { target, .. } | Insn::IfFalse { target, .. } | Insn::Jump(target) = self.find(insn) {
+                    num_in_edges[target.target.0] += 1;
+                }
+            }
+        }
+        let mut changed = false;
+        loop {
+            let mut iter_changed = false;
+            for block in self.rpo() {
+                // Ignore transient empty blocks
+                if self.blocks[block.0].insns.is_empty() { continue; }
+                loop {
+                    let absorbed = self.absorb_dst_block(&num_in_edges, block);
+                    if !absorbed { break; }
+                    iter_changed = true;
+                }
+            }
+            if !iter_changed { break; }
+            changed = true;
+        }
+        if changed {
+            self.infer_types();
+        }
+    }
+
     /// Return a traversal of the `Function`'s `BlockId`s in reverse post-order.
     pub fn rpo(&self) -> Vec<BlockId> {
         let mut result = self.po_from(self.entry_block);
@@ -1807,6 +1892,7 @@ impl Function {
         self.optimize_direct_sends();
         self.optimize_c_calls();
         self.fold_constants();
+        self.clean_cfg();
         self.eliminate_dead_code();
 
         // Dump HIR after optimization
@@ -2405,6 +2491,14 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                 }
                 YARVINSN_pop => { state.stack_pop()?; }
                 YARVINSN_dup => { state.stack_push(state.stack_top()?); }
+                YARVINSN_dupn => {
+                    // Duplicate the top N element of the stack. As we push, n-1 naturally
+                    // points higher in the original stack.
+                    let n = get_arg(pc, 0).as_usize();
+                    for _ in 0..n {
+                        state.stack_push(state.stack_topn(n-1)?);
+                    }
+                }
                 YARVINSN_swap => {
                     let right = state.stack_pop()?;
                     let left = state.stack_pop()?;
@@ -2669,6 +2763,26 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                     let insn_id = fun.push_insn(block, Insn::InvokeBuiltin { bf, args, state: exit_id });
                     state.stack_push(insn_id);
                 }
+                YARVINSN_objtostring => {
+                    let cd: *const rb_call_data = get_arg(pc, 0).as_ptr();
+                    let call_info = unsafe { rb_get_call_data_ci(cd) };
+
+                    if unknown_call_type(unsafe { rb_vm_ci_flag(call_info) }) {
+                        assert!(false, "objtostring should not have unknown call type");
+                    }
+                    let argc = unsafe { vm_ci_argc((*cd).ci) };
+                    assert_eq!(0, argc, "objtostring should not have args");
+
+                    let method_name: String = unsafe {
+                        let mid = rb_vm_ci_mid(call_info);
+                        mid.contents_lossy().into_owned()
+                    };
+
+                    let recv = state.stack_pop()?;
+                    let exit_id = fun.push_insn(block, Insn::Snapshot { state: exit_state });
+                    let objtostring = fun.push_insn(block, Insn::ObjToString { val: recv, call_info: CallInfo { method_name }, cd, state: exit_id });
+                    state.stack_push(objtostring)
+                }
                 _ => {
                     // Unknown opcode; side-exit into the interpreter
                     let exit_id = fun.push_insn(block, Insn::Snapshot { state: exit_state });
@@ -4314,6 +4428,43 @@ mod tests {
               Return v8
         "#]]);
     }
+
+    #[test]
+    fn dupn() {
+        eval("
+            def test(x) = (x[0, 1] ||= 2)
+        ");
+        assert_method_hir_with_opcode("test", YARVINSN_dupn, expect![[r#"
+            fn test:
+            bb0(v0:BasicObject, v1:BasicObject):
+              v3:NilClassExact = Const Value(nil)
+              v4:Fixnum[0] = Const Value(0)
+              v5:Fixnum[1] = Const Value(1)
+              v7:BasicObject = SendWithoutBlock v1, :[], v4, v5
+              v8:CBool = Test v7
+              IfTrue v8, bb1(v0, v1, v3, v1, v4, v5, v7)
+              v10:Fixnum[2] = Const Value(2)
+              v12:BasicObject = SendWithoutBlock v1, :[]=, v4, v5, v10
+              Return v10
+            bb1(v14:BasicObject, v15:BasicObject, v16:NilClassExact, v17:BasicObject, v18:Fixnum[0], v19:Fixnum[1], v20:BasicObject):
+              Return v20
+        "#]]);
+    }
+
+    #[test]
+    fn test_objtostring() {
+        eval("
+            def test = \"#{1}\"
+        ");
+        assert_method_hir_with_opcode("test", YARVINSN_objtostring, expect![[r#"
+            fn test:
+            bb0(v0:BasicObject):
+              v2:StringExact[VALUE(0x1000)] = Const Value(VALUE(0x1000))
+              v3:Fixnum[1] = Const Value(1)
+              v5:BasicObject = ObjToString v3
+              SideExit
+        "#]]);
+    }
 }
 
 #[cfg(test)]
@@ -4366,9 +4517,6 @@ mod opt_tests {
         assert_optimized_method_hir("test", expect![[r#"
             fn test:
             bb0(v0:BasicObject):
-              v3:FalseClassExact = Const Value(false)
-              Jump bb1(v0, v3)
-            bb1(v8:BasicObject, v9:FalseClassExact):
               v11:Fixnum[4] = Const Value(4)
               Return v11
         "#]]);
@@ -4545,8 +4693,6 @@ mod opt_tests {
             fn test:
             bb0(v0:BasicObject):
               PatchPoint BOPRedefined(INTEGER_REDEFINED_OP_FLAG, BOP_EQ)
-              Jump bb1(v0)
-            bb1(v10:BasicObject):
               v12:Fixnum[4] = Const Value(4)
               Return v12
         "#]]);
@@ -4609,8 +4755,6 @@ mod opt_tests {
             bb0(v0:BasicObject):
               PatchPoint BOPRedefined(INTEGER_REDEFINED_OP_FLAG, BOP_EQ)
               PatchPoint BOPRedefined(INTEGER_REDEFINED_OP_FLAG, BOP_NEQ)
-              Jump bb1(v0)
-            bb1(v10:BasicObject):
               v12:Fixnum[4] = Const Value(4)
               Return v12
         "#]]);
@@ -5555,12 +5699,8 @@ mod opt_tests {
               PatchPoint StableConstantNames(0x1000, C)
               v20:BasicObject[VALUE(0x1008)] = Const Value(VALUE(0x1008))
               v4:NilClassExact = Const Value(nil)
-              Jump bb1(v0, v4, v20)
-            bb1(v6:BasicObject, v7:NilClassExact, v8:BasicObject[VALUE(0x1008)]):
-              v11:BasicObject = SendWithoutBlock v8, :new
-              Jump bb2(v6, v11, v7)
-            bb2(v13:BasicObject, v14:BasicObject, v15:NilClassExact):
-              Return v14
+              v11:BasicObject = SendWithoutBlock v20, :new
+              Return v11
         "#]]);
     }
 
@@ -5583,12 +5723,8 @@ mod opt_tests {
               v22:BasicObject[VALUE(0x1008)] = Const Value(VALUE(0x1008))
               v4:NilClassExact = Const Value(nil)
               v5:Fixnum[1] = Const Value(1)
-              Jump bb1(v0, v4, v22, v5)
-            bb1(v7:BasicObject, v8:NilClassExact, v9:BasicObject[VALUE(0x1008)], v10:Fixnum[1]):
-              v13:BasicObject = SendWithoutBlock v9, :new, v10
-              Jump bb2(v7, v13, v8)
-            bb2(v15:BasicObject, v16:BasicObject, v17:NilClassExact):
-              Return v16
+              v13:BasicObject = SendWithoutBlock v22, :new, v5
+              Return v13
         "#]]);
     }
 
@@ -5872,4 +6008,34 @@ mod opt_tests {
               Return v7
         "#]]);
     }
+
+    #[test]
+    fn test_objtostring_string() {
+        eval(r##"
+            def test = "#{('foo')}"
+        "##);
+        assert_optimized_method_hir("test", expect![[r#"
+            fn test:
+            bb0(v0:BasicObject):
+              v2:StringExact[VALUE(0x1000)] = Const Value(VALUE(0x1000))
+              v3:StringExact[VALUE(0x1008)] = Const Value(VALUE(0x1008))
+              v4:StringExact = StringCopy v3
+              SideExit
+        "#]]);
+    }
+
+    #[test]
+    fn test_objtostring_with_non_string() {
+        eval(r##"
+            def test = "#{1}"
+        "##);
+        assert_optimized_method_hir("test", expect![[r#"
+            fn test:
+            bb0(v0:BasicObject):
+              v2:StringExact[VALUE(0x1000)] = Const Value(VALUE(0x1000))
+              v3:Fixnum[1] = Const Value(1)
+              v8:BasicObject = SendWithoutBlock v3, :to_s
+              SideExit
+        "#]]);
+    }
 }