jit: journal the portal coordinate and snapshot the locals the f_locals fold writes

pyre/bench/synth/getframe_flocals_mirror_survives_residual.cranelift.jitstats

@@ -0,0 +1,15 @@
+bridges_compiled=0
+descr_set_absent=0
+descr_set_ambiguous=0
+descr_set_stale_absent=0
+fbw_blackhole_adopted_multi_frame=0
+fbw_blackhole_adopted_single_frame=0
+fbw_rolled_back_with_effects=0
+fbw_store_journal_rollback_failed=0
+field_pos_attached_misplaced=0
+field_pos_spec_misplaced=0
+guard_failures=1
+internal_compile_panics=0
+loops_aborted=0
+loops_compiled=1
+retraces_compiled=0

pyre/bench/synth/getframe_flocals_mirror_survives_residual.dynasm.jitstats

@@ -0,0 +1,15 @@
+bridges_compiled=0
+descr_set_absent=0
+descr_set_ambiguous=0
+descr_set_stale_absent=0
+fbw_blackhole_adopted_multi_frame=0
+fbw_blackhole_adopted_single_frame=0
+fbw_rolled_back_with_effects=0
+fbw_store_journal_rollback_failed=0
+field_pos_attached_misplaced=0
+field_pos_spec_misplaced=0
+guard_failures=1
+internal_compile_panics=0
+loops_aborted=0
+loops_compiled=1
+retraces_compiled=0

pyre/bench/synth/getframe_flocals_mirror_survives_residual.py

@@ -0,0 +1,34 @@
+# pyre-check: max-pypy-ratio=12
+# pypy's exec time here sits at the startup-subtraction floor, so the printed
+# ratio is not a measurement; the ceiling is fitted to the slowest of nine
+# local readings across the three backends plus headroom.
+#
+# Folding `frame.f_locals` mirrors the virtualizable shadow into the live
+# `locals_cells_stack_w` array, because pyre answers the attribute with a 3.14
+# `FrameLocalsProxy` that reads that array lazily rather than copying out of it
+# at the call. The mirror is an eager walk-time write, so it needs an undo for
+# the walk that does not commit — and the undo it is recorded in decides how
+# long it survives.
+#
+# Recording it in the residual force's escape-flush capture is wrong: the tail
+# of `try_execute_residual_call_via_executor` restores that capture after EVERY
+# non-forcing residual call, so the next call in the same walk reverts the
+# mirror and the proxy answers from before the fold. That is what this fixture
+# reads — `str(i)` between the fold and the subscript — and the failure is one
+# wrong answer on the trace-recording iteration alone, not every iteration, so
+# a short loop passes either way.
+import sys
+
+
+def f(n):
+    bad = 0
+    for i in range(n):
+        x = i * 2
+        loc = sys._getframe(0).f_locals
+        s = str(i)                     # a residual call between fold and read
+        if loc['x'] != x:              # must be this iteration's value
+            bad += 1
+    return bad
+
+
+print("mismatches:", f(3000))

pyre/bench/synth/getframe_flocals_mirror_survives_residual.wasm.jitstats

@@ -0,0 +1,15 @@
+bridges_compiled=0
+descr_set_absent=0
+descr_set_ambiguous=0
+descr_set_stale_absent=0
+fbw_blackhole_adopted_multi_frame=0
+fbw_blackhole_adopted_single_frame=0
+fbw_rolled_back_with_effects=0
+fbw_store_journal_rollback_failed=0
+field_pos_attached_misplaced=0
+field_pos_spec_misplaced=0
+guard_failures=1
+internal_compile_panics=0
+loops_aborted=0
+loops_compiled=1
+retraces_compiled=0

pyre/pyre-jit-trace/src/jitcode_dispatch/fbw_state.rs

