ZJIT: Simplify invokesuper specialization to most common case

nirvdrum · nirvdrum · commit daa357865679 · 2026-01-08T12:52:51.000-05:00
Looking at ruby-bench, most `super` calls don't pass a block, which means we can use the already optimized `SendWithoutBlockDirect`.
diff --git a/test/ruby/test_zjit.rb b/test/ruby/test_zjit.rb
@@ -1241,7 +1241,7 @@ def foo
       end
 
       def test
-        B.new.f
+        B.new.foo
       end
 
       test  # profile invokesuper (super -> A#foo)
@@ -1259,6 +1259,35 @@ def foo
     }, call_threshold: 2
   end
 
+  # Test super with positional and keyword arguments (pattern from chunky_png)
+  def test_invokesuper_with_keyword_args
+    assert_compiles '["png", {content: "image data"}]', %q{
+      class Base
+        def initialize(type, attributes = {})
+          @type = type
+          @attributes = attributes
+        end
+
+        def info
+          [@type, @attributes]
+        end
+      end
+
+      class Generic < Base
+        def initialize(type, content = "")
+          super(type, content: content)
+        end
+      end
+
+      def test
+        Generic.new("png", "image data").info
+      end
+
+      test
+      test
+    }, call_threshold: 2
+  end
+
   def test_invokebuiltin
     # Not using assert_compiles due to register spill
     assert_runs '["."]', %q{
diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs
@@ -406,8 +406,6 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
             gen_send_without_block(jit, asm, *cd, &function.frame_state(*state), SendFallbackReason::TooManyArgsForLir),
         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)
@@ -457,6 +455,8 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         Insn::GuardNotShared { recv, state } => gen_guard_not_shared(jit, asm, opnd!(recv), &function.frame_state(*state)),
         &Insn::GuardLess { left, right, state } => gen_guard_less(jit, asm, opnd!(left), opnd!(right), &function.frame_state(state)),
         &Insn::GuardGreaterEq { left, right, state } => gen_guard_greater_eq(jit, asm, opnd!(left), opnd!(right), &function.frame_state(state)),
+        &Insn::GuardSuperMethodEntry { cme, state } => no_output!(gen_guard_super_method_entry(jit, asm, cme, &function.frame_state(state))),
+        Insn::GetBlockHandler => gen_get_block_handler(jit, asm),
         Insn::PatchPoint { invariant, state } => no_output!(gen_patch_point(jit, asm, invariant, &function.frame_state(*state))),
         Insn::CCall { cfunc, recv, args, name, return_type: _, elidable: _ } => gen_ccall(asm, *cfunc, *name, opnd!(recv), opnds!(args)),
         // Give up CCallWithFrame for 7+ args since asm.ccall() supports at most 6 args (recv + args).
@@ -719,6 +719,29 @@ fn gen_guard_greater_eq(jit: &JITState, asm: &mut Assembler, left: Opnd, right:
     left
 }
 
