YJIT: Fix not reading locals from `cfp->ep` after `YJIT.enable` and exceptional entry

[Backport #21941]

In case of `--yjit-disable`, YJIT only starts to record environment
escapes after `RubyVM::YJIT.enable`. Previously we falsely assumed that
we always have a full history all the way back to VM boot. This had YJIT
install and run code that assume EP=BP when EP≠BP for some exceptional
entry into the middle of a running frame, if the environment escaped
before `YJIT.enable`.

The fix is to reject exceptional entry with an escaped environment.
Rename things and explain in more detail how the predicate for deciding
to assume EP=BP works. It's quite subtle since it reasons about all
parties in the system that push a control frame and then run JIT code.

Note that while can_assume_on_stack_env() checks the currently running
environment if it so happens to be the one YJIT is compiling against, it
can return true for any ISEQ. The check isn't necessary for fixing the
bug, and the load bearing part of this patch is the change to
exceptional entries.

This fix is flat on speed and space on ruby-bench headline benchmarks.

Many thanks for the community effort to create a small test case for
this bug.

7 files changed, 73 insertions, 15 deletions

diff --git a/test/ruby/test_yjit.rb b/test/ruby/test_yjit.rb
index 2096585451..7d8d92475a 100644
--- a/test/ruby/test_yjit.rb
+++ b/test/ruby/test_yjit.rb
@@ -1784,6 +1784,27 @@ class TestYJIT < Test::Unit::TestCase
     RUBY
   end
 
+  def test_exceptional_entry_into_env_escaped_before_yjit_enablement
+    threshold = 2
+    assert_separately(["--disable-all", "--yjit-disable", "--yjit-call-threshold=#{threshold}"], <<~RUBY)
+      def run
+        @captured_env = ->{}
+        RubyVM::YJIT.enable
+
+        i = 0
+        while i < #{threshold}
+          next_i = i + 1
+          from_break = tap { break i + 1 } # break from the block generates an exceptional entry
+          assert_equal(from_break, next_i, '[Bug #21941]')
+          i = next_i
+        end
+      end
+
+      run
+      assert_equal(#{threshold}, @captured_env.binding.local_variable_get(:i))
+    RUBY
+  end
+
   private
 
   def code_gc_helpers
diff --git a/yjit.rb b/yjit.rb
index 9e9a79fa73..9fb4d6876b 100644
--- a/yjit.rb
+++ b/yjit.rb
@@ -353,6 +353,9 @@ module RubyVM::YJIT
       # Number of failed compiler invocations
       compilation_failure = stats[:compilation_failure]
 
+      # Number of refused exceptional entries with an escaped environment
+      exceptional_entry_escaped_env = stats[:exceptional_entry_escaped_env]
+
       code_region_overhead = stats[:code_region_size] - (stats[:inline_code_size] + stats[:outlined_code_size])
 
       out.puts "num_send:              " + format_number(13, stats[:num_send])
@@ -389,6 +392,7 @@ module RubyVM::YJIT
       out.puts "bindings_allocations:  " + format_number(13, stats[:binding_allocations])
       out.puts "bindings_set:          " + format_number(13, stats[:binding_set])
       out.puts "compilation_failure:   " + format_number(13, compilation_failure) if compilation_failure != 0
+      out.puts "exceptional_entry_escaped_env:" + format_number(6, exceptional_entry_escaped_env) if exceptional_entry_escaped_env != 0
       out.puts "live_iseq_count:       " + format_number(13, stats[:live_iseq_count])
       out.puts "iseq_alloc_count:      " + format_number(13, stats[:iseq_alloc_count])
       out.puts "compiled_iseq_entry:   " + format_number(13, stats[:compiled_iseq_entry])
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index 61de6a32da..4e4d58d5df 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -234,20 +234,36 @@ impl<'a> JITState<'a> {
         result
     }
 
-    /// Return true if the current ISEQ could escape an environment.
+    /// Return true if the JIT code can use [`Self::assume_no_ep_escape`]
+    /// and will run with an on-stack (`!VM_ENV_ESCAPED_P`) environment.
     ///
-    /// As of vm_push_frame(), EP is always equal to BP. However, after pushing
-    /// a frame, some ISEQ setups call vm_bind_update_env(), which redirects EP.
-    /// Also, some method calls escape the environment to the heap.
-    fn escapes_ep(&self) -> bool {
+    /// ## Reasoning about ISEQs that are not currently running
+    ///
+    /// As of vm_push_frame() and its JIT code equivalent, EP is always equal to BP (the
+    /// environment is on-stack and has not escaped). We can usually assume this is the starting
+    /// condition upon entry into JIT code. However, after pushing a frame and before entry into
+    /// JIT code, some ISEQ setups call vm_bind_update_env(), which redirects EP.
+    ///
+    /// ## After making the assumption
+    ///
+    /// After JIT code entry, many ruby operations can have the environment escape to heap. These
+    /// are handled by [`crate::invariants`].
+    ///
+    /// Exceptional entry through jit_exec_exception() is an extreme case of the environment state
+    /// changing between vm_push_frame() and entry into JIT code. The frame could have been pushed
+    /// before YJIT is enabled. The exception entry point refuses entry with an escaped environment.
+    fn can_assume_on_stack_env(&self) -> bool {
         match unsafe { get_iseq_body_type(self.iseq) } {
             // <main> frame is always associated to TOPLEVEL_BINDING.
             ISEQ_TYPE_MAIN |
             // Kernel#eval uses a heap EP when a Binding argument is not nil.
-            ISEQ_TYPE_EVAL => true,
-            // If this ISEQ has previously escaped EP, give up the optimization.
-            _ if iseq_escapes_ep(self.iseq) => true,
-            _ => false,
+            ISEQ_TYPE_EVAL => false,
+            // Check the running environment if compiling for it
+            _ if unsafe { self.iseq == get_cfp_iseq(self.get_cfp()) && cfp_env_has_escaped(self.get_cfp()) } => false,
+            // If we've seen this ISEQ run with an escaped environment, give up the optimization
+            // to avoid excessive invalidations (even though it may be fine for soundness).
+            _ if seen_escaped_env(self.iseq) => false,
+            _ => true,
         }
     }
 
@@ -376,8 +392,8 @@ impl<'a> JITState<'a> {
         if jit_ensure_block_entry_exit(self, asm).is_none() {
             return false; // out of space, give up
         }
-        if self.escapes_ep() {
-            return false; // EP has been escaped in this ISEQ. disable the optimization to avoid an invalidation loop.
+        if !self.can_assume_on_stack_env() {
+            return false; // Unsound or unprofitable to make the assumption
         }
         self.no_ep_escape = true;
         true
@@ -2448,7 +2464,7 @@ fn gen_getlocal_generic(
     level: u32,
 ) -> Option<CodegenStatus> {
     // Split the block if we need to invalidate this instruction when EP escapes
-    if level == 0 && !jit.escapes_ep() && !jit.at_compile_target() {
+    if level == 0 && jit.can_assume_on_stack_env() && !jit.at_compile_target() {
         return jit.defer_compilation(asm);
     }
 
@@ -2549,7 +2565,7 @@ fn gen_setlocal_generic(
     }
 
     // Split the block if we need to invalidate this instruction when EP escapes
-    if level == 0 && !jit.escapes_ep() && !jit.at_compile_target() {
+    if level == 0 && jit.can_assume_on_stack_env() && !jit.at_compile_target() {
         return jit.defer_compilation(asm);
     }
 
diff --git a/yjit/src/cruby.rs b/yjit/src/cruby.rs
index d34b049a45..62fc9d8f06 100644
--- a/yjit/src/cruby.rs
+++ b/yjit/src/cruby.rs
@@ -598,6 +598,13 @@ impl From<VALUE> for u16 {
     }
 }
 
+/// Check whether a control frame has an escaped environment
+pub unsafe fn cfp_env_has_escaped(cfp: *mut rb_control_frame_struct) -> bool {
+    use crate::utils::IntoUsize;
+    let ep = get_cfp_ep(cfp);
+    0 != ep.offset(VM_ENV_DATA_INDEX_FLAGS as isize).read().0 & VM_ENV_FLAG_ESCAPED.as_usize()
+}
+
 /// Produce a Ruby string from a Rust string slice
 pub fn rust_str_to_ruby(str: &str) -> VALUE {
     unsafe { rb_utf8_str_new(str.as_ptr() as *const _, str.len() as i64) }
diff --git a/yjit/src/invariants.rs b/yjit/src/invariants.rs
index 0f22fba6b8..5f319c9495 100644
--- a/yjit/src/invariants.rs
+++ b/yjit/src/invariants.rs
@@ -168,8 +168,8 @@ pub fn track_no_ep_escape_assumption(uninit_block: BlockRef, iseq: IseqPtr) {
         .insert(uninit_block);
 }
 
-/// Returns true if a given ISEQ has previously escaped an environment.
-pub fn iseq_escapes_ep(iseq: IseqPtr) -> bool {
+/// Returns true if a given ISEQ has escaped an environment since YJIT boot.
+pub fn seen_escaped_env(iseq: IseqPtr) -> bool {
     Invariants::get_instance()
         .no_ep_escape_iseqs
         .get(&iseq)
diff --git a/yjit/src/stats.rs b/yjit/src/stats.rs
index 84549fa5d3..e0dc257a31 100644
--- a/yjit/src/stats.rs
+++ b/yjit/src/stats.rs
@@ -248,6 +248,7 @@ pub const DEFAULT_COUNTERS: &'static [Counter] = &[
     Counter::compiled_blockid_count,
     Counter::compiled_block_count,
     Counter::deleted_defer_block_count,
+    Counter::exceptional_entry_escaped_env,
     Counter::compiled_branch_count,
     Counter::compile_time_ns,
     Counter::compilation_failure,
@@ -599,6 +600,8 @@ make_counters! {
     iseq_stack_too_large,
     iseq_too_long,
 
+    exceptional_entry_escaped_env,
+
     temp_reg_opnd,
     temp_mem_opnd,
     temp_spill,
diff --git a/yjit/src/yjit.rs b/yjit/src/yjit.rs
index 517a0daae5..e12706b79a 100644
--- a/yjit/src/yjit.rs
+++ b/yjit/src/yjit.rs
@@ -157,6 +157,13 @@ pub extern "C" fn rb_yjit_iseq_gen_entry_point(iseq: IseqPtr, ec: EcPtr, jit_exc
         return std::ptr::null();
     }
 
+    // In case of exceptional entry, reject escaped environment.
+    // This allows us to use the fact that new frames generally start with an on-stack environment.
+    if jit_exception && unsafe { cfp_env_has_escaped(get_ec_cfp(ec)) } {
+        incr_counter!(exceptional_entry_escaped_env);
+        return std::ptr::null();
+    }
+
     // If a custom call threshold was not specified on the command-line and
     // this is a large application (has very many ISEQs), switch to
     // using the call threshold for large applications after this entry point