ZJIT: Resurrect string key of opt_aref_with

The immediate problem this fixes is the frozen status of the key object
in the `recv[key]` call.

This also repairs an invariant that `FrameState` should always be at a
YARV instruction boundary. It's important to have this because if a
`SendWithoutBlock` is reduced to `SendWithoutBlockDirect`, the direct
variant may want to side exit, which needs a `FrameState` as described.

Author: XrXr

test/ruby/test_zjit.rb

@@ -1058,10 +1058,16 @@ def foo(#{'_,' * 39} n40) = n40
 
   def test_opt_aref_with
     assert_compiles ':ok', %q{
-      def aref_with(hash) = hash["key"]
+      def test(hash) = hash["key"]
 
-      aref_with({ "key" => :ok })
-    }
+      test({ "key" => :ok })
+    }, insns: [:opt_aref_with]
+
+    assert_compiles 'nil', %q{
+      def self.[](key) = (raise if key.frozen?)
+      def test = self["frozen string literal: false"]
+      test
+    }, insns: [:opt_aref_with]
   end
 
   def test_putself

zjit/bindgen/src/main.rs

@@ -246,6 +246,7 @@ fn main() {
         .allowlist_function("rb_obj_as_string_result")
         .allowlist_function("rb_str_byte_substr")
         .allowlist_function("rb_str_substr_two_fixnums")
+        .allowlist_function("rb_str_resurrect")
 
         // From include/ruby/internal/intern/parse.h
         .allowlist_function("rb_backref_get")

zjit/src/codegen.rs

@@ -370,6 +370,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         Insn::InvokeBuiltin { bf, .. } if bf.argc + 2 > (C_ARG_OPNDS.len() as i32) => return None,
         Insn::InvokeBuiltin { bf, args, state, .. } => gen_invokebuiltin(jit, asm, &function.frame_state(*state), bf, opnds!(args)),
         Insn::Return { val } => no_output!(gen_return(asm, opnd!(val))),
+        &Insn::ArefWith { receiver: _, key, cd, send_state } => gen_opt_aref_with(jit, asm, cd, opnd!(key), &function.frame_state(send_state)),
         Insn::FixnumAdd { left, right, state } => gen_fixnum_add(jit, asm, opnd!(left), opnd!(right), &function.frame_state(*state)),
         Insn::FixnumSub { left, right, state } => gen_fixnum_sub(jit, asm, opnd!(left), opnd!(right), &function.frame_state(*state)),
         Insn::FixnumMult { left, right, state } => gen_fixnum_mult(jit, asm, opnd!(left), opnd!(right), &function.frame_state(*state)),
@@ -861,16 +862,17 @@ fn gen_if_false(jit: &mut JITState, asm: &mut Assembler, val: lir::Opnd, branch:
     asm.write_label(if_true);
 }
 
+unsafe extern "C" {
+    fn rb_vm_opt_send_without_block(ec: EcPtr, cfp: CfpPtr, cd: VALUE) -> VALUE;
+}
+
 /// Compile a dynamic dispatch without block
 fn gen_send_without_block(
     jit: &mut JITState,
     asm: &mut Assembler,
     cd: *const rb_call_data,
     state: &FrameState,
 ) -> lir::Opnd {
-    // Note that it's incorrect to use this frame state to side exit because
-    // the state might not be on the boundary of an interpreter instruction.
-    // For example, `opt_aref_with` pushes to the stack and then sends.
     asm_comment!(asm, "spill frame state");
 
     // Save PC and SP
@@ -882,9 +884,6 @@ fn gen_send_without_block(
     gen_spill_stack(jit, asm, state);
 
     asm_comment!(asm, "call #{} with dynamic dispatch", ruby_call_method_name(cd));
-    unsafe extern "C" {
-        fn rb_vm_opt_send_without_block(ec: EcPtr, cfp: CfpPtr, cd: VALUE) -> VALUE;
-    }
     let ret = asm.ccall(
         rb_vm_opt_send_without_block as *const u8,
         vec![EC, CFP, (cd as usize).into()],
@@ -960,6 +959,31 @@ fn gen_send_without_block_direct(
     ret
 }
 
+/// Compile the equivalent of the general path of `opt_aref_with`
+fn gen_opt_aref_with(jit: &mut JITState, asm: &mut Assembler, cd: *const rb_call_data, key: Opnd, state: &FrameState) -> Opnd {
+    // First, resurrect the string key. Save PC before allocating.
+    gen_save_pc(asm, state);
+    let key = asm_ccall!(asm, rb_str_resurrect, key);
+
+    // Put resurrected key at the top of the stack
+    asm.mov(Opnd::mem(VALUE_BITS, SP, state.stack().len() as i32 * SIZEOF_VALUE_I32), key);
+
+    // Spill the rest of the frame
+    gen_spill_stack(jit, asm, state);
+    gen_save_sp(asm, state.stack().len() + 1); // +1 for the key
+    gen_spill_locals(jit, asm, state);
+
+    asm_comment!(asm, "call [] with dynamic dispatch");
+    let ret = asm.ccall(
+        rb_vm_opt_send_without_block as *const u8,
+        vec![EC, CFP, (cd as usize).into()],
+    );
+    // TODO(max): Add a PatchPoint here that can side-exit the function if the callee messed with
+    // the frame's locals
+
+    ret
+}
+
 /// Compile a string resurrection
 fn gen_string_copy(asm: &mut Assembler, recv: Opnd, chilled: bool, state: &FrameState) -> Opnd {
     // TODO: split rb_ec_str_resurrect into separate functions

zjit/src/cruby_bindings.inc.rs

@@ -564,6 +564,9 @@ pub enum Insn {
         return_type: Option<Type>,  // None for unannotated builtins
     },
 
+    /// Parallel of `opt_aref_with` in the interpreter.
+    ArefWith { receiver: InsnId, key: InsnId, send_state: InsnId, cd: *const rb_call_data },
+
     /// Control flow instructions
     Return { val: InsnId },
     /// Non-local control flow. See the throw YARV instruction
@@ -849,6 +852,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::ArefWith { receiver, key, .. } => write!(f, "ArefWith {receiver}, {key}"),
             Insn::ObjToString { val, .. } => { write!(f, "ObjToString {val}") },
             Insn::StringIntern { val, .. } => { write!(f, "StringIntern {val}") },
             Insn::AnyToString { val, str, .. } => { write!(f, "AnyToString {val}, str: {str}") },
@@ -1289,6 +1293,7 @@ impl Function {
                 NewHash { elements: found_elements, state: find!(state) }
             }
             &NewRange { low, high, flag, state } => NewRange { low: find!(low), high: find!(high), flag, state: find!(state) },
+            &ArefWith { receiver: self_val, key, send_state, cd } => ArefWith { receiver: find!(self_val), key: find!(key), send_state: find!(send_state), cd },
             &ArrayMax { ref elements, state } => ArrayMax { elements: find_vec!(elements), state: find!(state) },
             &SetGlobal { id, val, state } => SetGlobal { id, val: find!(val), state },
             &GetIvar { self_val, id, state } => GetIvar { self_val: find!(self_val), id, state },
@@ -1382,6 +1387,7 @@ impl Function {
             Insn::DefinedIvar { .. } => types::BasicObject,
             Insn::GetConstantPath { .. } => types::BasicObject,
             Insn::ArrayMax { .. } => types::BasicObject,
+            Insn::ArefWith { .. } => types::BasicObject,
             Insn::GetGlobal { .. } => types::BasicObject,
             Insn::GetIvar { .. } => types::BasicObject,
             Insn::GetSpecialSymbol { .. } => types::BasicObject,
@@ -2018,6 +2024,7 @@ impl Function {
             | &Insn::FixnumDiv { left, right, state }
             | &Insn::FixnumMod { left, right, state }
             | &Insn::ArrayExtend { left, right, state }
+            | &Insn::ArefWith { receiver: left, key: right, send_state: state, cd: _ }
             => {
                 worklist.push_back(left);
                 worklist.push_back(right);
@@ -3193,23 +3200,18 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                     // NB: opt_aref_with has an instruction argument for the call at get_arg(0)
                     let cd: *const rb_call_data = get_arg(pc, 1).as_ptr();
                     let call_info = unsafe { rb_get_call_data_ci(cd) };
+                    let exit_id = fun.push_insn(block, Insn::Snapshot { state: exit_state });
                     if unknown_call_type(unsafe { rb_vm_ci_flag(call_info) }) {
                         // Unknown call type; side-exit into the interpreter
-                        let exit_id = fun.push_insn(block, Insn::Snapshot { state: exit_state });
                         fun.push_insn(block, Insn::SideExit { state: exit_id, reason: SideExitReason::UnknownCallType });
                         break;  // End the block
                     }
                     let argc = unsafe { vm_ci_argc((*cd).ci) };
 
                     assert_eq!(1, argc, "opt_aref_with should only be emitted for argc=1");
-                    let aref_arg = fun.push_insn(block, Insn::Const { val: Const::Value(get_arg(pc, 0)) });
-                    let args = vec![aref_arg];
-
-                    let mut send_state = state.clone();
-                    send_state.stack_push(aref_arg);
-                    let send_state = fun.push_insn(block, Insn::Snapshot { state: send_state });
+                    let key = fun.push_insn(block, Insn::Const { val: Const::Value(get_arg(pc, 0)) });
                     let recv = state.stack_pop()?;
-                    let send = fun.push_insn(block, Insn::SendWithoutBlock { self_val: recv, cd, args, state: send_state });
+                    let send = fun.push_insn(block, Insn::ArefWith { receiver: recv, key, cd, send_state: exit_id });
                     state.stack_push(send);
                 }
                 YARVINSN_opt_neq => {
@@ -5229,8 +5231,8 @@ mod tests {
         assert_snapshot!(hir_string("test"), @r"
         fn test@<compiled>:2:
         bb0(v0:BasicObject, v1:BasicObject):
-          v3:StringExact[VALUE(0x1000)] = Const Value(VALUE(0x1000))
-          v5:BasicObject = SendWithoutBlock v1, :[], v3
+          v4:StringExact[VALUE(0x1000)] = Const Value(VALUE(0x1000))
+          v5:BasicObject = ArefWith v1, v4
           CheckInterrupts
           Return v5
         ");