jit: correct the vstack mirror's executed-store precedence, and record the exact opcode segmentation

pyre/gate-triage.md

@@ -974,7 +974,7 @@ the folds it selects, not before them.
 `PYRE_FBW_SPEC_CENSUS` in §6c is its read-only half: the per-fold
 consulted/fired tallies.
 
-### §6c — Default-OFF diagnostics, censuses and probes (62): keep, cost nothing
+### §6c — Default-OFF diagnostics, censuses and probes (65): keep, cost nothing
 
 Each is inert unless set, so none is a removal target by this file's
 already-ON criterion. They are listed so they cannot be missed again.
@@ -1000,6 +1000,7 @@ already-ON criterion. They are listed so they cannot be missed again.
 `PYRE_PORTAL_RCA`, `PYRE_PROBE_BH_STARTUP`, `PYRE_PROBE_SNAPSHOT`,
 `PYRE_PROBE_SUBSCR`, `PYRE_PROFILE_PIPELINE`, `PYRE_QMUT_MAPDICT_FORCE`,
 `PYRE_RERAISE_DIAG`, `PYRE_SIZE_SHELL_OWNERS`, `PYRE_SNAPSHOT_DIAG`,
+`PYRE_VSTACK_EXACT_AUDIT`, `PYRE_VSTACK_KEEP_REORDER`, `PYRE_VSTACK_NO_EXACT`,
 `PYRE_WASM_DUMP_BAD_TRACE`, `PYRE_WASM_EXEC_TRACE`, `PYRE_WASM_FBW_CENSUS`,
 `PYRE_WASM_GUARD_CENSUS`, `PYRE_WASM_JIT_STATS`, `PYRE_WASM_CALL_HIST`,
 `PYRE_WASM_NO_CACHE`, `PYRE_WASM_STARTUP_TRACE`.
@@ -1010,6 +1011,17 @@ value knobs bound the capture window, sampling rate, and report size. This is a
 diagnostic tool rather than a temporary runtime experiment, so it retires only
 if the example itself is removed.
 
+`PYRE_VSTACK_NO_EXACT` and `PYRE_VSTACK_KEEP_REORDER` are A/B switches over the
+walk-level operand-stack mirror, each restoring the behaviour its default
+replaced: resolving the mirror's Python-PC coordinate from the floor tier rather
+than from the per-emission segmentation, and leaving an armed out-of-order
+region in place across a mirror re-seed. They exist so each switchover stays
+measurable on one binary — any env var of any name moves the allocation layout,
+so a knob and its control must live in the same build. Each goes when the
+behaviour it restores has no plausible reader left to compare against.
+`PYRE_VSTACK_EXACT_AUDIT` dumps that segmentation table at build time and, on
+the walk side, the coordinate it yields; it is a report with no ON behaviour.
+
 `PYRE_CHECK_INHERIT_ENV` is the other odd one: an A/B switch, not a report.
 `check.py` starts a benchmark child from an allowlisted environment because the
 inherited one is startup allocation and moves `guard_failures`; setting the

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

@@ -28,6 +28,17 @@ pub(crate) fn pcmap_containing_audit_enabled() -> bool {
     *ENABLED.get_or_init(|| std::env::var_os("PYRE_PCMAP_CONTAINING_AUDIT").is_some())
 }
 
+/// `PYRE_VSTACK_EXACT_AUDIT`: report where the per-emission exact segmentation
+/// (`py_exact_by_jit_pc`) disagrees with the floor tier the vstack mirror
+/// currently resolves its boundary from. The mirror INFERS whether a Python
+/// opcode retired; the exact table states it. Every disagreement is a case the
+/// inference must special-case, so this is the gate that has to be quiet
+/// before any consumer switches over. Off in production.
+pub(crate) fn vstack_exact_audit_enabled() -> bool {
+    static ENABLED: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
+    *ENABLED.get_or_init(|| std::env::var_os("PYRE_VSTACK_EXACT_AUDIT").is_some())
+}
+
 /// `PYRE_PCMAP_AFTERRESIDUAL_AUDIT`: assert the Slice-C after-residual depth
 /// twin (`depth_after_residual_for_jitcode_pc`) equals the raw
 /// `depth_at_py_pc[semantic_fallthrough_pc(containing_py_pc_for_jitcode_pc(jit_pc))]`

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

@@ -1142,6 +1142,7 @@ pub(crate) fn try_walker_call_assembler_self_recursive<Sym: WalkSym>(
         ctx.vstack_cur_pypc = resume_py_pc;
         ctx.vstack_depth = resume_depth;
         ctx.vstack_last_ref = OpRef::NONE;
+        super::vstack_mirror::disarm_vstack_reorder_region(ctx);
     }
 
     // pyjitpl.py: GUARD_NOT_FORCED + resume snapshot advanced past

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

