ZJIT: Compile invokesuper with dynamic dispatch (ruby#14444)

Co-authored-by: Max Bernstein <[email protected]>

Author: k0kubun

test/ruby/test_zjit.rb

@@ -250,7 +250,7 @@ def test_nested_local_access
     }, call_threshold: 3, insns: [:getlocal, :setlocal, :getlocal_WC_0, :setlocal_WC_1]
   end
 
-  def test_read_local_written_by_children_iseqs
+  def test_send_with_local_written_by_blockiseq
     assert_compiles '[1, 2]', %q{
       def test
         l1 = nil
@@ -343,6 +343,46 @@ def test(a, b = 2) = [a, b]
     }
   end
 
+  def test_invokesuper
+    assert_compiles '[6, 60]', %q{
+      class Foo
+        def foo(a) = a + 1
+        def bar(a) = a + 10
+      end
+
+      class Bar < Foo
+        def foo(a) = super(a) + 2
+        def bar(a) = super + 20
+      end
+
+      bar = Bar.new
+      [bar.foo(3), bar.bar(30)]
+    }
+  end
+
+  def test_invokesuper_with_local_written_by_blockiseq
+    # Using `assert_runs` because we don't compile invokeblock yet
+    assert_runs '3', %q{
+      class Foo
+        def test
+          yield
+        end
+      end
+
+      class Bar < Foo
+        def test
+          a = 1
+          super do
+            a += 2
+          end
+          a
+        end
+      end
+
+      Bar.new.test
+    }
+  end
+
   def test_invokebuiltin
     omit 'Test fails at the moment due to not handling optional parameters'
     assert_compiles '["."]', %q{

zjit/src/codegen.rs

@@ -367,6 +367,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         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)),
         Insn::SendWithoutBlockDirect { cme, iseq, self_val, args, state, .. } => gen_send_without_block_direct(cb, jit, asm, *cme, *iseq, opnd!(self_val), opnds!(args), &function.frame_state(*state)),
+        &Insn::InvokeSuper { cd, blockiseq, state, .. } => gen_invokesuper(jit, asm, cd, blockiseq, &function.frame_state(state)),
         // 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)
         Insn::InvokeBuiltin { bf, state, .. } if bf.argc + 2 > (C_ARG_OPNDS.len() as i32) => return Err(*state),
@@ -1076,6 +1077,38 @@ fn gen_send_without_block_direct(
     ret
 }
 
+/// Compile a dynamic dispatch for `super`
+fn gen_invokesuper(
+    jit: &mut JITState,
+    asm: &mut Assembler,
+    cd: *const rb_call_data,
+    blockiseq: IseqPtr,
+    state: &FrameState,
+) -> lir::Opnd {
+    gen_incr_counter(asm, Counter::dynamic_send_count);
+
+    // Save PC and SP
+    gen_prepare_call_with_gc(asm, state);
+    gen_save_sp(asm, state.stack().len());
+
+    // Spill locals and stack
+    gen_spill_locals(jit, asm, state);
+    gen_spill_stack(jit, asm, state);
+
+    asm_comment!(asm, "call super with dynamic dispatch");
+    unsafe extern "C" {
+        fn rb_vm_invokesuper(ec: EcPtr, cfp: CfpPtr, cd: VALUE, blockiseq: IseqPtr) -> VALUE;
+    }
+    let ret = asm.ccall(
+        rb_vm_invokesuper as *const u8,
+        vec![EC, CFP, (cd as usize).into(), VALUE(blockiseq as usize).into()],
+    );
+    // TODO: Add a PatchPoint here that can side-exit the function if the callee messed with
+    // the frame's locals
+
+    ret
+}
+
 /// 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

