Add InstructionMetadata::stack_effect_jump for branch stack effects (#7585)

* Add InstructionMetadata::stack_effect_jump for branch stack effects

CPython's compile.c provides stack_effect(opcode, oparg, jump) where the
jump parameter selects between fallthrough and branch effects. The existing
stack_effect() only returns the fallthrough effect.

Add stack_effect_jump() that returns the branch effect. Most instructions
have identical fallthrough/branch effects; ForIter and Send are the
exceptions (ForIter: fallthrough=+1, branch=-1; Send: fallthrough=0,
branch=-1).

* apply review

crates/codegen/src/ir.rs

@@ -2207,40 +2207,18 @@ impl CodeInfo {
                 }
                 // Process target blocks for branching instructions
                 if ins.target != BlockIdx::NULL {
-                    if instr.is_block_push() {
-                        // SETUP_* pseudo ops: target is a handler block.
-                        // Handler entry depth uses the jump-path stack effect:
-                        //   SETUP_FINALLY:  +1  (pushes exc)
-                        //   SETUP_CLEANUP:  +2  (pushes lasti + exc)
-                        //   SETUP_WITH:     +1  (pops __enter__ result, pushes lasti + exc)
-                        let handler_effect: u32 = match instr.pseudo() {
-                            Some(PseudoInstruction::SetupCleanup { .. }) => 2,
-                            _ => 1, // SetupFinally and SetupWith
-                        };
-                        let handler_depth = depth + handler_effect;
-                        if handler_depth > maxdepth {
-                            maxdepth = handler_depth;
+                    let jump_effect = instr.stack_effect_jump(ins.arg.into());
+                    let target_depth = depth.checked_add_signed(jump_effect).ok_or({
+                        if jump_effect < 0 {
+                            InternalError::StackUnderflow
+                        } else {
+                            InternalError::StackOverflow
                         }
-                        stackdepth_push(&mut stack, &mut start_depths, ins.target, handler_depth);
-                    } else {
-                        // SEND jumps to END_SEND with receiver still on stack.
-                        // END_SEND performs the receiver pop.
-                        let jump_effect = match instr.real() {
-                            Some(Instruction::Send { .. }) => 0i32,
-                            _ => effect,
-                        };
-                        let target_depth = depth.checked_add_signed(jump_effect).ok_or({
-                            if jump_effect < 0 {
-                                InternalError::StackUnderflow
-                            } else {
-                                InternalError::StackOverflow
-                            }
-                        })?;
-                        if target_depth > maxdepth {
-                            maxdepth = target_depth
-                        }
-                        stackdepth_push(&mut stack, &mut start_depths, ins.target, target_depth);
+                    })?;
+                    if target_depth > maxdepth {
+                        maxdepth = target_depth;
                     }
+                    stackdepth_push(&mut stack, &mut start_depths, ins.target, target_depth);
                 }
                 depth = new_depth;
                 if instr.is_scope_exit() || instr.is_unconditional_jump() {

crates/compiler-core/src/bytecode/instruction.rs

@@ -1095,6 +1095,10 @@ impl InstructionMetadata for Instruction {
         StackEffect::new(pushed as u32, popped as u32)
     }
 
+    // In CPython 3.14 the metadata-based stack_effect is the same for both
+    // fallthrough and branch paths for all real instructions.
+    // Only pseudo-instructions (SETUP_*) differ — see PseudoInstruction.
+
     #[allow(clippy::too_many_arguments)]
     fn fmt_dis(
         &self,
@@ -1502,6 +1506,21 @@ impl InstructionMetadata for PseudoInstruction {
         StackEffect::new(pushed as u32, popped as u32)
     }
 
+    /// Handler entry effect for SETUP_* pseudo ops.
+    ///
+    /// Fallthrough effect is 0 (NOPs), but when the branch is taken the
+    /// handler block starts with extra values on the stack:
+    ///   SETUP_FINALLY:  +1  (exc)
+    ///   SETUP_CLEANUP:  +2  (lasti + exc)
+    ///   SETUP_WITH:     +1  (pops __enter__ result, pushes lasti + exc)
+    fn stack_effect_jump(&self, _oparg: u32) -> i32 {
+        match self {
+            Self::SetupFinally { .. } | Self::SetupWith { .. } => 1,
+            Self::SetupCleanup { .. } => 2,
+            _ => self.stack_effect(_oparg),
+        }
+    }
+
     fn is_unconditional_jump(&self) -> bool {
         matches!(self, Self::Jump { .. } | Self::JumpNoInterrupt { .. })
     }
@@ -1576,6 +1595,8 @@ impl InstructionMetadata for AnyInstruction {
 
     inst_either!(fn stack_effect(&self, oparg: u32) -> i32);
 
+    inst_either!(fn stack_effect_jump(&self, oparg: u32) -> i32);
+
     inst_either!(fn stack_effect_info(&self, oparg: u32) -> StackEffect);
 
     inst_either!(fn fmt_dis(
@@ -1692,6 +1713,16 @@ pub trait InstructionMetadata {
         self.stack_effect_info(oparg).effect()
     }
 
+    /// Stack effect when the instruction takes its branch (jump=true).
+    ///
+    /// CPython equivalent: `stack_effect(opcode, oparg, jump=True)`.
+    /// For most instructions this equals the fallthrough effect.
+    /// Override for instructions where branch and fallthrough differ
+    /// (e.g. `FOR_ITER`: fallthrough = +1, branch = −1).
+    fn stack_effect_jump(&self, oparg: u32) -> i32 {
+        self.stack_effect(oparg)
+    }
+
     #[allow(clippy::too_many_arguments)]
     fn fmt_dis(
         &self,

crates/stdlib/src/_opcode.rs

@@ -186,14 +186,18 @@ mod _opcode {
             })
             .unwrap_or(Ok(0))?;
 
-        let jump = args
-            .jump
-            .map(|v| {
-                v.try_to_bool(vm).map_err(|_| {
-                    vm.new_value_error("stack_effect: jump must be False, True or None")
-                })
-            })
-            .unwrap_or(Ok(false))?;
+        let jump: Option<bool> = match args.jump {
+            Some(v) => {
+                if vm.is_none(&v) {
+                    None
+                } else {
+                    Some(v.try_to_bool(vm).map_err(|_| {
+                        vm.new_value_error("stack_effect: jump must be False, True or None")
+                    })?)
+                }
+            }
+            None => None,
+        };
 
         let opcode = Opcode::try_from_pyint(args.opcode, vm)?;
 
@@ -202,8 +206,15 @@ mod _opcode {
             return Err(vm.new_value_error("invalid opcode or oparg"));
         }
 
-        let _ = jump; // Python API accepts jump but it's not used
-        Ok(opcode.stack_effect(oparg))
+        let effect = match jump {
+            Some(true) => opcode.stack_effect_jump(oparg),
+            Some(false) => opcode.stack_effect(oparg),
+            // jump=None: max of both paths (CPython convention)
+            None => opcode
+                .stack_effect(oparg)
+                .max(opcode.stack_effect_jump(oparg)),
+        };
+        Ok(effect)
     }
 
     #[pyfunction]