ZJIT: Inline << and push for Array in single arg case (ruby#14926)

Fixes Shopify#813

Author: aidenfoxivey

zjit/src/cruby_methods.rs

@@ -203,6 +203,8 @@ pub fn init() -> Annotations {
     annotate!(rb_cArray, "reverse", types::ArrayExact, leaf, elidable);
     annotate!(rb_cArray, "join", types::StringExact);
     annotate!(rb_cArray, "[]", inline_array_aref);
+    annotate!(rb_cArray, "<<", inline_array_push);
+    annotate!(rb_cArray, "push", inline_array_push);
     annotate!(rb_cHash, "[]", inline_hash_aref);
     annotate!(rb_cHash, "size", types::Fixnum, no_gc, leaf, elidable);
     annotate!(rb_cHash, "empty?", types::BoolExact, no_gc, leaf, elidable);
@@ -266,6 +268,15 @@ fn inline_array_aref(fun: &mut hir::Function, block: hir::BlockId, recv: hir::In
     None
 }
 
+fn inline_array_push(fun: &mut hir::Function, block: hir::BlockId, recv: hir::InsnId, args: &[hir::InsnId], state: hir::InsnId) -> Option<hir::InsnId> {
+    // Inline only the case of `<<` or `push` when called with a single argument.
+    if let &[val] = args {
+        let _ = fun.push_insn(block, hir::Insn::ArrayPush { array: recv, val, state });
+        return Some(recv);
+    }
+    None
+}
+
 fn inline_hash_aref(fun: &mut hir::Function, block: hir::BlockId, recv: hir::InsnId, args: &[hir::InsnId], state: hir::InsnId) -> Option<hir::InsnId> {
     if let &[key] = args  {
         let result = fun.push_insn(block, hir::Insn::HashAref { hash: recv, key, state });

zjit/src/hir.rs

@@ -14008,7 +14008,69 @@ mod opt_tests {
           PatchPoint MethodRedefined(Array@0x1000, <<@0x1008, cme:0x1010)
           PatchPoint NoSingletonClass(Array@0x1000)
           v26:ArrayExact = GuardType v9, ArrayExact
-          v27:BasicObject = CCallWithFrame <<@0x1038, v26, v13
+          ArrayPush v26, v13
+          IncrCounter inline_cfunc_optimized_send_count
+          CheckInterrupts
+          Return v26
+        ");
+    }
+
+    #[test]
+    fn test_optimize_array_push_single_arg() {
+        eval("
+            def test(arr)
+              arr.push(1)
+            end
+            test([])
+        ");
+        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):
+          v13:Fixnum[1] = Const Value(1)
+          PatchPoint MethodRedefined(Array@0x1000, push@0x1008, cme:0x1010)
+          PatchPoint NoSingletonClass(Array@0x1000)
+          v24:ArrayExact = GuardType v9, ArrayExact
+          ArrayPush v24, v13
+          IncrCounter inline_cfunc_optimized_send_count
+          CheckInterrupts
+          Return v24
+        ");
+    }
+
+    #[test]
+    fn test_do_not_optimize_array_push_multi_arg() {
+        eval("
+            def test(arr)
+              arr.push(1,2,3)
+            end
+            test([])
+        ");
+        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):
+          v13:Fixnum[1] = Const Value(1)
+          v14:Fixnum[2] = Const Value(2)
+          v15:Fixnum[3] = Const Value(3)
+          PatchPoint MethodRedefined(Array@0x1000, push@0x1008, cme:0x1010)
+          PatchPoint NoSingletonClass(Array@0x1000)
+          v26:ArrayExact = GuardType v9, ArrayExact
+          v27:BasicObject = CCallVariadic push@0x1038, v26, v13, v14, v15
           CheckInterrupts
           Return v27
         ");