ZJIT: Implement concatstrings

XrXr · XrXr · commit 9b7432c0bce3 · 2025-07-15T16:32:14.000-04:00
Quite tricky to use the native stack to pass a C `VALUE[]`.
diff --git a/test/ruby/test_zjit.rb b/test/ruby/test_zjit.rb
@@ -38,6 +38,13 @@ def test = "#{""}"
     }, insns: [:putstring]
   end
 
+  def test_concatstrings
+    assert_compiles '"Object is not Kernel"', %q{
+      def test = "#{Object} is not #{Kernel}"
+      test
+    }, insns: [:concatstrings]
+  end
+
   def test_putchilldedstring
     assert_compiles '""', %q{
       def test = ""
diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs
@@ -315,6 +315,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         Insn::SideExit { state, reason } => return gen_side_exit(jit, asm, reason, &function.frame_state(*state)),
         Insn::PutSpecialObject { value_type } => gen_putspecialobject(asm, *value_type),
         Insn::AnyToString { val, str, state } => gen_anytostring(asm, opnd!(val), opnd!(str), &function.frame_state(*state))?,
+        Insn::ConcatStrings { strings, state } => gen_concatstrings(jit, asm, strings, &function.frame_state(*state))?,
         Insn::Defined { op_type, obj, pushval, v } => gen_defined(jit, asm, *op_type, *obj, *pushval, opnd!(v))?,
         _ => {
             debug!("ZJIT: gen_function: unexpected insn {insn}");
@@ -330,6 +331,43 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
     Some(())
 }
 
+fn gen_concatstrings(jit: &JITState, asm: &mut Assembler, strings: &[InsnId], state: &FrameState) -> Option<Opnd> {
+    // In practice, concatstrings with 0 strings doesn't seem to happen.
+    if strings.len() == 0 {
+        return None;
+    }
+
+    // Put the strings onto the stack
+    for (i, string) in strings.iter().enumerate() {
+        let dest = Opnd::mem(64, NATIVE_STACK_PTR, i32::try_from(strings.len() - i).ok()? * -SIZEOF_VALUE_I32);
+        asm.mov(dest, jit.get_opnd(*string)?);
+    }
+    // Write, then move the SP, since the source might be from the stack
+    let stack_growth = aligned_stack_bytes(strings.len());
+    let new_sp = asm.sub(NATIVE_STACK_PTR, stack_growth.into());
+    asm.mov(NATIVE_STACK_PTR, new_sp);
+
+    // Save PC since the call allocates and can raise
+    gen_save_pc(asm, state);
+
+    // We've moved the SP, so the array is at SP or one element above
+    // if we grew the stack with extra space for alignment.
+    let str_array = if strings.len() % 2 == 0 {
+        NATIVE_STACK_PTR
+    } else {
+        asm.add(NATIVE_STACK_PTR, SIZEOF_VALUE.into())
+    };
+
+    // Call rb_str_concat_literals()
+    let new_string = asm.ccall(rb_str_concat_literals as _, vec![strings.len().into(), str_array]);
+
+    // Undo C SP change
+    let prev_sp = asm.add(NATIVE_STACK_PTR, stack_growth.into());
+    asm.mov(NATIVE_STACK_PTR, prev_sp);
+
+    Some(new_string)
+}
+
 /// Gets the EP of the ISeq of the containing method, or "local level".
 /// Equivalent of GET_LEP() macro.
 fn gen_get_lep(jit: &JITState, asm: &mut Assembler) -> Opnd {
diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
@@ -550,6 +550,7 @@ pub enum Insn {
     // 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 },
     AnyToString { val: InsnId, str: InsnId, state: InsnId },
+    ConcatStrings { strings: Vec<InsnId>, state: InsnId },
 
     /// Side-exit if val doesn't have the expected type.
     GuardType { val: InsnId, guard_type: Type, state: InsnId },
@@ -730,7 +731,14 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> {
                     write!(f, ", {arg}")?;
                 }
                 Ok(())
-            },
+            }
+            Insn::ConcatStrings { strings, state: _ } => {
+                write!(f, "ConcatStrings {}", strings[0])?;
+                for string in strings.iter().skip(1) {
+                    write!(f, ", {string}")?;
+                }
+                Ok(())
+            }
             Insn::Snapshot { state } => write!(f, "Snapshot {}", state),
             Insn::Defined { op_type, v, .. } => {
                 // op_type (enum defined_type) printing logic from iseq.c.
@@ -1138,6 +1146,7 @@ impl Function {
                 str: find!(*str),
                 state: *state,
             },
+            ConcatStrings { state, strings } => ConcatStrings { state: *state, strings: find_vec!(strings) },
             SendWithoutBlock { self_val, call_info, cd, args, state } => SendWithoutBlock {
                 self_val: find!(*self_val),
                 call_info: call_info.clone(),
@@ -1271,6 +1280,7 @@ impl Function {
             Insn::ToArray { .. } => types::ArrayExact,
             Insn::ObjToString { .. } => types::BasicObject,
             Insn::AnyToString { .. } => types::String,
+            Insn::ConcatStrings { .. } => types::StringExact,
             Insn::GetLocal { .. } => types::BasicObject,
             // The type of Snapshot doesn't really matter; it's never materialized. It's used only
             // as a reference for FrameState, which we use to generate side-exit code.
@@ -1902,7 +1912,8 @@ impl Function {
                 worklist.extend(args);
                 worklist.push_back(state);
             }
-            &Insn::InvokeBuiltin { ref args, state, .. } => {
+            &Insn::InvokeBuiltin { ref args, state, .. }
+            | &Insn::ConcatStrings { strings: ref args, state } => {
                 worklist.extend(args);
                 worklist.push_back(state)
             }
@@ -3162,6 +3173,16 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                     let anytostring = fun.push_insn(block, Insn::AnyToString { val, str, state: exit_id });
                     state.stack_push(anytostring);
                 }
+                YARVINSN_concatstrings => {
+                    let num = get_arg(pc, 0).as_u32();
+                    let mut strings = Vec::with_capacity(num.as_usize());
+                    for _ in 0..num {
+                        strings.push(state.stack_pop()?);
+                    }
+                    strings.reverse();
+                    let snapshot = fun.push_insn(block, Insn::Snapshot { state: exit_state });
+                    state.stack_push(fun.push_insn(block, Insn::ConcatStrings { state: snapshot, strings }));
+                }
                 _ => {
                     // Unknown opcode; side-exit into the interpreter
                     let exit_id = fun.push_insn(block, Insn::Snapshot { state: exit_state });
@@ -5046,7 +5067,8 @@ mod tests {
               v3:Fixnum[1] = Const Value(1)
               v5:BasicObject = ObjToString v3
               v7:String = AnyToString v3, str: v5
-              SideExit UnknownOpcode(concatstrings)
+              v9:StringExact = ConcatStrings v2, v7
+              Return v9
         "#]]);
     }
 
@@ -6827,7 +6849,8 @@ mod opt_tests {
               v2:StringExact[VALUE(0x1000)] = Const Value(VALUE(0x1000))
               v3:StringExact[VALUE(0x1008)] = Const Value(VALUE(0x1008))
               v4:StringExact = StringCopy v3
-              SideExit UnknownOpcode(concatstrings)
+              v10:StringExact = ConcatStrings v2, v4
+              Return v10
         "#]]);
     }
 
@@ -6841,9 +6864,10 @@ mod opt_tests {
             bb0(v0:BasicObject):
               v2:StringExact[VALUE(0x1000)] = Const Value(VALUE(0x1000))
               v3:Fixnum[1] = Const Value(1)
-              v10:BasicObject = SendWithoutBlock v3, :to_s
-              v7:String = AnyToString v3, str: v10
-              SideExit UnknownOpcode(concatstrings)
+              v11:BasicObject = SendWithoutBlock v3, :to_s
+              v7:String = AnyToString v3, str: v11
+              v9:StringExact = ConcatStrings v2, v7
+              Return v9
         "#]]);
     }
 
