ZJIT: Pass self through basic block params (#13529)

* ZJIT: Pass self through basic block params

Co-authored-by: Max Bernstein <tekknolagi@gmail.com>

* Add comments for self

* Use self_param for ivar

* Make self_param a loop local

* Fix rest parameter type check

* Push self_param first

* Add a test case for putself

* Use SELF_PARAM_IDX

Co-authored-by: Max Bernstein <tekknolagi@gmail.com>

* Fix test_unknown

---------

Co-authored-by: Max Bernstein <tekknolagi@gmail.com>

test/ruby/test_zjit.rb

@@ -534,6 +534,17 @@ class TestZJIT < Test::Unit::TestCase
     }
   end
 
+  def test_putself
+    assert_compiles '3', %q{
+      class Integer
+        def minus(a)
+          self - a
+        end
+      end
+      5.minus(2)
+    }
+  end
+
   # tool/ruby_vm/views/*.erb relies on the zjit instructions a) being contiguous and
   # b) being reliably ordered after all the other instructions.
   def test_instruction_order

zjit/src/backend/arm64/mod.rs

@@ -641,7 +641,7 @@ impl Assembler
                         },
                         // If we're loading a memory operand into a register, then
                         // we'll switch over to the load instruction.
-                        (Opnd::Reg(_), Opnd::Mem(_)) => {
+                        (Opnd::Reg(_) | Opnd::VReg { .. }, Opnd::Mem(_)) => {
                             let value = split_memory_address(asm, *src);
                             asm.load_into(*dest, value);
                         },
@@ -654,7 +654,7 @@ impl Assembler
                             };
                             asm.mov(*dest, value);
                         },
-                        _ => unreachable!()
+                        _ => unreachable!("unexpected combination of operands in Insn::Mov: {dest:?}, {src:?}")
                     };
                 },
                 Insn::Not { opnd, .. } => {

zjit/src/backend/lir.rs

@@ -77,7 +77,7 @@ impl fmt::Debug for Opnd {
         match self {
             Self::None => write!(fmt, "None"),
             Value(val) => write!(fmt, "Value({val:?})"),
-            VReg { idx, num_bits } => write!(fmt, "Out{num_bits}({idx})"),
+            VReg { idx, num_bits } => write!(fmt, "VReg{num_bits}({idx})"),
             Imm(signed) => write!(fmt, "{signed:x}_i64"),
             UImm(unsigned) => write!(fmt, "{unsigned:x}_u64"),
             // Say Mem and Reg only once
@@ -1122,7 +1122,7 @@ impl RegisterPool {
     fn take_reg(&mut self, reg: &Reg, vreg_idx: usize) -> Reg {
         let reg_idx = self.regs.iter().position(|elem| elem.reg_no == reg.reg_no)
             .unwrap_or_else(|| panic!("Unable to find register: {}", reg.reg_no));
-        assert_eq!(self.pool[reg_idx], None, "register already allocated");
+        assert_eq!(self.pool[reg_idx], None, "register already allocated for VReg({:?})", self.pool[reg_idx]);
         self.pool[reg_idx] = Some(vreg_idx);
         self.live_regs += 1;
         *reg

zjit/src/codegen.rs

@@ -7,7 +7,7 @@ use crate::state::ZJITState;
 use crate::{asm::CodeBlock, cruby::*, options::debug, virtualmem::CodePtr};
 use crate::invariants::{iseq_escapes_ep, track_no_ep_escape_assumption};
 use crate::backend::lir::{self, asm_comment, Assembler, Opnd, Target, CFP, C_ARG_OPNDS, C_RET_OPND, EC, SP};
-use crate::hir::{iseq_to_hir, Block, BlockId, BranchEdge, CallInfo, RangeType};
+use crate::hir::{iseq_to_hir, Block, BlockId, BranchEdge, CallInfo, RangeType, SELF_PARAM_IDX};
 use crate::hir::{Const, FrameState, Function, Insn, InsnId};
 use crate::hir_type::{types::Fixnum, Type};
 use crate::options::get_option;
@@ -248,7 +248,6 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
     }
 
     let out_opnd = match insn {
-        Insn::PutSelf => gen_putself(),
         Insn::Const { val: Const::Value(val) } => gen_const(*val),
         Insn::NewArray { elements, state } => gen_new_array(jit, asm, elements, &function.frame_state(*state)),
         Insn::NewRange { low, high, flag, state } => gen_new_range(asm, opnd!(low), opnd!(high), *flag, &function.frame_state(*state)),
@@ -324,13 +323,16 @@ fn gen_entry_prologue(asm: &mut Assembler, iseq: IseqPtr) {
 
 /// Assign method arguments to basic block arguments at JIT entry
 fn gen_method_params(asm: &mut Assembler, iseq: IseqPtr, entry_block: &Block) {
+    let self_param = gen_param(asm, SELF_PARAM_IDX);
+    asm.mov(self_param, Opnd::mem(VALUE_BITS, CFP, RUBY_OFFSET_CFP_SELF));
+
     let num_params = entry_block.params().len();
     if num_params > 0 {
         asm_comment!(asm, "set method params: {num_params}");
 
         // Allocate registers for basic block arguments
         let params: Vec<Opnd> = (0..num_params).map(|idx|
-            gen_param(asm, idx)
+            gen_param(asm, idx + 1) // +1 for self
         ).collect();
 
         // Assign local variables to the basic block arguments
@@ -374,11 +376,6 @@ fn gen_getlocal(asm: &mut Assembler, iseq: IseqPtr, local_idx: usize) -> lir::Op
     }
 }
 
-/// Compile self in the current frame
-fn gen_putself() -> lir::Opnd {
-    Opnd::mem(VALUE_BITS, CFP, RUBY_OFFSET_CFP_SELF)
-}
-
 /// Compile a constant
 fn gen_const(val: VALUE) -> lir::Opnd {
     // Just propagate the constant value and generate nothing
@@ -482,9 +479,6 @@ fn gen_send_without_block_direct(
     recv: Opnd,
     args: &Vec<InsnId>,
 ) -> Option<lir::Opnd> {
-    // Set up the new frame
-    gen_push_frame(asm, recv);
-
     asm_comment!(asm, "switch to new CFP");
     let new_cfp = asm.sub(CFP, RUBY_SIZEOF_CONTROL_FRAME.into());
     asm.mov(CFP, new_cfp);
@@ -492,6 +486,7 @@ fn gen_send_without_block_direct(
 
     // Set up arguments
     let mut c_args: Vec<Opnd> = vec![];
+    c_args.push(recv);
     for &arg in args.iter() {
         c_args.push(jit.get_opnd(arg)?);
     }
@@ -714,18 +709,6 @@ fn gen_save_sp(asm: &mut Assembler, stack_size: usize) {
     asm.mov(cfp_sp, sp_addr);
 }
 
-/// Compile an interpreter frame
-fn gen_push_frame(asm: &mut Assembler, recv: Opnd) {
-    // Write to a callee CFP
-    fn cfp_opnd(offset: i32) -> Opnd {
-        Opnd::mem(64, CFP, offset - (RUBY_SIZEOF_CONTROL_FRAME as i32))
-    }
-
-    asm_comment!(asm, "push callee control frame");
-    asm.mov(cfp_opnd(RUBY_OFFSET_CFP_SELF), recv);
-    // TODO: Write more fields as needed
-}
-
 /// Return a register we use for the basic block argument at a given index
 fn param_reg(idx: usize) -> Reg {
     // To simplify the implementation, allocate a fixed register for each basic block argument for now.