From 67c2cdc59a5fc8a82804c72a30a2d23ba36f2366 Mon Sep 17 00:00:00 2001
From: John Hawthorn <john@hawthorn.email>
Date: Tue, 8 Jun 2021 07:52:57 -0700
Subject: Implement gen_getlocal

This extracts the generation code from getlocal_wc1, since this is the
same just with more loops inside vm_get_ep.
---
 bootstraptest/test_yjit.rb | 14 ++++++++++++++
 yjit_codegen.c             | 33 ++++++++++++++++++++++++---------
 2 files changed, 38 insertions(+), 9 deletions(-)

diff --git a/bootstraptest/test_yjit.rb b/bootstraptest/test_yjit.rb
index 5bf415889e..51c82b55b5 100644
--- a/bootstraptest/test_yjit.rb
+++ b/bootstraptest/test_yjit.rb
@@ -1021,3 +1021,17 @@ assert_equal 'foo123', %q{
   make_str("foo", 123)
   make_str("foo", 123)
 }
+
+# getlocal with 2 levels
+assert_equal '7', %q{
+  def foo(foo, bar)
+    while foo > 0
+      while bar > 0
+        return foo + bar
+      end
+    end
+  end
+
+  foo(5,2)
+  foo(5,2)
+}
diff --git a/yjit_codegen.c b/yjit_codegen.c
index 87c5b12043..4870911404 100644
--- a/yjit_codegen.c
+++ b/yjit_codegen.c
@@ -787,22 +787,21 @@ gen_getlocal_wc0(jitstate_t* jit, ctx_t* ctx)
 }
 
 static codegen_status_t
-gen_getlocal_wc1(jitstate_t* jit, ctx_t* ctx)
+gen_getlocal_generic(ctx_t* ctx, uint32_t local_idx, uint32_t level)
 {
-    //fprintf(stderr, "gen_getlocal_wc1\n");
-
     // Load environment pointer EP from CFP
     mov(cb, REG0, member_opnd(REG_CFP, rb_control_frame_t, ep));
 
-    // 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))
-    mov(cb, REG0, mem_opnd(64, REG0, SIZEOF_VALUE * VM_ENV_DATA_INDEX_SPECVAL));
-    and(cb, REG0, imm_opnd(~0x03));
+    while (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))
+        mov(cb, REG0, mem_opnd(64, REG0, SIZEOF_VALUE * VM_ENV_DATA_INDEX_SPECVAL));
+        and(cb, REG0, imm_opnd(~0x03));
+    }
 
     // Load the local from the block
     // val = *(vm_get_ep(GET_EP(), level) - idx);
-    int32_t local_idx = (int32_t)jit_get_arg(jit, 0);
     const int32_t offs = -(SIZEOF_VALUE * local_idx);
     mov(cb, REG0, mem_opnd(64, REG0, offs));
 
@@ -813,6 +812,21 @@ gen_getlocal_wc1(jitstate_t* jit, ctx_t* ctx)
     return YJIT_KEEP_COMPILING;
 }
 
+static codegen_status_t
+gen_getlocal(jitstate_t* jit, ctx_t* ctx)
+{
+    int32_t idx = (int32_t)jit_get_arg(jit, 0);
+    int32_t level = (int32_t)jit_get_arg(jit, 1);
+    return gen_getlocal_generic(ctx, idx, level);
+}
+
+static codegen_status_t
+gen_getlocal_wc1(jitstate_t* jit, ctx_t* ctx)
+{
+    int32_t idx = (int32_t)jit_get_arg(jit, 0);
+    return gen_getlocal_generic(ctx, idx, 1);
+}
+
 static codegen_status_t
 gen_setlocal_wc0(jitstate_t* jit, ctx_t* ctx)
 {
@@ -3067,6 +3081,7 @@ yjit_init_codegen(void)
     yjit_reg_op(BIN(putobject_INT2FIX_0_), gen_putobject_int2fix);
     yjit_reg_op(BIN(putobject_INT2FIX_1_), gen_putobject_int2fix);
     yjit_reg_op(BIN(putself), gen_putself);
+    yjit_reg_op(BIN(getlocal), gen_getlocal);
     yjit_reg_op(BIN(getlocal_WC_0), gen_getlocal_wc0);
     yjit_reg_op(BIN(getlocal_WC_1), gen_getlocal_wc1);
     yjit_reg_op(BIN(setlocal_WC_0), gen_setlocal_wc0);
-- 
cgit v1.2.3