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/backend/arm64/mod.rs

@@ -32,6 +32,7 @@ pub const _C_RET_OPND: Opnd = Opnd::Reg(X0_REG);
 // These constants define the way we work with Arm64's stack pointer. The stack
 // pointer always needs to be aligned to a 16-byte boundary.
 pub const C_SP_REG: A64Opnd = X31;
+pub const C_SP_OPND: Opnd = Opnd::Reg(XZR_REG);
 pub const C_SP_STEP: i32 = 16;
 
 impl CodeBlock {

zjit/src/backend/lir.rs

@@ -16,7 +16,7 @@ pub const SP: Opnd = _SP;
 
 pub const C_ARG_OPNDS: [Opnd; 6] = _C_ARG_OPNDS;
 pub const C_RET_OPND: Opnd = _C_RET_OPND;
-pub use crate::backend::current::{Reg, C_RET_REG};
+pub use crate::backend::current::{Reg, C_RET_REG, C_SP_OPND};
 
 // Memory operand base
 #[derive(Clone, Copy, PartialEq, Eq, Debug)]

zjit/src/backend/x86_64/mod.rs

@@ -28,6 +28,7 @@ pub const _C_ARG_OPNDS: [Opnd; 6] = [
 // C return value register on this platform
 pub const C_RET_REG: Reg = RAX_REG;
 pub const _C_RET_OPND: Opnd = Opnd::Reg(RAX_REG);
+pub const C_SP_OPND: Opnd = Opnd::Reg(RSP_REG);
 
 impl CodeBlock {
     // The number of bytes that are generated by jmp_ptr

zjit/src/codegen.rs

@@ -6,7 +6,7 @@ use crate::backend::current::{Reg, ALLOC_REGS};
 use crate::profile::get_or_create_iseq_payload;
 use crate::state::ZJITState;
 use crate::{asm::CodeBlock, cruby::*, options::debug, virtualmem::CodePtr};
-use crate::backend::lir::{self, asm_comment, Assembler, Opnd, Target, CFP, C_ARG_OPNDS, C_RET_OPND, EC, SP};
+use crate::backend::lir::{self, asm_comment, Assembler, Opnd, Target, CFP, C_ARG_OPNDS, C_SP_OPND, C_RET_OPND, EC, SP};
 use crate::hir::{iseq_to_hir, Block, BlockId, BranchEdge, CallInfo, RangeType, SELF_PARAM_IDX, SpecialObjectType};
 use crate::hir::{Const, FrameState, Function, Insn, InsnId};
 use crate::hir_type::{types::Fixnum, Type};
@@ -278,6 +278,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         Insn::SideExit { state } => 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}");
@@ -293,6 +294,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, C_SP_OPND, 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(C_SP_OPND, stack_growth.into());
+    asm.mov(C_SP_OPND, 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 {
+        C_SP_OPND
+    } else {
+        asm.add(C_SP_OPND, 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(C_SP_OPND, stack_growth.into());
+    asm.mov(C_SP_OPND, 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

@@ -507,6 +507,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 },
@@ -683,7 +684,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.
@@ -1066,6 +1074,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(),
@@ -1197,6 +1206,7 @@ impl Function {
             Insn::ToArray { .. } => types::ArrayExact,
             Insn::ObjToString { .. } => types::BasicObject,
             Insn::AnyToString { .. } => types::String,
+            Insn::ConcatStrings { .. } => types::StringExact,
             Insn::GetLocal { .. } => types::BasicObject,
         }
     }
@@ -1756,7 +1766,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);
                 }
@@ -2917,6 +2928,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 });
@@ -4642,7 +4663,8 @@ mod tests {
               v3:Fixnum[1] = Const Value(1)
               v5:BasicObject = ObjToString v3
               v7:String = AnyToString v3, str: v5
-              SideExit
+              v9:StringExact = ConcatStrings v2, v7
+              Return v9
         "#]]);
     }
 
@@ -6220,7 +6242,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
+              v10:StringExact = ConcatStrings v2, v4
+              Return v10
         "#]]);
     }
 
@@ -6234,9 +6257,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
+              v11:BasicObject = SendWithoutBlock v3, :to_s
+              v7:String = AnyToString v3, str: v11
+              v9:StringExact = ConcatStrings v2, v7
+              Return v9
         "#]]);
     }