@@ -828,9 +828,30 @@ pub(crate) fn walker_capture_snapshot_for_last_guard_impl<Sym: WalkSym>(
                             // (`resume.py`). The orthodox fix is
                             // push-time boxing plus optimizer virtualization,
                             // beyond this capture hook.
+                            // Decline wherever the shadow already carries a
+                            // live Ref for this slot.  Omitting a slot leaves
+                            // the shadow's own value standing in the snapshot,
+                            // and the shadow is the executed-store record —
+                            // every push and pop writes it, including a NULL
+                            // push (`emit_pushvalue_ref_const!`) and a pop's
+                            // clear (`emit_popvalue_ref!`).  The per-PC color
+                            // map is a program-point LABEL rather than a
+                            // binding, so it must fill a slot neither the
+                            // mirror nor the shadow can source, not outrank a
+                            // slot the walk demonstrably stored: the label is
+                            // fixed per PC while the stored box varies per
+                            // capture, so overriding published one invariant
+                            // box over several different live ones.
+                            let shadow_has_live_ref =
+                                ctx.trace_ctx.virtualizable_box_at(vidx).is_some_and(|b| {
+                                    b != OpRef::NONE
+                                        && !opref_is_null_const_ptr(b)
+                                        && b.ty() == Some(majit_ir::Type::Ref)
+                                });
                             if box_op != OpRef::NONE
                                 && !opref_is_null_const_ptr(box_op)
                                 && box_op.ty() == Some(majit_ir::Type::Ref)
+                                && !shadow_has_live_ref
                             {
                                 augmented.push((vidx, box_op));
                                 covered.insert(vidx);

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

@@ -870,22 +870,32 @@ pub(crate) fn setarrayitem_vable_via_metainterp<Sym: WalkSym>(
                         _ => false,
                     })
             });
+            let stack_slot = (index_value - nlocals) as usize;
             if method_load {
-                let stack_slot = (index_value - nlocals) as usize;
                 if ctx.vstack_boxes.len() <= stack_slot {
                     ctx.vstack_boxes.resize(stack_slot + 1, OpRef::NONE);
                 }
                 ctx.vstack_boxes[stack_slot] = value;
-                // Mark the slot as execution-derived while an out-of-order
-                // region is armed, so the boundary restore does not put its
-                // pc-derived snapshot back over it.  Outside a region the
-                // mask does not exist and this costs one `Option` test.
-                if let Some((_, _, _, mask)) = ctx.vstack_reorder_saved.as_mut() {
-                    if mask.len() <= stack_slot {
-                        mask.resize(stack_slot + 1, false);
-                    }
-                    mask[stack_slot] = true;
+            }
+            // Mark the slot as execution-derived while an out-of-order region
+            // is armed, so the boundary restore does not put its pc-derived
+            // snapshot back over it.  Outside a region the mask does not exist
+            // and this costs one `Option` test.
+            //
+            // This is a property of the STORE, not of the opcode that issued
+            // it: the mask records which slot the walk demonstrably wrote, and
+            // that is the same fact whatever the surrounding opcode is.  Only
+            // the positional mirror write above is method-form specific, and
+            // for the reason its own comment gives.  Scoping the mask to that
+            // opcode too left every other push shape unprotected, so a window
+            // whose only executed store came from, say, `LOAD_NAME` or
+            // `BINARY_OP` reported an empty mask and the restore reinstated the
+            // pre-window mirror wholesale over it.
+            if let Some((_, _, _, mask)) = ctx.vstack_reorder_saved.as_mut() {
+                if mask.len() <= stack_slot {
+                    mask.resize(stack_slot + 1, false);
                 }
+                mask[stack_slot] = true;
             }
             ctx.vstack_last_ref = value;
         }

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

