Skip to content

Commit 7e5e026

Browse files
youknowoneyouknowone
authored
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
1 parent 27aed85 commit 7e5e026

File tree

3 files changed

+62
-42
lines changed

3 files changed

+62
-42
lines changed

‎crates/codegen/src/ir.rs‎

Lines changed: 10 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -2207,40 +2207,18 @@ impl CodeInfo {
22072207
}
22082208
// Process target blocks for branching instructions
22092209
if ins.target != BlockIdx::NULL {
2210-
if instr.is_block_push() {
2211-
// SETUP_* pseudo ops: target is a handler block.
2212-
// Handler entry depth uses the jump-path stack effect:
2213-
// SETUP_FINALLY: +1 (pushes exc)
2214-
// SETUP_CLEANUP: +2 (pushes lasti + exc)
2215-
// SETUP_WITH: +1 (pops __enter__ result, pushes lasti + exc)
2216-
let handler_effect: u32 = match instr.pseudo() {
2217-
Some(PseudoInstruction::SetupCleanup { .. }) => 2,
2218-
_ => 1, // SetupFinally and SetupWith
2219-
};
2220-
let handler_depth = depth + handler_effect;
2221-
if handler_depth > maxdepth {
2222-
maxdepth = handler_depth;
2223-
}
2224-
stackdepth_push(&mut stack, &mut start_depths, ins.target, handler_depth);
2225-
} else {
2226-
// SEND jumps to END_SEND with receiver still on stack.
2227-
// END_SEND performs the receiver pop.
2228-
let jump_effect = match instr.real() {
2229-
Some(Instruction::Send { .. }) => 0i32,
2230-
_ => effect,
2231-
};
2232-
let target_depth = depth.checked_add_signed(jump_effect).ok_or({
2233-
if jump_effect < 0 {
2234-
InternalError::StackUnderflow
2235-
} else {
2236-
InternalError::StackOverflow
2237-
}
2238-
})?;
2239-
if target_depth > maxdepth {
2240-
maxdepth = target_depth
2210+
let jump_effect = instr.stack_effect_jump(ins.arg.into());
2211+
let target_depth = depth.checked_add_signed(jump_effect).ok_or({
2212+
if jump_effect < 0 {
2213+
InternalError::StackUnderflow
2214+
} else {
2215+
InternalError::StackOverflow
22412216
}
2242-
stackdepth_push(&mut stack, &mut start_depths, ins.target, target_depth);
2217+
})?;
2218+
if target_depth > maxdepth {
2219+
maxdepth = target_depth;
22432220
}
2221+
stackdepth_push(&mut stack, &mut start_depths, ins.target, target_depth);
22442222
}
22452223
depth = new_depth;
22462224
if instr.is_scope_exit() || instr.is_unconditional_jump() {

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

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1095,6 +1095,10 @@ impl InstructionMetadata for Instruction {
10951095
StackEffect::new(pushed as u32, popped as u32)
10961096
}
10971097

1098+
// In CPython 3.14 the metadata-based stack_effect is the same for both
1099+
// fallthrough and branch paths for all real instructions.
1100+
// Only pseudo-instructions (SETUP_*) differ — see PseudoInstruction.
1101+
10981102
#[allow(clippy::too_many_arguments)]
10991103
fn fmt_dis(
11001104
&self,
@@ -1502,6 +1506,21 @@ impl InstructionMetadata for PseudoInstruction {
15021506
StackEffect::new(pushed as u32, popped as u32)
15031507
}
15041508

1509+
/// Handler entry effect for SETUP_* pseudo ops.
1510+
///
1511+
/// Fallthrough effect is 0 (NOPs), but when the branch is taken the
1512+
/// handler block starts with extra values on the stack:
1513+
/// SETUP_FINALLY: +1 (exc)
1514+
/// SETUP_CLEANUP: +2 (lasti + exc)
1515+
/// SETUP_WITH: +1 (pops __enter__ result, pushes lasti + exc)
1516+
fn stack_effect_jump(&self, _oparg: u32) -> i32 {
1517+
match self {
1518+
Self::SetupFinally { .. } | Self::SetupWith { .. } => 1,
1519+
Self::SetupCleanup { .. } => 2,
1520+
_ => self.stack_effect(_oparg),
1521+
}
1522+
}
1523+
15051524
fn is_unconditional_jump(&self) -> bool {
15061525
matches!(self, Self::Jump { .. } | Self::JumpNoInterrupt { .. })
15071526
}
@@ -1576,6 +1595,8 @@ impl InstructionMetadata for AnyInstruction {
15761595

15771596
inst_either!(fn stack_effect(&self, oparg: u32) -> i32);
15781597

1598+
inst_either!(fn stack_effect_jump(&self, oparg: u32) -> i32);
1599+
15791600
inst_either!(fn stack_effect_info(&self, oparg: u32) -> StackEffect);
15801601

15811602
inst_either!(fn fmt_dis(
@@ -1692,6 +1713,16 @@ pub trait InstructionMetadata {
16921713
self.stack_effect_info(oparg).effect()
16931714
}
16941715

1716+
/// Stack effect when the instruction takes its branch (jump=true).
1717+
///
1718+
/// CPython equivalent: `stack_effect(opcode, oparg, jump=True)`.
1719+
/// For most instructions this equals the fallthrough effect.
1720+
/// Override for instructions where branch and fallthrough differ
1721+
/// (e.g. `FOR_ITER`: fallthrough = +1, branch = −1).
1722+
fn stack_effect_jump(&self, oparg: u32) -> i32 {
1723+
self.stack_effect(oparg)
1724+
}
1725+
16951726
#[allow(clippy::too_many_arguments)]
16961727
fn fmt_dis(
16971728
&self,

‎crates/stdlib/src/_opcode.rs‎

Lines changed: 21 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -186,14 +186,18 @@ mod _opcode {
186186
})
187187
.unwrap_or(Ok(0))?;
188188

189-
let jump = args
190-
.jump
191-
.map(|v| {
192-
v.try_to_bool(vm).map_err(|_| {
193-
vm.new_value_error("stack_effect: jump must be False, True or None")
194-
})
195-
})
196-
.unwrap_or(Ok(false))?;
189+
let jump: Option<bool> = match args.jump {
190+
Some(v) => {
191+
if vm.is_none(&v) {
192+
None
193+
} else {
194+
Some(v.try_to_bool(vm).map_err(|_| {
195+
vm.new_value_error("stack_effect: jump must be False, True or None")
196+
})?)
197+
}
198+
}
199+
None => None,
200+
};
197201

198202
let opcode = Opcode::try_from_pyint(args.opcode, vm)?;
199203

@@ -202,8 +206,15 @@ mod _opcode {
202206
return Err(vm.new_value_error("invalid opcode or oparg"));
203207
}
204208

205-
let _ = jump; // Python API accepts jump but it's not used
206-
Ok(opcode.stack_effect(oparg))
209+
let effect = match jump {
210+
Some(true) => opcode.stack_effect_jump(oparg),
211+
Some(false) => opcode.stack_effect(oparg),
212+
// jump=None: max of both paths (CPython convention)
213+
None => opcode
214+
.stack_effect(oparg)
215+
.max(opcode.stack_effect_jump(oparg)),
216+
};
217+
Ok(effect)
207218
}
208219

209220
#[pyfunction]

0 commit comments

Comments
 (0)