ZJIT: Introduce the InvokeSuperDirect HIR instruction

The new instruction is an optimized version of `InvokeSuper` when we know the `super` target is an ISEQ.

Author: nirvdrum

zjit/src/codegen.rs

@@ -397,8 +397,10 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         // Give up SendWithoutBlockDirect for 6+ args since asm.ccall() doesn't support it.
         Insn::SendWithoutBlockDirect { cd, state, args, .. } if args.len() + 1 > C_ARG_OPNDS.len() => // +1 for self
             gen_send_without_block(jit, asm, *cd, &function.frame_state(*state), SendFallbackReason::TooManyArgsForLir),
-        Insn::SendWithoutBlockDirect { cme, iseq, recv, args, state, .. } => gen_send_without_block_direct(cb, jit, asm, *cme, *iseq, opnd!(recv), opnds!(args), &function.frame_state(*state)),
+        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::InvokeSuper { cd, blockiseq, state, reason, .. } => gen_invokesuper(jit, asm, cd, blockiseq, &function.frame_state(state), reason),
+        Insn::InvokeSuperDirect { recv, current_cme, super_cme, args, state, .. } =>
+            gen_invokesuper_direct(cb, jit, asm, *current_cme, *super_cme, opnd!(recv), opnds!(args), &function.frame_state(*state)),
         &Insn::InvokeBlock { cd, state, reason, .. } => gen_invokeblock(jit, asm, cd, &function.frame_state(state), reason),
         // Ensure we have enough room fit ec, self, and arguments
         // TODO remove this check when we have stack args (we can use Time.new to test it)
@@ -1299,8 +1301,10 @@ fn gen_send_without_block(
     )
 }
 
