ZJIT: Pull GetEP out of GetLocal (with level > 0) (ruby#16215)

Closes: Shopify#944

As a first step toward enabling CSE/LVN-based deduplication of frame EP loads, 
splitting `GetLocal` handling for `level > 0` into an explicit `GetEP + LoadField` path.

A follow-up PR will handle the `level == 0` case.

Author: nozomemein

zjit/src/hir.rs

@@ -854,6 +854,7 @@ pub enum Insn {
     /// Get a local variable from a higher scope or the heap.
     /// If `use_sp` is true, it uses the SP register to optimize the read.
     /// `rest_param` is used by infer_types to infer the ArrayExact type.
+    /// TODO: Replace the level == 0 + use_sp path with LoadSP + LoadField.
     GetLocal { level: u32, ep_offset: u32, use_sp: bool, rest_param: bool },
     /// Check whether VM_FRAME_FLAG_MODIFIED_BLOCK_PARAM is set in the environment flags.
     /// Returns CBool (0/1).
@@ -1169,7 +1170,7 @@ impl Insn {
             Insn::GetEP { .. } => effects::Empty,
             Insn::GetLEP { .. } => effects::Empty,
             Insn::LoadSelf { .. } => Effect::read_write(abstract_heaps::Frame, abstract_heaps::Empty),
-            Insn::LoadField { .. } => Effect::read_write(abstract_heaps::Other, abstract_heaps::Empty),
+            Insn::LoadField { .. } => Effect::read_write(abstract_heaps::Memory, abstract_heaps::Empty),
             Insn::StoreField { .. } => effects::Any,
             Insn::WriteBarrier { .. } => effects::Any,
             Insn::GetLocal   { .. } => Effect::read_write(abstract_heaps::Locals, abstract_heaps::Empty),
@@ -1272,6 +1273,15 @@ pub struct InsnPrinter<'a> {
     iseq: Option<IseqPtr>,
 }
 
+fn get_local_var_id(iseq: IseqPtr, level: u32, ep_offset: u32) -> ID {
+    let mut current_iseq = iseq;
+    for _ in 0..level {
+        current_iseq = unsafe { rb_get_iseq_body_parent_iseq(current_iseq) };
+    }
+    let local_idx = ep_offset_to_local_idx(current_iseq, ep_offset);
+    unsafe { rb_zjit_local_id(current_iseq, local_idx.try_into().unwrap()) }
+}
+
 /// Get the name of a local variable given iseq, level, and ep_offset.
 /// Returns
 /// - `":name"` if iseq is available and name is a real identifier,
@@ -1281,12 +1291,7 @@ pub struct InsnPrinter<'a> {
 ///
 /// This mimics local_var_name() from iseq.c.
 fn get_local_var_name_for_printer(iseq: Option<IseqPtr>, level: u32, ep_offset: u32) -> Option<String> {
-    let mut current_iseq = iseq?;
-    for _ in 0..level {
-        current_iseq = unsafe { rb_get_iseq_body_parent_iseq(current_iseq) };
-    }
-    let local_idx = ep_offset_to_local_idx(current_iseq, ep_offset);
-    let id: ID = unsafe { rb_zjit_local_id(current_iseq, local_idx.try_into().unwrap()) };
+    let id = get_local_var_id(iseq?, level, ep_offset);
 
     if id.0 == 0 || unsafe { rb_id2str(id) } == Qfalse {
         return Some(String::from("<empty>"));
@@ -3057,6 +3062,29 @@ impl Function {
         self.push_insn(block, Insn::GuardNoBitsSet { val: flags, mask: Const::CUInt64(RUBY_ELTS_SHARED as u64), mask_name: Some(ID!(RUBY_ELTS_SHARED)), reason: SideExitReason::GuardNotShared, state });
     }
 
+    // TODO: This helper is currently used for level>0 local reads only.
+    // Split GetLocal(level==0) into explicit SP/EP helpers in a follow-up.
+    fn get_local_from_ep(
+        &mut self,
+        block: BlockId,
+        ep_offset: u32,
+        level: u32,
+        return_type: Type,
+    ) -> InsnId {
+        let ep = self.push_insn(block, Insn::GetEP { level });
+        let local_id = get_local_var_id(self.iseq, level, ep_offset);
+        let ep_offset = i32::try_from(ep_offset)
+            .unwrap_or_else(|_| panic!("Could not convert ep_offset {ep_offset} to i32"));
+        let offset = -(SIZEOF_VALUE_I32 * ep_offset);
+
+        self.push_insn(block, Insn::LoadField {
+            recv: ep,
+            id: local_id,
+            offset,
+            return_type,
+        })
+    }
+
     /// Rewrite eligible Send opcodes into SendDirect
     /// opcodes if we know the target ISEQ statically. This removes run-time method lookups and
     /// opens the door for inlining.
@@ -6836,7 +6864,7 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                 }
                 YARVINSN_getlocal_WC_1 => {
                     let ep_offset = get_arg(pc, 0).as_u32();
-                    state.stack_push(fun.push_insn(block, Insn::GetLocal { ep_offset, level: 1, use_sp: false, rest_param: false }));
+                    state.stack_push(fun.get_local_from_ep(block, ep_offset, 1, types::BasicObject));
                 }
                 YARVINSN_setlocal_WC_1 => {
                     let ep_offset = get_arg(pc, 0).as_u32();
@@ -6845,7 +6873,11 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                 YARVINSN_getlocal => {
                     let ep_offset = get_arg(pc, 0).as_u32();
                     let level = get_arg(pc, 1).as_u32();
-                    state.stack_push(fun.push_insn(block, Insn::GetLocal { ep_offset, level, use_sp: false, rest_param: false }));
+                    if level == 0 {
+                        state.stack_push(fun.push_insn(block, Insn::GetLocal { ep_offset, level, use_sp: false, rest_param: false }));
+                    } else {
+                        state.stack_push(fun.get_local_from_ep(block, ep_offset, level, types::BasicObject));
+                    }
                 }
                 YARVINSN_setlocal => {
                     let ep_offset = get_arg(pc, 0).as_u32();
@@ -6956,12 +6988,11 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                     }));
 
                     // Push modified block: read Proc from EP.
-                    let modified_val = fun.push_insn(modified_block, Insn::GetLocal {
-                        ep_offset,
-                        level,
-                        use_sp: false,
-                        rest_param: false,
-                    });
+                    let modified_val = if level == 0 {
+                        fun.push_insn(modified_block, Insn::GetLocal { ep_offset, level, use_sp: false, rest_param: false })
+                    } else {
+                        fun.get_local_from_ep(modified_block, ep_offset, level, types::BasicObject)
+                    };
                     finish_getblockparam_branch(
                         &mut fun,
                         modified_block,

zjit/src/hir/opt_tests.rs

@@ -4461,14 +4461,15 @@ mod hir_opt_tests {
         bb3(v6:BasicObject):
           v10:CBool = IsBlockParamModified l1
           IfTrue v10, bb4(v6)
-          v19:BasicObject = GetBlockParam :block, l1, EP@3
-          Jump bb6(v6, v19)
+          v20:BasicObject = GetBlockParam :block, l1, EP@3
+          Jump bb6(v6, v20)
         bb4(v11:BasicObject):
-          v17:BasicObject = GetLocal :block, l1, EP@3
-          Jump bb6(v11, v17)
-        bb6(v21:BasicObject, v22:BasicObject):
+          v17:CPtr = GetEP 1
+          v18:BasicObject = LoadField v17, :block@0x1000
+          Jump bb6(v11, v18)
+        bb6(v22:BasicObject, v23:BasicObject):
           CheckInterrupts
-          Return v22
+          Return v23
         ");
     }
 

zjit/src/hir/tests.rs

@@ -946,18 +946,23 @@ pub mod hir_build_tests {
           v4:BasicObject = LoadArg :self@0
           Jump bb3(v4)
         bb3(v6:BasicObject):
-          v10:BasicObject = GetLocal :l2, l2, EP@4
-          SetLocal :l1, l1, EP@3, v10
-          v15:BasicObject = GetLocal :l1, l1, EP@3
-          v17:BasicObject = GetLocal :l2, l2, EP@4
-          v20:BasicObject = Send v15, :+, v17 # SendFallbackReason: Uncategorized(opt_plus)
-          SetLocal :l2, l2, EP@4, v20
-          v25:BasicObject = GetLocal :l2, l2, EP@4
-          v27:BasicObject = GetLocal :l3, l3, EP@5
-          v30:BasicObject = Send v25, :+, v27 # SendFallbackReason: Uncategorized(opt_plus)
-          SetLocal :l3, l3, EP@5, v30
+          v10:CPtr = GetEP 2
+          v11:BasicObject = LoadField v10, :l2@0x1000
+          SetLocal :l1, l1, EP@3, v11
+          v16:CPtr = GetEP 1
+          v17:BasicObject = LoadField v16, :l1@0x1001
+          v19:CPtr = GetEP 2
+          v20:BasicObject = LoadField v19, :l2@0x1000
+          v23:BasicObject = Send v17, :+, v20 # SendFallbackReason: Uncategorized(opt_plus)
+          SetLocal :l2, l2, EP@4, v23
+          v28:CPtr = GetEP 2
+          v29:BasicObject = LoadField v28, :l2@0x1000
+          v31:CPtr = GetEP 3
+          v32:BasicObject = LoadField v31, :l3@0x1002
+          v35:BasicObject = Send v29, :+, v32 # SendFallbackReason: Uncategorized(opt_plus)
+          SetLocal :l3, l3, EP@5, v35
           CheckInterrupts
-          Return v30
+          Return v35
         "
         );
     }
@@ -3088,14 +3093,15 @@ pub mod hir_build_tests {
           IfTrue v10, bb4(v6)
           Jump bb5(v6)
         bb4(v11:BasicObject):
-          v17:BasicObject = GetLocal :block, l1, EP@3
-          Jump bb6(v11, v17)
+          v17:CPtr = GetEP 1
+          v18:BasicObject = LoadField v17, :block@0x1000
+          Jump bb6(v11, v18)
         bb5(v13:BasicObject):
-          v19:BasicObject = GetBlockParam :block, l1, EP@3
-          Jump bb6(v13, v19)
-        bb6(v21:BasicObject, v22:BasicObject):
+          v20:BasicObject = GetBlockParam :block, l1, EP@3
+          Jump bb6(v13, v20)
+        bb6(v22:BasicObject, v23:BasicObject):
           CheckInterrupts
-          Return v22
+          Return v23
         ");
     }