Skip to content

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

Merged
youknowone merged 1 commit into
mainfrom
fix-foriter-review
Aug 23, 2026
Merged

jit: journal the portal coordinate and snapshot the locals the f_locals fold writes#1434
youknowone merged 1 commit into
mainfrom
fix-foriter-review

Conversation

@youknowone

@youknowone youknowone commented Aug 22, 2026

Copy link
Copy Markdown
Owner

Follow-up to #1401. Two eager walk-time writes had no undo, so a walk that does
not commit — which replays the frame from its pre-walk state — replayed on top
of them.

1. The portal frame's last_instr

try_walker_specialize_sys_getframe publishes the portal frame's coordinate
when a positive-depth _getframe fold lands on it. The per-opcode publisher
next to it already journals, but through a single Cell, so whichever write
came first in a walk was the only one recorded — and a walk can displace both
frames.

FBW_EXIT_LAST_INSTR_UNDO becomes one entry per frame, fbw_note_last_instr_undo
is the shared recorder, and the rollback walks entries newest first.

2. The locals the f_locals fold mirrors

walker_write_back_standard_frame_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 the array lazily instead of copying out of it
at the call.

The first version of this fix recorded that image in the residual force's
ESCAPE_FLUSH_UNDO, and that was wrong.
The tail of
try_execute_residual_call_via_executor calls restore_escape_flush_undo
after every non-forcing residual, so the next call in the same walk reverted
the mirror and the proxy answered from before the fold. Measured with an
env-gated A/B on one binary:

x = i * 2
loc = sys._getframe(0).f_locals
s = str(i)                 # residual call between the fold and the read
loc['x']                   # -> 2078 at i == 1040, the previous iteration's value

one wrong answer on the trace-recording iteration, identical on dynasm and
cranelift; check.py reports it as wrong output.

So the locals get their own journal, FBW_LOCALS_MIRROR_UNDO:

  • consumed only by the walk-end legs, beside the coordinate undo;
  • rolled back after restore_escape_flush_undo, so where a force flushed the
    same frame the older pre-fold image is the one that stands;
  • snapshots only the nlocals region the mirror writes, leaving the operand
    stack to whoever owns it;
  • its entries are forwarded by the store-journal root area — a displaced value
    can be the only reference left once the live slot holds the mirrored one.

mark_escape_flush_undo_pending is still deliberately not armed anywhere here:
that flag means "this walk withdrew its commit", which a fold cannot know.

Regression fixture

synth/getframe_flocals_mirror_survives_residual is the snippet above, and it
discriminates: mismatches: 1 / FAIL … wrong output with the escape-flush
capture, mismatches: 0 with the journal. Baselines recorded for all three
backends (loops_compiled=1, loops_aborted=0).

GC write barrier (Codex parity review)

Both restore paths now re-arm frame_array_write_barrier after writing the
saved image back. A restored value can be nursery-young while the frame and its
array are old-gen, and nothing re-traces it without a remembered-set entry. The
forward writes arm it per store because boxing allocates between them; a restore
writes values that are already boxed, so nothing can collect mid-loop and one
arming after it covers the image.

restore_escape_flush_undo was missing the same arming before this change —
Codex flagged it as pre-existing. It is fixed here too: arming a barrier is
monotonic, so the only cost is an entry the collector would otherwise have
skipped, and the mechanism is the one this PR is about.

Frame relocation (Codex P1)

Both journals name their 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
the restore would write the abandoned copy and leave the live frame at its
walk-time state. PyFrame::live_mut documents the same hazard for the
interpreter's own field writes.

All three restores now resolve the address through gc_current_object_address,
at record time so entries key on one identity and at restore time so an entry
recorded before the move still lands. One application is enough: the drag-out
promotes the frame out of the nursery, and nothing outside the nursery moves.

Resolving beats rooting the address here — the query answers addr unchanged
for anything the GC does not own, so it needs no ownership gate, and it does not
turn a journal entry into a strong root for a frame that is already reachable.

FBW_EXIT_LAST_INSTR_UNDO and ESCAPE_FLUSH_UNDO both carried this on main;
they are fixed here alongside the new journal.

