ZJIT: Implement include_p for opt_(new|dup)array_send YARV insns (#14885)

These just call to the C functions that do the optimized test but this avoids the side exit.

See #12123 for the original CRuby/YJIT implementation.

Author: rwstauner

test/ruby/test_zjit.rb

@@ -957,6 +957,24 @@ def test
     }, insns: [:opt_new], call_threshold: 2
   end
 
+  def test_opt_newarray_send_include_p
+    assert_compiles '[true, false]', %q{
+      def test(x)
+        [:y, 1, Object.new].include?(x)
+      end
+      [test(1), test("n")]
+    }, insns: [:opt_newarray_send], call_threshold: 1
+  end
+
+  def test_opt_duparray_send_include_p
+    assert_compiles '[true, false]', %q{
+      def test(x)
+        [:y, 1].include?(x)
+      end
+      [test(1), test("n")]
+    }, insns: [:opt_duparray_send], call_threshold: 1
+  end
+
   def test_new_hash_empty
     assert_compiles '{}', %q{
       def test = {}

zjit/src/codegen.rs

@@ -450,6 +450,8 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         Insn::LoadSelf => gen_load_self(),
         &Insn::LoadField { recv, id, offset, return_type: _ } => gen_load_field(asm, opnd!(recv), id, offset),
         &Insn::IsBlockGiven => gen_is_block_given(jit, asm),
+        Insn::ArrayInclude { elements, target, state } => gen_array_include(jit, asm, opnds!(elements), opnd!(target), &function.frame_state(*state)),
+        &Insn::DupArrayInclude { ary, target, state } => gen_dup_array_include(jit, asm, ary, opnd!(target), &function.frame_state(state)),
         &Insn::ArrayMax { state, .. }
         | &Insn::FixnumDiv { state, .. }
         | &Insn::Throw { state, .. }
@@ -1328,6 +1330,50 @@ fn gen_array_length(asm: &mut Assembler, array: Opnd) -> lir::Opnd {
     asm_ccall!(asm, rb_jit_array_len, array)
 }
 
+fn gen_array_include(
+    jit: &JITState,
+    asm: &mut Assembler,
+    elements: Vec<Opnd>,
+    target: Opnd,
+    state: &FrameState,
+) -> lir::Opnd {
+    gen_prepare_non_leaf_call(jit, asm, state);
+
+    let num: c_long = elements.len().try_into().expect("Unable to fit length of elements into c_long");
+
+    // After gen_prepare_non_leaf_call, the elements are spilled to the Ruby stack.
+    // The elements are at the bottom of the virtual stack, followed by the target.
+    // Get a pointer to the first element on the Ruby stack.
+    let stack_bottom = state.stack().len() - elements.len() - 1;
+    let elements_ptr = asm.lea(Opnd::mem(64, SP, stack_bottom as i32 * SIZEOF_VALUE_I32));
+
+    unsafe extern "C" {
+        fn rb_vm_opt_newarray_include_p(ec: EcPtr, num: c_long, elts: *const VALUE, target: VALUE) -> VALUE;
+    }
+    asm.ccall(
+        rb_vm_opt_newarray_include_p as *const u8,
+        vec![EC, num.into(), elements_ptr, target],
+    )
+}
+
+fn gen_dup_array_include(
+    jit: &JITState,
+    asm: &mut Assembler,
+    ary: VALUE,
+    target: Opnd,
+    state: &FrameState,
+) -> lir::Opnd {
+    gen_prepare_non_leaf_call(jit, asm, state);
+
+    unsafe extern "C" {
+        fn rb_vm_opt_duparray_include_p(ec: EcPtr, ary: VALUE, target: VALUE) -> VALUE;
+    }
+    asm.ccall(
+        rb_vm_opt_duparray_include_p as *const u8,
+        vec![EC, ary.into(), target],
+    )
+}
+
 /// Compile a new hash instruction
 fn gen_new_hash(
     jit: &mut JITState,

zjit/src/hir.rs

@@ -457,6 +457,7 @@ impl PtrPrintMap {
 #[derive(Debug, Clone, Copy)]
 pub enum SideExitReason {
     UnknownNewarraySend(vm_opt_newarray_send_type),
+    UnknownDuparraySend(u64),
     UnknownSpecialVariable(u64),
     UnhandledHIRInsn(InsnId),
     UnhandledYARVInsn(u32),
@@ -548,6 +549,7 @@ impl std::fmt::Display for SideExitReason {
             SideExitReason::UnknownNewarraySend(VM_OPT_NEWARRAY_SEND_PACK) => write!(f, "UnknownNewarraySend(PACK)"),
             SideExitReason::UnknownNewarraySend(VM_OPT_NEWARRAY_SEND_PACK_BUFFER) => write!(f, "UnknownNewarraySend(PACK_BUFFER)"),
             SideExitReason::UnknownNewarraySend(VM_OPT_NEWARRAY_SEND_INCLUDE_P) => write!(f, "UnknownNewarraySend(INCLUDE_P)"),
+            SideExitReason::UnknownDuparraySend(method_id) => write!(f, "UnknownDuparraySend({})", method_id),
             SideExitReason::GuardType(guard_type) => write!(f, "GuardType({guard_type})"),
             SideExitReason::GuardTypeNot(guard_type) => write!(f, "GuardTypeNot({guard_type})"),
             SideExitReason::GuardBitEquals(value) => write!(f, "GuardBitEquals({})", value.print(&PtrPrintMap::identity())),
@@ -616,6 +618,8 @@ pub enum Insn {
     NewRangeFixnum { low: InsnId, high: InsnId, flag: RangeType, state: InsnId },
     ArrayDup { val: InsnId, state: InsnId },
     ArrayMax { elements: Vec<InsnId>, state: InsnId },
+    ArrayInclude { elements: Vec<InsnId>, target: InsnId, state: InsnId },
+    DupArrayInclude { ary: VALUE, target: InsnId, state: InsnId },
     /// Extend `left` with the elements from `right`. `left` and `right` must both be `Array`.
     ArrayExtend { left: InsnId, right: InsnId, state: InsnId },
     /// Push `val` onto `array`, where `array` is already `Array`.
@@ -988,6 +992,18 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> {
                 }
                 Ok(())
             }
+            Insn::ArrayInclude { elements, target, .. } => {
+                write!(f, "ArrayInclude")?;
+                let mut prefix = " ";
+                for element in elements {
+                    write!(f, "{prefix}{element}")?;
+                    prefix = ", ";
+                }
+                write!(f, " | {target}")
+            }
+            Insn::DupArrayInclude { ary, target, .. } => {
+                write!(f, "DupArrayInclude {} | {}", ary.print(self.ptr_map), target)
+            }
             Insn::ArrayDup { val, .. } => { write!(f, "ArrayDup {val}") }
             Insn::HashDup { val, .. } => { write!(f, "HashDup {val}") }
             Insn::HashAref { hash, key, .. } => { write!(f, "HashAref {hash}, {key}")}
@@ -1789,6 +1805,8 @@ impl Function {
             &ArrayPop { array, state } => ArrayPop { array: find!(array), state: find!(state) },
             &ArrayLength { array } => ArrayLength { array: find!(array) },
             &ArrayMax { ref elements, state } => ArrayMax { elements: find_vec!(elements), state: find!(state) },
+            &ArrayInclude { ref elements, target, state } => ArrayInclude { elements: find_vec!(elements), target: find!(target), state: find!(state) },
+            &DupArrayInclude { ary, target, state } => DupArrayInclude { ary, target: find!(target), 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 },
             &LoadField { recv, id, offset, return_type } => LoadField { recv: find!(recv), id, offset, return_type },
@@ -1923,6 +1941,8 @@ impl Function {
             Insn::GetConstantPath { .. } => types::BasicObject,
             Insn::IsBlockGiven => types::BoolExact,
             Insn::ArrayMax { .. } => types::BasicObject,
+            Insn::ArrayInclude { .. } => types::BoolExact,
+            Insn::DupArrayInclude { .. } => types::BoolExact,
             Insn::GetGlobal { .. } => types::BasicObject,
             Insn::GetIvar { .. } => types::BasicObject,
             Insn::LoadPC => types::CPtr,
@@ -3211,6 +3231,15 @@ impl Function {
                 worklist.extend(elements);
                 worklist.push_back(state);
             }
+            &Insn::ArrayInclude { ref elements, target, state } => {
+                worklist.extend(elements);
+                worklist.push_back(target);
+                worklist.push_back(state);
+            }
+            &Insn::DupArrayInclude { target, state, .. } => {
+                worklist.push_back(target);
+                worklist.push_back(state);
+            }
             &Insn::NewRange { low, high, state, .. }
             | &Insn::NewRangeFixnum { low, high, state, .. } => {
                 worklist.push_back(low);
@@ -4448,6 +4477,11 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                     let exit_id = fun.push_insn(block, Insn::Snapshot { state: exit_state });
                     let (bop, insn) = match method {
                         VM_OPT_NEWARRAY_SEND_MAX => (BOP_MAX, Insn::ArrayMax { elements, state: exit_id }),
+                        VM_OPT_NEWARRAY_SEND_INCLUDE_P => {
+                            let target = elements[elements.len() - 1];
+                            let array_elements = elements[..elements.len() - 1].to_vec();
+                            (BOP_INCLUDE_P, Insn::ArrayInclude { elements: array_elements, target, state: exit_id })
+                        },
                         _ => {
                             // Unknown opcode; side-exit into the interpreter
                             fun.push_insn(block, Insn::SideExit { state: exit_id, reason: SideExitReason::UnknownNewarraySend(method) });
@@ -4468,6 +4502,30 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                     let insn_id = fun.push_insn(block, Insn::ArrayDup { val, state: exit_id });
                     state.stack_push(insn_id);
                 }
+                YARVINSN_opt_duparray_send => {
+                    let ary = get_arg(pc, 0);
+                    let method_id = get_arg(pc, 1).as_u64();
+                    let argc = get_arg(pc, 2).as_usize();
+                    if argc != 1 {
+                        break;
+                    }
+                    let target = state.stack_pop()?;
+                    let exit_id = fun.push_insn(block, Insn::Snapshot { state: exit_state });
+                    let bop = match method_id {
+                        x if x == ID!(include_p).0 => BOP_INCLUDE_P,
+                        _ => {
+                            fun.push_insn(block, Insn::SideExit { state: exit_id, reason: SideExitReason::UnknownDuparraySend(method_id) });
+                            break;
+                        },
+                    };
+                    if !unsafe { rb_BASIC_OP_UNREDEFINED_P(bop, ARRAY_REDEFINED_OP_FLAG) } {
+                        fun.push_insn(block, Insn::SideExit { state: exit_id, reason: SideExitReason::PatchPoint(Invariant::BOPRedefined { klass: ARRAY_REDEFINED_OP_FLAG, bop }) });
+                        break;
+                    }
+                    fun.push_insn(block, Insn::PatchPoint { invariant: Invariant::BOPRedefined { klass: ARRAY_REDEFINED_OP_FLAG, bop }, state: exit_id });
+                    let insn_id = fun.push_insn(block, Insn::DupArrayInclude { ary, target, state: exit_id });
+                    state.stack_push(insn_id);
+                }
                 YARVINSN_newhash => {
                     let count = get_arg(pc, 0).as_usize();
                     assert!(count % 2 == 0, "newhash count should be even");

zjit/src/hir/tests.rs

@@ -2023,7 +2023,70 @@ pub mod hir_build_tests {
           Jump bb2(v8, v9, v10, v11, v12)
         bb2(v14:BasicObject, v15:BasicObject, v16:BasicObject, v17:NilClass, v18:NilClass):
           v25:BasicObject = SendWithoutBlock v15, :+, v16
-          SideExit UnknownNewarraySend(INCLUDE_P)
+          PatchPoint BOPRedefined(ARRAY_REDEFINED_OP_FLAG, 33)
+          v30:BoolExact = ArrayInclude v15, v16 | v16
+          PatchPoint NoEPEscape(test)
+          v35:ArrayExact[VALUE(0x1000)] = Const Value(VALUE(0x1000))
+          v37:ArrayExact = ArrayDup v35
+          v39:BasicObject = SendWithoutBlock v14, :puts, v37
+          PatchPoint NoEPEscape(test)
+          CheckInterrupts
+          Return v30
+        ");
+    }
+
+    #[test]
+    fn test_opt_duparray_send_include_p() {
+        eval("
+            def test(x)
+              [:a, :b].include?(x)
+            end
+        ");
+        assert_contains_opcode("test", YARVINSN_opt_duparray_send);
+        assert_snapshot!(hir_string("test"), @r"
+        fn test@<compiled>:3:
+        bb0():
+          EntryPoint interpreter
+          v1:BasicObject = LoadSelf
+          v2:BasicObject = GetLocal l0, SP@4
+          Jump bb2(v1, v2)
+        bb1(v5:BasicObject, v6:BasicObject):
+          EntryPoint JIT(0)
+          Jump bb2(v5, v6)
+        bb2(v8:BasicObject, v9:BasicObject):
+          PatchPoint BOPRedefined(ARRAY_REDEFINED_OP_FLAG, 33)
+          v15:BoolExact = DupArrayInclude VALUE(0x1000) | v9
+          CheckInterrupts
+          Return v15
+        ");
+    }
+
+    #[test]
+    fn test_opt_duparray_send_include_p_redefined() {
+        eval("
+            class Array
+              alias_method :old_include?, :include?
+              def include?(x)
+                old_include?(x)
+              end
+            end
+            def test(x)
+              [:a, :b].include?(x)
+            end
+        ");
+        assert_contains_opcode("test", YARVINSN_opt_duparray_send);
+        assert_snapshot!(hir_string("test"), @r"
+        fn test@<compiled>:9:
+        bb0():
+          EntryPoint interpreter
+          v1:BasicObject = LoadSelf
+          v2:BasicObject = GetLocal l0, SP@4
+          Jump bb2(v1, v2)
+        bb1(v5:BasicObject, v6:BasicObject):
+          EntryPoint JIT(0)
+          Jump bb2(v5, v6)
+        bb2(v8:BasicObject, v9:BasicObject):
+          SideExit PatchPoint(BOPRedefined(ARRAY_REDEFINED_OP_FLAG, 33))
         ");
     }
 

zjit/src/stats.rs

@@ -128,6 +128,7 @@ make_counters! {
         // exit_: Side exits reasons
         exit_compile_error,
         exit_unknown_newarray_send,
+        exit_unknown_duparray_send,
         exit_unhandled_tailcall,
         exit_unhandled_splat,
         exit_unhandled_kwarg,
@@ -366,6 +367,7 @@ pub fn side_exit_counter(reason: crate::hir::SideExitReason) -> Counter {
     use crate::stats::Counter::*;
     match reason {
         UnknownNewarraySend(_)        => exit_unknown_newarray_send,
+        UnknownDuparraySend(_)        => exit_unknown_duparray_send,
         UnhandledCallType(Tailcall)   => exit_unhandled_tailcall,
         UnhandledCallType(Splat)      => exit_unhandled_splat,
         UnhandledCallType(Kwarg)      => exit_unhandled_kwarg,