ZJIT: Support optional keyword args in direct send

Author: rwstauner

test/ruby/test_zjit.rb

@@ -774,6 +774,52 @@ def test = [target(f: 6), target(10, 20, 30, f: 6), target(10, 20, 30, 40, 50, f
     }, call_threshold: 2
   end
 
+  def test_send_kwarg_partial_optional
+    assert_compiles '[[1, 2, 3], [1, 20, 3], [10, 2, 30]]', %q{
+      def test(a: 1, b: 2, c: 3) = [a, b, c]
+      def entry = [test, test(b: 20), test(c: 30, a: 10)]
+      entry
+      entry
+    }, call_threshold: 2
+  end
+
+  def test_send_kwarg_optional_a_lot
+    assert_compiles '[[1, 2, 3, 4, 5, 6], [1, 2, 3, 7, 8, 9], [2, 4, 6, 8, 10, 12]]', %q{
+      def test(a: 1, b: 2, c: 3, d: 4, e: 5, f: 6) = [a, b, c, d, e, f]
+      def entry = [test, test(d: 7, f: 9, e: 8), test(f: 12, e: 10, d: 8, c: 6, b: 4, a: 2)]
+      entry
+      entry
+    }, call_threshold: 2
+  end
+
+  def test_send_kwarg_non_constant_default
+    assert_compiles '[[1, 2], [10, 2]]', %q{
+      def make_default = 2
+      def test(a: 1, b: make_default) = [a, b]
+      def entry = [test, test(a: 10)]
+      entry
+      entry
+    }, call_threshold: 2
+  end
+
+  def test_send_kwarg_optional_static_with_side_exit
+    # verify frame reconstruction with synthesized keyword defaults is correct
+    assert_compiles '[10, 2, 10]', %q{
+      def callee(a: 1, b: 2)
+        # use binding to force side-exit
+        x = binding.local_variable_get(:a)
+        [a, b, x]
+      end
+
+      def entry
+        callee(a: 10) # b should get default value
+      end
+
+      entry
+      entry
+    }, call_threshold: 2
+  end
+
   def test_send_all_arg_types
     assert_compiles '[:req, :opt, :post, :kwr, :kwo, true]', %q{
       def test(a, b = :opt, c, d:, e: :kwo) = [a, b, c, d, e, block_given?]
@@ -1329,6 +1375,23 @@ def test
     }, call_threshold: 2
   end
 
+  def test_invokesuper_with_optional_keyword_args
+    assert_compiles '[1, 2, 3]', %q{
+      class Parent
+        def foo(a, b: 2, c: 3) = [a, b, c]
+      end
+
+      class Child < Parent
+        def foo(a) = super(a)
+      end
+
+      def test = Child.new.foo(1)
+
+      test
+      test
+    }, call_threshold: 2
+  end
+
   def test_invokebuiltin
     # Not using assert_compiles due to register spill
     assert_runs '["."]', %q{

zjit/src/codegen.rs

@@ -1403,11 +1403,14 @@ fn gen_send_iseq_direct(
     // Write "keyword_bits" to the callee's frame if the callee accepts keywords.
     // This is a synthetic local/parameter that the callee reads via checkkeyword to determine
     // which optional keyword arguments need their defaults evaluated.
+    // We write this to the local table slot at bits_start so that:
+    // 1. The interpreter can read it via checkkeyword if we side-exit
+    // 2. The JIT entry can read it via GetLocal
     if unsafe { rb_get_iseq_flags_has_kw(iseq) } {
         let keyword = unsafe { rb_get_iseq_body_param_keyword(iseq) };
         let bits_start = unsafe { (*keyword).bits_start } as usize;
-        // Currently we only support required keywords, so all bits are 0 (all keywords specified).
-        // TODO: When supporting optional keywords, calculate actual unspecified_bits here.
+        // kw_bits is always 0 because constant defaults are inlined directly,
+        // and non-constant defaults (Qundef) cause fallback to VM dispatch.
         let unspecified_bits = VALUE::fixnum_from_usize(0);
         let bits_offset = (state.stack().len() - args.len() + bits_start) * SIZEOF_VALUE;
         asm_comment!(asm, "write keyword bits to callee frame");
@@ -1434,10 +1437,11 @@ fn gen_send_iseq_direct(
         let lead_num = params.lead_num as u32;
         let opt_num = params.opt_num as u32;
         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;
+        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);
+        // For computing optional positional entry point, only count positional args
+        let positional_argc = args.len() as u32 - kw_total_num;
+        let num_optionals_passed = positional_argc.saturating_sub(lead_num);
         num_optionals_passed
     } else {
         0