diff --git a/crates/codegen/src/compile.rs b/crates/codegen/src/compile.rs index 2f7510f0d..5a3d98d54 100644 --- a/crates/codegen/src/compile.rs +++ b/crates/codegen/src/compile.rs @@ -12190,9 +12190,9 @@ def f(node): .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:?}" + assert_eq!( + return_count, 5, + "expected cloned return sites for each shared return edge, got ops={ops:?}" ); } @@ -13116,6 +13116,78 @@ def f(names, cls): assert_eq!(return_count, 1); } + #[test] + fn test_listcomp_cleanup_tail_keeps_split_store_fast_pair() { + let code = compile_exec( + "\ +def f(escaped_string, quote_types): + possible_quotes = [q for q in quote_types if q not in escaped_string] + return possible_quotes +", + ); + 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_idx = ops + .iter() + .position(|op| matches!(op, Instruction::PopIter)) + .expect("missing POP_ITER"); + let tail = &ops[pop_iter_idx + 1..]; + + assert!( + matches!( + tail, + [ + Instruction::StoreFast { .. }, + Instruction::StoreFast { .. }, + Instruction::LoadFastBorrow { .. }, + Instruction::ReturnValue, + .. + ] + ), + "expected split STORE_FAST pair after listcomp cleanup, got ops={ops:?}" + ); + } + + #[test] + fn test_with_suppress_tail_duplicates_final_return_none() { + let code = compile_exec( + "\ +def f(cm, cond): + if cond: + with cm(): + pass +", + ); + 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_eq!( + return_count, 3, + "expected duplicated return-none epilogues, got ops={ops:?}" + ); + assert!( + !ops.iter() + .any(|op| matches!(op, Instruction::JumpBackwardNoInterrupt { .. })), + "with suppress tail should not jump back to shared return block, got ops={ops:?}" + ); + } + #[test] fn test_fstring_adjacent_literals_are_merged() { let code = compile_exec( diff --git a/crates/codegen/src/ir.rs b/crates/codegen/src/ir.rs index 01e6971f6..d4cb210da 100644 --- a/crates/codegen/src/ir.rs +++ b/crates/codegen/src/ir.rs @@ -271,6 +271,8 @@ impl CodeInfo { reorder_jump_over_exception_cleanup_blocks(&mut self.blocks); self.eliminate_unreachable_blocks(); remove_redundant_nops_and_jumps(&mut self.blocks); + inline_with_suppress_return_blocks(&mut self.blocks); + self.eliminate_unreachable_blocks(); // Late CFG cleanup can create new same-line STORE_FAST/LOAD_FAST and // STORE_FAST/STORE_FAST adjacencies in match/capture code paths that // did not exist during the earlier flowgraph-like pass. @@ -301,6 +303,7 @@ impl CodeInfo { self.deoptimize_store_fast_store_fast_after_cleanup(); self.apply_static_swaps(); self.insert_superinstructions(); + self.deoptimize_store_fast_store_fast_after_cleanup(); self.optimize_load_global_push_null(); self.reorder_entry_prefix_cell_setup(); self.remove_unused_consts(); @@ -5319,7 +5322,7 @@ fn duplicate_end_returns(blocks: &mut Vec) { while current != BlockIdx::NULL { let block = &blocks[current.idx()]; let next = next_nonempty_block(blocks, block.next); - if current != last_block && !block.cold && !block.except_handler { + if current != last_block && !block.cold { let last_ins = block.instructions.last(); let has_fallthrough = last_ins .map(|ins| !ins.instr.is_scope_exit() && !ins.instr.is_unconditional_jump()) @@ -5335,20 +5338,28 @@ fn duplicate_end_returns(blocks: &mut Vec) { AnyInstruction::Real(Instruction::ReturnValue) ) }; - if next == last_block + if !block.except_handler + && next == last_block && has_fallthrough && trailing_conditional_jump_index(block).is_none() && !already_has_return { fallthrough_blocks_to_fix.push(current); } - if predecessors[last_block.idx()] > 1 - && let Some(last) = block.instructions.last() - && last.instr.is_unconditional_jump() - && last.target != BlockIdx::NULL - && next_nonempty_block(blocks, last.target) == last_block - { - jump_targets_to_fix.push((current, block.instructions.len() - 1)); + let jump_idx = trailing_conditional_jump_index(block).or_else(|| { + block.instructions.last().and_then(|last| { + (last.instr.is_unconditional_jump() && last.target != BlockIdx::NULL) + .then_some(block.instructions.len() - 1) + }) + }); + if let Some(jump_idx) = jump_idx { + let jump = &block.instructions[jump_idx]; + if jump.target != BlockIdx::NULL + && next_nonempty_block(blocks, jump.target) == last_block + && (is_conditional_jump(&jump.instr) || predecessors[last_block.idx()] > 1) + { + jump_targets_to_fix.push((current, jump_idx)); + } } } current = blocks[current.idx()].next; @@ -5371,7 +5382,7 @@ fn duplicate_end_returns(blocks: &mut Vec) { // Clone the final return block for jump predecessors so their target layout // matches CPython's duplicated exit blocks. - for (block_idx, instr_idx) in jump_targets_to_fix { + for (block_idx, instr_idx) in jump_targets_to_fix.into_iter().rev() { let jump = blocks[block_idx.idx()].instructions[instr_idx]; let mut cloned_return = return_insts.clone(); if let Some(first) = cloned_return.first_mut() { @@ -5404,6 +5415,64 @@ fn duplicate_end_returns(blocks: &mut Vec) { } } +fn inline_with_suppress_return_blocks(blocks: &mut [Block]) { + fn is_return_block(block: &Block) -> bool { + block.instructions.len() == 2 + && matches!( + block.instructions[0].instr.real(), + Some(Instruction::LoadConst { .. }) + ) + && matches!( + block.instructions[1].instr.real(), + Some(Instruction::ReturnValue) + ) + } + + fn has_with_suppress_prefix(block: &Block, jump_idx: usize) -> bool { + let tail: Vec<_> = block.instructions[..jump_idx] + .iter() + .filter_map(|info| info.instr.real()) + .rev() + .take(5) + .collect(); + matches!( + tail.as_slice(), + [ + Instruction::PopTop, + Instruction::PopTop, + Instruction::PopTop, + Instruction::PopExcept, + Instruction::PopTop, + ] + ) + } + + for block_idx in 0..blocks.len() { + let Some(jump_idx) = blocks[block_idx].instructions.len().checked_sub(1) else { + continue; + }; + let jump = blocks[block_idx].instructions[jump_idx]; + if !jump.instr.is_unconditional_jump() || jump.target == BlockIdx::NULL { + continue; + } + if !has_with_suppress_prefix(&blocks[block_idx], jump_idx) { + continue; + } + + let target = next_nonempty_block(blocks, jump.target); + if target == BlockIdx::NULL || !is_return_block(&blocks[target.idx()]) { + continue; + } + + let mut cloned_return = blocks[target.idx()].instructions.clone(); + for instr in &mut cloned_return { + overwrite_location(instr, jump.location, jump.end_location); + } + blocks[block_idx].instructions.pop(); + blocks[block_idx].instructions.extend(cloned_return); + } +} + /// Label exception targets: walk CFG with except stack, set per-instruction /// handler info and block preserve_lasti flag. Converts POP_BLOCK to NOP. /// flowgraph.c label_exception_targets + push_except_block