jit: drop the WITH_EXCEPT_START gate + with-frame exception-correctness fixes (#389)

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

@@ -417,7 +417,6 @@ pub struct InlineFrame {
 /// Per-trace-attempt walk session, owned by the walk driver and threaded
 /// through [`WalkContext`] — `MetaInterp.framestack` (`pyjitpl.py:2475`,
 /// `:2487`; depth scan `:1390`). Innermost level last.
-#[derive(Default)]
 pub struct WalkSession {
     /// Inlined callee levels. Parent snapshots are outermost-first, matching
     /// `Snapshot.frames`; a caller is pushed at its inline CALL and popped
@@ -428,6 +427,30 @@ pub struct WalkSession {
     /// the outer snapshot root's py_pc→jitcode translation, so abort-point
     /// flushing must decline after the sub-walk unwinds.
     pub abort_in_subwalk: bool,
+    /// Blackhole `tmpreg_r`/`tmpreg_i`/`tmpreg_f` (`blackhole.py:661-679`):
+    /// the single-slot scratch that `insert_renamings` (`flatten.py:154`)
+    /// routes a cyclic parallel move through via `*_push`/`*_pop` pairs.
+    /// `ref_push/r` writes the source Ref (+ its concrete shadow) here;
+    /// `ref_pop/>r` reads it back into a dst register.  Persisted on the
+    /// per-walk session because a push and its matching pop straddle
+    /// intervening `*_copy` ops within one trampoline.
+    pub tmpreg_r: OpRef,
+    pub tmpreg_r_concrete: ConcreteValue,
+    pub tmpreg_i: OpRef,
+    pub tmpreg_f: OpRef,
+}
+
+impl Default for WalkSession {
+    fn default() -> Self {
+        Self {
+            framestack: Vec::new(),
+            abort_in_subwalk: false,
+            tmpreg_r: OpRef::NONE,
+            tmpreg_r_concrete: ConcreteValue::Null,
+            tmpreg_i: OpRef::NONE,
+            tmpreg_f: OpRef::NONE,
+        }
+    }
 }
 
 /// Compile-time-constant frame fields of an inlined callee.
@@ -8419,6 +8442,7 @@ struct FbwStoreJournalRootArea {
     appends: *const std::cell::RefCell<Vec<(pyre_object::PyObjectRef, usize)>>,
     abort_overrides: *const std::cell::RefCell<Vec<(usize, pyre_object::PyObjectRef)>>,
     cell_stores: *const std::cell::RefCell<Vec<(pyre_object::PyObjectRef, i64)>>,
+    sys_exc: *const std::cell::RefCell<Vec<pyre_object::PyObjectRef>>,
     foriter: *const std::cell::RefCell<Vec<InflightForiter>>,
     abort_resume: *const std::cell::RefCell<Option<InlineAbortCarrier>>,
     active_session: *const std::cell::Cell<*const std::cell::RefCell<WalkSession>>,
@@ -8430,6 +8454,7 @@ thread_local! {
         appends: FBW_APPEND_JOURNAL.with(|value| value as *const _),
         abort_overrides: FBW_ABORT_OUTER_STACK_OVERRIDES.with(|value| value as *const _),
         cell_stores: FBW_CELL_STORE_JOURNAL.with(|value| value as *const _),
+        sys_exc: FBW_SYS_EXC_JOURNAL.with(|value| value as *const _),
         foriter: FBW_FORITER_INFLIGHT.with(|value| value as *const _),
         abort_resume: FBW_ABORT_CALL_RESUME.with(|value| value as *const _),
         active_session: ACTIVE_WALK_SESSION.with(|value| value as *const _),
@@ -9237,13 +9262,12 @@ pub unsafe fn fbw_store_journal_root_walker_area(
     // nursery-resident and no longer referenced elsewhere once the eager store
     // overwrote the EC slot; forward each so a minor collection during the rest
     // of the walk cannot free/move the value the rollback restores.
-    FBW_SYS_EXC_JOURNAL.with(|j| {
-        for displaced in j.borrow_mut().iter_mut() {
-            // SAFETY: `PyObjectRef` and `GcRef` share the usize repr; the
-            // borrow keeps the Vec storage alive for the visit.
-            visitor(unsafe { &mut *(displaced as *mut pyre_object::PyObjectRef).cast() });
-        }
-    });
+    let sys_exc = unsafe { &mut *(*area.sys_exc).as_ptr() };
+    for displaced in sys_exc.iter_mut() {
+        // SAFETY: `PyObjectRef` and `GcRef` share the usize repr; the
+        // borrowed area keeps the Vec storage alive for the visit.
+        visitor(unsafe { &mut *(displaced as *mut pyre_object::PyObjectRef).cast() });
+    }
     // #57 Option C: each captured in-flight FOR_ITER item is nursery-resident
     // across the rest of the walk (subsequent residual calls allocate and a
     // minor collection moves nursery objects), so forward every entry's item
@@ -17392,13 +17416,28 @@ fn try_walker_specialize_load_attr(
     // guard_value(getfield_gc_i(obj, map), C_map): `jit.promote(self.map)`
     // (`mapdict.py:905-906`).  The map nodes are interned + immortal, so the
     // pointer is a stable identity guarded as an opaque word (object_map_descr
-    // is Int-typed).  Skipped when the field read already const-folds to the map
-    // (heapcache hit).
+    // is Int-typed).
+    //
+    // The guard may only be elided when the map read is ALREADY a compile-time
+    // constant — i.e. a prior promotion in this trace pinned it via
+    // `replace_box`.  It must NOT be elided merely because `box_value(map_op)`
+    // reports the concrete map: every traced getfield op carries its live value
+    // (`opimpl_getfield_gc_i` → `set_opref_concrete`, `history.py:803`
+    // FrontendOp), so `box_value == map` holds for the very first read and
+    // would drop the guard on the trace's entry.  A trace whose map guard is
+    // dropped reads `storage[storageindex]` off any same-class receiver whose
+    // map differs (GuardClass alone does not pin the layout), returning a wild
+    // slot value.  Pin the map with `replace_box` after guarding so a later
+    // fold on the same receiver correctly elides (matching the trait
+    // `implement_guard_value`, `trace_opcode.rs:4631-4633`).
     let map_op =
         crate::state::opimpl_getfield_gc_i(ctx.trace_ctx, obj, crate::descr::object_map_descr());
-    if ctx.trace_ctx.box_value(map_op) != Some(majit_ir::Value::Int(map as i64)) {
+    if !map_op.is_constant() {
         let map_const = ctx.trace_ctx.const_int(map as i64);
         walker_emit_fold_guard_with_snapshot(ctx, op_pc, OpCode::GuardValue, &[map_op, map_const])?;
+        ctx.trace_ctx
+            .heap_cache_mut()
+            .replace_box(map_op, map_const);
     }
 
     // getfield_gc_r(obj, storage) + getarrayitem_gc_r(block, C_storageindex):
@@ -23252,6 +23291,88 @@ fn handle(
             *slot = src_val;
             Ok((DispatchOutcome::Continue, op.next_pc))
         }
+        "ref_push/r" => {
+            // Blackhole `bhimpl_ref_push` (`blackhole.py:664-666`):
+            // `self.tmpreg_r = a`.  `insert_renamings` (`flatten.py:154`)
+            // emits `*_push`/`*_pop` around a cyclic parallel move — the
+            // swap `r_a <-> r_b` lowers to `push r_b; copy r_b<-r_a;
+            // pop r_a` so the overwritten value survives in the tmpreg.
+            // Pure SSA-level scratch move, no IR op recorded.  Operand
+            // layout `r`: 1B src.
+            let src_val = read_ref_reg(code, op, 0, ctx)?;
+            let src_concrete = read_ref_reg_concrete(code, op, 0, ctx);
+            {
+                let mut sess = ctx.session.borrow_mut();
+                sess.tmpreg_r = src_val;
+                sess.tmpreg_r_concrete = src_concrete;
+            }
+            Ok((DispatchOutcome::Continue, op.next_pc))
+        }
+        "ref_pop/>r" => {
+            // Blackhole `bhimpl_ref_pop` (`blackhole.py:674-676`):
+            // `return self.get_tmpreg_r()`.  Reads the value stashed by
+            // the matching `ref_push/r` back into a dst register, in
+            // lock-step with its concrete shadow.  Operand layout `>r`:
+            // 1B dst.
+            let (val, concrete) = {
+                let sess = ctx.session.borrow();
+                (sess.tmpreg_r, sess.tmpreg_r_concrete)
+            };
+            let dst = code[op.pc + 1] as usize;
+            write_ref_reg(ctx, op.pc, dst, val, concrete)?;
+            Ok((DispatchOutcome::Continue, op.next_pc))
+        }
+        "int_push/i" => {
+            // Blackhole `bhimpl_int_push` (`blackhole.py:661-663`):
+            // `self.tmpreg_i = a`.  Int-bank sibling of `ref_push/r`.
+            // Operand layout `i`: 1B src.
+            let src_val = read_int_reg(code, op, 0, ctx)?;
+            ctx.session.borrow_mut().tmpreg_i = src_val;
+            Ok((DispatchOutcome::Continue, op.next_pc))
+        }
+        "int_pop/>i" => {
+            // Blackhole `bhimpl_int_pop` (`blackhole.py:671-673`).
+            // Operand layout `>i`: 1B dst.
+            let val = ctx.session.borrow().tmpreg_i;
+            let dst = code[op.pc + 1] as usize;
+            let len = ctx.registers_i.len();
+            let slot = ctx
+                .registers_i
+                .get_mut(dst)
+                .ok_or(DispatchError::RegisterOutOfRange {
+                    pc: op.pc,
+                    reg: dst,
+                    len,
+                    bank: "i",
+                })?;
+            *slot = val;
+            Ok((DispatchOutcome::Continue, op.next_pc))
+        }
+        "float_push/f" => {
+            // Blackhole `bhimpl_float_push` (`blackhole.py:667-669`).
+            // Operand layout `f`: 1B src.
+            let src_val = read_float_reg(code, op, 0, ctx)?;
+            ctx.session.borrow_mut().tmpreg_f = src_val;
+            Ok((DispatchOutcome::Continue, op.next_pc))
+        }
+        "float_pop/>f" => {
+            // Blackhole `bhimpl_float_pop` (`blackhole.py:677-679`).
+            // Operand layout `>f`: 1B dst.
+            let val = ctx.session.borrow().tmpreg_f;
+            let dst = code[op.pc + 1] as usize;
+            let len = ctx.registers_f.len();
+            let slot = ctx
+                .registers_f
+                .get_mut(dst)
+                .ok_or(DispatchError::RegisterOutOfRange {
+                    pc: op.pc,
+                    reg: dst,
+                    len,
+                    bank: "f",
+                })?;
+            *slot = val;
+            Ok((DispatchOutcome::Continue, op.next_pc))
+        }
         "ref_copy/r>r" => {
             // Ref-bank sibling of `int_copy/i>i`. Same RPython
             // `_opimpl_any_copy` body — the `>r` suffix only changes

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

@@ -1621,7 +1621,7 @@ fn run_perfn_walk(
                                 reserved_red_colors.first().copied() == Some(color);
                             let bridge_names_operand =
                                 !is_frame_color && opref != ec_box && opref != frame_box;
-                            if pyre_object::tagged_int::CAN_BE_TAGGED && bridge_names_operand {
+                            if bridge_names_operand {
                                 seed(color, opref);
                             }
                             continue;
@@ -1822,6 +1822,7 @@ fn run_perfn_walk(
                 ) => Some((*pc, false)),
                 _ => None,
             };
+            let mut entry_carrier_call_py_pc = None;
             if let Some((abort_jit_pc, is_marker_abort)) = call_forward_abort {
                 // gh#467: a supported abort fired inside a TOP-level inline
                 // sub-walk whose callee executed no concrete effect
@@ -1843,6 +1844,7 @@ fn run_perfn_walk(
                         call_py_pc,
                         call_stack,
                     }) => {
+                        entry_carrier_call_py_pc = Some(*call_py_pc);
                         if crate::state::flush_walk_end_state_at_outer_call(
                             ctx,
                             cf_addr,
@@ -1961,35 +1963,46 @@ fn run_perfn_walk(
                 } else if let Some(resume_py_pc) =
                     crate::jitcode_dispatch::fbw_abort_outer_resume_take()
                 {
-                    // Flush while the overrides stay rooted in
-                    // FBW_ABORT_OUTER_STACK_OVERRIDES (the flush boxes Int/Float
-                    // locals — an allocation that can move the nursery-resident
-                    // override refs; the area walker forwards them in place),
-                    // then clear the cell.
-                    let committed = crate::jitcode_dispatch::fbw_abort_outer_stack_overrides_with(
-                        |stack_overrides| {
-                            crate::state::flush_walk_end_state_to_frame_with_stack_overrides(
-                                ctx,
-                                cf_addr,
-                                resume_py_pc,
-                                stack_overrides,
-                            )
-                        },
-                    );
-                    crate::jitcode_dispatch::fbw_abort_outer_stack_overrides_clear();
-                    if committed {
+                    if entry_carrier_call_py_pc == Some(resume_py_pc) {
+                        crate::jitcode_dispatch::fbw_abort_outer_stack_overrides_clear();
                         if crate::jitcode_dispatch::fbw_debug_abort_enabled() {
                             eprintln!(
-                                "[fbw-abort-flush] COMMIT abort_jit_pc={abort_jit_pc} \
-                                 resume_py_pc={resume_py_pc} (nested inline decline)"
+                                "[fbw-abort-flush] skipped at resume_py_pc={resume_py_pc} \
+                                 (entry carrier already handled same resume)"
+                            );
+                        }
+                    } else {
+                        // Flush while the overrides stay rooted in
+                        // FBW_ABORT_OUTER_STACK_OVERRIDES (the flush boxes Int/Float
+                        // locals — an allocation that can move the nursery-resident
+                        // override refs; the area walker forwards them in place),
+                        // then clear the cell.
+                        let committed =
+                            crate::jitcode_dispatch::fbw_abort_outer_stack_overrides_with(
+                                |stack_overrides| {
+                                    crate::state::flush_walk_end_state_to_frame_with_stack_overrides(
+                                        ctx,
+                                        cf_addr,
+                                        resume_py_pc,
+                                        stack_overrides,
+                                    )
+                                },
+                            );
+                        crate::jitcode_dispatch::fbw_abort_outer_stack_overrides_clear();
+                        if committed {
+                            if crate::jitcode_dispatch::fbw_debug_abort_enabled() {
+                                eprintln!(
+                                    "[fbw-abort-flush] COMMIT abort_jit_pc={abort_jit_pc} \
+                                     resume_py_pc={resume_py_pc} (nested inline decline)"
+                                );
+                            }
+                            WALK_END_FLUSH_COMMITTED.with(|c| c.set(true));
+                        } else if crate::jitcode_dispatch::fbw_debug_abort_enabled() {
+                            eprintln!(
+                                "[fbw-abort-flush] declined at resume_py_pc={resume_py_pc} \
+                                 (shadow slot without concrete / depth / lastblock) — legacy replay kept"
                             );
                         }
-                        WALK_END_FLUSH_COMMITTED.with(|c| c.set(true));
-                    } else if crate::jitcode_dispatch::fbw_debug_abort_enabled() {
-                        eprintln!(
-                            "[fbw-abort-flush] declined at resume_py_pc={resume_py_pc} \
-                             (shadow slot without concrete / depth / lastblock) — legacy replay kept"
-                        );
                     }
                 } else if crate::jitcode_dispatch::fbw_debug_abort_enabled() {
                     eprintln!(

pyre/pyre-jit/src/eval.rs

@@ -4108,12 +4108,8 @@ fn unsupported_jit_shape(code: &pyre_interpreter::CodeObject) -> UnsupportedJitS
     let mut arg_state = pyre_interpreter::OpArgState::default();
     let mut has_for_iter = false;
     for unit in code.instructions.iter().copied() {
-        match arg_state.get(unit).0 {
-            pyre_interpreter::Instruction::WithExceptStart => {
-                return UnsupportedJitShape::StructuralRegion;
-            }
-            pyre_interpreter::Instruction::ForIter { .. } => has_for_iter = true,
-            _ => {}
+        if let pyre_interpreter::Instruction::ForIter { .. } = arg_state.get(unit).0 {
+            has_for_iter = true;
         }
     }
     if has_for_iter {
@@ -8946,24 +8942,20 @@ mod tests {
     }
 
     #[test]
-    fn with_block_frame_is_structural_region() {
-        // A `with` block compiles to `WITH_EXCEPT_START` for its exception link.
-        // The codewriter still residualizes that lowering, so tracing the frame
-        // (or its nested callees) miscompiles — a raw SIGSEGV in the exception
-        // path plus guard-failure storms (test_strftime/test_shlex/test_textwrap,
-        // #389). The frame must classify as `StructuralRegion` so it and its
-        // callees stay interpreted; dropping this gate regressed the CPython
-        // suite (commit 4daa5c517e, reverted).
+    fn with_block_frame_is_admitted() {
+        // A `with` block compiles to `WITH_EXCEPT_START` for its exception link,
+        // lowered as a residual with the exception disposition preserved across
+        // the guard-failure bridge. The frame is admitted for tracing.
+        // The body is kept free of `FOR_ITER` so the only classification axis
+        // is the `WITH_EXCEPT_START` shape; a `for` loop whose body is not
+        // allow-listed declines independently via `for_iter_bodies_all_jit_safe`.
         use pyre_interpreter::compile_exec;
         let module = compile_exec(
-            "def wf(cm):\n    total = 0\n    for _ in range(3):\n        with cm:\n            total += 1\n    return total\n",
+            "def wf(cm):\n    total = 0\n    with cm:\n        total += 1\n    return total\n",
         )
         .expect("test code should compile");
         let code = function_code_from_module(&module, "wf");
-        assert_eq!(
-            unsupported_jit_shape(&code),
-            UnsupportedJitShape::StructuralRegion
-        );
+        assert_eq!(unsupported_jit_shape(&code), UnsupportedJitShape::None);
     }
 
     fn ensure_test_jit_callbacks() {

pyre/pyre-jit/src/jit/codewriter.rs

@@ -9092,13 +9092,6 @@ impl CodeWriter {
                             // on top. Preserve the net `+1` stack effect in the
                             // shadow graph and fall back to the interpreter for
                             // the actual helper call semantics.
-                            //
-                            // Portable (flowspace records a `direct_call` /
-                            // `indirect_call`) but latent: a `with` block's
-                            // exception table prevents the enclosing loop/callee
-                            // from ever reaching a JIT token, so this abort is
-                            // never reached in practice — no residual yet.
-                            emit_abort_permanent!(py_pc);
                             push_fresh_ref(&mut current_state, &mut graph);
                             current_depth += 1;
                             emit_vsd!(current_depth, py_pc);