jit: PR#457 parity-review follow-ups — built_as_portal fold gate; callee-locals shadow TLS→WalkContext

[youknowone]

Follow-ups to the parity review on #457 (review comment: #457 (comment)). Both flagged orthodox gaps are fixed.

1. Gate strict fresh-frame fold on callee jitcode built_as_portal metadata (10a3e209aab)

Review §2: a mutable env var (PYRE_ALWAYS_PORTAL) was read at walk time to decide how a jitcode is interpreted. A JitCode's shape must come from the JitCode artifact itself.

• The fold gate now reads the callee jitcode's own built_as_portal metadata (stamped at build time in codewriter from frame_is_portal) via a new state::built_as_portal_at() accessor.
• fbw_always_portal_enabled() (the walk-time env read, single call site) is deleted. The build-time env read in codewriter stays — that is where the shape is legitimately decided and recorded.

2. Move the FBW inline-callee locals shadow from TLS onto WalkContext (0f32240831e)

Review §1: per-inlined-callee locals lived in TLS Vec-stacks (FBW_CALLEE_LOCALS_CONCRETE/_OPREF, FBW_STRICT_FOLD_FRAME_REG). RPython keeps per-callee state on MIFrame.registers_i/r/f (pyjitpl.py:177-234); pyre's MIFrame analog is WalkContext, which already owns registers_r/i/f.

• All three TLS stacks were accessed innermost-only, and innermost ≡ the current sub-walk's WalkContext. They are replaced by a callee_shadow: Option<CalleeLocalsShadow> field on WalkContext (slot-keyed OpRef/concrete maps + fold frame reg).
• Some only at the two inlined-callee sub-walk constructions (try_walker_inline_user_call, drive_bridge_carrier_subwalk); None everywhere else, preserving the previous empty-stack no-op semantics for top-level walks.
• The TLS statics, their helper fns, and the InlineRecursionGuard push/pop plumbing are deleted — nesting is now carried by the Rust call stack itself. This also structurally removes a latent hazard where a guard-less sub-walk could observe the caller's innermost maps through TLS.
• FBW_INLINE_CODE_STACK stays (genuine cross-level recursion-depth scan).

Verification

• python3 pyre/check.py: 161/161 on dynasm, cranelift, wasm (run for each commit).
• cargo test -p pyre-jit-trace --release --no-default-features --features dynasm: 265 passed.
• calls_closures.py output unchanged.

Out of scope

The review's deeper point — callee locals absent from resume data inside an inlined callee (currently handled by CALL-boundary replay) — belongs to the ongoing exact-resume epics (#467/#126/#215-P2) and is not touched here. FBW_INLINE_CALLEE_CONSTS follows the same TLS pattern and is a candidate for the same migration as a follow-up.

🤖 Generated with Claude Code

Summary by CodeRabbit

• New Features

  • Improved strict inline folding for more accurate handling of callee-local values and frame state.
  • Added portal-status detection for JIT-compiled code paths.

• Bug Fixes

  • Improved local variable reads and writes during inlined execution.
  • Ensured portal-based execution decisions consistently reflect compiled code metadata.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: ASSERTIVE

Plan: Pro

Run ID: 1a451a59-4728-4798-b4cc-1acc4453ffc6

📥 Commits

Reviewing files that changed from the base of the PR and between 4f5ea69 and 0f32240.

📒 Files selected for processing (2)

• pyre/pyre-jit-trace/src/jitcode_dispatch.rs
• pyre/pyre-jit-trace/src/state.rs

---

Walkthrough

Per-walk callee-local shadowing replaces thread-local strict-fold state during inline traversal. WalkContext now carries shadow state, callee reconstruction and local folding update it, and portal-frame selection reads initialized JIT metadata.

Changes

Strict inline folding

Layer / File(s)	Summary
Portal metadata lookup pyre/pyre-jit-trace/src/jitcode_dispatch.rs, pyre/pyre-jit-trace/src/state.rs	Adds built_as_portal_at and uses it to determine callee portal-frame registers.
Callee shadow context lifecycle pyre/pyre-jit-trace/src/jitcode_dispatch.rs	Adds CalleeLocalsShadow, extends WalkContext, initializes shadows for inline sub-walks, and updates test contexts.
Strict folding state propagation pyre/pyre-jit-trace/src/jitcode_dispatch.rs	Seeds reconstructed callee values and parameters into per-walk shadow state, and uses it for frame registers, local reads, and local writes instead of thread-local state.

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

