use correct svar even if env is escaped

This patch is follo-up of 0a82bfe.
Without this patch, if env is escaped (Proc'ed), strange svar
can be touched.

This patch tracks escaped env and use it.

6 files changed, 122 insertions, 20 deletions

diff --git a/internal/imemo.h b/internal/imemo.h
index 8d506dda7d..57a705ce77 100644
--- a/internal/imemo.h
+++ b/internal/imemo.h
@@ -82,7 +82,7 @@ struct vm_ifunc_argc {
 /*! IFUNC (Internal FUNCtion) */
 struct vm_ifunc {
     VALUE flags;
-    struct rb_control_frame_struct *owner_cfp;
+    VALUE *svar_lep;
     rb_block_call_func_t func;
     const void *data;
     struct vm_ifunc_argc argc;
diff --git a/proc.c b/proc.c
index 3e2d8be6cf..3d26fa9623 100644
--- a/proc.c
+++ b/proc.c
@@ -712,7 +712,8 @@ rb_vm_ifunc_new(rb_block_call_func_t func, const void *data, int min_argc, int m
     }
     arity.argc.min = min_argc;
     arity.argc.max = max_argc;
-    VALUE ret = rb_imemo_new(imemo_ifunc, (VALUE)func, (VALUE)data, arity.packed, (VALUE)GET_EC()->cfp);
+    rb_execution_context_t *ec = GET_EC();
+    VALUE ret = rb_imemo_new(imemo_ifunc, (VALUE)func, (VALUE)data, arity.packed, (VALUE)rb_vm_svar_lep(ec, ec->cfp));
     return (struct vm_ifunc *)ret;
 }
 
diff --git a/test/ruby/test_variable.rb b/test/ruby/test_variable.rb
index b7892f0cdc..1158313776 100644
--- a/test/ruby/test_variable.rb
+++ b/test/ruby/test_variable.rb
@@ -253,6 +253,84 @@ class TestVariable < Test::Unit::TestCase
     assert_include(gv, :$12)
   end
 
+  def prepare_klass_for_test_svar_with_ifunc
+    Class.new do
+      include Enumerable
+      def each(&b)
+        @b = b
+      end
+
+      def check1
+        check2.merge({check1: $1})
+      end
+
+      def check2
+        @b.call('foo')
+        {check2: $1}
+      end
+    end
+  end
+
+  def test_svar_with_ifunc
+    c = prepare_klass_for_test_svar_with_ifunc
+
+    expected_check1_result = {
+      check1: nil, check2: nil
+    }.freeze
+
+    obj = c.new
+    result = nil
+    obj.grep(/(f..)/){
+      result = $1
+    }
+    assert_equal nil, result
+    assert_equal nil, $1
+    assert_equal expected_check1_result, obj.check1
+    assert_equal 'foo', result
+    assert_equal 'foo', $1
+
+    # this frame was escaped so try it again
+    $~ = nil
+    obj = c.new
+    result = nil
+    obj.grep(/(f..)/){
+      result = $1
+    }
+    assert_equal nil, result
+    assert_equal nil, $1
+    assert_equal expected_check1_result, obj.check1
+    assert_equal 'foo', result
+    assert_equal 'foo', $1
+
+    # different context
+    result = nil
+    Fiber.new{
+      obj = c.new
+      obj.grep(/(f..)/){
+        result = $1
+      }
+    }.resume # obj is created in antoher Fiber
+    assert_equal nil, result
+    assert_equal expected_check1_result, obj.check1
+    assert_equal 'foo', result
+    assert_equal 'foo', $1
+
+    # different thread context
+    result = nil
+    Thread.new{
+      obj = c.new
+      obj.grep(/(f..)/){
+        result = $1
+      }
+    }.join # obj is created in another Thread
+
+    assert_equal nil, result
+    assert_equal expected_check1_result, obj.check1
+    assert_equal 'foo', result
+    assert_equal 'foo', $1
+  end
+
+
   def test_global_variable_0
     assert_in_out_err(["-e", "$0='t'*1000;print $0"], "", /\At+\z/, [])
   end
diff --git a/vm.c b/vm.c
index 60150d019f..528901ed8b 100644
--- a/vm.c
+++ b/vm.c
@@ -1291,17 +1291,41 @@ MJIT_FUNC_EXPORTED VALUE
 rb_vm_make_proc_lambda(const rb_execution_context_t *ec, const struct rb_captured_block *captured, VALUE klass, int8_t is_lambda)
 {
     VALUE procval;
+    enum imemo_type code_type = imemo_type(captured->code.val);
 
     if (!VM_ENV_ESCAPED_P(captured->ep)) {
         rb_control_frame_t *cfp = VM_CAPTURED_BLOCK_TO_CFP(captured);
         vm_make_env_object(ec, cfp);
     }
+
     VM_ASSERT(VM_EP_IN_HEAP_P(ec, captured->ep));
-    VM_ASSERT(imemo_type_p(captured->code.val, imemo_iseq) ||
-              imemo_type_p(captured->code.val, imemo_ifunc));
+    VM_ASSERT(code_type == imemo_iseq || code_type == imemo_ifunc);
 
     procval = vm_proc_create_from_captured(klass, captured,
-                                           imemo_type(captured->code.val) == imemo_iseq ? block_type_iseq : block_type_ifunc, FALSE, is_lambda);
+                                           code_type == imemo_iseq ? block_type_iseq : block_type_ifunc,
+                                           FALSE, is_lambda);
+
+    if (code_type ==  imemo_ifunc) {
+        struct vm_ifunc *ifunc = (struct vm_ifunc *)captured->code.val;
+        if (ifunc->svar_lep) {
+            VALUE ep0 = ifunc->svar_lep[0];
+            if (RB_TYPE_P(ep0, T_IMEMO) && imemo_type_p(ep0, imemo_env)) {
+                // `ep0 == imemo_env` means this ep is escaped to heap (in env object).
+                const rb_env_t *env = (const rb_env_t *)ep0;
+                ifunc->svar_lep = (VALUE *)env->ep;
+            }
+            else {
+                VM_ASSERT(FIXNUM_P(ep0));
+                if (ep0 & VM_ENV_FLAG_ESCAPED) {
+                    // ok. do nothing
+                }
+                else {
+                    ifunc->svar_lep = NULL;
+                }
+            }
+        }
+    }
+
     return procval;
 }
 
@@ -1620,19 +1644,13 @@ rb_vm_invoke_proc_with_self(rb_execution_context_t *ec, rb_proc_t *proc, VALUE s
 
 /* special variable */
 
-static rb_control_frame_t *
-vm_svar_frame(const rb_execution_context_t *ec, rb_control_frame_t *cfp)
+VALUE *
+rb_vm_svar_lep(const rb_execution_context_t *ec, const rb_control_frame_t *cfp)
 {
     while (cfp->pc == 0) {
         if (VM_FRAME_TYPE(cfp) ==  VM_FRAME_MAGIC_IFUNC) {
             struct vm_ifunc *ifunc = (struct vm_ifunc *)cfp->iseq;
-            rb_control_frame_t *owner_cfp = ifunc->owner_cfp;
-            if (cfp < owner_cfp) {
-                cfp = owner_cfp;
-            }
-            else {
-                return NULL;
-            }
+            return ifunc->svar_lep;
         }
         else {
             cfp = RUBY_VM_PREVIOUS_CONTROL_FRAME(cfp);
@@ -1642,21 +1660,25 @@ vm_svar_frame(const rb_execution_context_t *ec, rb_control_frame_t *cfp)
             return NULL;
         }
     }
-    return cfp;
+
+    if (cfp) {
+        return (VALUE *)VM_CF_LEP(cfp);
+    }
+    else {
+        return NULL;
+    }
 }
 
 static VALUE
 vm_cfp_svar_get(const rb_execution_context_t *ec, rb_control_frame_t *cfp, VALUE key)
 {
-    cfp = vm_svar_frame(ec, cfp);
-    return lep_svar_get(ec, cfp ? VM_CF_LEP(cfp) : 0, key);
+    return lep_svar_get(ec, rb_vm_svar_lep(ec, cfp), key);
 }
 
 static void
 vm_cfp_svar_set(const rb_execution_context_t *ec, rb_control_frame_t *cfp, VALUE key, const VALUE val)
 {
-    cfp = vm_svar_frame(ec, cfp);
-    lep_svar_set(ec, cfp ? VM_CF_LEP(cfp) : 0, key, val);
+    lep_svar_set(ec, rb_vm_svar_lep(ec, cfp), key, val);
 }
 
 static VALUE
diff --git a/vm_core.h b/vm_core.h
index e07b49077a..9d7c54b4d1 100644
--- a/vm_core.h
+++ b/vm_core.h
@@ -1768,6 +1768,7 @@ rb_vm_living_threads_init(rb_vm_t *vm)
 typedef int rb_backtrace_iter_func(void *, VALUE, int, VALUE);
 rb_control_frame_t *rb_vm_get_ruby_level_next_cfp(const rb_execution_context_t *ec, const rb_control_frame_t *cfp);
 rb_control_frame_t *rb_vm_get_binding_creatable_next_cfp(const rb_execution_context_t *ec, const rb_control_frame_t *cfp);
+VALUE *rb_vm_svar_lep(const rb_execution_context_t *ec, const rb_control_frame_t *cfp);
 int rb_vm_get_sourceline(const rb_control_frame_t *);
 void rb_vm_stack_to_heap(rb_execution_context_t *ec);
 void ruby_thread_init_stack(rb_thread_t *th);
diff --git a/yjit/src/cruby_bindings.inc.rs b/yjit/src/cruby_bindings.inc.rs
index 08ff196f05..989a0eb279 100644
--- a/yjit/src/cruby_bindings.inc.rs
+++ b/yjit/src/cruby_bindings.inc.rs
@@ -570,7 +570,7 @@ pub struct vm_ifunc_argc {
 #[repr(C)]
 pub struct vm_ifunc {
     pub flags: VALUE,
-    pub owner_cfp: *mut rb_control_frame_struct,
+    pub svar_lep: *mut VALUE,
     pub func: rb_block_call_func_t,
     pub data: *const ::std::os::raw::c_void,
     pub argc: vm_ifunc_argc,