@@ -565,10 +565,13 @@ pub enum Insn {
     /// `name` is for printing purposes only
     CCall { cfun: *const u8, args: Vec<InsnId>, name: ID, return_type: Type, elidable: bool },
 
-    /// Send without block with dynamic dispatch
+    /// Un-optimized fallback implementation (dynamic dispatch) for send-ish instructions
     /// Ignoring keyword arguments etc for now
     SendWithoutBlock { self_val: InsnId, cd: *const rb_call_data, args: Vec<InsnId>, state: InsnId },
     Send { self_val: InsnId, cd: *const rb_call_data, blockiseq: IseqPtr, args: Vec<InsnId>, state: InsnId },
+    InvokeSuper { self_val: InsnId, cd: *const rb_call_data, blockiseq: IseqPtr, args: Vec<InsnId>, state: InsnId },
+
+    /// Optimized ISEQ call
     SendWithoutBlockDirect {
         self_val: InsnId,
         cd: *const rb_call_data,
@@ -817,6 +820,13 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> {
                 }
                 Ok(())
             }
+            Insn::InvokeSuper { self_val, blockiseq, args, .. } => {
+                write!(f, "InvokeSuper {self_val}, {:p}", self.ptr_map.map_ptr(blockiseq))?;
+                for arg in args {
+                    write!(f, ", {arg}")?;
+                }
+                Ok(())
+            }
             Insn::InvokeBuiltin { bf, args, .. } => {
                 write!(f, "InvokeBuiltin {}", unsafe { CStr::from_ptr(bf.name) }.to_str().unwrap())?;
                 for arg in args {
@@ -1310,6 +1320,13 @@ impl Function {
                 args: find_vec!(args),
                 state,
             },
+            &InvokeSuper { self_val, cd, blockiseq, ref args, state } => InvokeSuper {
+                self_val: find!(self_val),
+                cd,
+                blockiseq,
+                args: find_vec!(args),
+                state,
+            },
             &InvokeBuiltin { bf, ref args, state, return_type } => InvokeBuiltin { bf, args: find_vec!(args), state, return_type },
             &ArrayDup { val, state } => ArrayDup { val: find!(val), state },
             &HashDup { val, state } => HashDup { val: find!(val), state },
@@ -1418,6 +1435,7 @@ impl Function {
             Insn::SendWithoutBlock { .. } => types::BasicObject,
             Insn::SendWithoutBlockDirect { .. } => types::BasicObject,
             Insn::Send { .. } => types::BasicObject,
+            Insn::InvokeSuper { .. } => types::BasicObject,
             Insn::InvokeBuiltin { return_type, .. } => return_type.unwrap_or(types::BasicObject),
             Insn::Defined { pushval, .. } => Type::from_value(*pushval).union(types::NilClass),
             Insn::DefinedIvar { .. } => types::BasicObject,
@@ -2206,7 +2224,8 @@ impl Function {
             }
             &Insn::Send { self_val, ref args, state, .. }
             | &Insn::SendWithoutBlock { self_val, ref args, state, .. }
-            | &Insn::SendWithoutBlockDirect { self_val, ref args, state, .. } => {
+            | &Insn::SendWithoutBlockDirect { self_val, ref args, state, .. }
+            | &Insn::InvokeSuper { self_val, ref args, state, .. } => {
                 worklist.push_back(self_val);
                 worklist.extend(args);
                 worklist.push_back(state);
@@ -2830,14 +2849,14 @@ fn insn_idx_at_offset(idx: u32, offset: i64) -> u32 {
 
 struct BytecodeInfo {
     jump_targets: Vec<u32>,
-    has_send: bool,
+    has_blockiseq: bool,
 }
 
 fn compute_bytecode_info(iseq: *const rb_iseq_t) -> BytecodeInfo {
     let iseq_size = unsafe { get_iseq_encoded_size(iseq) };
     let mut insn_idx = 0;
     let mut jump_targets = HashSet::new();
-    let mut has_send = false;
+    let mut has_blockiseq = false;
     while insn_idx < iseq_size {
         // Get the current pc and opcode
         let pc = unsafe { rb_iseq_pc_at_idx(iseq, insn_idx) };
@@ -2861,13 +2880,18 @@ fn compute_bytecode_info(iseq: *const rb_iseq_t) -> BytecodeInfo {
                     jump_targets.insert(insn_idx);
                 }
             }
-            YARVINSN_send => has_send = true,
+            YARVINSN_send | YARVINSN_invokesuper => {
+                let blockiseq: IseqPtr = get_arg(pc, 1).as_iseq();
+                if !blockiseq.is_null() {
+                    has_blockiseq = true;
+                }
+            }
             _ => {}
         }
     }
     let mut result = jump_targets.into_iter().collect::<Vec<_>>();
     result.sort();
-    BytecodeInfo { jump_targets: result, has_send }
+    BytecodeInfo { jump_targets: result, has_blockiseq }
 }
 
 #[derive(Debug, PartialEq, Clone, Copy)]
@@ -2984,7 +3008,7 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
     let mut profiles = ProfileOracle::new(payload);
     let mut fun = Function::new(iseq);
     // Compute a map of PC->Block by finding jump targets
-    let BytecodeInfo { jump_targets, has_send } = compute_bytecode_info(iseq);
+    let BytecodeInfo { jump_targets, has_blockiseq } = compute_bytecode_info(iseq);
     let mut insn_idx_to_block = HashMap::new();
     for insn_idx in jump_targets {
         if insn_idx == 0 {
@@ -3321,7 +3345,7 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                 }
                 YARVINSN_getlocal_WC_0 => {
                     let ep_offset = get_arg(pc, 0).as_u32();
-                    if iseq_type == ISEQ_TYPE_EVAL || has_send {
+                    if iseq_type == ISEQ_TYPE_EVAL || has_blockiseq {
                         // On eval, the locals are always on the heap, so read the local using EP.
                         let val = fun.push_insn(block, Insn::GetLocal { ep_offset, level: 0 });
                         state.setlocal(ep_offset, val);
@@ -3341,7 +3365,7 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                     let ep_offset = get_arg(pc, 0).as_u32();
                     let val = state.stack_pop()?;
                     state.setlocal(ep_offset, val);
-                    if iseq_type == ISEQ_TYPE_EVAL || has_send {
+                    if iseq_type == ISEQ_TYPE_EVAL || has_blockiseq {
                         // On eval, the locals are always on the heap, so write the local using EP.
                         fun.push_insn(block, Insn::SetLocal { val, ep_offset, level: 0 });
                     }
@@ -3521,6 +3545,34 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                         state.setlocal(ep_offset, val);
                     }
                 }
+                YARVINSN_invokesuper => {
+                    let cd: *const rb_call_data = get_arg(pc, 0).as_ptr();
+                    let call_info = unsafe { rb_get_call_data_ci(cd) };
+                    if let Err(call_type) = unknown_call_type(unsafe { rb_vm_ci_flag(call_info) } & !VM_CALL_SUPER & !VM_CALL_ZSUPER) {
+                        // Unknown call type; side-exit into the interpreter
+                        let exit_id = fun.push_insn(block, Insn::Snapshot { state: exit_state });
+                        fun.push_insn(block, Insn::SideExit { state: exit_id, reason: SideExitReason::UnhandledCallType(call_type) });
+                        break;  // End the block
+                    }
+                    let argc = unsafe { vm_ci_argc((*cd).ci) };
+                    let args = state.stack_pop_n(argc as usize)?;
+                    let recv = state.stack_pop()?;
+                    let blockiseq: IseqPtr = get_arg(pc, 1).as_ptr();
+                    let exit_id = fun.push_insn(block, Insn::Snapshot { state: exit_state });
+                    let result = fun.push_insn(block, Insn::InvokeSuper { self_val: recv, cd, blockiseq, args, state: exit_id });
+                    state.stack_push(result);
+
+                    if !blockiseq.is_null() {
+                        // Reload locals that may have been modified by the blockiseq.
+                        // TODO: Avoid reloading locals that are not referenced by the blockiseq
+                        // or not used after this. Max thinks we could eventually DCE them.
+                        for local_idx in 0..state.locals.len() {
+                            let ep_offset = local_idx_to_ep_offset(iseq, local_idx) as u32;
+                            let val = fun.push_insn(block, Insn::GetLocal { ep_offset, level: 0 });
+                            state.setlocal(ep_offset, val);
+                        }
+                    }
+                }
                 YARVINSN_getglobal => {
                     let id = ID(get_arg(pc, 0).as_u64());
                     let exit_id = fun.push_insn(block, Insn::Snapshot { state: exit_state });
@@ -4971,7 +5023,6 @@ mod tests {
         assert_snapshot!(hir_string("test"), @r"
         fn test@<compiled>:2:
         bb0(v0:BasicObject, v1:BasicObject):
-          v5:BasicObject = GetLocal l0, EP@3
           SideExit UnhandledCallType(BlockArg)
         ");
     }
@@ -5004,26 +5055,43 @@ mod tests {
     // TODO(max): Figure out how to generate a call with TAILCALL flag
 
     #[test]
-    fn test_cant_compile_super() {
+    fn test_compile_super() {
         eval("
             def test = super()
         ");
         assert_snapshot!(hir_string("test"), @r"
         fn test@<compiled>:2:
         bb0(v0:BasicObject):
-          SideExit UnhandledYARVInsn(invokesuper)
+          v5:BasicObject = InvokeSuper v0, 0x1000
+          CheckInterrupts
+          Return v5
         ");
     }
 
     #[test]
-    fn test_cant_compile_zsuper() {
+    fn test_compile_zsuper() {
         eval("
             def test = super
         ");
         assert_snapshot!(hir_string("test"), @r"
         fn test@<compiled>:2:
         bb0(v0:BasicObject):
-          SideExit UnhandledYARVInsn(invokesuper)
+          v5:BasicObject = InvokeSuper v0, 0x1000
+          CheckInterrupts
+          Return v5
+        ");
+    }
+
+    #[test]
+    fn test_cant_compile_super_nil_blockarg() {
+        eval("
+            def test = super(&nil)
+        ");
+        assert_snapshot!(hir_string("test"), @r"
+        fn test@<compiled>:2:
+        bb0(v0:BasicObject):
+          v4:NilClass = Const Value(nil)
+          SideExit UnhandledCallType(BlockArg)
         ");
     }