@@ -958,10 +958,34 @@ pub(crate) fn vstack_step_py_pc(
     // block (for example its RETURN_VALUE); applying that opcode here corrupts
     // the live caller stack before the destination block is entered.
     if metadata_block_head_py_pc(metadata, jit_pc).is_some() {
-        current_py_pc
-    } else {
-        vstack_containing_py_pc(metadata, jit_pc)
+        return current_py_pc;
+    }
+    // Prefer the per-emission segmentation over the floor tier.  The floor
+    // table keys each Python PC to its FIRST jitcode offset, so a PC that
+    // emits in two disjoint regions is collapsed onto the earlier one and the
+    // later region reads as belonging to whichever PC last opened a segment.
+    // The mirror then sees a boundary the walk never crossed and replays that
+    // opcode's stack effect against a stack it does not describe — the same
+    // shape as applying a SWAP at a boundary that did not retire it.
+    //
+    // `py_exact_by_jit_pc` records the owning PC at every point the emitted
+    // stream changes owner, so it answers directly.  Empty for skeleton /
+    // fixture metadata, where the floor tier remains the only source.
+    if let Some(exact) =
+        crate::pyjitcode::exact_py_pc_for_jitcode_pc(&metadata.py_exact_by_jit_pc, jit_pc)
+        && !exact_segmentation_disabled()
+    {
+        return exact;
     }
+    vstack_containing_py_pc(metadata, jit_pc)
+}
+
+/// `PYRE_VSTACK_NO_EXACT`: fall back to the floor tier for the mirror's
+/// coordinate. The escape hatch for A/B-ing the segmentation switchover
+/// against the tier it replaces.
+fn exact_segmentation_disabled() -> bool {
+    static OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
+    *OFF.get_or_init(|| std::env::var_os("PYRE_VSTACK_NO_EXACT").is_some())
 }
 
 /// #73: step the walk-level operand-stack box mirror at
@@ -1047,6 +1071,30 @@ pub(crate) fn step_vstack_mirror<Sym: WalkSym>(ctx: &mut WalkContext<'_, '_, Sym
             } else {
                 raw_depth()
             };
+            // Stage gate for the per-emission exact segmentation.  The
+            // boundary the mirror acts on is currently INFERRED — `py_pc` is
+            // the floor segment's owner, and whether an opcode retired is then
+            // guessed from (prev, new, depth).  `py_exact_by_jit_pc` answers it
+            // directly.  Compare the two at every step before any consumer
+            // trusts the new table: a disagreement is a boundary the floor tier
+            // reports differently from the emission record, and each one is a
+            // case the inference has to special-case.
+            if vstack_exact_audit_enabled() {
+                if let Some(exact) = crate::pyjitcode::exact_py_pc_for_jitcode_pc(
+                    &jc.payload.metadata.py_exact_by_jit_pc,
+                    jit_pc,
+                ) && exact != py_pc
+                {
+                    let code = &*jc.payload.code_ptr;
+                    eprintln!(
+                        "[vstack-exact] code={} jit_pc={jit_pc} floor_py={py_pc} \
+                         exact_py={exact} cur_py={} boundary={}",
+                        code.obj_name,
+                        ctx.vstack_cur_pypc,
+                        py_pc != ctx.vstack_cur_pypc,
+                    );
+                }
+            }
             (py_pc, jc.payload.code_ptr, depth)
         }
     };
@@ -1243,6 +1291,31 @@ fn py_stack_depth(code_ptr: *const pyre_interpreter::CodeObject, py: u32) -> usi
         .unwrap_or(0) as usize
 }
 
+/// Drop any armed out-of-order region.
+///
+/// The region's snapshot is keyed on the `(py_pc, depth)` coordinate the walk
+/// left, and it is restored when the walk returns to that same coordinate on
+/// the premise that no Python opcode retired in between.  A caller that
+/// rewrites `vstack_boxes` wholesale and moves the coordinate — an exception
+/// handler entry, a call-assembler return — breaks that premise: the stack the
+/// snapshot describes no longer exists, so a later boundary that happens to
+/// land on the same `(py_pc, depth)` would restore it over the live mirror.
+pub(crate) fn disarm_vstack_reorder_region<Sym: WalkSym>(ctx: &mut WalkContext<'_, '_, Sym>) {
+    if keep_reorder_region_across_reseed() {
+        return;
+    }
+    ctx.vstack_reorder_ceiling = u32::MAX;
+    ctx.vstack_reorder_saved = None;
+}
+
+/// `PYRE_VSTACK_KEEP_REORDER`: leave an armed region in place across a mirror
+/// re-seed. The escape hatch for A/B-ing the disarm against the behaviour it
+/// replaces.
+fn keep_reorder_region_across_reseed() -> bool {
+    static KEEP: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
+    *KEEP.get_or_init(|| std::env::var_os("PYRE_VSTACK_KEEP_REORDER").is_some())
+}
+
 /// The handler-entry coordinate the mirror adopted, under `PYRE_VSTACK_DIAG`.
 /// `floor_py` is reported next to it: the two disagreeing is what an
 /// out-of-line catch target looks like from the walk's side, and without both
