ZJIT: Specialize direct sends to methods with post-required positional parameters

Author: nirvdrum

zjit/src/codegen.rs

@@ -1671,12 +1671,14 @@ fn gen_send_iseq_direct(
         // See vm_call_iseq_setup_normal_opt_start in vm_inshelper.c
         let lead_num = params.lead_num as u32;
         let opt_num = params.opt_num as u32;
+        let post_num = params.post_num as u32;
         let keyword = params.keyword;
         let kw_total_num = if keyword.is_null() { 0 } else { unsafe { (*keyword).num } } as u32;
-        assert!(args.len() as u32 <= lead_num + opt_num + kw_total_num);
+        assert!(args.len() as u32 <= lead_num + opt_num + post_num + kw_total_num);
         // For computing optional positional entry point, only count positional args
+        // and exclude the always-present lead and post slots.
         let positional_argc = args.len() as u32 - kw_total_num;
-        let num_optionals_passed = positional_argc.saturating_sub(lead_num);
+        let num_optionals_passed = positional_argc.saturating_sub(lead_num + post_num);
         num_optionals_passed
     } else {
         0

zjit/src/hir.rs

@@ -2402,7 +2402,6 @@ fn can_direct_send(function: &mut Function, block: BlockId, iseq: *const rb_iseq
 
     use Counter::*;
     if 0 != params.flags.has_rest()    { count_failure(complex_arg_pass_param_rest) }
-    if 0 != params.flags.has_post()    { count_failure(complex_arg_pass_param_post) }
     if 0 != params.flags.forwardable() { count_failure(complex_arg_pass_param_forwardable) }
     if callee_has_block_param && caller_passes_block_arg
                                        { count_failure(complex_arg_pass_param_block) }
@@ -2423,9 +2422,9 @@ fn can_direct_send(function: &mut Function, block: BlockId, iseq: *const rb_iseq
         return false;
     }
 
-    // Because we exclude e.g. post parameters above, they are also excluded from the checks below.
     let lead_num = params.lead_num;
     let opt_num = params.opt_num;
+    let post_num = params.post_num;
     let keyword = params.keyword;
     let kw_req_num = if keyword.is_null() { 0 } else { unsafe { (*keyword).required_num } };
     let kw_total_num = if keyword.is_null() { 0 } else { unsafe { (*keyword).num } };
@@ -2441,7 +2440,7 @@ fn can_direct_send(function: &mut Function, block: BlockId, iseq: *const rb_iseq
 
     let positional_ok = c_int::try_from(caller_positional)
         .as_ref()
-        .map(|argc| (lead_num..=lead_num + opt_num).contains(argc))
+        .map(|argc| (lead_num + post_num..=lead_num + opt_num + post_num).contains(argc))
         .unwrap_or(false);
     let keyword_ok = c_int::try_from(caller_kw_count)
         .as_ref()

zjit/src/hir/opt_tests.rs

@@ -3762,12 +3762,12 @@ mod hir_opt_tests {
     }
 
     #[test]
-    fn dont_specialize_call_to_post_param_iseq() {
-        enable_zjit_stats();
+    fn specialize_call_to_post_param_iseq() {
         eval("
             def foo(opt=80, post) = post
             def test = foo(10)
             test
+            test
         ");
         assert_snapshot!(hir_string("test"), @"
         fn test@<compiled>:3:
@@ -3778,18 +3778,44 @@ mod hir_opt_tests {
         bb2():
           EntryPoint JIT(0)
           v4:BasicObject = LoadArg :self@0
-          IncrCounterPtr
           Jump bb3(v4)
-        bb3(v7:BasicObject):
-          IncrCounter zjit_insn_count
-          IncrCounter zjit_insn_count
-          v14:Fixnum[10] = Const Value(10)
-          IncrCounter zjit_insn_count
-          IncrCounter complex_arg_pass_param_post
-          v17:BasicObject = Send v7, :foo, v14 # SendFallbackReason: Complex argument passing
-          IncrCounter zjit_insn_count
+        bb3(v6:BasicObject):
+          v11:Fixnum[10] = Const Value(10)
+          PatchPoint MethodRedefined(Object@0x1000, foo@0x1008, cme:0x1010)
+          v20:ObjectSubclass[class_exact*:Object@VALUE(0x1000)] = GuardType v6, ObjectSubclass[class_exact*:Object@VALUE(0x1000)]
+          v21:BasicObject = SendDirect v20, 0x1038, :foo (0x1048), v11
           CheckInterrupts
-          Return v17
+          Return v21
+        ");
+    }
+
+    #[test]
+    fn specialize_call_to_iseq_with_optional_between_required_params() {
+        let result = eval("
+            def foo(lead, opt=80, post) = lead + opt + post
+            def test = foo(10, 20)
+            test
+            test
+        ");
+        assert_eq!(VALUE::fixnum_from_usize(110), result);
+        assert_snapshot!(hir_string("test"), @"
+        fn test@<compiled>:3:
+        bb1():
+          EntryPoint interpreter
+          v1:BasicObject = LoadSelf
+          Jump bb3(v1)
+        bb2():
+          EntryPoint JIT(0)
+          v4:BasicObject = LoadArg :self@0
+          Jump bb3(v4)
+        bb3(v6:BasicObject):
+          v11:Fixnum[10] = Const Value(10)
+          v13:Fixnum[20] = Const Value(20)
+          PatchPoint MethodRedefined(Object@0x1000, foo@0x1008, cme:0x1010)
+          v22:ObjectSubclass[class_exact*:Object@VALUE(0x1000)] = GuardType v6, ObjectSubclass[class_exact*:Object@VALUE(0x1000)]
+          v23:BasicObject = SendDirect v22, 0x1038, :foo (0x1048), v11, v13
+          CheckInterrupts
+          Return v23
         ");
     }