local reasoning fix

Author: XrXr

yjit/src/codegen.rs

@@ -234,20 +234,21 @@ impl<'a> JITState<'a> {
         result
     }
 
-    /// Return true if the current ISEQ could escape an 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 {
-        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,
+    // TODO
+    fn can_assume_on_stack_env(&self) -> bool {
+        // In principle, all ISEQs can run with an escaped environment. But,
+        // if we're compiling the currently running iseq, and the environment
+        // is currently on-stack, we can say our code will run as such. The
+        // compiler never causes an environment escape and maintains the
+        // current state.
+        if unsafe { self.iseq == get_cfp_iseq(self.get_cfp()) && !cfp_env_has_escaped(self.get_cfp()) } {
+            //
+            if iseq_escapes_ep(self.iseq) {
+                return false;
+            }
+            true
+        } else {
+            false
         }
     }
 
@@ -370,15 +371,13 @@ impl<'a> JITState<'a> {
         true
     }
 
-    /// Assume that base pointer is equal to environment pointer in the current ISEQ.
-    /// Return true if it's safe to assume so.
     fn assume_no_ep_escape(&mut self, asm: &mut Assembler) -> bool {
+        if !self.can_assume_on_stack_env() {
+            return false; // Unsound to rely on TODO
+        }
         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.
-        }
         self.no_ep_escape = true;
         true
     }
@@ -2509,7 +2508,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);
     }
 
@@ -2610,7 +2609,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);
     }
 

yjit/src/cruby.rs

@@ -176,7 +176,7 @@ pub use rb_iseq_encoded_size as get_iseq_encoded_size;
 pub use rb_get_iseq_body_local_iseq as get_iseq_body_local_iseq;
 pub use rb_get_iseq_body_iseq_encoded as get_iseq_body_iseq_encoded;
 pub use rb_get_iseq_body_stack_max as get_iseq_body_stack_max;
-pub use rb_get_iseq_body_type as get_iseq_body_type;
+//pub use rb_get_iseq_body_type as get_iseq_body_type;
 pub use rb_get_iseq_flags_has_lead as get_iseq_flags_has_lead;
 pub use rb_get_iseq_flags_has_opt as get_iseq_flags_has_opt;
 pub use rb_get_iseq_flags_has_kw as get_iseq_flags_has_kw;
@@ -598,6 +598,12 @@ impl From<VALUE> for u16 {
     }
 }
 
+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) }