-/// Compile a direct jump to an ISEQ call without block
-fn gen_send_without_block_direct(
+/// Compile a direct call to an ISEQ method.
+/// If `block_handler` is provided, it's used as the specval for the new frame (for forwarding blocks).
+/// Otherwise, `VM_BLOCK_HANDLER_NONE` is used.
+fn gen_send_iseq_direct(
     cb: &mut CodeBlock,
     jit: &mut JITState,
     asm: &mut Assembler,
@@ -1309,6 +1313,7 @@ fn gen_send_without_block_direct(
     recv: Opnd,
     args: Vec<Opnd>,
     state: &FrameState,
+    block_handler: Option<Opnd>,
 ) -> lir::Opnd {
     gen_incr_counter(asm, Counter::iseq_optimized_send_count);
 
@@ -1335,7 +1340,8 @@ fn gen_send_without_block_direct(
         let bmethod_specval = (capture.ep.addr() | 1).into();
         (bmethod_frame_type, bmethod_specval)
     } else {
-        (VM_FRAME_MAGIC_METHOD | VM_ENV_FLAG_LOCAL, VM_BLOCK_HANDLER_NONE.into())
+        let specval = block_handler.unwrap_or_else(|| VM_BLOCK_HANDLER_NONE.into());
+        (VM_FRAME_MAGIC_METHOD | VM_ENV_FLAG_LOCAL, specval)
     };
 
     // Set up the new frame
@@ -1459,6 +1465,37 @@ fn gen_invokesuper(
     )
 }
 
+/// Compile an optimized super call to an ISEQ method
+fn gen_invokesuper_direct(
+    cb: &mut CodeBlock,
+    jit: &mut JITState,
+    asm: &mut Assembler,
+    current_cme: *const rb_callable_method_entry_t,
+    super_cme: *const rb_callable_method_entry_t,
+    recv: lir::Opnd,
+    args: Vec<lir::Opnd>,
+    state: &FrameState,
+) -> lir::Opnd {
+    // Guard that ep[VM_ENV_DATA_INDEX_ME_CREF] matches current_cme, ensuring that we're calling
+    // `super` from the expected method context.
+    asm_comment!(asm, "guard super method entry");
+    let lep = gen_get_lep(jit, asm);
+    let ep_me_opnd = Opnd::mem(64, lep, SIZEOF_VALUE_I32 * VM_ENV_DATA_INDEX_ME_CREF);
+    let ep_me = asm.load(ep_me_opnd);
+    asm.cmp(ep_me, Opnd::UImm(current_cme as u64));
+    asm.jne(side_exit(jit, state, SideExitReason::GuardSuperMethodEntry));
+
+    // Get the super method's ISEQ since we'll be dispatching to that.
+    let super_iseq = unsafe { get_def_iseq_ptr((*super_cme).def) };
+
+    // Read the block handler from the LEP to forward the caller's block.
+    asm_comment!(asm, "load block handler from LEP");
+    let block_handler = asm.load(Opnd::mem(64, lep, SIZEOF_VALUE_I32 * VM_ENV_DATA_INDEX_SPECVAL));
+
+    // Generate the code to invoke the `super` target method.
+    gen_send_iseq_direct(cb, jit, asm, super_cme, super_iseq, recv, args, state, Some(block_handler))
+}
+
 /// Compile a string resurrection
 fn gen_string_copy(asm: &mut Assembler, recv: Opnd, chilled: bool, state: &FrameState) -> Opnd {
     // TODO: split rb_ec_str_resurrect into separate functions

zjit/src/hir.rs

@@ -493,6 +493,7 @@ pub enum SideExitReason {
     GuardNotFrozen,
     GuardLess,
     GuardGreaterEq,
+    GuardSuperMethodEntry,
     PatchPoint(Invariant),
     CalleeSideExit,
     ObjToStringFallback,
@@ -884,6 +885,17 @@ pub enum Insn {
         state: InsnId,
         reason: SendFallbackReason,
     },
+    /// Optimized super call to an ISEQ method
+    InvokeSuperDirect {
+        recv: InsnId,
+        cd: *const rb_call_data,
+        /// The CME of the method containing the super call (for runtime guard)
+        current_cme: *const rb_callable_method_entry_t,
+        /// The resolved super method's CME (for PatchPoint and to get target ISEQ)
+        super_cme: *const rb_callable_method_entry_t,
+        args: Vec<InsnId>,
+        state: InsnId,
+    },
     InvokeBlock {
         cd: *const rb_call_data,
         args: Vec<InsnId>,
@@ -1279,6 +1291,13 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> {
                 write!(f, " # SendFallbackReason: {reason}")?;
                 Ok(())
             }
+            Insn::InvokeSuperDirect { recv, args, .. } => {
+                write!(f, "InvokeSuperDirect {recv}")?;
+                for arg in args {
+                    write!(f, ", {arg}")?;
+                }
+                Ok(())
+            }
             Insn::InvokeBlock { args, reason, .. } => {
                 write!(f, "InvokeBlock")?;
                 for arg in args {
@@ -2018,6 +2037,14 @@ impl Function {
                 state,
                 reason,
             },
+            &InvokeSuperDirect { recv, cd, current_cme, super_cme, ref args, state } => InvokeSuperDirect {
+                recv: find!(recv),
+                cd,
+                current_cme,
+                super_cme,
+                args: find_vec!(args),
+                state,
+            },
             &InvokeBlock { cd, ref args, state, reason } => InvokeBlock {
                 cd,
                 args: find_vec!(args),
@@ -2203,6 +2230,7 @@ impl Function {
             Insn::Send { .. } => types::BasicObject,
             Insn::SendForward { .. } => types::BasicObject,
             Insn::InvokeSuper { .. } => types::BasicObject,
+            Insn::InvokeSuperDirect { .. } => types::BasicObject,
             Insn::InvokeBlock { .. } => types::BasicObject,
             Insn::InvokeBuiltin { return_type, .. } => return_type.unwrap_or(types::BasicObject),
             Insn::Defined { pushval, .. } => Type::from_value(*pushval).union(types::NilClass),
@@ -2968,6 +2996,74 @@ impl Function {
                             self.push_insn_id(block, insn_id);
                         };
                     }
+                    Insn::InvokeSuper { recv, cd, blockiseq, args, state, .. } => {
+                        // Don't handle calls with literal blocks (e.g., super { ... })
+                        if !blockiseq.is_null() {
+                            self.push_insn_id(block, insn_id); continue;
+                        }
+
+                        // Don't handle calls with complex arguments (kwarg, splat, kw_splat, blockarg, forwarding)
+                        let ci = unsafe { get_call_data_ci(cd) };
+                        let flags = unsafe { rb_vm_ci_flag(ci) };
+                        if unspecializable_call_type(flags) {
+                            self.push_insn_id(block, insn_id); continue;
+                        }
+
+                        let frame_state = self.frame_state(state);
+
+                        // Get the profiled CME from the current method
+                        let Some(profiles) = self.profiles.as_ref() else {
+                            self.push_insn_id(block, insn_id); continue;
+                        };
+                        let Some(current_cme) = profiles.payload.profile.get_super_method_entry(frame_state.insn_idx) else {
+                            // No profile or polymorphic
+                            self.push_insn_id(block, insn_id); continue;
+                        };
+
+                        // Get defined_class and method ID from the profiled CME
+                        let current_defined_class = unsafe { (*current_cme).defined_class };
+                        let mid = unsafe { get_def_original_id((*current_cme).def) };
+
+                        // Compute superclass: RCLASS_SUPER(RCLASS_ORIGIN(defined_class))
+                        let superclass = unsafe { rb_class_get_superclass(RCLASS_ORIGIN(current_defined_class)) };
+                        if superclass.nil_p() {
+                            self.push_insn_id(block, insn_id); continue;
+                        }
+
+                        // Look up the super method
+                        let super_cme = unsafe { rb_callable_method_entry(superclass, mid) };
+                        if super_cme.is_null() {
+                            self.push_insn_id(block, insn_id); continue;
+                        }
+
+                        // Check if it's an ISEQ method
+                        let def_type = unsafe { get_cme_def_type(super_cme) };
+                        if def_type != VM_METHOD_TYPE_ISEQ {
+                            // TODO: Handle CFUNCs
+                            self.push_insn_id(block, insn_id); continue;
+                        }
+
+                        // Add PatchPoints for method redefinition
+                        // TODO: Add guard that ep[-2] matches current_cme
+                        self.push_insn(block, Insn::PatchPoint {
+                            invariant: Invariant::MethodRedefined {
+                                klass: unsafe { (*super_cme).defined_class },
+                                method: mid,
+                                cme: super_cme
+                            },
+                            state
+                        });
+
+                        let send_direct = self.push_insn(block, Insn::InvokeSuperDirect {
+                            recv,
+                            cd,
+                            current_cme,
+                            super_cme,
+                            args,
+                            state
+                        });
+                        self.make_equal_to(insn_id, send_direct);
+                    }
                     _ => { self.push_insn_id(block, insn_id); }
                 }
             }
@@ -3974,7 +4070,8 @@ impl Function {
             | &Insn::CCallWithFrame { recv, ref args, state, .. }
             | &Insn::SendWithoutBlockDirect { recv, ref args, state, .. }
             | &Insn::InvokeBuiltin { recv, ref args, state, .. }
-            | &Insn::InvokeSuper { recv, ref args, state, .. } => {
+            | &Insn::InvokeSuper { recv, ref args, state, .. }
+            | &Insn::InvokeSuperDirect { recv, ref args, state, .. } => {
                 worklist.push_back(recv);
                 worklist.extend(args);
                 worklist.push_back(state);
@@ -4590,6 +4687,7 @@ impl Function {
             | Insn::Send { recv, ref args, .. }
             | Insn::SendForward { recv, ref args, .. }
             | Insn::InvokeSuper { recv, ref args, .. }
+            | Insn::InvokeSuperDirect { recv, ref args, .. }
             | Insn::CCallWithFrame { recv, ref args, .. }
             | Insn::CCallVariadic { recv, ref args, .. }
             | Insn::InvokeBuiltin { recv, ref args, .. }