YJIT: Support iseq sends with mixed kwargs (#5082)

* YJIT: Support iseq sends with mixed kwargs

Co-authored-by: Kevin Newton <kddnewton@gmail.com>

* Add additional comments to iseq sends

Co-authored-by: Kevin Newton <kddnewton@gmail.com>

2 files changed, 106 insertions, 48 deletions

diff --git a/bootstraptest/test_yjit.rb b/bootstraptest/test_yjit.rb
index 13d63b521c..33e9a35462 100644
--- a/bootstraptest/test_yjit.rb
+++ b/bootstraptest/test_yjit.rb
@@ -2187,6 +2187,60 @@ assert_equal '[3]', %q{
   5.times.map { default_expression }.uniq
 }
 
+# reordered optional kwargs
+assert_equal '[[100, 1]]', %q{
+  def foo(capacity: 100, max: nil)
+    [capacity, max]
+  end
+
+  5.times.map { foo(max: 1) }.uniq
+}
+
+# invalid lead param
+assert_equal 'ok', %q{
+  def bar(baz: 2)
+    baz
+  end
+
+  def foo
+    bar(1, baz: 123)
+  end
+
+  begin
+    foo
+    foo
+  rescue ArgumentError => e
+    print "ok"
+  end
+}
+
+# reordered required kwargs
+assert_equal '[[1, 2, 3, 4]]', %q{
+  def foo(default1: 1, required1:, default2: 3, required2:)
+    [default1, required1, default2, required2]
+  end
+
+  5.times.map { foo(required1: 2, required2: 4) }.uniq
+}
+
+# reordered default expression kwargs
+assert_equal '[[:one, :two, 3]]', %q{
+  def foo(arg1: (1+0), arg2: (2+0), arg3: (3+0))
+    [arg1, arg2, arg3]
+  end
+
+  5.times.map { foo(arg2: :two, arg1: :one) }.uniq
+}
+
+# complex kwargs
+assert_equal '[[1, 2, 3, 4]]', %q{
+  def foo(required:, specified: 999, simple_default: 3, complex_default: "4".to_i)
+    [required, specified, simple_default, complex_default]
+  end
+
+  5.times.map { foo(specified: 2, required: 1) }.uniq
+}
+
 # attr_reader on frozen object
 assert_equal 'false', %q{
   class Foo
diff --git a/yjit_codegen.c b/yjit_codegen.c
index db27c52ed9..337ccb5765 100644
--- a/yjit_codegen.c
+++ b/yjit_codegen.c
@@ -3527,11 +3527,9 @@ gen_send_iseq(jitstate_t *jit, ctx_t *ctx, const struct rb_callinfo *ci, const r
         }
 
         if (vm_ci_flag(ci) & VM_CALL_KWARG) {
-            if ((kw_arg->keyword_len != keyword->num) || (lead_num != argc - kw_arg->keyword_len)) {
-                // Here the method being called specifies optional and required
-                // keyword arguments and the callee is not specifying every one
-                // of them.
-                GEN_COUNTER_INC(cb, send_iseq_kwargs_req_and_opt_missing);
+            // Check that the size of non-keyword arguments matches
+            if (lead_num != argc - kw_arg->keyword_len) {
+                GEN_COUNTER_INC(cb, send_iseq_complex_callee);
                 return YJIT_CANT_COMPILE;
             }
 
@@ -3656,33 +3654,65 @@ gen_send_iseq(jitstate_t *jit, ctx_t *ctx, const struct rb_callinfo *ci, const r
         // keyword parameters.
         const struct rb_iseq_param_keyword *keyword = iseq->body->param.keyword;
 
-        // Note: we are about to do argument shuffling for a keyword argument
-        // call. Assert that we are in one of currently supported cases: either
-        // all keywords are provided or they are all the default values.
-        // The various checks for whether we can do it happened earlier in this
-        // function.
-        RUBY_ASSERT(
-                ((caller_keyword_len == keyword->num) &&
-                 (lead_num == argc - caller_keyword_len)) ||
-                (caller_keyword_len == 0 && lead_num == argc));
+        ADD_COMMENT(cb, "keyword args");
 
         // This is the list of keyword arguments that the callee specified
         // in its initial declaration.
         const ID *callee_kwargs = keyword->table;
 
+        int total_kwargs = keyword->num;
+
         // Here we're going to build up a list of the IDs that correspond to
         // the caller-specified keyword arguments. If they're not in the
         // same order as the order specified in the callee declaration, then
         // we're going to need to generate some code to swap values around
         // on the stack.
-        ID *caller_kwargs = ALLOCA_N(VALUE, caller_keyword_len);
-        for (int kwarg_idx = 0; kwarg_idx < caller_keyword_len; kwarg_idx++)
-            caller_kwargs[kwarg_idx] = SYM2ID(caller_keywords[kwarg_idx]);
+        ID *caller_kwargs = ALLOCA_N(VALUE, total_kwargs);
+        int kwarg_idx;
+        for (kwarg_idx = 0; kwarg_idx < caller_keyword_len; kwarg_idx++) {
+                caller_kwargs[kwarg_idx] = SYM2ID(caller_keywords[kwarg_idx]);
+        }
+
+        int unspecified_bits = 0;
+
+        for (int callee_idx = keyword->required_num; callee_idx < total_kwargs; callee_idx++) {
+            bool already_passed = false;
+            ID callee_kwarg = callee_kwargs[callee_idx];
+
+            for (int caller_idx = 0; caller_idx < caller_keyword_len; caller_idx++) {
+                if (caller_kwargs[caller_idx] == callee_kwarg) {
+                    already_passed = true;
+                    break;
+                }
+            }
+
+            if (!already_passed) {
+                // Reserve space on the stack for each default value we'll be
+                // filling in (which is done in the next loop). Also increments
+                // argc so that the callee's SP is recorded correctly.
+                argc++;
+                x86opnd_t default_arg = ctx_stack_push(ctx, TYPE_UNKNOWN);
+                VALUE default_value = keyword->default_values[callee_idx - keyword->required_num];
+
+                if (default_value == Qundef) {
+                    // Qundef means that this value is not constant and must be
+                    // recalculated at runtime, so we record it in unspecified_bits
+                    // (Qnil is then used as a placeholder instead of Qundef).
+                    unspecified_bits |= 0x01 << (callee_idx - keyword->required_num);
+                    default_value = Qnil;
+                }
+
+                mov(cb, default_arg, imm_opnd(default_value));
+
+                caller_kwargs[kwarg_idx++] = callee_kwarg;
+            }
+        }
+        RUBY_ASSERT(kwarg_idx == total_kwargs);
 
         // Next, we're going to loop through every keyword that was
         // specified by the caller and make sure that it's in the correct
         // place. If it's not we're going to swap it around with another one.
-        for (int kwarg_idx = 0; kwarg_idx < caller_keyword_len; kwarg_idx++) {
+        for (kwarg_idx = 0; kwarg_idx < total_kwargs; kwarg_idx++) {
             ID callee_kwarg = callee_kwargs[kwarg_idx];
 
             // If the argument is already in the right order, then we don't
@@ -3693,7 +3723,7 @@ gen_send_iseq(jitstate_t *jit, ctx_t *ctx, const struct rb_callinfo *ci, const r
             // In this case the argument is not in the right place, so we
             // need to find its position where it _should_ be and swap with
             // that location.
-            for (int swap_idx = kwarg_idx + 1; swap_idx < caller_keyword_len; swap_idx++) {
+            for (int swap_idx = kwarg_idx + 1; swap_idx < total_kwargs; swap_idx++) {
                 if (callee_kwarg == caller_kwargs[swap_idx]) {
                     // First we're going to generate the code that is going
                     // to perform the actual swapping at runtime.
@@ -3711,35 +3741,6 @@ gen_send_iseq(jitstate_t *jit, ctx_t *ctx, const struct rb_callinfo *ci, const r
             }
         }
 
-        int unspecified_bits = 0;
-
-        if (caller_keyword_len == 0) {
-            ADD_COMMENT(cb, "default kwarg values");
-
-            for (int callee_idx = 0; callee_idx < keyword->num; callee_idx++) {
-                // Reserve space on the stack for each default value we'll be
-                // filling in (which is done in the next loop). Also increments
-                // argc so that the callee's SP is recorded correctly.
-                argc++;
-                ctx_stack_push(ctx, TYPE_UNKNOWN);
-            }
-
-            for (int callee_idx = 0; callee_idx < keyword->num; callee_idx++) {
-                VALUE default_value = keyword->default_values[callee_idx];
-
-                if (default_value == Qundef) {
-                    // Qundef means that this value is not constant and must be
-                    // recalculated at runtime, so we record it in unspecified_bits
-                    // (Qnil is then used as a placeholder instead of Qundef).
-                    unspecified_bits |= 0x01 << callee_idx;
-                    default_value = Qnil;
-                }
-
-                x86opnd_t stack_arg = ctx_stack_opnd(ctx, keyword->num - callee_idx - 1);
-                mov(cb, stack_arg, imm_opnd(default_value));
-            }
-        }
-
         // Keyword arguments cause a special extra local variable to be
         // pushed onto the stack that represents the parameters that weren't
         // explicitly given a value and have a non-constant default.
@@ -3749,6 +3750,7 @@ gen_send_iseq(jitstate_t *jit, ctx_t *ctx, const struct rb_callinfo *ci, const r
     x86opnd_t recv = ctx_stack_opnd(ctx, argc);
 
     // Store the updated SP on the current frame (pop arguments and receiver)
+    ADD_COMMENT(cb, "store caller sp");
     lea(cb, REG0, ctx_sp_opnd(ctx, sizeof(VALUE) * -(argc + 1)));
     mov(cb, member_opnd(REG_CFP, rb_control_frame_t, sp), REG0);
 
@@ -3771,6 +3773,7 @@ gen_send_iseq(jitstate_t *jit, ctx_t *ctx, const struct rb_callinfo *ci, const r
         mov(cb, mem_opnd(64, REG0, sizeof(VALUE) * (i - num_locals - 3)), imm_opnd(Qnil));
     }
 
+    ADD_COMMENT(cb, "push env");
     // Put compile time cme into REG1. It's assumed to be valid because we are notified when
     // any cme we depend on become outdated. See rb_yjit_method_lookup_change().
     jit_mov_gc_ptr(jit, cb, REG1, (VALUE)cme);
@@ -3795,6 +3798,7 @@ gen_send_iseq(jitstate_t *jit, ctx_t *ctx, const struct rb_callinfo *ci, const r
     uint64_t frame_type = VM_FRAME_MAGIC_METHOD | VM_ENV_FLAG_LOCAL;
     mov(cb, mem_opnd(64, REG0, 8 * -1), imm_opnd(frame_type));
 
+    ADD_COMMENT(cb, "push callee CFP");
     // Allocate a new CFP (ec->cfp--)
     sub(cb, REG_CFP, imm_opnd(sizeof(rb_control_frame_t)));
     mov(cb, member_opnd(REG_EC, rb_execution_context_t, cfp), REG_CFP);