Skip to content

Commit abf9495

Browse files
youknowoneyouknowone
committed
Align nested with cleanup bytecode
1 parent 597c759 commit abf9495

1 file changed

Lines changed: 132 additions & 4 deletions

File tree

‎crates/codegen/src/compile.rs‎

Lines changed: 132 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -28,9 +28,9 @@ use rustpython_compiler_core::{
2828
Mode, OneIndexed, PositionEncoding, SourceFile, SourceLocation,
2929
bytecode::{
3030
self, AnyInstruction, Arg as OpArgMarker, BinaryOperator, BuildSliceArgCount, CodeObject,
31-
ComparisonOperator, ConstantData, ConvertValueOparg, Instruction, IntrinsicFunction1,
32-
Invert, LoadAttr, LoadSuperAttr, OpArg, OpArgType, PseudoInstruction, SpecialMethod,
33-
UnpackExArgs, oparg,
31+
ComparisonOperator, ConstantData, ConvertValueOparg, Instruction, InstructionMetadata,
32+
IntrinsicFunction1, Invert, LoadAttr, LoadSuperAttr, OpArg, OpArgType, PseudoInstruction,
33+
SpecialMethod, UnpackExArgs, oparg,
3434
},
3535
};
3636
use rustpython_wtf8::Wtf8Buf;
@@ -608,6 +608,16 @@ impl Compiler {
608608
}
609609
}
610610

611+
fn statements_end_with_with_cleanup_scope_exit(body: &[ast::Stmt]) -> bool {
612+
body.last().is_some_and(|stmt| match stmt {
613+
ast::Stmt::With(ast::StmtWith { body, .. }) => {
614+
Self::statements_end_with_scope_exit(body)
615+
|| Self::statements_end_with_with_cleanup_scope_exit(body)
616+
}
617+
_ => false,
618+
})
619+
}
620+
611621
fn preserves_finally_entry_nop(body: &[ast::Stmt]) -> bool {
612622
body.last().is_some_and(|stmt| match stmt {
613623
ast::Stmt::Try(ast::StmtTry {
@@ -5818,6 +5828,10 @@ impl Compiler {
58185828
self.compile_with(items, body, is_async)?;
58195829
}
58205830

5831+
let preserve_outer_cleanup_target_nop = !is_async
5832+
&& (Self::statements_end_with_with_cleanup_scope_exit(body)
5833+
|| self.current_block_has_terminal_with_suppress_exit_predecessor());
5834+
58215835
// Pop fblock before normal exit. CPython emits this POP_BLOCK with
58225836
// no location for sync with, but with the with-item location for
58235837
// async with.
@@ -5827,7 +5841,11 @@ impl Compiler {
58275841
emit!(self, PseudoInstruction::PopBlock);
58285842
if !is_async {
58295843
self.set_no_location();
5830-
self.remove_last_no_location_nop();
5844+
if preserve_outer_cleanup_target_nop {
5845+
self.preserve_last_redundant_nop();
5846+
} else {
5847+
self.remove_last_no_location_nop();
5848+
}
58315849
self.set_source_range(with_range);
58325850
}
58335851
self.pop_fblock(if is_async {
@@ -10003,6 +10021,44 @@ impl Compiler {
1000310021
}
1000410022
}
1000510023

10024+
fn current_block_has_terminal_with_suppress_exit_predecessor(&self) -> bool {
10025+
let code = self.code_stack.last().expect("no code on stack");
10026+
let target = code.current_block;
10027+
let mut has_suppress_exit = false;
10028+
let mut has_normal_exit = false;
10029+
10030+
for block in &code.blocks {
10031+
let Some((last, prefix)) = block.instructions.split_last() else {
10032+
continue;
10033+
};
10034+
if last.target != target {
10035+
continue;
10036+
}
10037+
match last.instr.pseudo() {
10038+
Some(PseudoInstruction::JumpNoInterrupt { .. }) => {
10039+
let real_instrs: Vec<_> =
10040+
prefix.iter().filter_map(|info| info.instr.real()).collect();
10041+
has_suppress_exit |= matches!(
10042+
real_instrs.as_slice(),
10043+
[
10044+
Instruction::PopTop,
10045+
Instruction::PopExcept,
10046+
Instruction::PopTop,
10047+
Instruction::PopTop,
10048+
Instruction::PopTop,
10049+
]
10050+
);
10051+
}
10052+
Some(PseudoInstruction::Jump { .. }) => {
10053+
has_normal_exit |= !prefix.iter().any(|info| info.instr.is_scope_exit());
10054+
}
10055+
_ => {}
10056+
}
10057+
}
10058+
10059+
has_suppress_exit && !has_normal_exit
10060+
}
10061+
1000610062
fn remove_last_no_location_nop(&mut self) {
1000710063
if let Some(info) = self.current_block().instructions.last_mut() {
1000810064
info.remove_no_location_nop = true;
@@ -20181,6 +20237,78 @@ async def foo():
2018120237
);
2018220238
}
2018320239

20240+
#[test]
20241+
fn test_nested_terminal_with_keeps_outer_cleanup_target_nop() {
20242+
let code = compile_exec(
20243+
"\
20244+
def f():
20245+
with a():
20246+
with b():
20247+
raise E()
20248+
",
20249+
);
20250+
let f = find_code(&code, "f").expect("missing f code");
20251+
let ops: Vec<_> = f
20252+
.instructions
20253+
.iter()
20254+
.map(|unit| unit.op)
20255+
.filter(|op| !matches!(op, Instruction::Cache))
20256+
.collect();
20257+
20258+
assert!(
20259+
ops.windows(6).any(|window| {
20260+
matches!(
20261+
window,
20262+
[
20263+
Instruction::Copy { .. },
20264+
Instruction::PopExcept,
20265+
Instruction::Reraise { .. },
20266+
Instruction::Nop,
20267+
Instruction::LoadConst { .. },
20268+
Instruction::LoadConst { .. },
20269+
]
20270+
)
20271+
}),
20272+
"expected CPython-style outer with-exit target NOP after terminal nested with cleanup, got ops={ops:?}"
20273+
);
20274+
}
20275+
20276+
#[test]
20277+
fn test_nested_nonterminal_with_drops_outer_cleanup_target_nop() {
20278+
let code = compile_exec(
20279+
"\
20280+
def f():
20281+
with a():
20282+
with b():
20283+
x()
20284+
",
20285+
);
20286+
let f = find_code(&code, "f").expect("missing f code");
20287+
let ops: Vec<_> = f
20288+
.instructions
20289+
.iter()
20290+
.map(|unit| unit.op)
20291+
.filter(|op| !matches!(op, Instruction::Cache))
20292+
.collect();
20293+
20294+
assert!(
20295+
!ops.windows(6).any(|window| {
20296+
matches!(
20297+
window,
20298+
[
20299+
Instruction::Copy { .. },
20300+
Instruction::PopExcept,
20301+
Instruction::Reraise { .. },
20302+
Instruction::Nop,
20303+
Instruction::LoadConst { .. },
20304+
Instruction::LoadConst { .. },
20305+
]
20306+
)
20307+
}),
20308+
"unexpected outer with-exit target NOP for nested with with normal fallthrough, got ops={ops:?}"
20309+
);
20310+
}
20311+
2018420312
#[test]
2018520313
fn test_try_loop_elif_places_return_before_orelse_tail() {
2018620314
let code = compile_exec(

0 commit comments

Comments
 (0)