Not fixed

The FrameLocalsProxy write-through vs. virtualizable-shadow coherence gap is
real but pre-existing on main — gating the fold to the inline receiver alone
still reproduces it — so it stays out of this PR.

FLAG_ESCAPED is still not rolled back, on risk asymmetry: a stale flag costs
an unnecessary force, while clearing a flag another unjournaled site set costs a
skipped force and a dangling vref.

Verification

All three backends built from this head:

pattern dynasm cranelift wasm
getframe_* 29/29 29/29 28/28
*inlined_callee* 14/14 14/14 14/14
*locals* 8/8 8/8 7/7
gc_* 10/10 10/10 7/7
traceback_* 1/1 1/1 1/1

Every pre-existing jit-stat unmoved. cargo test --all --no-default-features --features dynasm,cpyext clean; cargo fmt --all --check clean;
check-new-line-citations.py --base origin/main reports no new line-number
citations across 168 added Rust lines.

@coderabbitai

coderabbitai Bot commented Aug 22, 2026

Copy link
Copy Markdown

Review Change Stack

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Organization UI

Review profile: ASSERTIVE

Plan: Pro Plus

Run ID: f0b1b67f-d5e5-4c51-8e06-a310ea9dbe3a

📥 Commits

Reviewing files that changed from the base of the PR and between 9be6297 and 1ecfab2.

📒 Files selected for processing (9)
  • pyre/bench/synth/getframe_flocals_mirror_survives_residual.cranelift.jitstats
  • pyre/bench/synth/getframe_flocals_mirror_survives_residual.dynasm.jitstats
  • pyre/bench/synth/getframe_flocals_mirror_survives_residual.py
  • pyre/bench/synth/getframe_flocals_mirror_survives_residual.wasm.jitstats
  • pyre/pyre-jit-trace/src/jitcode_dispatch/fbw_state.rs
  • pyre/pyre-jit-trace/src/jitcode_dispatch/mod.rs
  • pyre/pyre-jit-trace/src/jitcode_dispatch/residual_call.rs
  • pyre/pyre-jit-trace/src/jitcode_dispatch/specialize.rs
  • pyre/pyre-jit-trace/src/trace.rs

Included review availability: Your plan provides up to 2 included reviews per hour; 0 remain after this review.


Walkthrough

The JIT now uses per-frame journals to preserve last_instr and mirrored locals across multi-frame walks. Walk completion commits or rolls back these values. GC traversal and residual restoration now handle displaced references. A PyPy benchmark validates the behavior.

Changes

Frame locals mirror lifecycle

Layer / File(s) Summary
Per-frame undo journals
pyre/pyre-jit-trace/src/jitcode_dispatch/fbw_state.rs, pyre/pyre-jit-trace/src/jitcode_dispatch/mod.rs
FBW state now stores per-frame instruction and mirrored-local undo entries. Rollback, commit, reset, and GC-root traversal process these entries.
Walk journaling and completion
pyre/pyre-jit-trace/src/jitcode_dispatch/specialize.rs, pyre/pyre-jit-trace/src/trace.rs
Specialization records frame changes before applying them. Walk completion commits mirrored locals for committed walks and restores them for non-committed walks.
Residual restoration and benchmark validation
pyre/pyre-jit-trace/src/jitcode_dispatch/residual_call.rs, pyre/bench/synth/getframe_flocals_mirror_survives_residual.*
Residual restoration applies a frame-array write barrier. The benchmark checks f_locals consistency and records JIT statistics for Cranelift, DynASM, and WASM.

Estimated code review effort: 4 (Complex) | ~45 minutes

Merge Risk: ⚪ Minimal · up to 1ecfa

The change preserves frame coordinates and mirrored locals across aborted walks and residual calls, with targeted coverage across all supported backends; no actionable merge-blocking risk remains after normal checks and review.

Poem

