ZJIT: Fn stub: Move args to create appropriate unfilled optional param gap

Fixes the included test case. It returned junk previously. To keep the
IseqCall struct 24 bytes, I've imposed a ~65K limit on the number of
arguments that can be codegen compiles. Should be fine.

Author: XrXr

zjit/src/codegen.rs

@@ -612,7 +612,10 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         &Insn::Send { cd, blockiseq: None, state, reason, .. } => gen_send_without_block(jit, asm, cd, &function.frame_state(state), reason),
         &Insn::Send { cd, blockiseq: Some(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::SendDirect { cme, iseq, recv, args, kw_bits, blockiseq, state, .. } => gen_send_iseq_direct(cb, jit, asm, *cme, *iseq, opnd!(recv), opnds!(args), *kw_bits, &function.frame_state(*state), *blockiseq),
+        &Insn::SendDirect { cme, iseq, recv, ref args, kw_bits, blockiseq, state, .. } => gen_send_iseq_direct(
+                cb, jit, asm, cme, iseq, opnd!(recv), opnds!(args),
+                kw_bits, &function.frame_state(state), blockiseq,
+            ).ok_or(state)?,
         &Insn::InvokeSuper { cd, blockiseq, state, reason, .. } => gen_invokesuper(jit, asm, cd, blockiseq, &function.frame_state(state), reason),
         &Insn::InvokeSuperForward { cd, blockiseq, state, reason, .. } => gen_invokesuperforward(jit, asm, cd, blockiseq, &function.frame_state(state), reason),
         &Insn::InvokeBlock { cd, state, reason, .. } => gen_invokeblock(jit, asm, cd, &function.frame_state(state), reason),
@@ -1535,7 +1538,7 @@ fn gen_send_iseq_direct(
     kw_bits: u32,
     state: &FrameState,
     blockiseq: Option<IseqPtr>,
-) -> lir::Opnd {
+) -> Option<lir::Opnd> {
     gen_incr_counter(asm, Counter::iseq_optimized_send_count);
 
     let local_size = unsafe { get_iseq_body_local_table_size(iseq) }.to_usize();
@@ -1651,7 +1654,7 @@ fn gen_send_iseq_direct(
     };
 
     // Make a method call. The target address will be rewritten once compiled.
-    let iseq_call = IseqCall::new(iseq, num_optionals_passed);
+    let iseq_call = IseqCall::new(iseq, num_optionals_passed.try_into().ok()?, args.len().try_into().ok()?);
     let dummy_ptr = cb.get_write_ptr().raw_ptr(cb);
     jit.iseq_calls.push(iseq_call.clone());
     let ret = asm.ccall_with_iseq_call(dummy_ptr, c_args, &iseq_call);
@@ -1668,7 +1671,7 @@ fn gen_send_iseq_direct(
     let new_sp = asm.sub(SP, sp_offset.into());
     asm.mov(SP, new_sp);
 
-    ret
+    Some(ret)
 }
 
 /// Compile for invokeblock
@@ -2937,22 +2940,39 @@ c_callable! {
             // function_stub_hit_body() may allocate and call gc_validate_pc(), so we always set PC.
             let iseq_call = unsafe { Rc::from_raw(iseq_call_ptr as *const IseqCall) };
             let iseq = iseq_call.iseq.get();
+            let argc = iseq_call.argc;
+            let num_opts_filled = iseq_call.jit_entry_idx;
             let entry_insn_idxs = crate::hir::jit_entry_insns(iseq);
             let pc = unsafe { rb_iseq_pc_at_idx(iseq, entry_insn_idxs[iseq_call.jit_entry_idx.to_usize()]) };
             unsafe { rb_set_cfp_pc(cfp, pc) };
 
             // JIT-to-JIT calls don't eagerly fill nils to non-parameter locals.
             // If we side-exit from function_stub_hit (before JIT code runs), we need to set them here.
-            fn prepare_for_exit(iseq: IseqPtr, cfp: CfpPtr, sp: *mut VALUE, compile_error: &CompileError) {
+            fn prepare_for_exit(iseq: IseqPtr, cfp: CfpPtr, sp: *mut VALUE, argc: u16, num_opts_filled: u16, compile_error: &CompileError) {
                 unsafe {
                     // Set SP which gen_push_frame() doesn't set
                     rb_set_cfp_sp(cfp, sp);
 
-                    // Fill nils to uninitialized (non-argument) locals
                     let local_size = get_iseq_body_local_table_size(iseq).to_usize();
-                    let num_params = iseq.params().size.to_usize();
-                    let base = sp.offset(-local_size_and_idx_to_bp_offset(local_size, num_params) as isize);
-                    slice::from_raw_parts_mut(base, local_size - num_params).fill(Qnil);
+                    let params = iseq.params();
+                    let params_size = params.size.to_usize();
+                    let frame_base = sp.offset(-local_size_and_idx_to_bp_offset(local_size, 0) as isize);
+                    let locals = slice::from_raw_parts_mut(frame_base, local_size);
+                    // Fill nils to uninitialized (non-parameter) locals
+                    locals.get_mut(params_size..).unwrap_or_default().fill(Qnil);
+
+                    // When there are unfilled optionals
+                    let opt_num: usize = params.opt_num.try_into().expect("ISEQ opt_num should be non-negative");
+                    let opts_unfilled = opt_num - num_opts_filled.to_usize();
+                    if opts_unfilled > 0 {
+                        let lead_num: usize = params.lead_num.try_into().expect("should be non-negative");
+                        let after_opts = lead_num + opt_num;
+                        let param_locals = &mut locals[..params_size];
+                        // Right-justify params that are after the unfilled optionals
+                        param_locals.copy_within(num_opts_filled.to_usize()..argc.to_usize(), after_opts);
+                        // Nil-fill unfilled optionals
+                        (&mut param_locals[num_opts_filled.to_usize()..after_opts]).fill(Qnil);
+                    }
                 }
 
                 // Increment a compile error counter for --zjit-stats
@@ -2976,7 +2996,7 @@ c_callable! {
                 // We'll use this Rc again, so increment the ref count decremented by from_raw.
                 unsafe { Rc::increment_strong_count(iseq_call_ptr as *const IseqCall); }
 
-                prepare_for_exit(iseq, cfp, sp, compile_error);
+                prepare_for_exit(iseq, cfp, sp, argc, num_opts_filled, compile_error);
                 return ZJITState::get_exit_trampoline_with_counter().raw_ptr(cb);
             }
 
@@ -2991,7 +3011,7 @@ c_callable! {
                 // We'll use this Rc again, so increment the ref count decremented by from_raw.
                 unsafe { Rc::increment_strong_count(iseq_call_ptr as *const IseqCall); }
 
-                prepare_for_exit(iseq, cfp, sp, &compile_error);
+                prepare_for_exit(iseq, cfp, sp, argc, num_opts_filled, &compile_error);
                 ZJITState::get_exit_trampoline_with_counter()
             });
             cb.mark_all_executable();
@@ -3316,7 +3336,10 @@ pub struct IseqCall {
     pub iseq: Cell<IseqPtr>,
 
     /// Index that corresponds to [crate::hir::jit_entry_insns]
-    jit_entry_idx: u32,
+    jit_entry_idx: u16,
+
+    /// Argument count passing to the HIR function
+    argc: u16,
 
     /// Position where the call instruction starts
     start_addr: Cell<Option<CodePtr>>,
@@ -3329,12 +3352,13 @@ pub type IseqCallRef = Rc<IseqCall>;
 
 impl IseqCall {
     /// Allocate a new IseqCall
-    fn new(iseq: IseqPtr, jit_entry_idx: u32) -> IseqCallRef {
+    fn new(iseq: IseqPtr, jit_entry_idx: u16, argc: u16) -> IseqCallRef {
         let iseq_call = IseqCall {
             iseq: Cell::new(iseq),
             start_addr: Cell::new(None),
             end_addr: Cell::new(None),
             jit_entry_idx,
+            argc,
         };
         Rc::new(iseq_call)
     }

zjit/src/hir/opt_tests.rs

@@ -14857,6 +14857,28 @@ mod hir_opt_tests {
         ");
     }
 
+    #[test]
+    fn test_exit_from_function_stub_for_opt_keyword_callee() {
+        // We have a SendDirect to a callee that fails to compile,
+        // so the function stub has to take care of exiting to
+        // interpreter.
+        eval("
+            def target(a = binding.local_variable_get(:a), b: nil)
+              ::RubyVM::ZJIT.induce_compile_failure!
+              [a, b]
+            end
+
+            def entry = target(b: -1)
+
+            raise 'wrong' unless [nil, -1] == entry
+            raise 'wrong' unless [nil, -1] == entry
+        ");
+
+        crate::hir::tests::hir_build_tests::assert_compile_fails("target", ParseError::DirectiveInduced);
+        let hir = hir_string("entry");
+        assert!(hir.contains("SendDirect"), "{hir}");
+    }
+
     #[test]
     fn test_recompile_no_profile_send() {
         // Define a callee method and a test method that calls it

zjit/src/hir/tests.rs

@@ -207,7 +207,7 @@ mod snapshot_tests {
 }
 
 #[cfg(test)]
-pub mod hir_build_tests {
+pub(crate) mod hir_build_tests {
     use super::*;
     use crate::options::set_call_threshold;
     use insta::assert_snapshot;
@@ -275,7 +275,7 @@ pub mod hir_build_tests {
     }
 
     #[track_caller]
-    fn assert_compile_fails(method: &str, reason: ParseError) {
+    pub fn assert_compile_fails(method: &str, reason: ParseError) {
         let iseq = crate::cruby::with_rubyvm(|| get_method_iseq("self", method));
         unsafe { crate::cruby::rb_zjit_profile_disable(iseq) };
         let result = iseq_to_hir(iseq);