ZJIT: Implement concatstrings

Quite tricky to use the native stack to pass a C `VALUE[]`.

Author: XrXr

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 = ""

zjit/src/codegen.rs

@@ -299,6 +299,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         Insn::SideExit { state, reason: _ } => return gen_side_exit(jit, asm, &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}");
@@ -314,6 +315,48 @@ 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;
+    }
+
+    // Find how much to grow the stack, aligning to 16 for calling convention
+    let mut stack_growth = strings.len() * SIZEOF_VALUE;
+    if stack_growth % 16 != 0 {
+        stack_growth += 16 - stack_growth % 16;
+    }
+
+    // Grow the stack and 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 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 {

zjit/src/hir.rs

@@ -529,6 +529,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 },
@@ -705,7 +706,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.
@@ -1088,6 +1096,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(),
@@ -1219,6 +1228,7 @@ impl Function {
             Insn::ToArray { .. } => types::ArrayExact,
             Insn::ObjToString { .. } => types::BasicObject,
             Insn::AnyToString { .. } => types::String,
+            Insn::ConcatStrings { .. } => types::StringExact,
             Insn::GetLocal { .. } => types::BasicObject,
         }
     }
@@ -1778,7 +1788,8 @@ impl Function {
                     worklist.push_back(state);
                 }
                 Insn::ArrayMax { elements, state }
-                | Insn::NewArray { elements, state } => {
+                | Insn::NewArray { elements, state }
+                | Insn::ConcatStrings { strings: elements, state } => {
                     worklist.extend(elements);
                     worklist.push_back(state);
                 }
@@ -2939,6 +2950,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 });
@@ -4664,7 +4685,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
         "#]]);
     }
 
@@ -6305,7 +6327,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
         "#]]);
     }
 
@@ -6319,9 +6342,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
         "#]]);
     }