Update gen_send_iseq for CRuby PRs ruby#5285 and ruby#5379. (ruby#220)

Noah Gibbs · web-flow · commit c63042e8ee2b · 2022-03-07T10:46:06.000-05:00
diff --git a/yjit.c b/yjit.c
@@ -393,6 +393,11 @@ rb_get_def_iseq_ptr(rb_method_definition_t *def)
     return def_iseq_ptr(def);
 }
 
+rb_iseq_t *
+rb_get_iseq_body_local_iseq(rb_iseq_t * iseq) {
+    return iseq->body->local_iseq;
+}
+
 unsigned int
 rb_get_iseq_body_local_table_size(rb_iseq_t* iseq) {
     return iseq->body->local_table_size;
@@ -413,11 +418,41 @@ rb_get_iseq_body_stack_max(rb_iseq_t* iseq) {
     return iseq->body->stack_max;
 }
 
-int
+bool
 rb_get_iseq_flags_has_opt(rb_iseq_t* iseq) {
     return iseq->body->param.flags.has_opt;
 }
 
+bool
+rb_get_iseq_flags_has_kw(rb_iseq_t* iseq) {
+    return iseq->body->param.flags.has_kw;
+}
+
+bool
+rb_get_iseq_flags_has_post(rb_iseq_t* iseq) {
+    return iseq->body->param.flags.has_post;
+}
+
+bool
+rb_get_iseq_flags_has_kwrest(rb_iseq_t* iseq) {
+    return iseq->body->param.flags.has_kwrest;
+}
+
+bool
+rb_get_iseq_flags_has_rest(rb_iseq_t *iseq) {
+    return iseq->body->param.flags.has_rest;
+}
+
+bool
+rb_get_iseq_flags_has_block(rb_iseq_t* iseq) {
+    return iseq->body->param.flags.has_block;
+}
+
+bool
+rb_get_iseq_flags_has_accepts_no_kwarg(rb_iseq_t* iseq) {
+    return iseq->body->param.flags.accepts_no_kwarg;
+}
+
 const rb_seq_param_keyword_struct*
 rb_get_iseq_body_param_keyword(rb_iseq_t* iseq) {
     return iseq->body->param.keyword;
@@ -443,22 +478,6 @@ rb_get_iseq_body_param_opt_table(rb_iseq_t* iseq) {
     return iseq->body->param.opt_table;
 }
 
-// Returns whether the iseq only needs positional (lead) argument setup.
-bool
-rb_iseq_needs_lead_args_only(const rb_iseq_t *iseq)
-{
-    // When iseq->body->local_iseq == iseq, setup_parameters_complex()
-    // doesn't do anything to setup the block parameter.
-    bool takes_block = iseq->body->param.flags.has_block;
-    return (!takes_block || iseq->body->local_iseq == iseq) &&
-        iseq->body->param.flags.has_opt          == false &&
-        iseq->body->param.flags.has_rest         == false &&
-        iseq->body->param.flags.has_post         == false &&
-        iseq->body->param.flags.has_kw           == false &&
-        iseq->body->param.flags.has_kwrest       == false &&
-        iseq->body->param.flags.accepts_no_kwarg == false;
-}
-
 // If true, the iseq is leaf and it can be replaced by a single C call.
 bool
 rb_leaf_invokebuiltin_iseq_p(const rb_iseq_t *iseq)
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
@@ -676,7 +676,7 @@ pub fn gen_entry_prologue(cb: &mut CodeBlock, iseq: IseqPtr) -> Option<CodePtr>
     // has optional parameters, we'll add a runtime check that the PC we've
     // compiled for is the same PC that the interpreter wants us to run with.
     // If they don't match, then we'll take a side exit.
-    if unsafe { get_iseq_flags_has_opt(iseq) } != 0 {
+    if unsafe { get_iseq_flags_has_opt(iseq) } {
         gen_pc_guard(cb, iseq);
     }
 
@@ -3696,80 +3696,100 @@ fn gen_send_iseq(jit: &mut JITState, ctx: &mut Context, cb: &mut CodeBlock, ocb:
     // specified at the call site. We need to keep track of the fact that this
     // value is present on the stack in order to properly set up the callee's
     // stack pointer.
-    let mut doing_kw_call = false;
+    let doing_kw_call = unsafe { get_iseq_flags_has_kw(iseq) };
+    let supplying_kws = unsafe { vm_ci_flag(ci) & VM_CALL_KWARG } != 0;
 
     if unsafe{ vm_ci_flag(ci) } & VM_CALL_TAILCALL != 0 {
         // We can't handle tailcalls
         gen_counter_incr!(cb, send_iseq_tailcall);
         return CantCompile;
     }
 
-    // Arity handling and optional parameter setup
-    let mut num_params = unsafe { get_iseq_body_param_size(iseq) as i32 };
-    let mut start_pc_offset:u32 = 0;
-
-    if unsafe { iseq_needs_lead_args_only(iseq) } {
-        // If we have keyword arguments being passed to a callee that only takes
-        // positionals, then we need to allocate a hash. For now we're going to
-        // call that too complex and bail.
-        if unsafe { vm_ci_flag(ci) } & VM_CALL_KWARG != 0 {
+    // No support for callees with these parameters yet as they require allocation
+    // or complex handling.
+    if unsafe { get_iseq_flags_has_rest(iseq) ||
+        get_iseq_flags_has_post(iseq) ||
+        get_iseq_flags_has_kwrest(iseq) } {
             gen_counter_incr!(cb, send_iseq_complex_callee);
             return CantCompile;
-        }
+    }
 
-        num_params = unsafe { get_iseq_body_param_lead_num(iseq) };
+    // If we have keyword arguments being passed to a callee that only takes
+    // positionals, then we need to allocate a hash. For now we're going to
+    // call that too complex and bail.
+    if supplying_kws && ! unsafe { get_iseq_flags_has_kw(iseq) } {
+        gen_counter_incr!(cb, send_iseq_complex_callee);
+        return CantCompile;
+    }
 
-        if num_params != argc {
-            gen_counter_incr!(cb, send_iseq_arity_error);
-            return CantCompile;
-        }
+    // If we have a method accepting no kwargs (**nil), exit if we have passed
+    // it any kwargs.
+    if supplying_kws && unsafe { get_iseq_flags_has_accepts_no_kwarg(iseq) } {
+        gen_counter_incr!(cb, send_iseq_complex_callee);
+        return CantCompile;
     }
-    else if unsafe { rb_iseq_only_optparam_p(iseq) } {
-        // If we have keyword arguments being passed to a callee that only takes
-        // positionals and optionals, then we need to allocate a hash. For now
-        // we're going to call that too complex and bail.
-        if unsafe { vm_ci_flag(ci) } & VM_CALL_KWARG != 0 {
+
+    // For computing number of locals to set up for the callee
+    let mut num_params = unsafe { get_iseq_body_param_size(iseq) };
+
+    // Block parameter handling. This mirrors setup_parameters_complex().
+    if unsafe { get_iseq_flags_has_block(iseq) } {
+        if unsafe { get_iseq_body_local_iseq(iseq) == iseq } {
+            num_params -= 1;
+        } else {
+            // In this case (param.flags.has_block && local_iseq != iseq),
+            // the block argument is setup as a local variable and requires
+            // materialization (allocation). Bail.
             gen_counter_incr!(cb, send_iseq_complex_callee);
             return CantCompile;
         }
+    }
 
-        // These are iseqs with 0 or more required parameters followed by 1
-        // or more optional parameters.
-        // We follow the logic of vm_call_iseq_setup_normal_opt_start()
-        // and these are the preconditions required for using that fast path.
-        //RUBY_ASSERT(vm_ci_markable(ci) && ((vm_ci_flag(ci) &
-        //                (VM_CALL_KW_SPLAT | VM_CALL_KWARG | VM_CALL_ARGS_SPLAT)) == 0));
+    let mut start_pc_offset = 0;
+    let required_num = unsafe { get_iseq_body_param_lead_num(iseq) };
 
-        let required_num = unsafe { get_iseq_body_param_lead_num(iseq) };
-        let opts_filled = argc - required_num;
-        let opt_num = unsafe { get_iseq_body_param_opt_num(iseq) };
+    // This struct represents the metadata about the caller-specified
+    // keyword arguments.
+    let kw_arg = unsafe { vm_ci_kwarg(ci) };
+    let kw_arg_num = if kw_arg.is_null() { 0 } else { unsafe { get_cikw_keyword_len(kw_arg) } };
 
-        if opts_filled < 0 || opts_filled > opt_num {
-            gen_counter_incr!(cb, send_iseq_arity_error);
-            return CantCompile;
-        }
+    // Arity handling and optional parameter setup
+    let opts_filled = argc - required_num - kw_arg_num;
+    let opt_num = unsafe { get_iseq_body_param_opt_num(iseq) };
+    let opts_missing:u32 = (opt_num - opts_filled).try_into().unwrap();
+
+    if opts_filled < 0 || opts_filled > opt_num {
+        gen_counter_incr!(cb, send_iseq_arity_error);
+        return CantCompile;
+    }
+
+    // If we have unfilled optional arguments and keyword arguments then we
+    // would need to move adjust the arguments location to account for that.
+    // For now we aren't handling this case.
+    if doing_kw_call && opts_missing > 0 {
+        gen_counter_incr!(cb, send_iseq_complex_callee);
+        return CantCompile;
+    }
 
-        num_params -= opt_num - opts_filled;
+    if opt_num > 0 {
+        num_params -= opts_missing;
         unsafe {
             let opt_table = get_iseq_body_param_opt_table(iseq);
             start_pc_offset = (*opt_table.offset(opts_filled as isize)).as_u32();
         }
     }
 
-    else if unsafe { rb_iseq_only_kwparam_p(iseq) } {
-        let lead_num = unsafe { get_iseq_body_param_lead_num(iseq) };
-
-        doing_kw_call = true;
-
+    if doing_kw_call {
         // Here we're calling a method with keyword arguments and specifying
         // keyword arguments at this call site.
 
-        // This struct represents the metadata about the caller-specified
-        // keyword arguments.
-        let kw_arg = unsafe { vm_ci_kwarg(ci) };
-
+        // This struct represents the metadata about the callee-specified
+        // keyword parameters.
         let keyword = unsafe { get_iseq_body_param_keyword(iseq) };
         let keyword_num = unsafe { (*keyword).num };
+        let keyword_required_num = unsafe { (*keyword).required_num };
+
+        let mut required_kwargs_filled = 0;
 
         if keyword_num > 30 {
             // We have so many keywords that (1 << num) encoded as a FIXNUM
@@ -3779,13 +3799,8 @@ fn gen_send_iseq(jit: &mut JITState, ctx: &mut Context, cb: &mut CodeBlock, ocb:
             return CantCompile;
         }
 
-        if unsafe { vm_ci_flag(ci) } & VM_CALL_KWARG != 0 {
-            // Check that the size of non-keyword arguments matches
-            if lead_num != argc - unsafe { get_cikw_keyword_len(kw_arg) } {
-                gen_counter_incr!(cb, send_iseq_complex_callee);
-                return CantCompile;
-            }
-
+        // Check that the kwargs being passed are valid
+        if supplying_kws {
             // This is the list of keyword arguments that the callee specified
             // in its initial declaration.
             let callee_kwargs = unsafe { (*keyword).table };
@@ -3822,35 +3837,22 @@ fn gen_send_iseq(jit: &mut JITState, ctx: &mut Context, cb: &mut CodeBlock, ocb:
                     gen_counter_incr!(cb, send_iseq_kwargs_mismatch);
                     return CantCompile;
                 }
-            }
-        }
-        else if argc == lead_num {
-            // Here we are calling a method that accepts keyword arguments
-            // (optional or required) but we're not passing any keyword
-            // arguments at this call site
-
-            if unsafe { (*keyword).required_num } != 0 {
-                // If any of the keywords are required this is a mismatch
-                gen_counter_incr!(cb, send_iseq_kwargs_mismatch);
-                return CantCompile;
-            }
 
-            doing_kw_call = true;
+                // Keep a count to ensure all required kwargs are specified
+                if callee_idx < keyword_required_num {
+                    required_kwargs_filled += 1;
+                }
+            }
         }
-        else {
-            gen_counter_incr!(cb, send_iseq_complex_callee);
+        assert!(required_kwargs_filled <= keyword_required_num);
+        if required_kwargs_filled != keyword_required_num {
+            gen_counter_incr!(cb, send_iseq_kwargs_mismatch);
             return CantCompile;
         }
     }
-    else {
-        // Only handle iseqs that have simple parameter setup.
-        // See vm_callee_setup_arg().
-        gen_counter_incr!(cb, send_iseq_complex_callee);
-        return CantCompile;
-    }
 
     // Number of locals that are not parameters
-    let num_locals = unsafe { get_iseq_body_local_table_size(iseq) as i32 } - num_params;
+    let num_locals = unsafe { get_iseq_body_local_table_size(iseq) as i32 } - (num_params as i32);
 
     // Create a side-exit to fall back to the interpreter
     let side_exit = get_side_exit(jit, ocb, ctx);
@@ -3906,7 +3908,10 @@ fn gen_send_iseq(jit: &mut JITState, ctx: &mut Context, cb: &mut CodeBlock, ocb:
     if doing_kw_call {
         // Here we're calling a method with keyword arguments and specifying
         // keyword arguments at this call site.
-        let lead_num = unsafe { get_iseq_body_param_lead_num(iseq) };
+
+        // Number of positional arguments the callee expects before the first
+        // keyword argument
+        let args_before_kw = required_num + opt_num;
 
         // This struct represents the metadata about the caller-specified
         // keyword arguments.
@@ -4013,8 +4018,8 @@ fn gen_send_iseq(jit: &mut JITState, ctx: &mut Context, cb: &mut CodeBlock, ocb:
                     // to perform the actual swapping at runtime.
                     let swap_idx_i32:i32 = swap_idx.try_into().unwrap();
                     let kwarg_idx_i32:i32 = kwarg_idx.try_into().unwrap();
-                    let offset0:u16 = (argc - 1 - swap_idx_i32 - lead_num).try_into().unwrap();
-                    let offset1:u16 = (argc - 1 - kwarg_idx_i32 - lead_num).try_into().unwrap();
+                    let offset0:u16 = (argc - 1 - swap_idx_i32 - args_before_kw).try_into().unwrap();
+                    let offset1:u16 = (argc - 1 - kwarg_idx_i32 - args_before_kw).try_into().unwrap();
                     stack_swap(ctx, cb, offset0, offset1, REG1, REG0);
 
                     // Next we're going to do some bookkeeping on our end so
diff --git a/yjit/src/cruby.rs b/yjit/src/cruby.rs
@@ -168,6 +168,9 @@ extern "C" {
     #[link_name = "rb_iseq_encoded_size"]
     pub fn get_iseq_encoded_size(iseq: IseqPtr) -> c_uint;
 
+    #[link_name = "rb_get_iseq_body_local_iseq"]
+    pub fn get_iseq_body_local_iseq(iseq: IseqPtr) -> IseqPtr;
+
     #[link_name = "rb_get_iseq_body_iseq_encoded"]
     pub fn get_iseq_body_iseq_encoded(iseq: IseqPtr) -> *mut VALUE;
 
@@ -178,7 +181,25 @@ extern "C" {
     pub fn get_iseq_body_stack_max(iseq:IseqPtr) -> c_uint;
 
     #[link_name = "rb_get_iseq_flags_has_opt"]
-    pub fn get_iseq_flags_has_opt(iseq: IseqPtr) -> c_int;
+    pub fn get_iseq_flags_has_opt(iseq: IseqPtr) -> bool;
+
+    #[link_name = "rb_get_iseq_flags_has_kw"]
+    pub fn get_iseq_flags_has_kw(iseq: IseqPtr) -> bool;
+
+    #[link_name = "rb_get_iseq_flags_has_rest"]
+    pub fn get_iseq_flags_has_rest(iseq: IseqPtr) -> bool;
+
+    #[link_name = "rb_get_iseq_flags_has_post"]
+    pub fn get_iseq_flags_has_post(iseq: IseqPtr) -> bool;
+
+    #[link_name = "rb_get_iseq_flags_has_kwrest"]
+    pub fn get_iseq_flags_has_kwrest(iseq: IseqPtr) -> bool;
+
+    #[link_name = "rb_get_iseq_flags_has_block"]
+    pub fn get_iseq_flags_has_block(iseq: IseqPtr) -> bool;
+
+    #[link_name = "rb_get_iseq_flags_has_accepts_no_kwarg"]
+    pub fn get_iseq_flags_has_accepts_no_kwarg(iseq: IseqPtr) -> bool;
 
     #[link_name = "rb_get_iseq_body_local_table_size"]
     pub fn get_iseq_body_local_table_size(iseq: IseqPtr) -> c_uint;
@@ -204,9 +225,6 @@ extern "C" {
     #[link_name = "rb_get_cikw_keywords_idx"]
     pub fn get_cikw_keywords_idx(cikw: *const rb_callinfo_kwarg, idx: c_int) -> VALUE;
 
-    #[link_name = "rb_iseq_needs_lead_args_only"]
-    pub fn iseq_needs_lead_args_only(iseq: *const rb_iseq_t) -> bool;
-
     #[link_name = "rb_get_call_data_ci"]
     pub fn get_call_data_ci(cd: * const rb_call_data) -> *const rb_callinfo;
 