+/// Guard that the method entry at ep[VM_ENV_DATA_INDEX_ME_CREF] matches the expected CME.
+/// This ensures we're calling super from the expected method context.
+fn gen_guard_super_method_entry(
+    jit: &JITState,
+    asm: &mut Assembler,
+    cme: *const rb_callable_method_entry_t,
+    state: &FrameState,
+) {
+    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(cme as u64));
+    asm.jne(side_exit(jit, state, SideExitReason::GuardSuperMethodEntry));
+}
+
+/// Get the block handler from ep[VM_ENV_DATA_INDEX_SPECVAL] at the local EP (LEP).
+fn gen_get_block_handler(jit: &JITState, asm: &mut Assembler) -> Opnd {
+    asm_comment!(asm, "get block handler from LEP");
+    let lep = gen_get_lep(jit, asm);
+    asm.load(Opnd::mem(64, lep, SIZEOF_VALUE_I32 * VM_ENV_DATA_INDEX_SPECVAL))
+}
+
 fn gen_get_constant_path(jit: &JITState, asm: &mut Assembler, ic: *const iseq_inline_constant_cache, state: &FrameState) -> Opnd {
     unsafe extern "C" {
         fn rb_vm_opt_getconstant_path(ec: EcPtr, cfp: CfpPtr, ic: *const iseq_inline_constant_cache) -> VALUE;
@@ -1499,37 +1522,6 @@ 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
diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
@@ -895,17 +895,6 @@ 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>,
@@ -982,6 +971,11 @@ pub enum Insn {
     GuardGreaterEq { left: InsnId, right: InsnId, state: InsnId },
     /// Side-exit if left is not less than right (both operands are C long).
     GuardLess { left: InsnId, right: InsnId, state: InsnId },
+    /// Side-exit if the method entry at ep[VM_ENV_DATA_INDEX_ME_CREF] doesn't match the expected CME.
+    /// Used to ensure super calls are made from the expected method context.
+    GuardSuperMethodEntry { cme: *const rb_callable_method_entry_t, state: InsnId },
+    /// Get the block handler from ep[VM_ENV_DATA_INDEX_SPECVAL] at the local EP (LEP).
+    GetBlockHandler,
 
     /// Generate no code (or padding if necessary) and insert a patch point
     /// that can be rewritten to a side exit when the Invariant is broken.
@@ -1010,7 +1004,7 @@ impl Insn {
             | Insn::PatchPoint { .. } | Insn::SetIvar { .. } | Insn::SetClassVar { .. } | Insn::ArrayExtend { .. }
             | Insn::ArrayPush { .. } | Insn::SideExit { .. } | Insn::SetGlobal { .. }
             | Insn::SetLocal { .. } | Insn::Throw { .. } | Insn::IncrCounter(_) | Insn::IncrCounterPtr { .. }
-            | Insn::CheckInterrupts { .. } | Insn::GuardBlockParamProxy { .. }
+            | Insn::CheckInterrupts { .. } | Insn::GuardBlockParamProxy { .. } | Insn::GuardSuperMethodEntry { .. }
             | Insn::StoreField { .. } | Insn::WriteBarrier { .. } | Insn::HashAset { .. }
             | Insn::ArrayAset { .. } => false,
             _ => true,
@@ -1310,13 +1304,6 @@ 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 {
@@ -1364,6 +1351,8 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> {
             Insn::GuardNotShared { recv, .. } => write!(f, "GuardNotShared {recv}"),
             Insn::GuardLess { left, right, .. } => write!(f, "GuardLess {left}, {right}"),
             Insn::GuardGreaterEq { left, right, .. } => write!(f, "GuardGreaterEq {left}, {right}"),
+            Insn::GuardSuperMethodEntry { cme, .. } => write!(f, "GuardSuperMethodEntry {:p}", self.ptr_map.map_ptr(cme)),
+            Insn::GetBlockHandler => write!(f, "GetBlockHandler"),
             Insn::PatchPoint { invariant, .. } => { write!(f, "PatchPoint {}", invariant.print(self.ptr_map)) },
             Insn::GetConstantPath { ic, .. } => { write!(f, "GetConstantPath {:p}", self.ptr_map.map_ptr(ic)) },
             Insn::IsBlockGiven => { write!(f, "IsBlockGiven") },
@@ -2001,6 +1990,8 @@ impl Function {
             &GuardNotShared { recv, state } => GuardNotShared { recv: find!(recv), state },
             &GuardGreaterEq { left, right, state } => GuardGreaterEq { left: find!(left), right: find!(right), state },
             &GuardLess { left, right, state } => GuardLess { left: find!(left), right: find!(right), state },
+            &GuardSuperMethodEntry { cme, state } => GuardSuperMethodEntry { cme, state },
+            &GetBlockHandler => GetBlockHandler,
             &FixnumAdd { left, right, state } => FixnumAdd { left: find!(left), right: find!(right), state },
             &FixnumSub { left, right, state } => FixnumSub { left: find!(left), right: find!(right), state },
             &FixnumMult { left, right, state } => FixnumMult { left: find!(left), right: find!(right), state },
@@ -2066,14 +2057,6 @@ 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),
@@ -2181,8 +2164,9 @@ impl Function {
             Insn::SetGlobal { .. } | Insn::Jump(_) | Insn::EntryPoint { .. }
             | Insn::IfTrue { .. } | Insn::IfFalse { .. } | Insn::Return { .. } | Insn::Throw { .. }
             | Insn::PatchPoint { .. } | Insn::SetIvar { .. } | Insn::SetClassVar { .. } | Insn::ArrayExtend { .. }
-            | Insn::ArrayPush { .. } | Insn::SideExit { .. } | Insn::SetLocal { .. } | Insn::IncrCounter(_)
-            | Insn::CheckInterrupts { .. } | Insn::GuardBlockParamProxy { .. } | Insn::IncrCounterPtr { .. }
+            | Insn::ArrayPush { .. } | Insn::SideExit { .. } | Insn::SetLocal { .. }
+            | Insn::IncrCounter(_) | Insn::IncrCounterPtr { .. }
+            | Insn::CheckInterrupts { .. } | Insn::GuardBlockParamProxy { .. } | Insn::GuardSuperMethodEntry { .. }
             | Insn::StoreField { .. } | Insn::WriteBarrier { .. } | Insn::HashAset { .. } | Insn::ArrayAset { .. } =>
                 panic!("Cannot infer type of instruction with no output: {}. See Insn::has_output().", self.insns[insn.0]),
             Insn::Const { val: Const::Value(val) } => Type::from_value(*val),
@@ -2262,7 +2246,6 @@ 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),
@@ -2290,6 +2273,7 @@ impl Function {
             Insn::AnyToString { .. } => types::String,
             Insn::GetLocal { rest_param: true, .. } => types::ArrayExact,
             Insn::GetLocal { .. } => types::BasicObject,
+            Insn::GetBlockHandler => types::RubyValue,
             // The type of Snapshot doesn't really matter; it's never materialized. It's used only
             // as a reference for FrameState, which we use to generate side-exit code.
             Insn::Snapshot { .. } => types::Any,
@@ -3044,16 +3028,17 @@ impl Function {
                             self.push_insn_id(block, insn_id); continue;
                         }
 
-                        // Don't handle calls with complex arguments (kwarg, splat, kw_splat, blockarg, forwarding)
+                        // Only handle bare `super` calls (i.e., with no arguments) so we can
+                        // forward the caller arguments without needing to modify them.
                         let ci = unsafe { get_call_data_ci(cd) };
                         let flags = unsafe { rb_vm_ci_flag(ci) };
-                        if unspecializable_call_type(flags) {
+                        if (flags & VM_CALL_ZSUPER) == 0 {
                             self.push_insn_id(block, insn_id); continue;
                         }
 
                         let frame_state = self.frame_state(state);
 
-                        // Get the profiled CME from the current method
+                        // Get the profiled CME from the current method.
                         let Some(profiles) = self.profiles.as_ref() else {
                             self.push_insn_id(block, insn_id); continue;
                         };
@@ -3062,7 +3047,7 @@ impl Function {
                             self.push_insn_id(block, insn_id); continue;
                         };
 
-                        // Get defined_class and method ID from the profiled CME
+                        // 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) };
 
@@ -3072,16 +3057,15 @@ impl Function {
                             self.push_insn_id(block, insn_id); continue;
                         }
 
-                        // Look up the super method
+                        // 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
+                        // Check if it's an ISEQ method; bail if it isn't.
                         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;
                         }
 
@@ -3092,8 +3076,7 @@ impl Function {
                             self.push_insn_id(block, insn_id); continue;
                         }
 
-                        // Add PatchPoints for method redefinition
-                        // TODO: Add guard that ep[-2] matches current_cme
+                        // Add PatchPoint for method redefinition.
                         self.push_insn(block, Insn::PatchPoint {
                             invariant: Invariant::MethodRedefined {
                                 klass: unsafe { (*super_cme).defined_class },
@@ -3103,11 +3086,23 @@ impl Function {
                             state
                         });
 
-                        let send_direct = self.push_insn(block, Insn::InvokeSuperDirect {
+                        // Guard that we're calling `super` from the expected method context.
+                        self.push_insn(block, Insn::GuardSuperMethodEntry { cme: current_cme, state });
+
+                        // Guard that no block is being passed (implicit or explicit).
+                        let block_handler = self.push_insn(block, Insn::GetBlockHandler);
+                        self.push_insn(block, Insn::GuardBitEquals {
+                            val: block_handler,
+                            expected: Const::Value(VALUE(VM_BLOCK_HANDLER_NONE as usize)),
+                            state
+                        });
+
+                        // Use SendWithoutBlockDirect with the super method's CME and ISEQ.
+                        let send_direct = self.push_insn(block, Insn::SendWithoutBlockDirect {
                             recv,
                             cd,
-                            current_cme,
-                            super_cme,
+                            cme: super_cme,
+                            iseq: super_iseq,
                             args,
                             state
                         });
@@ -3975,6 +3970,7 @@ impl Function {
             | &Insn::LoadEC
             | &Insn::LoadSelf
             | &Insn::GetLocal { .. }
+            | &Insn::GetBlockHandler
             | &Insn::PutSpecialObject { .. }
             | &Insn::IsBlockGiven
             | &Insn::IncrCounter(_)
@@ -4151,8 +4147,7 @@ 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::InvokeSuperDirect { recv, ref args, state, .. } => {
+            | &Insn::InvokeSuper { recv, ref args, state, .. } => {
                 worklist.push_back(recv);
                 worklist.extend(args);
                 worklist.push_back(state);
@@ -4204,6 +4199,7 @@ impl Function {
                 worklist.push_back(val);
             }
             &Insn::GuardBlockParamProxy { state, .. } |
+            &Insn::GuardSuperMethodEntry { state, .. } |
             &Insn::GetGlobal { state, .. } |
             &Insn::GetSpecialSymbol { state, .. } |
             &Insn::GetSpecialNumber { state, .. } |
@@ -4715,6 +4711,8 @@ impl Function {
             | Insn::Jump { .. }
             | Insn::EntryPoint { .. }
             | Insn::GuardBlockParamProxy { .. }
+            | Insn::GuardSuperMethodEntry { .. }
+            | Insn::GetBlockHandler
             | Insn::PatchPoint { .. }
             | Insn::SideExit { .. }
             | Insn::IncrCounter { .. }
@@ -4768,7 +4766,6 @@ 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, .. }
diff --git a/zjit/src/hir/opt_tests.rs b/zjit/src/hir/opt_tests.rs
