Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
106 changes: 97 additions & 9 deletions crates/codegen/src/compile.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1513,12 +1513,12 @@ impl Compiler {

FBlockType::ForLoop => {
// When returning from a for-loop, CPython swaps the preserved
// value with the iterator and uses POP_TOP for the iterator slot.
// value with the iterator and uses POP_TOP for loop cleanup.
if preserve_tos {
emit!(self, Instruction::Swap { i: 2 });
emit!(self, Instruction::PopTop);
} else {
emit!(self, Instruction::PopIter);
emit!(self, Instruction::PopTop);
}
}

Expand Down Expand Up @@ -5809,7 +5809,10 @@ impl Compiler {
emit!(self, Instruction::PopTop); // pop lasti
emit!(self, Instruction::PopTop); // pop self_exit
emit!(self, Instruction::PopTop); // pop exit_func
emit!(self, PseudoInstruction::Jump { delta: after_block });
emit!(
self,
PseudoInstruction::JumpNoInterrupt { delta: after_block }
);

// ===== Cleanup block (for nested exception during __exit__) =====
// Stack: [..., __exit__, lasti, prev_exc, lasti2, exc2]
Expand Down Expand Up @@ -9920,9 +9923,9 @@ impl Compiler {
}
}

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

// Jump to target
Expand Down Expand Up @@ -11155,12 +11158,12 @@ def f(base, cls, state):
fn test_loop_store_subscr_threads_direct_backedge() {
let code = compile_exec(
"\
def f(kwonlyargs, kwonlydefaults, arg2value):
def f(kwonlyargs, kw_only_defaults, arg2value):
missing = 0
for kwarg in kwonlyargs:
if kwarg not in arg2value:
if kwonlydefaults and kwarg in kwonlydefaults:
arg2value[kwarg] = kwonlydefaults[kwarg]
if kw_only_defaults and kwarg in kw_only_defaults:
arg2value[kwarg] = kw_only_defaults[kwarg]
else:
missing += 1
return missing
Expand Down Expand Up @@ -11244,6 +11247,91 @@ def f(obj):
);
}

#[test]
fn test_shared_final_return_is_cloned_for_jump_target() {
let code = compile_exec(
"\
def f(node):
if not isinstance(
node, (AsyncFunctionDef, FunctionDef, ClassDef, Module)
) or len(node.body) < 1:
return None
node = node.body[0]
if not isinstance(node, Expr):
return None
node = node.value
if isinstance(node, Constant) and isinstance(node.value, str):
return node
",
);
let f = find_code(&code, "f").expect("missing function code");
let ops: Vec<_> = f
.instructions
.iter()
.map(|unit| unit.op)
.filter(|op| !matches!(op, Instruction::Cache))
.collect();

let return_count = ops
.iter()
.filter(|op| matches!(op, Instruction::ReturnValue))
.count();
assert!(
return_count >= 3,
"expected multiple explicit return sites for shared final return case, got ops={ops:?}"
);
}

#[test]
fn test_for_break_uses_poptop_cleanup() {
let code = compile_exec(
"\
def f(parts):
for value in parts:
if value:
break
",
);
let f = find_code(&code, "f").expect("missing function code");
let ops: Vec<_> = f
.instructions
.iter()
.map(|unit| unit.op)
.filter(|op| !matches!(op, Instruction::Cache))
.collect();

let pop_iter_count = ops
.iter()
.filter(|op| matches!(op, Instruction::PopIter))
.count();
assert_eq!(
pop_iter_count, 1,
"expected only the loop-exhaustion POP_ITER, got ops={ops:?}"
);

let break_cleanup_idx = ops
.windows(3)
.position(|window| {
matches!(
window,
[
Instruction::PopTop,
Instruction::LoadConst { .. },
Instruction::ReturnValue
]
)
})
.expect("missing POP_TOP/LOAD_CONST/RETURN_VALUE break cleanup");
let end_for_idx = ops
.iter()
.position(|op| matches!(op, Instruction::EndFor))
.expect("missing END_FOR");
assert!(
break_cleanup_idx < end_for_idx,
"expected break cleanup before END_FOR, got ops={ops:?}"
);
}

#[test]
fn test_assert_without_message_raises_class_directly() {
let code = compile_exec(
Expand Down Expand Up @@ -11335,7 +11423,7 @@ def f(names, cls):
.filter(|unit| matches!(unit.op, Instruction::ReturnValue))
.count();

assert_eq!(return_count, 2);
assert_eq!(return_count, 1);
}

#[test]
Expand Down
Loading
Loading