@@ -166,11 +166,39 @@ thread_local! {
     /// recorded `setfield_vable_i` half needs no undo — it only reaches a frame
     /// on a compiled run.
     ///
-    /// Only the first publish of a walk is recorded, so the restore targets the
-    /// value the frame carried when the walk began rather than an intermediate
-    /// one.  Cleared with the store journal at the start of every walk.
-    static FBW_EXIT_LAST_INSTR_UNDO: std::cell::Cell<Option<(usize, isize)>> =
-        const { std::cell::Cell::new(None) };
+    /// Only the first publish of a walk is recorded PER FRAME, so the restore
+    /// targets the value each frame carried when the walk began rather than an
+    /// intermediate one.  More than one frame can be displaced in a walk — the
+    /// recording frame's per-opcode publication, and the portal frame's
+    /// coordinate when a positive-depth `_getframe` fold lands on it — so this
+    /// holds one entry each rather than one entry overall.  Cleared with the
+    /// store journal at the start of every walk.
+    static FBW_EXIT_LAST_INSTR_UNDO: std::cell::RefCell<Vec<(usize, isize)>> =
+        const { std::cell::RefCell::new(Vec::new()) };
+
+    /// The `locals_cells_stack_w` image a walk-time `f_locals` fold displaced,
+    /// per frame it mirrored into.
+    ///
+    /// The fold answers `frame.f_locals` by writing the virtualizable shadow
+    /// into the live array so the `FrameLocalsProxy` it hands out reads current
+    /// values (`virtualizable.py write_boxes`, which `pyframe.py fast2locals`
+    /// reaches through `force`).  That write lands whether or not the walk
+    /// commits, and a walk that does not commit replays the frame from its
+    /// pre-walk instruction — the replay would then re-derive these locals on
+    /// top of walk-time values instead of the ones the frame entered with.
+    ///
+    /// Kept separate from the residual force's `ESCAPE_FLUSH_UNDO` even though
+    /// both snapshot the same array: that capture is consumed by the tail
+    /// restore in `try_execute_residual_call_via_executor`, which runs after
+    /// EVERY non-forcing residual, so a mirror recorded there would be reverted
+    /// by the next call in the same walk.  This one is consumed only by the
+    /// walk-end legs, beside [`fbw_exit_last_instr_rollback`].
+    ///
+    /// First write per frame wins, for the same reason as the coordinate undo
+    /// above.  Rolled back AFTER `restore_escape_flush_undo`, so where both are
+    /// armed for one frame the older (pre-fold) image is the one that stands.
+    static FBW_LOCALS_MIRROR_UNDO: std::cell::RefCell<Vec<FbwLocalsMirrorUndo>> =
+        const { std::cell::RefCell::new(Vec::new()) };
 
     /// Armed by the bridge tracer (`call_jit::trace_and_compile_from_bridge`)
     /// before a single-frame, direct-return-capable guard-failure walk.  When
@@ -402,29 +430,147 @@ pub(crate) fn fbw_store_journal_reset() {
     // unrelated POP_EXCEPT in this walk.
     FBW_EXC_PREV.with(|s| s.borrow_mut().clear());
     FBW_EXC_PENDING_PUSH_SET.with(|c| c.set(false));
-    FBW_EXIT_LAST_INSTR_UNDO.with(|c| c.set(None));
+    FBW_EXIT_LAST_INSTR_UNDO.with(|c| c.borrow_mut().clear());
+    FBW_LOCALS_MIRROR_UNDO.with(|c| c.borrow_mut().clear());
+}
+
+/// The address a journalled frame lives at NOW.
+///
+/// A journal entry names its frame by raw address, and a JIT-created frame can
+/// be nursery-resident: a minor collection between the eager write and the
+/// walk-end restore drags it out and leaves a forwarding stub at the recorded
+/// address, so restoring through that address would write the abandoned copy
+/// and leave the live frame at its walk-time state.  `PyFrame::live_mut`
+/// carries the same reload for the interpreter's own field writes.
+///
+/// Applied at both ends: at record time so entries key on one identity, and at
+/// restore time so an entry recorded before the move still lands.  One
+/// application suffices — the drag-out promotes the frame out of the nursery,
+/// and nothing outside it moves (`gc_current_object_address`).
+fn live_frame_addr(frame: usize) -> usize {
+    pyre_object::gc_hook::try_gc_current_object_address(frame as *mut u8) as usize
+}
+
+/// Record the `last_instr` an eager walk-time write is about to displace, so
+/// [`fbw_exit_last_instr_rollback`] can put it back when the walk does not
+/// commit.  First write per frame wins: that is the coordinate the frame
+/// carried when the walk began.
+pub(crate) fn fbw_note_last_instr_undo(frame: usize) {
+    if frame == 0 {
+        return;
+    }
+    let frame = live_frame_addr(frame);
+    FBW_EXIT_LAST_INSTR_UNDO.with(|c| {
+        let mut undo = c.borrow_mut();
+        if undo.iter().any(|(f, _)| *f == frame) {
+            return;
+        }
+        // SAFETY: the caller holds a live `PyFrame` at this address, and
+        // `frame_layout` pins `last_instr` to this offset with a compile-time
+        // assertion against the interpreter's own constant.
+        let before =
+            unsafe { *((frame + crate::frame_layout::PYFRAME_LAST_INSTR_OFFSET) as *const isize) };
+        undo.push((frame, before));
+    });
 }
 
 /// Put `last_instr` back for a walk that did not commit its end state, so the
 /// replay resumes where the frame stood before the walk.  Runs beside
 /// [`fbw_store_journal_rollback`] on every non-committed exit; the commit side
 /// just drops the undo ([`fbw_exit_last_instr_commit`]).
 pub(crate) fn fbw_exit_last_instr_rollback() {
-    let Some((frame, before)) = FBW_EXIT_LAST_INSTR_UNDO.with(|c| c.take()) else {
-        return;
-    };
-    // SAFETY: the frame the publish wrote is the walk's live recording frame,
-    // which outlives the walk, and `frame_layout` pins `last_instr` to this
-    // offset with a compile-time assertion against the interpreter's constant.
-    unsafe {
-        *((frame + crate::frame_layout::PYFRAME_LAST_INSTR_OFFSET) as *mut isize) = before;
+    let undo = FBW_EXIT_LAST_INSTR_UNDO.with(|c| std::mem::take(&mut *c.borrow_mut()));
+    // Newest first, so a frame displaced twice — which the per-frame guard
+    // above already prevents — could never end on the later value.
+    for (frame, before) in undo.into_iter().rev() {
+        let frame = live_frame_addr(frame);
+        // SAFETY: every frame a publish wrote is live for the whole walk, and
+        // `frame_layout` pins `last_instr` to this offset with a compile-time
+        // assertion against the interpreter's constant.
+        unsafe {
+            *((frame + crate::frame_layout::PYFRAME_LAST_INSTR_OFFSET) as *mut isize) = before;
+        }
     }
 }
 
 /// Drop the undo: the walk's end state is kept, so the published exit
 /// coordinate is the one the frame should carry.
 pub(crate) fn fbw_exit_last_instr_commit() {
-    FBW_EXIT_LAST_INSTR_UNDO.with(|c| c.set(None));
+    FBW_EXIT_LAST_INSTR_UNDO.with(|c| c.borrow_mut().clear());
+}
+
+/// One frame's pre-fold `locals_cells_stack_w` image.  See
+/// [`FBW_LOCALS_MIRROR_UNDO`].
+pub(crate) struct FbwLocalsMirrorUndo {
+    frame: usize,
+    pub(crate) slots: Vec<pyre_object::PyObjectRef>,
+}
+
+/// Record the locals image a walk-time `f_locals` mirror is about to
+/// overwrite.  `nlocals` is the region the mirror writes — the operand stack
+/// above it is deliberately left out, so the rollback cannot revert a stack
+/// slot some other walk-time write owns.  First write per frame wins: that is
+/// the image the frame carried when the walk began.
+pub(crate) fn fbw_note_locals_mirror_undo(frame: usize, nlocals: usize) {
+    if frame == 0 {
+        return;
+    }
+    let frame = live_frame_addr(frame);
+    FBW_LOCALS_MIRROR_UNDO.with(|c| {
+        let mut undo = c.borrow_mut();
+        if undo.iter().any(|entry| entry.frame == frame) {
+            return;
+        }
+        // SAFETY: the caller holds a live `PyFrame` at this address for the
+        // rest of the walk — it is the frame whose `f_locals` it just folded.
+        let pf = unsafe { &*(frame as *const pyre_interpreter::PyFrame) };
+        let live = pyre_interpreter::locals_w!(pf).as_slice();
+        let slots = live[..nlocals.min(live.len())].to_vec();
+        undo.push(FbwLocalsMirrorUndo { frame, slots });
+    });
+}
+
+/// Non-commit epilogue: put each mirrored frame's pre-fold locals back, so the
+/// replay re-derives them from the frame's entry state instead of compounding
+/// onto the walk's.  Newest first, matching the store journal.
+pub(crate) fn fbw_locals_mirror_rollback() {
+    let undo = FBW_LOCALS_MIRROR_UNDO.with(|c| std::mem::take(&mut *c.borrow_mut()));
+    for entry in undo.into_iter().rev() {
+        let frame = live_frame_addr(entry.frame);
+        // SAFETY: as in the capture — the frame is live for the whole walk.
+        unsafe {
+            let pf = &mut *(frame as *mut pyre_interpreter::PyFrame);
+            let arr_ptr = pf.locals_cells_stack_w;
+            let dst = pyre_interpreter::locals_w_mut!(pf);
+            // The array cannot be reallocated mid-walk, but clamp anyway: a
+            // short write is recoverable, a write past the end is not.
+            let n = entry.slots.len().min(dst.as_slice().len());
+            for (i, &value) in entry.slots.iter().take(n).enumerate() {
+                dst[i] = value;
+            }
+            // A restored value can be nursery-young while the frame and its
+            // array are old-gen, and nothing re-traces it unless an owner is in
+            // the remembered set.  The forward write arms this per store
+            // because boxing allocates between them; the restore writes values
+            // that are already boxed, so nothing can collect mid-loop and one
+            // arming after it covers every slot.
+            crate::state::frame_array_write_barrier(frame as *mut u8, arr_ptr);
+        }
+    }
+}
+
+/// Drop the undo: the walk's end state is kept, so the mirrored locals are the
+/// ones the frame should carry (the resumed interpreter reads its fastlocals
+/// straight out of that array).
+pub(crate) fn fbw_locals_mirror_commit() {
+    FBW_LOCALS_MIRROR_UNDO.with(|c| c.borrow_mut().clear());
+}
+
+/// GC: a pre-fold image can be the ONLY reference to a value the mirror
+/// displaced, so the root area forwards them (see
+/// [`capture_fbw_store_journal_root_area`]).
+pub(crate) fn locals_mirror_undo_cell_ptr() -> *const std::cell::RefCell<Vec<FbwLocalsMirrorUndo>> {
+    FBW_LOCALS_MIRROR_UNDO.with(|c| c as *const _)
 }
 
 /// Record the element a walked eager list store displaces, for rollback
@@ -2188,11 +2334,7 @@ pub(crate) fn fbw_publish_exit_last_instr<Sym: WalkSym>(
         // compile-time assertion against the interpreter's own constant.
         let slot =
             (recording_frame_ptr + crate::frame_layout::PYFRAME_LAST_INSTR_OFFSET) as *mut isize;
-        FBW_EXIT_LAST_INSTR_UNDO.with(|c| {
-            if c.get().is_none() {
-                c.set(Some((recording_frame_ptr, unsafe { *slot })));
-            }
-        });
+        fbw_note_last_instr_undo(recording_frame_ptr);
         unsafe {
             *slot = py_pc as isize;
         }

pyre/pyre-jit-trace/src/jitcode_dispatch/mod.rs

@@ -6637,6 +6637,7 @@ struct FbwStoreJournalRootArea {
     abort_resume: *const std::cell::RefCell<Option<InlineAbortCarrier>>,
     active_session: *const std::cell::Cell<*const std::cell::RefCell<WalkSession>>,
     escape_flush_undo: *const std::cell::RefCell<Option<EscapeFlushUndo>>,
+    locals_mirror_undo: *const std::cell::RefCell<Vec<FbwLocalsMirrorUndo>>,
     single_frame_blackhole: *const std::cell::RefCell<Option<LatchedSingleFrameBlackhole>>,
     multi_frame_blackhole: *const std::cell::RefCell<Option<LatchedMultiFrameBlackhole>>,
 }
@@ -6655,6 +6656,7 @@ thread_local! {
         abort_resume: FBW_ABORT_CALL_RESUME.with(|value| value as *const _),
         active_session: ACTIVE_WALK_SESSION.with(|value| value as *const _),
         escape_flush_undo: escape_flush_undo_cell_ptr(),
+        locals_mirror_undo: locals_mirror_undo_cell_ptr(),
         single_frame_blackhole: single_frame_blackhole_cell_ptr(),
         multi_frame_blackhole: multi_frame_blackhole_cell_ptr(),
     };
@@ -7034,6 +7036,16 @@ pub unsafe fn fbw_store_journal_root_walker_area(
             visitor(unsafe { &mut *(slot as *mut pyre_object::PyObjectRef).cast() });
         }
     }
+    // The `f_locals` fold's pre-mirror images are armed the same way and live
+    // just as long — until the walk's commit or rollback leg consumes them —
+    // and the live-frame slots they shadow now hold the mirrored values, so
+    // each displaced entry needs forwarding too.
+    let locals_mirror = unsafe { &mut *(*area.locals_mirror_undo).as_ptr() };
+    for entry in locals_mirror.iter_mut() {
+        for slot in entry.slots.iter_mut() {
+            visitor(unsafe { &mut *(slot as *mut pyre_object::PyObjectRef).cast() });
+        }
+    }
     // The vable-force and trace-too-long MIFrame images survive the dispatch
     // unwind in TLS. Their Option<i64> Ref banks are not otherwise visible to
     // the collector, so forward every populated color until the walk-end

pyre/pyre-jit-trace/src/jitcode_dispatch/residual_call.rs

@@ -1585,15 +1585,28 @@ pub(crate) fn restore_escape_flush_undo() {
         let Some(undo) = slot.borrow_mut().take() else {
             return;
         };
+        // A JIT-created frame can be nursery-resident, and a minor collection
+        // between the flush and this restore drags it out, leaving a forwarding
+        // stub at the captured address.  Restoring through the stale address
+        // would write the abandoned copy and leave the live frame flushed.
+        let frame =
+            pyre_object::gc_hook::try_gc_current_object_address(undo.frame as *mut u8) as usize;
         unsafe {
-            let pf = &mut *(undo.frame as *mut pyre_interpreter::PyFrame);
+            let pf = &mut *(frame as *mut pyre_interpreter::PyFrame);
+            let arr_ptr = pf.locals_cells_stack_w;
             let dst = locals_w_mut!(pf);
             let n = undo.slots.len().min(dst.as_slice().len());
             for (i, &v) in undo.slots.iter().take(n).enumerate() {
                 dst[i] = v;
             }
             pf.last_instr = undo.last_instr;
             pf.valuestackdepth = undo.valuestackdepth;
+            // The forward flush arms this per store; the restore writes values
+            // that are already boxed, so nothing allocates between them and one
+            // arming covers the whole image.  Without it a pre-flush value that
+            // is nursery-young goes back into an old-gen array with no
+            // remembered-set entry to re-trace it.
+            crate::state::frame_array_write_barrier(frame as *mut u8, arr_ptr);
         }
     });
 }

