ZJIT: Support optional keyword arguments in direct send (ruby#15873)

This fills in constants when unspecified optional keyword args have static default values.
For complex defaults we calculate the kw_bits and utilize the checkkeyword logic we already had.

The following benchmarks used to register param_kw_opt.
Some of them (like graphql*) just trade that for some other complexity, or "too_many_args_for_lir".
Notable improvements include activerecord where the previous param_kw_opt count has a corresponding drop in complex args and dynamic_send_count and a nearly equal rise in optimized_send_count.
The gains are similar but not as complete in hexapdf, liquid-render, lobsters, railsbench, shipit.

| Benchmark | param_kw_opt | Δ one_or_more_complex | Δ too_many_args | Δ dynamic_send | Δ optimized_send |
|-----------|-------------:|----------------------:|----------------:|---------------:|-----------------:|
| activerecord | 6,307,141 | -6,253,823 | +4,084 | -6,306,223 | +6,279,766 |
| blurhash | 21 | -21 | +0 | -23 | +20 |
| chunky-png | 813,604 | -813,604 | +0 | -813,616 | +813,556 |
| erubi-rails | 1,590,395 | -590,274 | +35,578 | -552,914 | +550,826 |
| fluentd | 4,906 | -4,854 | +21 | -5,745 | +5,080 |
| graphql | 1,610,439 | -1,610,432 | +1,605,751 | -4,688 | +4,628 |
| graphql-native | 16,332,386 | -16,332,375 | +16,309,681 | -22,701 | +22,638 |
| hexapdf | 9,165,465 | -9,124,509 | +203,754 | -8,920,727 | +8,839,295 |
| liquid-compile | 14,817 | -14,792 | +0 | -14,705 | +15,045 |
| liquid-render | 3,994,905 | -3,994,901 | +0 | -3,994,868 | +3,020,779 |
| lobsters | 2,467,510 | -2,297,298 | +205,610 | -2,216,583 | +1,694,092 |
| protoboeuf | 11,521 | -11,521 | +0 | -11,523 | +11,520 |
| psych-load | 77,612 | -77,609 | +29,942 | -77,613 | -12,242 |
| rack | 2,743 | -2,742 | +0 | -2,750 | +2,668 |
| railsbench | 3,579,778 | -2,517,615 | +432,575 | -2,084,480 | +1,882,928 |
| ruby-lsp | 287,171 | -379,716 | +37 | -409,368 | -267,248 |
| rubyboy | 5,993,004 | -5,993,003 | +0 | -5,993,006 | +5,992,993 |
| sequel | 182,652 | -182,631 | +0 | -182,563 | +122,687 |
| shipit | 3,289,456 | -2,778,419 | +306,867 | -3,201,395 | +1,068,505 |
| tinygql | 2,732 | -2,732 | +1 | -2,734 | +2,729 |

Author: rwstauner

test/ruby/test_zjit.rb

@@ -833,6 +833,61 @@ def entry
     }, call_threshold: 2
   end
 