A rabbit journals each frame with care,
Mirrors local values, then checks them there.
Commit keeps the changes in line,
Rollback restores each state in time.
GC roots hop through the guarded trail.
f_locals stays true without fail.

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Docstring Coverage ✅ Passed Docstring check was indeterminate for this PR — some files could not be analyzed in time. Not blocking.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title clearly summarizes the main changes: journaling the portal coordinate and snapshotting locals written by the f_locals fold.
✨ Finishing Touches
📝 Generate docstrings
  • Create stacked PR
  • Commit on current branch
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch fix-foriter-review

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

https://github.com/youknowone/pyre/blob/c0689c199e3a2f88732e10cee660e8ec5446a546/pyre-jit-trace/src/jitcode_dispatch/specialize.rs#L2622
P1 Badge Separate fold snapshots from residual-force undo

When a standard-frame f_locals fold is followed by any non-escaping residual, this capture remains in the force-specific ESCAPE_FLUSH_UNDO slot, and execute_residual_call unconditionally calls restore_escape_flush_undo() after that residual. Consequently, code such as two successive proxy["x"] reads sees the current shadow value on the first call, then restores the pre-fold locals array and sees an old value or KeyError on the second, corrupting the recording-time result and potentially the compiled trace. The fold snapshot needs a lifecycle distinct from the temporary force capture so ordinary residual completion cannot consume it.

AGENTS.md reference: AGENTS.md:L26-L33


https://github.com/youknowone/pyre/blob/c0689c199e3a2f88732e10cee660e8ec5446a546/pyre-jit-trace/src/jitcode_dispatch/residual_call.rs#L1607-L1608
P2 Badge Clear the locals capture when a walk commits

When the folded f_locals operation is near the end of a successfully committed walk and no later residual consumes this capture, the normal commit epilogue never calls discard_escape_flush_undo; it is only cleared when another trace starts. Because the permanently registered extra-root walker visits every saved slot, objects that were present in the frame before the fold remain GC-reachable for the lifetime of the thread if no later trace occurs, potentially retaining an arbitrarily large locals graph. The committed walk path should explicitly discard this snapshot.

AGENTS.md reference: AGENTS.md:L98-L104

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

@github-actions

github-actions Bot commented Aug 23, 2026

Copy link
Copy Markdown

🤖 Codex parity review

Static analysis of this diff vs the local RPython/PyPy sources (commit dc52193).
Updated: 2026-08-23T07:10:00.176Z

Files in the reviewed diff
pyre/bench/synth/getframe_flocals_mirror_survives_residual.py
pyre/pyre-jit-trace/src/jitcode_dispatch/fbw_state.rs
pyre/pyre-jit-trace/src/jitcode_dispatch/mod.rs
pyre/pyre-jit-trace/src/jitcode_dispatch/residual_call.rs
pyre/pyre-jit-trace/src/jitcode_dispatch/specialize.rs
pyre/pyre-jit-trace/src/trace.rs

1. Regressions to PyPy parity introduced by this patch

None.

2. Other mismatches introduced by this patch

None.

3. Pre-existing mismatches (already present before this patch)

  • pyre/pyre-interpreter/src/pyframe.rs:70 ↔ pypy/interpreter/pyframe.py:525 — Pyre exposes optimized-frame locals through a CPython-style FrameLocalsProxy; PyPy calls fast2locals() and returns its cached dictionary. CPython 3.14 requires write-through locals behavior (lib-python/3/test/test_frame.py:325-330), but this cannot be filed as a CPython-spec structural adaptation: PyPy’s deciding fast2locals implementation is governed by @jit.unroll_safe at pypy/interpreter/pyframe.py:539. This mismatch predates the authoritative changed-file list; the patch only preserves the proxy’s already-selected semantics during an aborted walk.

4. Structural adaptations

  • pyre/pyre-jit-trace/src/jitcode_dispatch/fbw_state.rs:514 ↔ rpython/jit/metainterp/virtualizable.py:101 — Pyre snapshots/restores the eagerly mirrored per-frame locals image, roots it across collection, and applies a remembered-set barrier on restoration. RPython’s write_boxes directly writes the virtualizable during tracing. The undo/rooting machinery is required by Pyre’s Rust full-body walker, which speculatively executes and can abandon a walk; it is a fundamental implementation-architecture adaptation, not a PyPy semantic divergence.

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 364dffd33e

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

