YJIT: Refactor into gen_push_frame (#6412)

This refactors the "push frame" operation common to both gen_send_iseq
and gen_send_cfunc into its own method. This allows that logic to live
in one place.

1 files changed, 161 insertions, 102 deletions

diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index a61795f5d4..15ff2a467f 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -3998,6 +3998,135 @@ unsafe extern "C" fn build_kwhash(ci: *const rb_callinfo, sp: *const VALUE) -> V
     hash
 }
 
+enum BlockHandler {
+    None,
+    CurrentFrame,
+}
+
+struct ControlFrame {
+    recv: Opnd,
+    sp: Opnd,
+    iseq: Option<IseqPtr>,
+    pc: Option<u64>,
+    frame_type: u32,
+    block_handler: BlockHandler,
+    cme: *const rb_callable_method_entry_t,
+    local_size: i32
+}
+
+// Codegen performing a similar (but not identical) function to vm_push_frame
+//
+// This will generate the code to:
+//   * initialize locals to Qnil
+//   * push the environment (cme, block handler, frame type)
+//   * push a new CFP
+//   * save the new CFP to ec->cfp
+//
+// Notes:
+//   * Provided sp should point to the new frame's sp, immediately following locals and the environment
+//   * At entry, CFP points to the caller (not callee) frame
+//   * At exit, ec->cfp is updated to the pushed CFP
+//   * CFP and SP registers are updated only if switch_in_jit is set
+//   * Stack overflow is not checked (should be done by the caller)
+//   * Interrupts are not checked (should be done by the caller)
+fn gen_push_frame(
+    jit: &mut JITState,
+    ctx: &mut Context,
+    asm: &mut Assembler,
+    set_pc_cfp: bool, // if true CFP and SP will be switched to the callee
+    frame: ControlFrame,
+) {
+    assert!(frame.local_size >= 0);
+
+    let sp = frame.sp;
+
+    let num_locals = frame.local_size;
+    if num_locals > 0 {
+        asm.comment("initialize locals");
+
+        // Initialize local variables to Qnil
+        for i in 0..num_locals {
+            let offs = (SIZEOF_VALUE as i32) * (i - num_locals - 3);
+            asm.store(Opnd::mem(64, sp, offs), Qnil.into());
+        }
+    }
+
+    asm.comment("push cme, block handler, frame type");
+
+    // Write method entry at sp[-3]
+    // sp[-3] = me;
+    // Use compile time cme. It's assumed to be valid because we are notified when
+    // any cme we depend on become outdated. See yjit_method_lookup_change().
+    asm.store(Opnd::mem(64, sp, SIZEOF_VALUE_I32 * -3), VALUE::from(frame.cme).into());
+
+    // Write block handler at sp[-2]
+    // sp[-2] = block_handler;
+    let block_handler: Opnd = match frame.block_handler {
+        BlockHandler::None => {
+            VM_BLOCK_HANDLER_NONE.into()
+        },
+        BlockHandler::CurrentFrame => {
+            let cfp_self = asm.lea(Opnd::mem(64, CFP, RUBY_OFFSET_CFP_SELF));
+            asm.or(cfp_self, Opnd::Imm(1))
+        },
+    };
+    asm.store(Opnd::mem(64, sp, SIZEOF_VALUE_I32 * -2), block_handler);
+
+    // Write env flags at sp[-1]
+    // sp[-1] = frame_type;
+    asm.store(Opnd::mem(64, sp, SIZEOF_VALUE_I32 * -1), frame.frame_type.into());
+
+    // Allocate a new CFP (ec->cfp--)
+    fn cfp_opnd(offset: i32) -> Opnd {
+        Opnd::mem(64, CFP, offset - (RUBY_SIZEOF_CONTROL_FRAME as i32))
+    }
+
+    // Setup the new frame
+    // *cfp = (const struct rb_control_frame_struct) {
+    //    .pc         = <unset for iseq, 0 for cfunc>,
+    //    .sp         = sp,
+    //    .iseq       = <iseq for iseq, 0 for cfunc>,
+    //    .self       = recv,
+    //    .ep         = <sp - 1>,
+    //    .block_code = 0,
+    //    .__bp__     = sp,
+    // };
+    asm.comment("push callee control frame");
+
+    // For an iseq call PC may be None, in which case we will not set PC and will allow jitted code
+    // to set it as necessary.
+    let pc = if let Some(pc) = frame.pc {
+        asm.mov(cfp_opnd(RUBY_OFFSET_CFP_PC), pc.into());
+    };
+    asm.mov(cfp_opnd(RUBY_OFFSET_CFP_BP), sp);
+    asm.mov(cfp_opnd(RUBY_OFFSET_CFP_SP), sp);
+    let iseq: Opnd = if let Some(iseq) = frame.iseq {
+        VALUE::from(iseq).into()
+    } else {
+        0.into()
+    };
+    asm.mov(cfp_opnd(RUBY_OFFSET_CFP_ISEQ), iseq);
+    asm.mov(cfp_opnd(RUBY_OFFSET_CFP_SELF), frame.recv);
+    asm.mov(cfp_opnd(RUBY_OFFSET_CFP_BLOCK_CODE), 0.into());
+
+    if set_pc_cfp {
+        // Saving SP before calculating ep avoids a dependency on a register
+        // However this must be done after referencing frame.recv, which may be SP-relative
+        asm.mov(SP, sp);
+    }
+    let ep = asm.sub(sp, SIZEOF_VALUE.into());
+    asm.mov(cfp_opnd(RUBY_OFFSET_CFP_EP), ep);
+
+    asm.comment("switch to new CFP");
+    let new_cfp = asm.lea(cfp_opnd(0));
+    if set_pc_cfp {
+        asm.mov(CFP, new_cfp);
+        asm.store(Opnd::mem(64, EC, RUBY_OFFSET_EC_CFP), CFP);
+    } else {
+        asm.store(Opnd::mem(64, EC, RUBY_OFFSET_EC_CFP), new_cfp);
+    }
+}
+
 fn gen_send_cfunc(
     jit: &mut JITState,
     ctx: &mut Context,
@@ -4103,60 +4232,27 @@ fn gen_send_cfunc(
     // sp += 3
     let sp = asm.lea(ctx.sp_opnd((SIZEOF_VALUE as isize) * 3));
 
-    // Write method entry at sp[-3]
-    // sp[-3] = me;
-    // Put compile time cme into REG1. It's assumed to be valid because we are notified when
-    // any cme we depend on become outdated. See yjit_method_lookup_change().
-    asm.comment("push cme, block handler, frame type");
-    asm.mov(Opnd::mem(64, sp, SIZEOF_VALUE_I32 * -3), Opnd::UImm(cme as u64));
-
-    // Write block handler at sp[-2]
-    // sp[-2] = block_handler;
-    if let Some(_block_iseq) = block {
-        // reg1 = VM_BH_FROM_ISEQ_BLOCK(VM_CFP_TO_CAPTURED_BLOCK(reg_cfp));
-        let cfp_self = asm.lea(Opnd::mem(64, CFP, RUBY_OFFSET_CFP_SELF));
-        let block_handler = asm.or(cfp_self, Opnd::Imm(1));
-        asm.store(Opnd::mem(64, sp, SIZEOF_VALUE_I32 * -2), block_handler);
+    let frame_block_handler = if let Some(_block_iseq) = block {
+        BlockHandler::CurrentFrame
     } else {
-        let dst_opnd = Opnd::mem(64, sp, SIZEOF_VALUE_I32 * -2);
-        asm.store(dst_opnd, Opnd::UImm(VM_BLOCK_HANDLER_NONE.into()));
-    }
+        BlockHandler::None
+    };
 
-    // Write env flags at sp[-1]
-    // sp[-1] = frame_type;
     let mut frame_type = VM_FRAME_MAGIC_CFUNC | VM_FRAME_FLAG_CFRAME | VM_ENV_FLAG_LOCAL;
     if !kw_arg.is_null() {
         frame_type |= VM_FRAME_FLAG_CFRAME_KW
     }
-    asm.store(Opnd::mem(64, sp, SIZEOF_VALUE_I32 * -1), Opnd::UImm(frame_type.into()));
-
-    // Allocate a new CFP (ec->cfp--)
-    asm.comment("push callee control frame");
-    let ec_cfp_opnd = Opnd::mem(64, EC, RUBY_OFFSET_EC_CFP);
-    let new_cfp = asm.sub(ec_cfp_opnd, Opnd::UImm(RUBY_SIZEOF_CONTROL_FRAME as u64));
-    asm.mov(ec_cfp_opnd, new_cfp);
 
-    // Setup the new frame
-    // *cfp = (const struct rb_control_frame_struct) {
-    //    .pc         = 0,
-    //    .sp         = sp,
-    //    .iseq       = 0,
-    //    .self       = recv,
-    //    .ep         = sp - 1,
-    //    .block_code = 0,
-    //    .__bp__     = sp,
-    // };
-
-    // Can we re-use ec_cfp_opnd from above?
-    let ec_cfp_opnd = asm.load(Opnd::mem(64, EC, RUBY_OFFSET_EC_CFP));
-    asm.mov(Opnd::mem(64, ec_cfp_opnd, RUBY_OFFSET_CFP_PC), Opnd::Imm(0));
-    asm.mov(Opnd::mem(64, ec_cfp_opnd, RUBY_OFFSET_CFP_SP), sp);
-    asm.mov(Opnd::mem(64, ec_cfp_opnd, RUBY_OFFSET_CFP_ISEQ), Opnd::Imm(0));
-    asm.mov(Opnd::mem(64, ec_cfp_opnd, RUBY_OFFSET_CFP_BLOCK_CODE), Opnd::Imm(0));
-    asm.mov(Opnd::mem(64, ec_cfp_opnd, RUBY_OFFSET_CFP_BP), sp);
-    let ep = asm.sub(sp, Opnd::UImm(SIZEOF_VALUE as u64));
-    asm.mov(Opnd::mem(64, ec_cfp_opnd, RUBY_OFFSET_CFP_EP), ep);
-    asm.mov(Opnd::mem(64, ec_cfp_opnd, RUBY_OFFSET_CFP_SELF), recv);
+    gen_push_frame(jit, ctx, asm, false, ControlFrame {
+        frame_type,
+        block_handler: frame_block_handler,
+        cme,
+        recv,
+        sp,
+        pc: Some(0),
+        iseq: None,
+        local_size: 0,
+    });
 
     if !kw_arg.is_null() {
         // Build a hash from all kwargs passed
@@ -4219,10 +4315,10 @@ fn gen_send_cfunc(
     asm.mov(stack_ret, ret);
 
     // Pop the stack frame (ec->cfp++)
-    // Can we reuse ec_cfp_opnd from above?
+    // Instead of recalculating, we can reuse the previous CFP, which is stored in a callee-saved
+    // register
     let ec_cfp_opnd = Opnd::mem(64, EC, RUBY_OFFSET_EC_CFP);
-    let new_cfp = asm.add(ec_cfp_opnd, Opnd::UImm(RUBY_SIZEOF_CONTROL_FRAME as u64));
-    asm.store(ec_cfp_opnd, new_cfp);
+    asm.store(ec_cfp_opnd, CFP);
 
     // cfunc calls may corrupt types
     ctx.clear_local_types();
@@ -4766,62 +4862,25 @@ fn gen_send_iseq(
         (SIZEOF_VALUE as isize) * (3 + (num_locals as isize) + if doing_kw_call { 1 } else { 0 });
     let callee_sp = asm.lea(ctx.sp_opnd(offs));
 
-    // Initialize local variables to Qnil
-    for i in 0..num_locals {
-        let offs = (SIZEOF_VALUE as i32) * (i - num_locals - 3);
-        asm.store(Opnd::mem(64, callee_sp, offs), Qnil.into());
-    }
-
-    // Write the callee CME on the stack. It's assumed to be valid because we are notified when
-    // any cme we depend on become outdated. See yjit_method_lookup_change().
-    // Write method entry at sp[-3]
-    // sp[-3] = me;
-    asm.comment("push cme, block handler, frame type");
-    asm.store(Opnd::mem(64, callee_sp, SIZEOF_VALUE_I32 * -3), VALUE(cme as usize).into());
-
-    // Write block handler at sp[-2]
-    // sp[-2] = block_handler;
-    match block {
-        Some(_) => {
-            // reg1 = VM_BH_FROM_ISEQ_BLOCK(VM_CFP_TO_CAPTURED_BLOCK(reg_cfp));
-            let block_handler = asm.lea(Opnd::mem(64, CFP, RUBY_OFFSET_CFP_SELF));
-            let block_handler = asm.or(block_handler, 1.into());
-            asm.store(Opnd::mem(64, callee_sp, SIZEOF_VALUE_I32 * -2), block_handler);
-        }
-        None => {
-            asm.store(Opnd::mem(64, callee_sp, SIZEOF_VALUE_I32 * -2), VM_BLOCK_HANDLER_NONE.into());
-        }
-    }
+    let frame_block_handler = if let Some(_) = block {
+        BlockHandler::CurrentFrame
+    } else {
+        BlockHandler::None
+    };
 
-    // Write env flags at sp[-1]
-    // sp[-1] = frame_type;
     let frame_type = VM_FRAME_MAGIC_METHOD | VM_ENV_FLAG_LOCAL;
-    asm.store(Opnd::mem(64, callee_sp, SIZEOF_VALUE_I32 * -1), frame_type.into());
-
-    asm.comment("push callee control frame");
-    // Allocate a new CFP (ec->cfp--)
-    let new_cfp = asm.sub(CFP, (RUBY_SIZEOF_CONTROL_FRAME as u64).into());
-    asm.mov(CFP, new_cfp);
-    asm.store(Opnd::mem(64, EC, RUBY_OFFSET_EC_CFP), CFP);
 
     // Setup the new frame
-    // *cfp = (const struct rb_control_frame_struct) {
-    //    .pc         = pc,
-    //    .sp         = sp,
-    //    .iseq       = iseq,
-    //    .self       = recv,
-    //    .ep         = sp - 1,
-    //    .block_code = 0,
-    //    .__bp__     = sp,
-    // };
-    asm.mov(Opnd::mem(64, CFP, RUBY_OFFSET_CFP_SELF), recv);
-    asm.mov(SP, callee_sp);
-    asm.mov(Opnd::mem(64, CFP, RUBY_OFFSET_CFP_SP), callee_sp);
-    asm.mov(Opnd::mem(64, CFP, RUBY_OFFSET_CFP_BP), callee_sp);
-    let callee_ep = asm.sub(callee_sp, (SIZEOF_VALUE as u64).into());
-    asm.mov(Opnd::mem(64, CFP, RUBY_OFFSET_CFP_EP), callee_ep);
-    asm.mov(Opnd::mem(64, CFP, RUBY_OFFSET_CFP_ISEQ), VALUE(iseq as usize).into());
-    asm.mov(Opnd::mem(64, CFP, RUBY_OFFSET_CFP_BLOCK_CODE), 0.into());
+    gen_push_frame(jit, ctx, asm, true, ControlFrame {
+        frame_type,
+        block_handler: frame_block_handler,
+        cme,
+        recv,
+        sp: callee_sp,
+        iseq: Some(iseq),
+        pc: None, // We are calling into jitted code, which will set the PC as necessary
+        local_size: num_locals
+    });
 
     // No need to set cfp->pc since the callee sets it whenever calling into routines
     // that could look at it through jit_save_pc().