@@ -1343,6 +1416,7 @@ pub(crate) fn vstack_enter_exception_handler<Sym: WalkSym>(
     };
     ctx.vstack_boxes.clear();
     ctx.vstack_boxes.resize(handler_depth, OpRef::NONE);
+    disarm_vstack_reorder_region(ctx);
     // The unwinder pushes the caught exception onto the new TOS.
     if handler_depth >= 1 && exc != OpRef::NONE {
         ctx.vstack_boxes[handler_depth - 1] = exc;
@@ -1416,6 +1490,7 @@ fn vstack_enter_exception_handler_callee<Sym: WalkSym>(
     // mirror shallower than the handler depth pads with NONE holes, which
     // `mirror_covers_kept` declines per slot rather than latching invalid.
     ctx.vstack_boxes.resize(handler_depth, OpRef::NONE);
+    disarm_vstack_reorder_region(ctx);
     // The unwinder pushes the caught exception onto the new TOS.
     if handler_depth >= 1 && exc != OpRef::NONE {
         ctx.vstack_boxes[handler_depth - 1] = exc;

pyre/pyre-jit-trace/src/pyjitcode.rs

@@ -89,6 +89,22 @@ pub struct PyJitCodeMetadata {
     /// predecessor op-start tier: exact block-head marker precedence remains in
     /// `block_head_py_by_jit_pc`. Empty for skeleton / fixture metadata.
     pub py_floor_by_jit_pc: Vec<(u32, u32)>,
+    /// JitCode byte-offset → owning Python PC, one entry per contiguous
+    /// emission run, sorted ascending by offset.
+    ///
+    /// `py_floor_by_jit_pc` above answers the same question from the
+    /// FIRST-offset-per-PC table, so it collapses a PC that emits in two
+    /// disjoint regions and its segment extent is not one opcode's byte range.
+    /// This table is recorded at emission instead, so consecutive entries
+    /// bracket exactly one opcode's bytes even when a block is drained out of
+    /// source order, a mid-opcode split reopens the same PC, or a can-raise
+    /// trailing marker is re-keyed to the call's fallthrough.
+    ///
+    /// That makes "did the walk leave the opcode it was in" answerable rather
+    /// than inferable: the PC owning `jit_pc` differing from the mirror's
+    /// current coordinate IS the opcode boundary. Empty for skeleton / fixture
+    /// metadata, in which case consumers must fall back to the floor tier.
+    pub py_exact_by_jit_pc: Vec<(u32, u32)>,
     /// Exact `abort_permanent` jitcode offset → the Python PC whose lowering
     /// emitted it, sorted ascending by offset for binary search.
     ///
@@ -472,6 +488,14 @@ pub fn derive_resume_marker(
 /// Return the floor segment containing `jit_pc` in a codewriter-built JitCode
 /// PC pivot. An empty table is deliberately distinguishable from the `(0, 0)`
 /// fallback segment carried by every drained install.
+/// Exact owning Python PC for a JitCode byte offset, from the per-emission
+/// run table. `None` when the table is empty (skeleton / fixture metadata) so
+/// the caller can fall back to the floor tier rather than read a bogus 0.
+pub fn exact_py_pc_for_jitcode_pc(py_exact_by_jit_pc: &[(u32, u32)], jit_pc: usize) -> Option<u32> {
+    let end = py_exact_by_jit_pc.partition_point(|&(off, _)| (off as usize) <= jit_pc);
+    end.checked_sub(1).map(|idx| py_exact_by_jit_pc[idx].1)
+}
+
 pub fn floor_segment_for_jitcode_pc(
     py_floor_by_jit_pc: &[(u32, u32)],
     jit_pc: usize,
@@ -1063,6 +1087,7 @@ impl PyJitCodeMetadata {
             n_py_instrs: 0,
             block_head_py_by_jit_pc: Vec::new(),
             py_floor_by_jit_pc: Vec::new(),
+            py_exact_by_jit_pc: Vec::new(),
             abort_permanent_py_pc_by_jit_pc: Vec::new(),
             merge_entry_by_green: Vec::new(),
             pcdep_by_jit_pc: Vec::new(),

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

@@ -13655,6 +13655,7 @@ mod tests {
                 n_py_instrs: 0,
                 block_head_py_by_jit_pc: vec![(0, 0)],
                 py_floor_by_jit_pc: Vec::new(),
+                py_exact_by_jit_pc: Vec::new(),
                 abort_permanent_py_pc_by_jit_pc: Vec::new(),
                 merge_entry_by_green: Vec::new(),
                 pcdep_by_jit_pc: vec![(0, Vec::new())],