Load Param off of cfp->ep (https://github.com/Shopify/zjit/pull/31)

* Load Param off of cfp->ep

* Test with --zjit-call-threshold=1 as well

* Fix get_opnd's debug output

* Return Mem operand from gen_param

* Test both first and second calls

* Spell out the namespace for Opnd returns

* Update a comment about gen_param

* Explain why we take a lock

* Fix a typo

5 files changed, 117 insertions, 63 deletions

diff --git a/.github/workflows/zjit-macos.yml b/.github/workflows/zjit-macos.yml
index 0416981f7e..6829bbe113 100644
--- a/.github/workflows/zjit-macos.yml
+++ b/.github/workflows/zjit-macos.yml
@@ -30,19 +30,20 @@ jobs:
         include:
           - test_task: 'zjit-test'
             configure: '--enable-zjit=dev'
-            rust_version: 1.85.0
+
+          - test_task: 'btest'
+            zjit_opts: '--zjit-call-threshold=1'
+            configure: '--enable-zjit=dev'
+            btests: '../src/bootstraptest/test_zjit.rb'
 
           - test_task: 'btest'
             zjit_opts: '--zjit-call-threshold=2'
             configure: '--enable-zjit=dev'
             btests: '../src/bootstraptest/test_zjit.rb'
-            rust_version: 1.85.0
 
+          # Test without ZJIT for now
           - test_task: 'check'
-            # Test without ZJIT for now
-            #zjit_opts: '--zjit-call-threshold=2'
             configure: '--enable-zjit'
-            rust_version: 1.85.0
 
     env:
       GITPULLOPTIONS: --no-tags origin ${{ github.ref }}
@@ -85,8 +86,7 @@ jobs:
         if: ${{ matrix.test_task == 'zjit-test' }}
 
       - name: Install Rust # TODO(alan): remove when GitHub images catch up past 1.85.0
-        if: ${{ matrix.rust_version }}
-        run: rustup default ${{ matrix.rust_version }}
+        run: rustup default 1.85.0
 
       - name: Run configure
         run: ../src/configure -C --disable-install-doc ${{ matrix.configure }}
diff --git a/.github/workflows/zjit-ubuntu.yml b/.github/workflows/zjit-ubuntu.yml
index d77ac78283..2c72f2edc4 100644
--- a/.github/workflows/zjit-ubuntu.yml
+++ b/.github/workflows/zjit-ubuntu.yml
@@ -32,23 +32,23 @@ jobs:
             hint: 'To fix: use patch in logs'
             configure: '--enable-zjit=dev --with-gcc=clang-14'
             libclang_path: '/usr/lib/llvm-14/lib/libclang.so.1'
-            rust_version: 1.85.0 # TODO(alan): remove when GitHub's images catch up
 
           - test_task: 'zjit-test'
             configure: '--enable-zjit=dev'
-            rust_version: 1.85.0
+
+          - test_task: 'btest'
+            zjit_opts: '--zjit-call-threshold=1'
+            configure: '--enable-zjit=dev'
+            btests: '../src/bootstraptest/test_zjit.rb'
 
           - test_task: 'btest'
             zjit_opts: '--zjit-call-threshold=2'
             configure: '--enable-zjit=dev'
             btests: '../src/bootstraptest/test_zjit.rb'
-            rust_version: 1.85.0
 
+          # Test without ZJIT for now
           - test_task: 'check'
-            # Test without ZJIT for now
-            #zjit_opts: '--zjit-call-threshold=2'
             configure: '--enable-zjit'
-            rust_version: 1.85.0
 
     env:
       GITPULLOPTIONS: --no-tags origin ${{ github.ref }}
@@ -100,8 +100,7 @@ jobs:
           fetch-depth: 10
 
       - name: Install Rust
-        if: ${{ matrix.rust_version }}
-        run: rustup default ${{ matrix.rust_version }}
+        run: rustup default 1.85.0
 
       - name: Install rustfmt
         if: ${{ matrix.test_task == 'zjit-bindgen' }}
diff --git a/bootstraptest/test_zjit.rb b/bootstraptest/test_zjit.rb
index a94594d098..a8c6c85c0f 100644
--- a/bootstraptest/test_zjit.rb
+++ b/bootstraptest/test_zjit.rb
@@ -15,3 +15,8 @@ assert_equal '3', %q{
   def test = 1 + 2
   test; test
 }
+
+assert_equal '[6, 3]', %q{
+  def test(a, b) = a + b
+  [test(2, 4), test(1, 2)]
+}
diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs
index 2a6ffcc4ac..f3a9b68e4b 100644
--- a/zjit/src/codegen.rs
+++ b/zjit/src/codegen.rs
@@ -1,10 +1,10 @@
 use crate::{
-    asm::CodeBlock, backend::lir::{asm_comment, Assembler, Opnd, Target, CFP, C_ARG_OPNDS, EC, SP}, cruby::*, debug, hir::{Const, FrameState, Function, Insn, InsnId}, hir_type::{types::Fixnum, Type}, virtualmem::CodePtr
+    asm::CodeBlock, backend::lir, backend::lir::{asm_comment, Assembler, Opnd, Target, CFP, C_ARG_OPNDS, EC, SP}, cruby::*, debug, hir::{Const, FrameState, Function, Insn, InsnId}, hir_type::{types::Fixnum, Type}, virtualmem::CodePtr
 };
 
 /// Ephemeral code generation state
 struct JITState {
-    /// Instruction sequence for the compiling method
+    /// Instruction sequence for the method being compiled
     iseq: IseqPtr,
 
     /// Low-level IR Operands indexed by High-level IR's Instruction ID
@@ -12,12 +12,22 @@ struct JITState {
 }
 
 impl JITState {
+    /// Create a new JITState instance
     fn new(iseq: IseqPtr, insn_len: usize) -> Self {
         JITState {
             iseq,
             opnds: vec![None; insn_len],
         }
     }
+
+    /// Retrieve the output of a given instruction that has been compiled
+    fn get_opnd(&self, insn_id: InsnId) -> Option<lir::Opnd> {
+        let opnd = self.opnds[insn_id.0];
+        if opnd.is_none() {
+            debug!("Failed to get_opnd({insn_id})");
+        }
+        opnd
+    }
 }
 
 /// Compile High-level IR into machine code
@@ -29,23 +39,10 @@ pub fn gen_function(cb: &mut CodeBlock, function: &Function, iseq: IseqPtr) -> O
 
     // Compile each instruction in the IR
     for (insn_idx, insn) in function.insns.iter().enumerate() {
-        let insn_id = InsnId(insn_idx);
-        if !matches!(*insn, Insn::Snapshot { .. }) {
-            asm_comment!(asm, "Insn: {:04} {:?}", insn_idx, insn);
+        if gen_insn(&mut jit, &mut asm, function, InsnId(insn_idx), insn).is_none() {
+            debug!("Failed to compile insn: {:04} {:?}", insn_idx, insn);
+            return None;
         }
-        match insn {
-            Insn::Const { val: Const::Value(val) } => gen_const(&mut jit, insn_id, *val),
-            Insn::Snapshot { .. } => {}, // we don't need to do anything for this instruction at the moment
-            Insn::Return { val } => gen_return(&jit, &mut asm, *val)?,
-            Insn::FixnumAdd { left, right, state } => gen_fixnum_add(&mut jit, &mut asm, insn_id, *left, *right, function.frame_state(*state))?,
-            Insn::GuardType { val, guard_type, state } => gen_guard_type(&mut jit, &mut asm, insn_id, *val, *guard_type, function.frame_state(*state))?,
-            Insn::PatchPoint(_) => {}, // For now, rb_zjit_bop_redefined() panics. TODO: leave a patch point and fix rb_zjit_bop_redefined()
-            _ => {
-                debug!("ZJIT: gen_function: unexpected insn {:?}", insn);
-                return None;
-            }
-        }
-        debug!("Compiled insn: {:04} {:?}", insn_idx, insn);
     }
 
     // Generate code if everything can be compiled
@@ -55,6 +52,31 @@ pub fn gen_function(cb: &mut CodeBlock, function: &Function, iseq: IseqPtr) -> O
     start_ptr
 }
 
+/// Compile an instruction
+fn gen_insn(jit: &mut JITState, asm: &mut Assembler, function: &Function, insn_id: InsnId, insn: &Insn) -> Option<()> {
+    if !matches!(*insn, Insn::Snapshot { .. }) {
+        asm_comment!(asm, "Insn: {:04} {:?}", insn_id.0, insn);
+    }
+    let out_opnd = match insn {
+        Insn::Const { val: Const::Value(val) } => gen_const(*val),
+        Insn::Param { idx } => gen_param(jit, asm, *idx)?,
+        Insn::Snapshot { .. } => return Some(()), // we don't need to do anything for this instruction at the moment
+        Insn::Return { val } => return Some(gen_return(&jit, asm, *val)?),
+        Insn::FixnumAdd { left, right, state } => gen_fixnum_add(jit, asm, *left, *right, function.frame_state(*state))?,
+        Insn::GuardType { val, guard_type, state } => gen_guard_type(jit, asm, *val, *guard_type, function.frame_state(*state))?,
+        Insn::PatchPoint(_) => return Some(()), // For now, rb_zjit_bop_redefined() panics. TODO: leave a patch point and fix rb_zjit_bop_redefined()
+        _ => {
+            debug!("ZJIT: gen_function: unexpected insn {:?}", insn);
+            return None;
+        }
+    };
+
+    // If the instruction has an output, remember it in jit.opnds
+    jit.opnds[insn_id.0] = Some(out_opnd);
+
+    Some(())
+}
+
 /// Compile an interpreter entry block to be inserted into an ISEQ
 fn gen_entry_prologue(jit: &JITState, asm: &mut Assembler) {
     asm_comment!(asm, "YJIT entry point: {}", iseq_get_location(jit.iseq, 0));
@@ -76,9 +98,25 @@ fn gen_entry_prologue(jit: &JITState, asm: &mut Assembler) {
 }
 
 /// Compile a constant
-fn gen_const(jit: &mut JITState, insn_id: InsnId, val: VALUE) {
-    // Just remember the constant value and generate nothing
-    jit.opnds[insn_id.0] = Some(Opnd::Value(val));
+fn gen_const(val: VALUE) -> Opnd {
+    // Just propagate the constant value and generate nothing
+    Opnd::Value(val)
+}
+
+/// Compile a method/block paramter read. For now, it only supports method parameters.
+fn gen_param(jit: &JITState, asm: &mut Assembler, local_idx: usize) -> Option<lir::Opnd> {
+    // Get the EP of the current CFP
+    // TODO: Use the SP register and invalidate on EP escape
+    let ep_opnd = Opnd::mem(64, CFP, RUBY_OFFSET_CFP_EP);
+    let ep_reg = asm.load(ep_opnd);
+
+    // Load the local variable
+    // val = *(vm_get_ep(GET_EP(), level) - idx);
+    let ep_offset = local_idx_to_ep_offset(jit.iseq, local_idx);
+    let offs = -(SIZEOF_VALUE_I32 * ep_offset);
+    let local_opnd = Opnd::mem(64, ep_reg, offs);
+
+    Some(local_opnd)
 }
 
 /// Compile code that exits from JIT code with a return value
@@ -104,9 +142,9 @@ fn gen_return(jit: &JITState, asm: &mut Assembler, val: InsnId) -> Option<()> {
 }
 
 /// Compile Fixnum + Fixnum
-fn gen_fixnum_add(jit: &mut JITState, asm: &mut Assembler, insn_id: InsnId, left: InsnId, right: InsnId, state: &FrameState) -> Option<()> {
-    let left_opnd = jit.opnds[left.0]?;
-    let right_opnd = jit.opnds[right.0]?;
+fn gen_fixnum_add(jit: &mut JITState, asm: &mut Assembler, left: InsnId, right: InsnId, state: &FrameState) -> Option<lir::Opnd> {
+    let left_opnd = jit.get_opnd(left)?;
+    let right_opnd = jit.get_opnd(right)?;
 
     // Load left into a register if left is a constant. The backend doesn't support sub(imm, imm).
     let left_reg = match left_opnd {
@@ -119,13 +157,12 @@ fn gen_fixnum_add(jit: &mut JITState, asm: &mut Assembler, insn_id: InsnId, left
     let out_val = asm.add(left_untag, right_opnd);
     asm.jo(Target::SideExit(state.clone()));
 
-    jit.opnds[insn_id.0] = Some(out_val);
-    Some(())
+    Some(out_val)
 }
 
 /// Compile a type check with a side exit
-fn gen_guard_type(jit: &mut JITState, asm: &mut Assembler, insn_id: InsnId, val: InsnId, guard_type: Type, state: &FrameState) -> Option<()> {
-    let opnd = jit.opnds[val.0]?;
+fn gen_guard_type(jit: &mut JITState, asm: &mut Assembler, val: InsnId, guard_type: Type, state: &FrameState) -> Option<lir::Opnd> {
+    let opnd = jit.get_opnd(val)?;
     if guard_type.is_subtype(Fixnum) {
         // Load opnd into a register if opnd is a constant. The backend doesn't support test(imm, imm) yet.
         let opnd_reg = match opnd {
@@ -139,7 +176,13 @@ fn gen_guard_type(jit: &mut JITState, asm: &mut Assembler, insn_id: InsnId, val:
     } else {
         unimplemented!("unsupported type: {guard_type}");
     }
+    Some(opnd)
+}
 
-    jit.opnds[insn_id.0] = Some(opnd);
-    Some(())
+/// Inverse of ep_offset_to_local_idx(). See ep_offset_to_local_idx() for details.
+fn local_idx_to_ep_offset(iseq: IseqPtr, local_idx: usize) -> i32 {
+    let local_table_size: i32 = unsafe { get_iseq_body_local_table_size(iseq) }
+        .try_into()
+        .unwrap();
+    local_table_size - local_idx as i32 - 1 + VM_ENV_DATA_SIZE as i32
 }
diff --git a/zjit/src/lib.rs b/zjit/src/lib.rs
index c2102959dd..9e93f0f2be 100644
--- a/zjit/src/lib.rs
+++ b/zjit/src/lib.rs
@@ -89,23 +89,30 @@ fn rb_bug_panic_hook() {
 /// Generate JIT code for a given ISEQ, which takes EC and CFP as its arguments.
 #[unsafe(no_mangle)]
 pub extern "C" fn rb_zjit_iseq_gen_entry_point(iseq: IseqPtr, _ec: EcPtr) -> *const u8 {
-    // TODO: acquire the VM barrier
-
-    // Compile ISEQ into High-level IR
-    let ssa = match hir::iseq_to_hir(iseq) {
-        Ok(ssa) => ssa,
-        Err(err) => {
-            debug!("ZJIT: iseq_to_hir: {:?}", err);
-            return std::ptr::null();
-        }
-    };
-
-    // Compile High-level IR into machine code
-    let cb = ZJITState::get_code_block();
-    match gen_function(cb, &ssa, iseq) {
-        Some(start_ptr) => start_ptr.raw_ptr(cb),
-
-        // Compilation failed, continue executing in the interpreter only
-        None => std::ptr::null(),
+    // Do not test the JIT code in HIR tests
+    if cfg!(test) {
+        return std::ptr::null();
     }
+
+    // Take a lock to avoid writing to ISEQ in parallel with Ractors.
+    // with_vm_lock() does nothing if the program doesn't use Ractors.
+    with_vm_lock(src_loc!(), || {
+        // Compile ISEQ into High-level IR
+        let ssa = match hir::iseq_to_hir(iseq) {
+            Ok(ssa) => ssa,
+            Err(err) => {
+                debug!("ZJIT: iseq_to_hir: {:?}", err);
+                return std::ptr::null();
+            }
+        };
+
+        // Compile High-level IR into machine code
+        let cb = ZJITState::get_code_block();
+        match gen_function(cb, &ssa, iseq) {
+            Some(start_ptr) => start_ptr.raw_ptr(cb),
+
+            // Compilation failed, continue executing in the interpreter only
+            None => std::ptr::null(),
+        }
+    })
 }