ZJIT: Add codegen for GetLocal and SetLocal

They're only used when level≠0. Same EP hopping logic as interpreter and
YJIT. Change assert_compiles() to get ISeq by method name since the old
code didn't support methods defined using define_method() which I need
for the new test.

2 files changed, 61 insertions, 6 deletions

diff --git a/test/ruby/test_zjit.rb b/test/ruby/test_zjit.rb
index 6e0f274c30..be8b910fd6 100644
--- a/test/ruby/test_zjit.rb
+++ b/test/ruby/test_zjit.rb
@@ -62,6 +62,25 @@ class TestZJIT < Test::Unit::TestCase
     }
   end
 
+  def test_nested_local_access
+    assert_compiles '[1, 2, 3]', %q{
+      1.times do |l2|
+        1.times do |l1|
+          define_method(:test) do
+            l1 = 1
+            l2 = 2
+            l3 = 3
+            [l1, l2, l3]
+          end
+        end
+      end
+
+      test
+      test
+      test
+    }, call_threshold: 3, insns: [:getlocal, :setlocal, :getlocal_WC_0, :setlocal_WC_1]
+  end
+
   def test_send_without_block
     assert_compiles '[1, 2, 3]', %q{
       def foo = 1
@@ -791,7 +810,7 @@ class TestZJIT < Test::Unit::TestCase
       result = {
         ret_val:,
         #{ unless insns.empty?
-          'insns: RubyVM::InstructionSequence.of(_test_proc).enum_for(:each_child).map(&:to_a)'
+          'insns: RubyVM::InstructionSequence.of(method(:test)).to_a'
         end}
       }
       IO.open(#{pipe_fd}).write(Marshal.dump(result))
@@ -805,13 +824,12 @@ class TestZJIT < Test::Unit::TestCase
     assert status.success?, message
 
     result = Marshal.load(result)
-    assert_equal expected, result.fetch(:ret_val).inspect
+    assert_equal(expected, result.fetch(:ret_val).inspect)
 
     unless insns.empty?
-      iseqs = result.fetch(:insns)
-      iseqs.filter! { it[9] == :method } # ISeq type
-      assert_equal 1, iseqs.size, "Opcode assertions tests must define exactly one method"
-      iseq_insns = iseqs.first.last
+      iseq = result.fetch(:insns)
+      assert_equal("YARVInstructionSequence/SimpleDataFormat", iseq.first, "failed to get iseq disassembly")
+      iseq_insns = iseq.last
 
       expected_insns = Set.new(insns)
       iseq_insns.each do
diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs
index f4f1109134..915a1e93bd 100644
--- a/zjit/src/codegen.rs
+++ b/zjit/src/codegen.rs
@@ -1,5 +1,6 @@
 use std::cell::Cell;
 use std::rc::Rc;
+use std::num::NonZeroU32;
 
 use crate::backend::current::{Reg, ALLOC_REGS};
 use crate::profile::get_or_create_iseq_payload;
@@ -279,6 +280,8 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         Insn::GetIvar { self_val, id, state: _ } => gen_getivar(asm, opnd!(self_val), *id),
         Insn::SetGlobal { id, val, state: _ } => gen_setglobal(asm, *id, opnd!(val)),
         Insn::GetGlobal { id, state: _ } => gen_getglobal(asm, *id),
+        &Insn::GetLocal { ep_offset, level } => gen_nested_getlocal(asm, ep_offset, level)?,
+        Insn::SetLocal { val, ep_offset, level } => return gen_nested_setlocal(asm, opnd!(val), *ep_offset, *level),
         Insn::GetConstantPath { ic, state } => gen_get_constant_path(asm, *ic, &function.frame_state(*state)),
         Insn::SetIvar { self_val, id, val, state: _ } => return gen_setivar(asm, opnd!(self_val), *id, opnd!(val)),
         Insn::SideExit { state } => return gen_side_exit(jit, asm, &function.frame_state(*state)),
@@ -298,6 +301,40 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
     Some(())
 }
 
+// Get EP at `level` from CFP
+fn gen_get_ep(asm: &mut Assembler, level: u32) -> Opnd {
+    // Load environment pointer EP from CFP into a register
+    let ep_opnd = Opnd::mem(64, CFP, RUBY_OFFSET_CFP_EP);
+    let mut ep_opnd = asm.load(ep_opnd);
+
+    for _ in 0..level {
+        // Get the previous EP from the current EP
+        // See GET_PREV_EP(ep) macro
+        // VALUE *prev_ep = ((VALUE *)((ep)[VM_ENV_DATA_INDEX_SPECVAL] & ~0x03))
+        const UNTAGGING_MASK: Opnd = Opnd::Imm(!0x03);
+        let offset = SIZEOF_VALUE_I32 * VM_ENV_DATA_INDEX_SPECVAL;
+        ep_opnd = asm.load(Opnd::mem(64, ep_opnd, offset));
+        ep_opnd = asm.and(ep_opnd, UNTAGGING_MASK);
+    }
+
+    ep_opnd
+}
+
+/// Get a local variable from a higher scope. `local_ep_offset` is in number of VALUEs.
+fn gen_nested_getlocal(asm: &mut Assembler, local_ep_offset: u32, level: NonZeroU32) -> Option<lir::Opnd> {
+    let ep = gen_get_ep(asm, level.get());
+    let offset = -(SIZEOF_VALUE_I32 * i32::try_from(local_ep_offset).ok()?);
+    Some(asm.load(Opnd::mem(64, ep, offset)))
+}
+
+/// Set a local variable from a higher scope. `local_ep_offset` is in number of VALUEs.
+fn gen_nested_setlocal(asm: &mut Assembler, val: Opnd, local_ep_offset: u32, level: NonZeroU32) -> Option<()> {
+    let ep = gen_get_ep(asm, level.get());
+    let offset = -(SIZEOF_VALUE_I32 * i32::try_from(local_ep_offset).ok()?);
+    asm.mov(Opnd::mem(64, ep, offset), val);
+    Some(())
+}
+
 fn gen_get_constant_path(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;