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

Author: XrXr

zjit/src/codegen.rs

@@ -1630,7 +1630,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, args.len().try_into().expect("TODO: no panic here from downcast from usize"));
     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);
@@ -2872,22 +2872,40 @@ 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) };
 
             // Successful JIT-to-JIT calls fill nils to non-parameter locals in generated code.
             // 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: u32, num_opts_filled: u32, 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
+                    // Fill nils to uninitialized (non-parameter) locals
                     let local_size = get_iseq_body_local_table_size(iseq).to_usize();
-                    let num_params = iseq.params().size.to_usize();
+                    let params = iseq.params();
+                    let num_params = params.size.to_usize();
                     let base = sp.offset(-local_size_and_idx_to_bp_offset(local_size, num_params) as isize);
+                    // TODO one slice for everything
                     slice::from_raw_parts_mut(base, local_size - num_params).fill(Qnil);
+
+                    let opt_num: usize = params.opt_num.try_into().expect("ISEQ opt_num should be non-negative");
+                    let num_opts_unfilled = opt_num - num_opts_filled.to_usize();
+                    if num_opts_unfilled > 0 {
+                        let lead_num: usize = params.lead_num.try_into().expect("should be non-negative");
+                        let after_opts = lead_num + opt_num;
+
+                        // TODO rename
+                        let base = sp.offset(-local_size_and_idx_to_bp_offset(local_size, 0) as isize);
+                        let params = slice::from_raw_parts_mut(base, num_params);
+                        params.copy_within(num_opts_filled.to_usize()..argc.to_usize(), after_opts);
+
+                        (&mut params[num_opts_filled.to_usize()..after_opts]).fill(Qnil);
+                    }
                 }
 
                 // Increment a compile error counter for --zjit-stats
@@ -2911,7 +2929,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);
             }
 
@@ -2926,7 +2944,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();
@@ -3253,6 +3271,9 @@ pub struct IseqCall {
     /// Index that corresponds to [crate::hir::jit_entry_insns]
     jit_entry_idx: u32,
 
+    /// Argument count passing to the HIR function
+    argc: u32,
+
     /// Position where the call instruction starts
     start_addr: Cell<Option<CodePtr>>,
 
@@ -3264,12 +3285,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: u32, argc: u32) -> 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

@@ -14765,4 +14765,24 @@ mod hir_opt_tests {
           Return v13
         ");
     }
+
+    #[test]
+    fn test_exit_from_function_stub_for_opt_keyword_callee() {
+        set_call_threshold(2);
+        let result = eval("
+            def target(a = nil, b: nil)
+              ::RubyVM::ZJIT.induce_compile_failure!
+              b
+            end
+
+            def entry = target(b: -1)
+
+            entry
+            entry
+        ");
+        // TODO assert failure for target
+        // TODO it's not entering into jit code. is rb_zjit_enabled_p=false in harness?
+        //assert_snapshot!(hir_string("entry"), @"");
+        assert_eq!(VALUE::fixnum_from_isize(-1), result);
+    }
 }