rusty OuterVariables

Author: XrXr

zjit.c

@@ -29,7 +29,8 @@
 STATIC_ASSERT(pointer_tagging_scheme, USE_FLONUM);
 
 enum zjit_struct_offsets {
-    ISEQ_BODY_OFFSET_PARAM = offsetof(struct rb_iseq_constant_body, param)
+    ISEQ_BODY_OFFSET_PARAM = offsetof(struct rb_iseq_constant_body, param),
+    ISEQ_BODY_OFFSET_OUTER_VARIABLES = offsetof(struct rb_iseq_constant_body, outer_variables)
 };
 
 // Special JITFrame used by all C method calls. We don't control the native
@@ -220,23 +221,6 @@ rb_zjit_local_id(const rb_iseq_t *iseq, unsigned idx)
     return ISEQ_BODY(iseq)->local_table[idx];
 }
 
-// True if `blockiseq` (or any iseq nested within it) has bytecode that assigns
-// to the outer local variable named `id`.
-// The same table backs Ractor.shareable_proc's isolation checks.
-bool
-rb_zjit_iseq_writes_outer_local_p(const rb_iseq_t *blockiseq, ID id)
-{
-    struct rb_id_table *ovs = ISEQ_BODY(blockiseq)->outer_variables;
-    if (ovs == NULL) return false;
-    VALUE write = Qfalse;
-    if (rb_id_table_lookup(ovs, id, &write)) {
-        // Table entry precense means local is referenced.
-        // Truth entry means it's referenced through a setlocal.
-        return RTEST(write);
-    }
-    return false;
-}
-
 bool rb_zjit_cme_is_cfunc(const rb_callable_method_entry_t *me, const void *func);
 
 const struct rb_callable_method_entry_struct *

zjit/bindgen/src/main.rs

@@ -297,7 +297,7 @@ fn main() {
         .allowlist_function("rb_zjit_iseq_inspect")
         .allowlist_function("rb_zjit_iseq_insn_set")
         .allowlist_function("rb_zjit_local_id")
-        .allowlist_function("rb_zjit_iseq_writes_outer_local_p")
+        .allowlist_function("rb_id_table_lookup")
         .allowlist_function("rb_set_cfp_(pc|sp)")
         .allowlist_function("rb_c_method_tracing_currently_enabled")
         .allowlist_function("rb_zjit_method_tracing_currently_enabled")
@@ -451,6 +451,9 @@ fn main() {
         .blocklist_type("ID")
         .blocklist_type("rb_iseq_constant_body")
 
+        // We only need id_table as an opaque pointer to pass to its APIs
+        .opaque_type("rb_id_table")
+
         // Avoid binding to stuff we don't use
         .blocklist_item("rb_thread_struct.*")
         .opaque_type("rb_thread_struct.*")

zjit/src/cruby.rs

@@ -93,6 +93,7 @@ use std::ffi::{c_void, CString, CStr};
 use std::fmt::{Debug, Display, Formatter};
 use std::os::raw::{c_char, c_int, c_long, c_uint};
 use std::panic::{catch_unwind, UnwindSafe};
+use std::ptr::NonNull;
 
 use crate::cast::IntoUsize as _;
 
@@ -749,9 +750,42 @@ impl VALUE {
 
 pub type IseqParameters = rb_iseq_constant_body_rb_iseq_parameters;
 
+/// How a block iseq refers to a variable in an enclosing scope, as recorded in
+/// `ISEQ_BODY(blockiseq)->outer_variables`. `compile.c` aggregates accesses from
+/// nested blocks up the chain, and the same table backs `Ractor.shareable_proc`'s
+/// isolation checks.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum OuterLocalAccess {
+    /// The variable is read but never assigned to.
+    ReadOnly,
+    /// The variable is assigned to and maybe also read.
+    ReadWrite,
+}
+
+/// Wrapper over an iseq's `outer_variables` table, which describes
+/// how a block iseq refers to a variable in an enclosing scope.
+#[derive(Clone, Copy)]
+pub struct OuterVariables(Option<NonNull<rb_id_table>>);
+
+impl OuterVariables {
+    /// Look up how the enclosing-scope local `id` is accessed by the iseq (or any
+    /// iseq nested within it). Returns `None` when the variable isn't referenced.
+    pub fn local_access(self, id: ID) -> Option<OuterLocalAccess> {
+        let table = self.0?;
+        let mut write = Qfalse;
+        // Non-zero return means there's a table entry, i.e. the variable is referenced.
+        if unsafe { rb_id_table_lookup(table.as_ptr(), id, &mut write) } == 0 {
+            return None;
+        }
+        // Truthy means write
+        Some(if write.test() { OuterLocalAccess::ReadWrite } else { OuterLocalAccess::ReadOnly })
+    }
+}
+
 /// Extension trait to enable method calls on [`IseqPtr`]
 pub trait IseqAccess {
     unsafe fn params<'a>(self) -> &'a IseqParameters;
+    unsafe fn outer_variables(self) -> OuterVariables;
 }
 
 impl IseqAccess for IseqPtr {
@@ -760,6 +794,13 @@ impl IseqAccess for IseqPtr {
         use crate::cast::IntoUsize;
         unsafe { &*((*self).body.byte_add(ISEQ_BODY_OFFSET_PARAM.to_usize()) as *const IseqParameters) }
     }
+
+    /// The iseq's `outer_variables` table. See [`OuterVariables`].
+    unsafe fn outer_variables(self) -> OuterVariables {
+        use crate::cast::IntoUsize;
+        let field = unsafe { (*self).body.byte_add(ISEQ_BODY_OFFSET_OUTER_VARIABLES.to_usize()) } as *const *mut rb_id_table;
+        OuterVariables(NonNull::new(unsafe { *field }))
+    }
 }
 
 impl IseqParameters {

zjit/src/cruby_bindings.inc.rs

@@ -5058,21 +5058,26 @@ impl Function {
             &mut (0..state.locals.len())
         } else {
             let params = unsafe { iseq.params() };
-            let block_param_local_idx = if params.flags.has_block() != 0 {
-                Some(params.block_start)
+            let block_param_local_idx: Option<usize> = if params.flags.has_block() != 0 {
+                params.block_start.try_into().ok()
             } else {
                 None
             };
-            // When not escaped, only reload the ones syntactically written to
+            let outer_variables = unsafe { blockiseq.outer_variables() };
+            // When not escaped, only reload the locals the block can have modified.
             &mut (0..state.locals.len()).filter(move |&local_idx| {
-                if block_param_local_idx.and_then(|idx| idx.try_into().ok()).is_some_and(|idx: usize| idx == local_idx) {
-                    // An ostensibly read of the the block param, through `getblockparam` can
-                    // write to the local slot for it. TODO(alan): no reload when outer blocker
-                    // param not referenced in block.
-                    return true;
-                }
                 let id = unsafe { rb_zjit_local_id(iseq, local_idx.try_into().unwrap()) };
-                unsafe { rb_zjit_iseq_writes_outer_local_p(blockiseq, id) }
+                let access = outer_variables.local_access(id);
+                if block_param_local_idx == Some(local_idx) {
+                    // The block param slot is special: `getblockparam` is recorded as a
+                    // read, but it materializes the captured block into this slot. So
+                    // reload it whenever the block references it at all (read or write),
+                    // not just on a setlocal. When the block never references it, the
+                    // slot can't have changed, so skip the reload.
+                    access.is_some()
+                } else {
+                    access == Some(OuterLocalAccess::Write)
+                }
             })
         };
         let mut base: Option<InsnId> = None;