+  def test_pos_optional_with_maybe_too_many_args
+    assert_compiles '[[1, 2, 3, 4, 5, 6], [10, 20, 30, 4, 5, 6], [10, 20, 30, 40, 50, 60]]', %q{
+      def target(a = 1, b = 2, c = 3, d = 4, e = 5, f:) = [a, b, c, d, e, f]
+      def test = [target(f: 6), target(10, 20, 30, f: 6), target(10, 20, 30, 40, 50, f: 60)]
+      test
+      test
+    }, 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?]
@@ -1388,6 +1443,190 @@ 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_send_with_non_constant_keyword_default
+    assert_compiles '[[2, 4, 16], [10, 4, 16], [2, 20, 16], [2, 4, 30], [10, 20, 30]]', %q{
+      def dbl(x = 1) = x * 2
+
+      def foo(a: dbl, b: dbl(2), c: dbl(2 ** 3))
+        [a, b, c]
+      end
+
+      def test
+        [
+          foo,
+          foo(a: 10),
+          foo(b: 20),
+          foo(c: 30),
+          foo(a: 10, b: 20, c: 30)
+        ]
+      end
+
+      test
+      test
+    }, call_threshold: 2
+  end
+
+  def test_send_with_non_constant_keyword_default_not_evaluated_when_provided
+    assert_compiles '[1, 2, 3]', %q{
+      def foo(a: raise, b: raise, c: raise)
+        [a, b, c]
+      end
+
+      def test
+        foo(a: 1, b: 2, c: 3)
+      end
+
+      test
+      test
+    }, call_threshold: 2
+  end
+
+  def test_send_with_non_constant_keyword_default_evaluated_when_not_provided
+    assert_compiles '["a", "b", "c"]', %q{
+      def raise_a = raise "a"
+      def raise_b = raise "b"
+      def raise_c = raise "c"
+
+      def foo(a: raise_a, b: raise_b, c: raise_c)
+        [a, b, c]
+      end
+
+      def test_a
+        foo(b: 2, c: 3)
+      rescue RuntimeError => e
+        e.message
+      end
+
+      def test_b
+        foo(a: 1, c: 3)
+      rescue RuntimeError => e
+        e.message
+      end
+
+      def test_c
+        foo(a: 1, b: 2)
+      rescue RuntimeError => e
+        e.message
+      end
+
+      def test
+        [test_a, test_b, test_c]
+      end
+
+      test
+      test
+    }, call_threshold: 2
+  end
+
+  def test_send_with_non_constant_keyword_default_jit_to_jit
+    # Test that kw_bits passing works correctly in JIT-to-JIT calls
+    assert_compiles '[2, 4, 6]', %q{
+      def make_default(x) = x * 2
+
+      def callee(a: make_default(1), b: make_default(2), c: make_default(3))
+        [a, b, c]
+      end
+
+      def caller_method
+        callee
+      end
+
+      # Warm up callee first so it gets JITted
+      callee
+      callee
+
+      # Now warm up caller - this creates JIT-to-JIT call
+      caller_method
+      caller_method
+    }, call_threshold: 2
+  end
+
+  def test_send_with_non_constant_keyword_default_side_exit
+    # Verify frame reconstruction includes correct values for non-constant defaults
+    assert_compiles '[10, 2, 30]', %q{
+      def make_b = 2
+
+      def callee(a: 1, b: make_b, c: 3)
+        x = binding.local_variable_get(:a)
+        y = binding.local_variable_get(:b)
+        z = binding.local_variable_get(:c)
+        [x, y, z]
+      end
+
+      def test
+        callee(a: 10, c: 30)
+      end
+
+      test
+      test
+    }, call_threshold: 2
+  end
+
+  def test_send_with_non_constant_keyword_default_evaluation_order
+    # Verify defaults are evaluated left-to-right and only when not provided
+    assert_compiles '[["a", "b", "c"], ["b", "c"], ["a", "c"], ["a", "b"]]', %q{
+      def log(x)
+        $order << x
+        x
+      end
+
+      def foo(a: log("a"), b: log("b"), c: log("c"))
+        [a, b, c]
+      end
+
+      def test
+        results = []
+
+        $order = []
+        foo
+        results << $order.dup
+
+        $order = []
+        foo(a: "A")
+        results << $order.dup
+
+        $order = []
+        foo(b: "B")
+        results << $order.dup
+
+        $order = []
+        foo(c: "C")
+        results << $order.dup
+
+        results
+      end
+
+      test
+      test
+    }, call_threshold: 2
+  end
+
+  def test_send_with_too_many_non_constant_keyword_defaults
+    assert_compiles '35', %q{
+      def many_kwargs( k1: 1, k2: 2, k3: 3, k4: 4, k5: 5, k6: 6, k7: 7, k8: 8, k9: 9, k10: 10, k11: 11, k12: 12, k13: 13, k14: 14, k15: 15, k16: 16, k17: 17, k18: 18, k19: 19, k20: 20, k21: 21, k22: 22, k23: 23, k24: 24, k25: 25, k26: 26, k27: 27, k28: 28, k29: 29, k30: 30, k31: 31, k32: 32, k33: 33, k34: k33 + 1) = k1 + k34
+      def t = many_kwargs
+      t
+      t
+    }, call_threshold: 2
+  end
+
   def test_invokebuiltin
     # Not using assert_compiles due to register spill
     assert_runs '["."]', %q{

zjit/src/codegen.rs

@@ -401,7 +401,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         &Insn::Send { cd, blockiseq, state, reason, .. } => gen_send(jit, asm, cd, blockiseq, &function.frame_state(state), reason),
         &Insn::SendForward { cd, blockiseq, state, reason, .. } => gen_send_forward(jit, asm, cd, blockiseq, &function.frame_state(state), reason),
         &Insn::SendWithoutBlock { cd, state, reason, .. } => gen_send_without_block(jit, asm, cd, &function.frame_state(state), reason),
-        Insn::SendWithoutBlockDirect { cme, iseq, recv, args, state, .. } => gen_send_iseq_direct(cb, jit, asm, *cme, *iseq, opnd!(recv), opnds!(args), &function.frame_state(*state), None),
+        Insn::SendWithoutBlockDirect { cme, iseq, recv, args, kw_bits, state, .. } => gen_send_iseq_direct(cb, jit, asm, *cme, *iseq, opnd!(recv), opnds!(args), *kw_bits, &function.frame_state(*state), None),
         &Insn::InvokeSuper { cd, blockiseq, state, reason, .. } => gen_invokesuper(jit, asm, cd, blockiseq, &function.frame_state(state), reason),
         &Insn::InvokeBlock { cd, state, reason, .. } => gen_invokeblock(jit, asm, cd, &function.frame_state(state), reason),
         Insn::InvokeProc { recv, args, state, kw_splat } => gen_invokeproc(jit, asm, opnd!(recv), opnds!(args), *kw_splat, &function.frame_state(*state)),
@@ -1358,6 +1358,7 @@ fn gen_send_iseq_direct(
     iseq: IseqPtr,
     recv: Opnd,
     args: Vec<Opnd>,
+    kw_bits: u32,
     state: &FrameState,
     block_handler: Option<Opnd>,
 ) -> lir::Opnd {
@@ -1404,12 +1405,13 @@ 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.
-        let unspecified_bits = VALUE::fixnum_from_usize(0);
+        let unspecified_bits = VALUE::fixnum_from_usize(kw_bits as usize);
         let bits_offset = (state.stack().len() - args.len() + bits_start) * SIZEOF_VALUE;
         asm_comment!(asm, "write keyword bits to callee frame");
         asm.store(Opnd::mem(64, SP, bits_offset as i32), unspecified_bits.into());
@@ -1435,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