Visitar URL original
Bytecode parity - for-loop cleanup, return duplication, late jump thr… · RustPython/RustPython@a082b4c · GitHub
Skip to content

Commit a082b4c

Browse files
youknowoneyouknowone
committed
Bytecode parity - for-loop cleanup, return duplication, late jump threading fix
- Use POP_TOP instead of POP_ITER for for-loop break/return cleanup - Expand duplicate_end_returns to clone final return for jump predecessors - Restrict late jump threading pass to unconditional jumps only - Skip exception blocks in inline/reorder passes - Simplify threaded_jump_instr NoInterrupt handling
1 parent 9cf7bcd commit a082b4c

12 files changed

Lines changed: 579 additions & 178 deletions

‎crates/codegen/src/compile.rs‎

Lines changed: 211 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1513,12 +1513,12 @@ impl Compiler {
15131513

15141514
FBlockType::ForLoop => {
15151515
// When returning from a for-loop, CPython swaps the preserved
1516-
// value with the iterator and uses POP_TOP for the iterator slot.
1516+
// value with the iterator and uses POP_TOP for loop cleanup.
15171517
if preserve_tos {
15181518
emit!(self, Instruction::Swap { i: 2 });
15191519
emit!(self, Instruction::PopTop);
15201520
} else {
1521-
emit!(self, Instruction::PopIter);
1521+
emit!(self, Instruction::PopTop);
15221522
}
15231523
}
15241524

@@ -5809,7 +5809,10 @@ impl Compiler {
58095809
emit!(self, Instruction::PopTop); // pop lasti
58105810
emit!(self, Instruction::PopTop); // pop self_exit
58115811
emit!(self, Instruction::PopTop); // pop exit_func
5812-
emit!(self, PseudoInstruction::Jump { delta: after_block });
5812+
emit!(
5813+
self,
5814+
PseudoInstruction::JumpNoInterrupt { delta: after_block }
5815+
);
58135816

58145817
// ===== Cleanup block (for nested exception during __exit__) =====
58155818
// Stack: [..., __exit__, lasti, prev_exc, lasti2, exc2]
@@ -9920,9 +9923,9 @@ impl Compiler {
99209923
}
99219924
}
99229925

9923-
// For break in a for loop, pop the iterator
9926+
// CPython unwinds a for-loop break with POP_TOP rather than POP_ITER.
99249927
if is_break && is_for_loop {
9925-
emit!(self, Instruction::PopIter);
9928+
emit!(self, Instruction::PopTop);
99269929
}
99279930

99289931
// Jump to target
@@ -11155,12 +11158,12 @@ def f(base, cls, state):
1115511158
fn test_loop_store_subscr_threads_direct_backedge() {
1115611159
let code = compile_exec(
1115711160
"\
11158-
def f(kwonlyargs, kwonlydefaults, arg2value):
11161+
def f(kwonlyargs, kw_only_defaults, arg2value):
1115911162
missing = 0
1116011163
for kwarg in kwonlyargs:
1116111164
if kwarg not in arg2value:
11162-
if kwonlydefaults and kwarg in kwonlydefaults:
11163-
arg2value[kwarg] = kwonlydefaults[kwarg]
11165+
if kw_only_defaults and kwarg in kw_only_defaults:
11166+
arg2value[kwarg] = kw_only_defaults[kwarg]
1116411167
else:
1116511168
missing += 1
1116611169
return missing
@@ -11244,6 +11247,205 @@ def f(obj):
1124411247
);
1124511248
}
1124611249

11250+
#[test]
11251+
fn test_conditional_body_is_preserved_before_final_return() {
11252+
let code = compile_exec(
11253+
"\
11254+
def f(x, y):
11255+
if x == y:
11256+
print('then', flush=True)
11257+
",
11258+
);
11259+
let f = find_code(&code, "f").expect("missing function code");
11260+
let ops: Vec<_> = f
11261+
.instructions
11262+
.iter()
11263+
.map(|unit| unit.op)
11264+
.filter(|op| !matches!(op, Instruction::Cache))
11265+
.collect();
11266+
11267+
let cond_idx = ops
11268+
.iter()
11269+
.position(|op| matches!(op, Instruction::PopJumpIfFalse { .. }))
11270+
.expect("missing POP_JUMP_IF_FALSE");
11271+
let first_return_idx = ops
11272+
.iter()
11273+
.position(|op| matches!(op, Instruction::ReturnValue))
11274+
.expect("missing RETURN_VALUE");
11275+
11276+
assert!(
11277+
ops[cond_idx..first_return_idx]
11278+
.iter()
11279+
.any(|op| matches!(op, Instruction::CallKw { .. })),
11280+
"expected conditional body call before final return, got ops={ops:?}"
11281+
);
11282+
}
11283+
11284+
#[test]
11285+
fn test_nested_conditional_body_is_preserved_before_final_return() {
11286+
let code = compile_exec(
11287+
"\
11288+
def outer():
11289+
def side():
11290+
print('side', flush=True)
11291+
def cb():
11292+
flag = True
11293+
if flag:
11294+
side()
11295+
return cb
11296+
",
11297+
);
11298+
let cb = find_code(&code, "cb").expect("missing nested cb code");
11299+
let ops: Vec<_> = cb
11300+
.instructions
11301+
.iter()
11302+
.map(|unit| unit.op)
11303+
.filter(|op| !matches!(op, Instruction::Cache))
11304+
.collect();
11305+
11306+
let cond_idx = ops
11307+
.iter()
11308+
.position(|op| matches!(op, Instruction::PopJumpIfFalse { .. }))
11309+
.expect("missing POP_JUMP_IF_FALSE");
11310+
let first_return_idx = ops
11311+
.iter()
11312+
.position(|op| matches!(op, Instruction::ReturnValue))
11313+
.expect("missing RETURN_VALUE");
11314+
11315+
assert!(
11316+
ops[cond_idx..first_return_idx]
11317+
.iter()
11318+
.any(|op| matches!(op, Instruction::Call { .. })),
11319+
"expected nested conditional body call before final return, got ops={ops:?}"
11320+
);
11321+
}
11322+
11323+
#[test]
11324+
fn test_conditional_compare_uses_bool_compare_oparg() {
11325+
let code = compile_exec(
11326+
"\
11327+
def f(x, y):
11328+
if x == y:
11329+
return 1
11330+
return 0
11331+
",
11332+
);
11333+
let f = find_code(&code, "f").expect("missing function code");
11334+
let compare = f
11335+
.instructions
11336+
.iter()
11337+
.find(|unit| matches!(unit.op, Instruction::CompareOp { .. }))
11338+
.expect("missing COMPARE_OP");
11339+
11340+
assert_eq!(u8::from(compare.arg), 88);
11341+
}
11342+
11343+
#[test]
11344+
fn test_chained_conditional_compares_use_bool_compare_oparg() {
11345+
let code = compile_exec(
11346+
"\
11347+
def f(a, b, c):
11348+
if a < b < c:
11349+
return 1
11350+
return 0
11351+
",
11352+
);
11353+
let f = find_code(&code, "f").expect("missing function code");
11354+
let compare_args: Vec<_> = f
11355+
.instructions
11356+
.iter()
11357+
.filter(|unit| matches!(unit.op, Instruction::CompareOp { .. }))
11358+
.map(|unit| u8::from(unit.arg))
11359+
.collect();
11360+
11361+
assert_eq!(compare_args, vec![18, 18]);
11362+
}
11363+
11364+
#[test]
11365+
fn test_shared_final_return_is_cloned_for_jump_target() {
11366+
let code = compile_exec(
11367+
"\
11368+
def f(node):
11369+
if not isinstance(
11370+
node, (AsyncFunctionDef, FunctionDef, ClassDef, Module)
11371+
) or len(node.body) < 1:
11372+
return None
11373+
node = node.body[0]
11374+
if not isinstance(node, Expr):
11375+
return None
11376+
node = node.value
11377+
if isinstance(node, Constant) and isinstance(node.value, str):
11378+
return node
11379+
",
11380+
);
11381+
let f = find_code(&code, "f").expect("missing function code");
11382+
let ops: Vec<_> = f
11383+
.instructions
11384+
.iter()
11385+
.map(|unit| unit.op)
11386+
.filter(|op| !matches!(op, Instruction::Cache))
11387+
.collect();
11388+
11389+
let return_count = ops
11390+
.iter()
11391+
.filter(|op| matches!(op, Instruction::ReturnValue))
11392+
.count();
11393+
assert!(
11394+
return_count >= 3,
11395+
"expected multiple explicit return sites for shared final return case, got ops={ops:?}"
11396+
);
11397+
}
11398+
11399+
#[test]
11400+
fn test_for_break_uses_poptop_cleanup() {
11401+
let code = compile_exec(
11402+
"\
11403+
def f(parts):
11404+
for value in parts:
11405+
if value:
11406+
break
11407+
",
11408+
);
11409+
let f = find_code(&code, "f").expect("missing function code");
11410+
let ops: Vec<_> = f
11411+
.instructions
11412+
.iter()
11413+
.map(|unit| unit.op)
11414+
.filter(|op| !matches!(op, Instruction::Cache))
11415+
.collect();
11416+
11417+
let pop_iter_count = ops
11418+
.iter()
11419+
.filter(|op| matches!(op, Instruction::PopIter))
11420+
.count();
11421+
assert_eq!(
11422+
pop_iter_count, 1,
11423+
"expected only the loop-exhaustion POP_ITER, got ops={ops:?}"
11424+
);
11425+
11426+
let break_cleanup_idx = ops
11427+
.windows(3)
11428+
.position(|window| {
11429+
matches!(
11430+
window,
11431+
[
11432+
Instruction::PopTop,
11433+
Instruction::LoadConst { .. },
11434+
Instruction::ReturnValue
11435+
]
11436+
)
11437+
})
11438+
.expect("missing POP_TOP/LOAD_CONST/RETURN_VALUE break cleanup");
11439+
let end_for_idx = ops
11440+
.iter()
11441+
.position(|op| matches!(op, Instruction::EndFor))
11442+
.expect("missing END_FOR");
11443+
assert!(
11444+
break_cleanup_idx < end_for_idx,
11445+
"expected break cleanup before END_FOR, got ops={ops:?}"
11446+
);
11447+
}
11448+
1124711449
#[test]
1124811450
fn test_assert_without_message_raises_class_directly() {
1124911451
let code = compile_exec(
@@ -11335,7 +11537,7 @@ def f(names, cls):
1133511537
.filter(|unit| matches!(unit.op, Instruction::ReturnValue))
1133611538
.count();
1133711539

11338-
assert_eq!(return_count, 2);
11540+
assert_eq!(return_count, 1);
1133911541
}
1134011542

1134111543
#[test]

0 commit comments

Comments
 (0)