pyre/pyre-jit-trace/src/jitcode_dispatch/specialize.rs

@@ -2595,7 +2595,18 @@ fn walker_write_back_standard_frame_locals<Sym: WalkSym>(
             Some((value, _)) => slots.push((slot as i64, value)),
         }
     }
+    // The mirror below writes the live frame's locals array, and a walk that
+    // does not commit replays from its pre-walk instruction — so the pre-walk
+    // values have to be recoverable.  Journal them against the walk's own
+    // non-commit epilogue rather than the escape-flush capture: that capture is
+    // consumed by every non-forcing residual (`try_execute_residual_call_via_
+    // executor`'s tail restore), which would revert this mirror mid-walk and
+    // leave a live `FrameLocalsProxy` reading pre-fold values.
+    crate::jitcode_dispatch::fbw_note_locals_mirror_undo(concrete_frame, nlocals);
     if !crate::state::flush_locals_region_to_frame(ctx.trace_ctx, concrete_frame) {
+        // All-or-nothing decline: nothing was written.  The journal entry is
+        // harmless — restoring the values still in place is a no-op — and the
+        // first-per-frame rule means dropping it could discard a real one.
         return false;
     }
     ctx.trace_ctx
@@ -9602,6 +9613,10 @@ pub(crate) fn try_walker_specialize_sys_getframe<Sym: WalkSym>(
             .trace_ctx
             .virtualizable_entry_at(crate::virtualizable_spec::LAST_INSTR_VABLE_FIELD_INDEX)
     {
+        // Journaled like the per-opcode publication: this store lands whether
+        // or not the walk commits, and a walk that does not commit replays the
+        // frame from its pre-walk coordinate.
+        crate::jitcode_dispatch::fbw_note_last_instr_undo(cur_ptr as usize);
         unsafe { (*cur_ptr).last_instr = last_instr as isize };
     }