Sequence Diagram(s)

sequenceDiagram
  participant CallerWalk
  participant InlineWalk
  participant CalleeLocalsShadow
  participant JitMetadata
  CallerWalk->>JitMetadata: read built_as_portal_at(jitcode_index)
  CallerWalk->>InlineWalk: create WalkContext with callee_shadow
  InlineWalk->>CalleeLocalsShadow: seed reconstructed slots and parameters
  InlineWalk->>CalleeLocalsShadow: read and update local OpRef/concrete state
  CalleeLocalsShadow-->>InlineWalk: provide strict fold frame register

Loading

Possibly related PRs

• youknowone/pyre#440: Adjusts strict multiframe inline dispatch and CALL-site liveness alongside this PR’s callee shadowing changes.

Poem

A bunny hops through frames so bright,
Shadows track each slot just right.
Portal flags now guide the way,
Strict folds bloom through code all day.
No thread-local trails to chase—
Per-walk carrots win the race!

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly summarizes the two main changes: portal-based fold gating and moving callee-locals shadow state into WalkContext.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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.

✨ 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 nbody

---

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

• X
• Mastodon
• Reddit
• LinkedIn

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

[github-actions]

🤖 Codex parity review

Static analysis of this diff vs the local RPython/PyPy sources (commit 2350a54).

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-jit-trace/src/jitcode_dispatch.rs:357 ↔ rpython/jit/metainterp/pyjitpl.py:68-90 — “HashMap<i64, OpRef> / HashMap<i64, Value> shadow” remains a parallel per-local side table, whereas PyPy stores each frame’s execution values in that frame’s registers_i/r/f. The patch correctly moves ownership from thread-local stacks to the sub-walk, but it retains the non-parity slot-map representation.

• pyre/pyre-jit-trace/src/jitcode_dispatch.rs:4418 ↔ rpython/jit/metainterp/pyjitpl.py:1218-1230 — “fall back to … shadow.concrete after the metainterp vable read” has no PyPy equivalent: PyPy returns the live virtualizable_boxes[index] value. The fallback can preserve an independently cached concrete after heap-cache loss rather than deriving it from the frame’s live box state.

4. Structural adaptations

• pyre/pyre-jit-trace/src/jitcode_dispatch.rs:4324 ↔ rpython/jit/codewriter/jtransform.py:990-993 — “strict fresh-frame fold … emit NO GC op” is a CPython-compiler/JitCode-shape adaptation. RPython removes fresh-virtualizable accesses during codewriter translation; Pyre receives CPython-compatible frame-vable opcodes and suppresses them in the walker instead.

• pyre/pyre-jit-trace/src/jitcode_dispatch.rs:13227 ↔ rpython/jit/metainterp/pyjitpl.py:74-90 — “consult built_as_portal metadata before using portal_frame_reg” is a Rust-generated-JitCode adaptation. PyPy’s JitCode directly owns its register layout; Pyre must query generated metadata. This patch improves parity over the prior process-environment gate because it uses the actual callee JitCode shape.

• pyre/pyre-jit-trace/src/jitcode_dispatch.rs:2878 ↔ rpython/jit/metainterp/pyjitpl.py:2456-2475 — “callee_shadow: Some(Default::default()) per sub-walk” is an implementation-language adaptation toward PyPy’s one-MIFrame-per-call model. It removes the prior thread-local ownership collapse, though the remaining HashMap representation is the pre-existing mismatch above.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 0f32240831

ℹ️ 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".

[chatgpt-codex-connector]

Preserve callee shadow across inline_call sub-walks

When an FBW-inlined Python callee reaches an inline_call_* helper, this sub-WalkContext now drops the active callee locals shadow. Before this change, run_sub_jitcode_walk did not push a new TLS shadow frame, so helper sub-walks inside the inlined callee still saw the current frame's fold/concrete maps; with callee_shadow: None, any vable reads/writes in those helper jitcodes lose the fresh-frame fold and concrete fallback, leading to VableBoxNotSeeded aborts or stale local/stack shadows for callees whose body contains helper inline_call_* plus getarrayitem_vable/setarrayitem_vable.

Useful? React with 👍 / 👎.