ZJIT: Handle caller_kwarg in direct send when all keyword params are required

Author: rwstauner

test/ruby/test_zjit.rb

@@ -615,6 +615,15 @@ def entry = test(%(a\nb\nc))
     }, call_threshold: 2
   end
 
+  def test_send_kwarg_with_too_many_args_to_c_call
+    assert_compiles '"a b c d {kwargs: :e}"', %q{
+      def test(a:, b:, c:, d:, e:) = sprintf("%s %s %s %s %s", a, b, c, d, kwargs: e)
+      def entry = test(e: :e, d: :d, c: :c, a: :a, b: :b)
+      entry
+      entry
+    }, call_threshold: 2
+  end
+
   def test_send_kwrest
     assert_compiles '{a: 3}', %q{
       def test(**kwargs) = kwargs
@@ -633,6 +642,15 @@ def entry = test(1, c: 3)
     }, call_threshold: 2
   end
 
+  def test_send_req_opt_kwreq
+    assert_compiles '[[1, 2, 3], [-1, -2, -3]]', %q{
+      def test(a, b = 2, c:) = [a, b, c]
+      def entry = [test(1, c: 3), test(-1, -2, c: -3)] # specify all, change kw order
+      entry
+      entry
+    }, call_threshold: 2
+  end
+
   def test_send_req_opt_kwreq_kwopt
     assert_compiles '[[1, 2, 3, 4], [-1, -2, -3, -4]]', %q{
       def test(a, b = 2, c:, d: 4) = [a, b, c, d]

zjit/src/codegen.rs

@@ -383,6 +383,10 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         &Insn::SendWithoutBlock { cd, state, reason, .. } => gen_send_without_block(jit, asm, cd, &function.frame_state(state), reason),
         // Give up SendWithoutBlockDirect for 6+ args since asm.ccall() doesn't support it.
         Insn::SendWithoutBlockDirect { cd, state, args, .. } if args.len() + 1 > C_ARG_OPNDS.len() => // +1 for self
+            gen_send_without_block(jit, asm, *cd, &function.frame_state(*state), SendFallbackReason::TooManyArgsForLir),
+        // Give up SendWithoutBlockDirect for 5+ args (plus 1 arg for keyword bits) since asm.ccall() doesn't support it.
+        Insn::SendWithoutBlockDirect { cd, state, args, iseq, .. } if args.len() + 2 > C_ARG_OPNDS.len() && unsafe { rb_get_iseq_flags_has_kw(*iseq) } => // +1 for self +1 for keyword bits
+
             gen_send_without_block(jit, asm, *cd, &function.frame_state(*state), SendFallbackReason::TooManyArgsForLir),
         Insn::SendWithoutBlockDirect { cme, iseq, recv, args, state, .. } => gen_send_without_block_direct(cb, jit, asm, *cme, *iseq, opnd!(recv), opnds!(args), &function.frame_state(*state)),
         &Insn::InvokeSuper { cd, blockiseq, state, reason, .. } => gen_invokesuper(jit, asm, cd, blockiseq, &function.frame_state(state), reason),
@@ -1350,13 +1354,23 @@ fn gen_send_without_block_direct(
     let mut c_args = vec![recv];
     c_args.extend(&args);
 
+    if unsafe { rb_get_iseq_flags_has_kw(iseq) } {
+        // Currently we only get to this point if all the accepted keyword args are required.
+        let unspecified_bits = 0;
+        // For each optional keyword that isn't passed we would `unspecified_bits |= (0x01 << idx)`.
+        c_args.push(VALUE::fixnum_from_usize(unspecified_bits).into());
+    }
+
     let params = unsafe { iseq.params() };
     let num_optionals_passed = if params.flags.has_opt() != 0 {
         // 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;
-        assert!(args.len() as u32 <= lead_num + opt_num);
-        let num_optionals_passed = args.len() as u32 - lead_num;
+        let keyword = params.keyword;
+        let kw_req_num = if keyword.is_null() { 0 } else { unsafe { (*keyword).required_num } } as u32;
+        let req_num = lead_num + kw_req_num;
+        assert!(args.len() as u32 <= req_num + opt_num);
+        let num_optionals_passed = args.len() as u32 - req_num;
         num_optionals_passed
     } else {
         0

zjit/src/hir.rs

@@ -615,6 +615,10 @@ pub enum SendFallbackReason {
     SendWithoutBlockNotOptimizedMethodTypeOptimized(OptimizedMethodType),
     SendWithoutBlockBopRedefined,
     SendWithoutBlockOperandsNotFixnum,
+    SendWithoutBlockDirectKeywordMismatch,
+    SendWithoutBlockDirectOptionalKeywords,
+    SendWithoutBlockDirectKeywordCountMismatch,
+    SendWithoutBlockDirectMissingKeyword,
     SendPolymorphic,
     SendMegamorphic,
     SendNoProfiles,
@@ -631,6 +635,8 @@ pub enum SendFallbackReason {
     /// The call has at least one feature on the caller or callee side that the optimizer does not
     /// support.
     ComplexArgPass,
+    /// Caller has keyword arguments but callee doesn't expect them; need to convert to hash.
+    UnexpectedKeywordArgs,
     /// Initial fallback reason for every instruction, which should be mutated to
     /// a more actionable reason when an attempt to specialize the instruction fails.
     Uncategorized(ruby_vminsn_type),
@@ -1528,7 +1534,19 @@ 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.has_kw()      { count_failure(complex_arg_pass_param_kw) }
+
+    if 0 != params.flags.has_kw() {
+        let keyword = params.keyword;
+        if !keyword.is_null() {
+            let num = unsafe { (*keyword).num };
+            let required_num = unsafe { (*keyword).required_num };
+            // Only support required keywords for now (no optional keywords)
+            if num != required_num {
+                count_failure(complex_arg_pass_param_kw_opt)
+            }
+        }
+    }
+
     if 0 != params.flags.has_kwrest()  { count_failure(complex_arg_pass_param_kwrest) }
     if 0 != params.flags.has_block()   { count_failure(complex_arg_pass_param_block) }
     if 0 != params.flags.forwardable() { count_failure(complex_arg_pass_param_forwardable) }
@@ -1541,9 +1559,12 @@ fn can_direct_send(function: &mut Function, block: BlockId, iseq: *const rb_iseq
     // Because we exclude e.g. post parameters above, they are also excluded from the sum below.
     let lead_num = params.lead_num;
     let opt_num = params.opt_num;
+    let keyword = params.keyword;
+    let kw_req_num = if keyword.is_null() { 0 } else { unsafe { (*keyword).required_num } };
+    let req_num = lead_num + kw_req_num;
     can_send = c_int::try_from(args.len())
         .as_ref()
-        .map(|argc| (lead_num..=lead_num + opt_num).contains(argc))
+        .map(|argc| (req_num..=req_num + opt_num).contains(argc))
         .unwrap_or(false);
     if !can_send {
         function.set_dynamic_send_reason(send_insn, ArgcParamMismatch);
@@ -2255,6 +2276,72 @@ impl Function {
         }
     }
 
+    /// Reorder keyword arguments to match the callee's expectation.
+    ///
+    /// Returns Ok with reordered arguments if successful, or Err with the fallback reason if not.
+    fn reorder_keyword_arguments(
+        &self,
+        args: &[InsnId],
+        kwarg: *const rb_callinfo_kwarg,
+        iseq: IseqPtr,
+    ) -> Result<Vec<InsnId>, SendFallbackReason> {
+        let callee_keyword = unsafe { rb_get_iseq_body_param_keyword(iseq) };
+        if callee_keyword.is_null() {
+            // Caller is passing kwargs but callee doesn't expect them.
+            return Err(SendWithoutBlockDirectKeywordMismatch);
+        }
+
+        let caller_kw_count = unsafe { get_cikw_keyword_len(kwarg) } as usize;
+        let callee_kw_count = unsafe { (*callee_keyword).num } as usize;
+        let callee_kw_required = unsafe { (*callee_keyword).required_num } as usize;
+        let callee_kw_table = unsafe { (*callee_keyword).table };
+
+        // For now, only handle the case where all keywords are required.
+        if callee_kw_count != callee_kw_required {
+            return Err(SendWithoutBlockDirectOptionalKeywords);
+        }
+        if caller_kw_count != callee_kw_count {
+            return Err(SendWithoutBlockDirectKeywordCountMismatch);
+        }
+
+        // The keyword arguments are the last arguments in the args vector.
+        let kw_args_start = args.len() - caller_kw_count;
+
+        // Build a mapping from caller keywords to their positions.
+        let mut caller_kw_order: Vec<ID> = Vec::with_capacity(caller_kw_count);
+        for i in 0..caller_kw_count {
+            let sym = unsafe { get_cikw_keywords_idx(kwarg, i as i32) };
+            let id = unsafe { rb_sym2id(sym) };
+            caller_kw_order.push(id);
+        }
+
+        // Reorder keyword arguments to match callee expectation.
+        let mut reordered_kw_args: Vec<InsnId> = Vec::with_capacity(callee_kw_count);
+        for i in 0..callee_kw_count {
+            let expected_id = unsafe { *callee_kw_table.add(i) };
+
+            // Find where this keyword is in the caller's order
+            let mut found = false;
+            for (j, &caller_id) in caller_kw_order.iter().enumerate() {
+                if caller_id == expected_id {
+                    reordered_kw_args.push(args[kw_args_start + j]);
+                    found = true;
+                    break;
+                }
+            }
+
+            if !found {
+                // Required keyword not provided by caller which will raise an ArgumentError.
+                return Err(SendWithoutBlockDirectMissingKeyword);
+            }
+        }
+
+        // Replace the keyword arguments with the reordered ones.
+        let mut processed_args = args[..kw_args_start].to_vec();
+        processed_args.extend(reordered_kw_args);
+        Ok(processed_args)
+    }
+
     /// Resolve the receiver type for method dispatch optimization.
     ///
     /// Takes the receiver's Type, receiver HIR instruction, and ISEQ instruction index.
@@ -2470,6 +2557,7 @@ impl Function {
                             cme = unsafe { rb_aliased_callable_method_entry(cme) };
                             def_type = unsafe { get_cme_def_type(cme) };
                         }
+
                         if def_type == VM_METHOD_TYPE_ISEQ {
                             // TODO(max): Allow non-iseq; cache cme
                             // Only specialize positional-positional calls
@@ -2485,7 +2573,29 @@ impl Function {
                             if let Some(profiled_type) = profiled_type {
                                 recv = self.push_insn(block, Insn::GuardType { val: recv, guard_type: Type::from_profiled_type(profiled_type), state });
                             }
-                            let send_direct = self.push_insn(block, Insn::SendWithoutBlockDirect { recv, cd, cme, iseq, args, state });
+
+                            // Check if caller is passing keywords but callee doesn't expect them.
+                            let kwarg = unsafe { rb_vm_ci_kwarg(ci) };
+                            if !kwarg.is_null() && !unsafe { rb_get_iseq_flags_has_kw(iseq) } {
+                                // Caller has keywords but callee doesn't; Need to convert to hash.
+                                self.set_dynamic_send_reason(insn_id, UnexpectedKeywordArgs);
+                                self.push_insn_id(block, insn_id); continue;
+                            }
+
+                            // Handle keyword argument reordering if present.
+                            let processed_args = if !kwarg.is_null() {
+                                match self.reorder_keyword_arguments(&args, kwarg, iseq) {
+                                    Ok(reordered) => reordered,
+                                    Err(reason) => {
+                                        self.set_dynamic_send_reason(insn_id, reason);
+                                        self.push_insn_id(block, insn_id); continue;
+                                    }
+                                }
+                            } else {
+                                args.clone()
+                            };
+
+                            let send_direct = self.push_insn(block, Insn::SendWithoutBlockDirect { recv, cd, cme, iseq, args: processed_args, state });
                             self.make_equal_to(insn_id, send_direct);
                         } else if def_type == VM_METHOD_TYPE_BMETHOD {
                             let procv = unsafe { rb_get_def_bmethod_proc((*cme).def) };
@@ -2520,7 +2630,29 @@ impl Function {
                             if let Some(profiled_type) = profiled_type {
                                 recv = self.push_insn(block, Insn::GuardType { val: recv, guard_type: Type::from_profiled_type(profiled_type), state });
                             }
-                            let send_direct = self.push_insn(block, Insn::SendWithoutBlockDirect { recv, cd, cme, iseq, args, state });
+
+                            // Check if caller is passing keywords but callee doesn't expect them.
+                            let kwarg = unsafe { rb_vm_ci_kwarg(ci) };
+                            if !kwarg.is_null() && !unsafe { rb_get_iseq_flags_has_kw(iseq) } {
+                                // Caller has keywords but callee doesn't; Need to convert to hash.
+                                self.set_dynamic_send_reason(insn_id, UnexpectedKeywordArgs);
+                                self.push_insn_id(block, insn_id); continue;
+                            }
+
+                            // Handle keyword argument reordering if present.
+                            let processed_args = if !kwarg.is_null() {
+                                match self.reorder_keyword_arguments(&args, kwarg, iseq) {
+                                    Ok(reordered) => reordered,
+                                    Err(reason) => {
+                                        self.set_dynamic_send_reason(insn_id, reason);
+                                        self.push_insn_id(block, insn_id); continue;
+                                    }
+                                }
+                            } else {
+                                args.clone()
+                            };
+
+                            let send_direct = self.push_insn(block, Insn::SendWithoutBlockDirect { recv, cd, cme, iseq, args: processed_args, state });
                             self.make_equal_to(insn_id, send_direct);
                         } else if def_type == VM_METHOD_TYPE_IVAR && args.is_empty() {
                             // Check if we're accessing ivars of a Class or Module object as they require single-ractor mode.
@@ -3062,7 +3194,7 @@ impl Function {
 
             // When seeing &block argument, fall back to dynamic dispatch for now
             // TODO: Support block forwarding
-            if unspecializable_call_type(ci_flags) {
+            if unspecializable_c_call_type(ci_flags) {
                 fun.count_complex_call_features(block, ci_flags);
                 fun.set_dynamic_send_reason(send_insn_id, ComplexArgPass);
                 return Err(());
@@ -4960,9 +5092,14 @@ fn unhandled_call_type(flags: u32) -> Result<(), CallType> {
     Ok(())
 }
 
+/// If a given call to a c func uses overly complex arguments, then we won't specialize.
+fn unspecializable_c_call_type(flags: u32) -> bool {
+    ((flags & VM_CALL_KWARG) != 0) ||
+    unspecializable_call_type(flags)
+}
+
 /// If a given call uses overly complex arguments, then we won't specialize.
 fn unspecializable_call_type(flags: u32) -> bool {
-    ((flags & VM_CALL_KWARG) != 0) ||
     ((flags & VM_CALL_ARGS_SPLAT) != 0) ||
     ((flags & VM_CALL_ARGS_BLOCKARG) != 0)
 }