Use find() in frame_state()

We want to return representatives for the stack and locals.

2 files changed, 20 insertions, 10 deletions

diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs
index 7fbe4f55dd..d431889f63 100644
--- a/zjit/src/codegen.rs
+++ b/zjit/src/codegen.rs
@@ -188,11 +188,11 @@ fn gen_insn(jit: &mut JITState, asm: &mut Assembler, function: &Function, insn_i
         Insn::IfTrue { val, target } => return gen_if_true(jit, asm, opnd!(val), target),
         Insn::IfFalse { val, target } => return gen_if_false(jit, asm, opnd!(val), target),
         Insn::SendWithoutBlock { call_info, cd, state, .. } | Insn::SendWithoutBlockDirect { call_info, cd, state, .. }
-            => gen_send_without_block(jit, asm, call_info, *cd, function.frame_state(*state))?,
+            => gen_send_without_block(jit, asm, call_info, *cd, &function.frame_state(*state))?,
         Insn::Return { val } => return Some(gen_return(asm, opnd!(val))?),
-        Insn::FixnumAdd { left, right, state } => gen_fixnum_add(asm, opnd!(left), opnd!(right), function.frame_state(*state))?,
-        Insn::FixnumSub { left, right, state } => gen_fixnum_sub(asm, opnd!(left), opnd!(right), function.frame_state(*state))?,
-        Insn::FixnumMult { left, right, state } => gen_fixnum_mult(asm, opnd!(left), opnd!(right), function.frame_state(*state))?,
+        Insn::FixnumAdd { left, right, state } => gen_fixnum_add(asm, opnd!(left), opnd!(right), &function.frame_state(*state))?,
+        Insn::FixnumSub { left, right, state } => gen_fixnum_sub(asm, opnd!(left), opnd!(right), &function.frame_state(*state))?,
+        Insn::FixnumMult { left, right, state } => gen_fixnum_mult(asm, opnd!(left), opnd!(right), &function.frame_state(*state))?,
         Insn::FixnumEq { left, right } => gen_fixnum_eq(asm, opnd!(left), opnd!(right))?,
         Insn::FixnumNeq { left, right } => gen_fixnum_neq(asm, opnd!(left), opnd!(right))?,
         Insn::FixnumLt { left, right } => gen_fixnum_lt(asm, opnd!(left), opnd!(right))?,
@@ -200,8 +200,8 @@ fn gen_insn(jit: &mut JITState, asm: &mut Assembler, function: &Function, insn_i
         Insn::FixnumGt { left, right } => gen_fixnum_gt(asm, opnd!(left), opnd!(right))?,
         Insn::FixnumGe { left, right } => gen_fixnum_ge(asm, opnd!(left), opnd!(right))?,
         Insn::Test { val } => gen_test(asm, opnd!(val))?,
-        Insn::GuardType { val, guard_type, state } => gen_guard_type(asm, opnd!(val), *guard_type, function.frame_state(*state))?,
-        Insn::GuardBitEquals { val, expected, state } => gen_guard_bit_equals(asm, opnd!(val), *expected, function.frame_state(*state))?,
+        Insn::GuardType { val, guard_type, state } => gen_guard_type(asm, opnd!(val), *guard_type, &function.frame_state(*state))?,
+        Insn::GuardBitEquals { val, expected, state } => gen_guard_bit_equals(asm, opnd!(val), *expected, &function.frame_state(*state))?,
         Insn::PatchPoint(_) => return Some(()), // For now, rb_zjit_bop_redefined() panics. TODO: leave a patch point and fix rb_zjit_bop_redefined()
         _ => {
             debug!("ZJIT: gen_function: unexpected insn {:?}", insn);
diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
index a592586ee7..40cc88b52d 100644
--- a/zjit/src/hir.rs
+++ b/zjit/src/hir.rs
@@ -663,9 +663,9 @@ impl Function {
         self.insns.len()
     }
 
-    /// Return a reference to the FrameState at the given instruction index.
-    pub fn frame_state(&self, insn_id: InsnId) -> &FrameState {
-        match &self.insns[insn_id.0] {
+    /// Return a FrameState at the given instruction index.
+    pub fn frame_state(&self, insn_id: InsnId) -> FrameState {
+        match self.find(insn_id) {
             Insn::Snapshot { state } => state,
             insn => panic!("Unexpected non-Snapshot {insn} when looking up FrameState"),
         }
@@ -712,8 +712,18 @@ impl Function {
         let insn_id = self.union_find.find_const(insn_id);
         use Insn::*;
         match &self.insns[insn_id.0] {
-            result@(PutSelf | Const {..} | Param {..} | NewArray {..} | GetConstantPath {..} | Snapshot {..}
+            result@(PutSelf | Const {..} | Param {..} | NewArray {..} | GetConstantPath {..}
                     | Jump(_) | PatchPoint {..}) => result.clone(),
+            Snapshot { state: FrameState { iseq, insn_idx, pc, stack, locals } } =>
+                Snapshot {
+                    state: FrameState {
+                        iseq: *iseq,
+                        insn_idx: *insn_idx,
+                        pc: *pc,
+                        stack: stack.iter().map(|v| find!(*v)).collect(),
+                        locals: locals.iter().map(|v| find!(*v)).collect(),
+                    }
+                },
             Return { val } => Return { val: find!(*val) },
             StringCopy { val } => StringCopy { val: find!(*val) },
             StringIntern { val } => StringIntern { val: find!(*val) },