/// One frame's pre-fold `locals_cells_stack_w` image. See
/// [`FBW_LOCALS_MIRROR_UNDO`].
pub(crate) struct FbwLocalsMirrorUndo {
frame: usize,

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Root the frame held by the locals undo journal

When f_locals is folded for a nursery-allocated standard virtualizable, such as a bridge or recursive JIT frame, a subsequent proxy or residual allocation can relocate that frame before the non-commit epilogue. The extra-root walker forwards entry.slots but never this raw usize, so fbw_locals_mirror_rollback later dereferences the abandoned nursery address instead of the live red frame, failing to restore its locals and potentially corrupting reused nursery memory. Store and visit this frame identity as a forwardable GC root rather than a raw address.

AGENTS.md reference: AGENTS.md:L26-L33

Useful? React with 👍 / 👎.

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 1ecfab2e38

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

// 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);

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Forward the frame stored in the coordinate undo

When this positive-depth _getframe fold lands on a nursery-allocated portal frame and a later residual allocation triggers a minor collection before a non-committing walk exit, CURRENT_FRAME is forwarded but the raw usize saved by fbw_note_last_instr_undo is not. fbw_exit_last_instr_rollback then restores last_instr through the abandoned address, leaving the live frame at the walk-time coordinate and potentially writing into reused nursery memory; store this identity in a forwardable root slot and visit it alongside the other journal roots.

AGENTS.md reference: AGENTS.md:L26-L33

Useful? React with 👍 / 👎.

…irrors

`try_walker_specialize_sys_getframe` wrote the portal frame's `last_instr`
without recording the value it displaced, and
`walker_write_back_standard_frame_locals` mirrored the virtualizable shadow
into the live locals array without recording the values it overwrote. Both
land whether or not the walk commits, and a walk that does not commit replays
the frame from its pre-walk state.

`FBW_EXIT_LAST_INSTR_UNDO` becomes one entry per frame rather than one entry
overall: a walk can displace both the recording frame's coordinate (the
per-opcode publication) and the portal frame's, and a single cell kept only
whichever came first. `fbw_note_last_instr_undo` is the shared recorder; the
rollback walks the entries newest first.

The locals get their own journal, `FBW_LOCALS_MIRROR_UNDO`, rather than the
residual force's `ESCAPE_FLUSH_UNDO`. That capture is consumed by the tail
`restore_escape_flush_undo` in `try_execute_residual_call_via_executor`, which
runs after every non-forcing residual, so a mirror recorded there is reverted
by the next call in the same walk: `f_locals['x']` then answers with the
previous iteration's value on the trace-recording iteration, measured on both
dynasm and cranelift. The new journal is consumed only by the walk-end legs,
beside the coordinate undo, and runs after `restore_escape_flush_undo` so the
older pre-fold image stands where both are armed for one frame. It snapshots
only the `nlocals` region the mirror writes, leaving the operand stack alone,
and the store-journal root area forwards its entries.

`synth/getframe_flocals_mirror_survives_residual` reads the proxy after a
`str(i)` call in the same body. It reports `mismatches: 1` when the mirror is
recorded in the escape-flush capture and `mismatches: 0` with the journal;
check.py calls the former `wrong output`.

Both restore paths re-arm the frame/array write barrier after writing the
saved image back. The forward writes arm it per store because boxing allocates
between them; a restore writes values that are already boxed, so one arming
after the loop covers the image. `restore_escape_flush_undo` was missing the
same arming before this change.

Both journals name their frame by raw address, and a JIT-created frame can be
nursery-resident, so a minor collection between the eager write and the restore
drags it out and leaves a forwarding stub at the recorded address.  All three
restores now resolve the address through `gc_current_object_address` — the same
reload `PyFrame::live_mut` performs for the interpreter's own field writes —
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.

Assisted-by: Claude
@youknowone
youknowone merged commit da122cc into main Aug 23, 2026
27 of 29 checks passed
@youknowone
youknowone deleted the fix-foriter-review branch August 23, 2026 13:24
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant