From 47068d7edd2d137ea932f4daeb3baf0111361dbd Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Sun, 19 Apr 2026 22:27:17 +0900 Subject: [PATCH 1/3] Bytecode parity - boolop, comprehension, CFG passes - Flatten nested same-op BoolOp and add IfExpr to jump_if - Simplify is_name_imported to module-level only - Enable inlined comprehensions in module/class scope - Add TweakInlinedComprehensionScopes with fast_hidden tracking - Track fb_range in FBlockInfo for with-statement line info - Add constant subscript folding and unary Not folding - Add emit_return_const_no_location for implicit returns - Use JumpNoInterrupt for ternary/except jumps - Add STACK_USE_GUIDELINE threshold for collection building - Reorder CFG pipeline: inline small blocks earlier, resolve line numbers before cold block extraction - Add redirect_empty_unconditional_jump_targets, reorder_conditional_chain_and_jump_back_blocks, materialize_empty_conditional_exit_targets, duplicate_shared_jump_back_targets passes - Add borrow deoptimization for multi-handler, named-except, protected conditional tail, and protected import joins - Run folding/optimization passes twice around peephole --- crates/codegen/src/compile.rs | 2909 ++++++++++++++--- crates/codegen/src/ir.rs | 1877 ++++++++--- ...degen__compile__tests__nested_bool_op.snap | 2 +- ...pile__tests__nested_double_async_with.snap | 8 +- crates/codegen/src/symboltable.rs | 27 +- crates/compiler-core/src/marshal.rs | 20 +- scripts/compare_bytecode.py | 126 +- scripts/dis_dump.py | 17 +- 8 files changed, 4052 insertions(+), 934 deletions(-) diff --git a/crates/codegen/src/compile.rs b/crates/codegen/src/compile.rs index 0b82f400dfb..34ea9aead72 100644 --- a/crates/codegen/src/compile.rs +++ b/crates/codegen/src/compile.rs @@ -134,6 +134,7 @@ pub struct FBlockInfo { pub fb_type: FBlockType, pub fb_block: BlockIdx, pub fb_exit: BlockIdx, + pub fb_range: TextRange, // additional data for fblock unwinding pub fb_datum: FBlockDatum, } @@ -443,6 +444,8 @@ enum CollectionType { Set, } +const STACK_USE_GUIDELINE: u32 = 30; + impl Compiler { fn constant_truthiness(constant: &ConstantData) -> bool { match constant { @@ -460,6 +463,18 @@ impl Compiler { } } + fn boolop_fast_fold_literal(expr: &ast::Expr) -> bool { + matches!( + expr, + ast::Expr::NumberLiteral(_) + | ast::Expr::StringLiteral(_) + | ast::Expr::BytesLiteral(_) + | ast::Expr::BooleanLiteral(_) + | ast::Expr::NoneLiteral(_) + | ast::Expr::EllipsisLiteral(_) + ) + } + fn constant_expr_truthiness(&mut self, expr: &ast::Expr) -> CompileResult> { Ok(self .try_fold_constant_expr(expr)? @@ -662,14 +677,11 @@ impl Compiler { pushed: u32, collection_type: CollectionType, ) -> CompileResult<()> { + let collection_range = self.current_source_range; let n = elts.len().to_u32(); let seen_star = elts.iter().any(|e| matches!(e, ast::Expr::Starred(_))); - // Determine collection size threshold for optimization - let big = match collection_type { - CollectionType::Set => n > 8, - _ => n > 4, - }; + let big = n + pushed > STACK_USE_GUIDELINE; let can_fold_const_collection = match collection_type { CollectionType::Tuple => n > 0, @@ -689,13 +701,17 @@ impl Compiler { self.emit_load_const(folded); } CollectionType::List => { + self.set_source_range(collection_range); emit!(self, Instruction::BuildList { count: 0 }); self.emit_load_const(folded); + self.set_source_range(collection_range); emit!(self, Instruction::ListExtend { i: 1 }); } CollectionType::Set => { + self.set_source_range(collection_range); emit!(self, Instruction::BuildSet { count: 0 }); self.emit_load_const(folded); + self.set_source_range(collection_range); emit!(self, Instruction::SetUpdate { i: 1 }); } } @@ -708,6 +724,7 @@ impl Compiler { self.compile_expression(elt)?; } let total_size = n + pushed; + self.set_source_range(collection_range); match collection_type { CollectionType::List => { emit!(self, Instruction::BuildList { count: total_size }); @@ -722,25 +739,11 @@ impl Compiler { return Ok(()); } - // Has stars or too big: use streaming approach - let stream_big_nonconst_collection = if !seen_star - && big - && matches!(collection_type, CollectionType::List | CollectionType::Set) - { - elts.iter().try_fold(false, |has_nonconst, elt| { - if has_nonconst { - return Ok(true); - } - Ok(self.try_fold_constant_expr(elt)?.is_none()) - })? - } else { - false - }; - + // Has stars or too big: use streaming approach. let mut sequence_built = false; let mut i = 0u32; - if stream_big_nonconst_collection { + if big { match collection_type { CollectionType::List => { emit!(self, Instruction::BuildList { count: pushed }); @@ -750,7 +753,10 @@ impl Compiler { emit!(self, Instruction::BuildSet { count: pushed }); sequence_built = true; } - CollectionType::Tuple => {} + CollectionType::Tuple => { + emit!(self, Instruction::BuildList { count: pushed }); + sequence_built = true; + } } } @@ -758,6 +764,7 @@ impl Compiler { if let ast::Expr::Starred(ast::ExprStarred { value, .. }) = elt { // When we hit first star, build sequence with elements so far if !sequence_built { + self.set_source_range(collection_range); match collection_type { CollectionType::List => { emit!(self, Instruction::BuildList { count: i + pushed }); @@ -774,6 +781,7 @@ impl Compiler { // Compile the starred expression and extend self.compile_expression_without_const_boolop_folding(value)?; + self.set_source_range(collection_range); match collection_type { CollectionType::List => { emit!(self, Instruction::ListExtend { i: 1 }); @@ -791,6 +799,7 @@ impl Compiler { if sequence_built { // Sequence already exists, append to it + self.set_source_range(collection_range); match collection_type { CollectionType::List => { emit!(self, Instruction::ListAppend { i: 1 }); @@ -811,6 +820,7 @@ impl Compiler { // If we never built sequence (all non-starred), build it now if !sequence_built { + self.set_source_range(collection_range); match collection_type { CollectionType::List => { emit!(self, Instruction::BuildList { count: i + pushed }); @@ -824,6 +834,7 @@ impl Compiler { } } else if collection_type == CollectionType::Tuple { // For tuples, convert the list to tuple + self.set_source_range(collection_range); emit!( self, Instruction::CallIntrinsic1 { @@ -858,29 +869,13 @@ impl Compiler { .expect("symbol_table_stack is empty! This is a compiler bug.") } - /// Check if a name is imported in current scope or any enclosing scope. + /// Match CPython's `is_import_originated()`: only imports recorded in the + /// module-level symbol table suppress method-call optimization. fn is_name_imported(&self, name: &str) -> bool { - let current = self.current_symbol_table(); - if let Some(sym) = current.symbols.get(name) { - if sym.flags.contains(SymbolFlags::IMPORTED) { - // Module/class scope imports use plain LOAD_ATTR - // Function-local imports use method mode (scope is Local) - return !matches!( - current.typ, - CompilerScope::Function | CompilerScope::AsyncFunction | CompilerScope::Lambda - ); - } - if sym.scope == SymbolScope::Local { - return false; - } - } - // Check enclosing scopes for module-level imports accessed as globals - self.symbol_table_stack.iter().rev().skip(1).any(|table| { - table - .symbols - .get(name) - .is_some_and(|sym| sym.flags.contains(SymbolFlags::IMPORTED)) - }) + self.symbol_table_stack + .first() + .and_then(|table| table.symbols.get(name)) + .is_some_and(|sym| sym.flags.contains(SymbolFlags::IMPORTED)) } /// Get the cell-relative index of a free variable. @@ -1162,8 +1157,6 @@ impl Compiler { } /// Check if this is an inlined comprehension context (PEP 709). - /// PEP 709: Inline comprehensions in function-like scopes. - /// TODO: Module/class scope inlining needs more work (Cell name resolution edge cases). /// Generator expressions are never inlined. fn is_inlined_comprehension_context( &self, @@ -1173,9 +1166,6 @@ impl Compiler { if comprehension_type == ComprehensionType::Generator { return false; } - if !self.ctx.in_func() { - return false; - } comp_table.comp_inlined } @@ -1609,6 +1599,7 @@ impl Compiler { fb_exit: BlockIdx, fb_datum: FBlockDatum, ) -> CompileResult<()> { + let fb_range = self.current_source_range; let code = self.current_code_info(); if code.fblock.len() >= MAXBLOCKS { return Err(self.error(CodegenErrorType::SyntaxError( @@ -1619,6 +1610,7 @@ impl Compiler { fb_type, fb_block, fb_exit, + fb_range, fb_datum, }); Ok(()) @@ -1680,6 +1672,7 @@ impl Compiler { FBlockType::With | FBlockType::AsyncWith => { // Stack: [..., exit_func, self_exit, return_value (if preserve_tos)] + self.set_source_range(info.fb_range); emit!(self, PseudoInstruction::PopBlock); if preserve_tos { @@ -2002,8 +1995,10 @@ impl Compiler { assert_eq!(self.code_stack.len(), size_before); - // Emit None at end: - self.emit_return_const(ConstantData::None); + // Match _PyCodegen_AddReturnAtEnd(): implicit scope epilogues start + // without a source location and receive one later via CFG line + // propagation. + self.emit_return_const_no_location(ConstantData::None); Ok(()) } @@ -2047,6 +2042,7 @@ impl Compiler { ); emit!(self, Instruction::PopTop); + self.set_no_location(); } else { self.compile_statement(statement)?; } @@ -2063,6 +2059,7 @@ impl Compiler { ); emit!(self, Instruction::PopTop); + self.set_no_location(); } else { self.compile_statement(last)?; self.emit_load_const(ConstantData::None); @@ -2270,7 +2267,14 @@ impl Compiler { SymbolScope::Local => { if module_global_from_nested_scope { NameOp::Global - } else if is_function_like { + } else if is_function_like + || self + .current_code_info() + .metadata + .fast_hidden + .get(name.as_ref()) + .is_some_and(|&hidden| hidden) + { NameOp::Fast } else { NameOp::Name @@ -2306,23 +2310,13 @@ impl Compiler { _ => unreachable!("Invalid scope for Deref operation"), }; - // Mark cell variables accessed inside inlined comprehensions as hidden - if self.current_code_info().in_inlined_comp { - let info = self.code_stack.last_mut().unwrap(); - if info - .metadata - .fast_hidden - .get(name.as_ref()) - .is_none_or(|&v| v) - { - info.metadata.fast_hidden.insert(name.to_string(), true); - } - } - match usage { NameUsage::Load => { // ClassBlock (not inlined comp): LOAD_LOCALS first, then LOAD_FROM_DICT_OR_DEREF - if self.ctx.in_class && !self.ctx.in_func() { + if self.ctx.in_class + && !self.ctx.in_func() + && !self.current_code_info().in_inlined_comp + { emit!(self, Instruction::LoadLocals); emit!(self, Instruction::LoadFromDictOrDeref { i }); // can_see_class_scope: LOAD_DEREF(__classdict__) first @@ -2340,18 +2334,6 @@ impl Compiler { } NameOp::Fast => { let var_num = self.get_local_var_index(&name)?; - // Mark variables accessed inside inlined comprehensions as hidden - if self.current_code_info().in_inlined_comp { - let info = self.code_stack.last_mut().unwrap(); - if info - .metadata - .fast_hidden - .get(name.as_ref()) - .is_none_or(|&v| v) - { - info.metadata.fast_hidden.insert(name.to_string(), true); - } - } match usage { NameUsage::Load => emit!(self, Instruction::LoadFast { var_num }), NameUsage::Store => emit!(self, Instruction::StoreFast { var_num }), @@ -2542,6 +2524,7 @@ impl Compiler { } emit!(self, Instruction::PopTop); + self.set_no_location(); } } ast::Stmt::Global(_) | ast::Stmt::Nonlocal(_) => { @@ -2706,6 +2689,8 @@ impl Compiler { ); } + let prev_source_range = self.current_source_range; + let stmt_range = statement.range(); match value { Some(v) => { if self.ctx.func == FunctionContext::AsyncFunction @@ -2723,19 +2708,36 @@ impl Compiler { let preserve_tos = folded_constant.is_none(); if preserve_tos { self.compile_expression(v)?; + } else { + self.set_source_range(v.range()); + emit!(self, Instruction::Nop); + } + + let source = self.source_file.to_source_code(); + if source.line_index(v.range().start()) + != source.line_index(stmt_range.start()) + { + self.set_source_range(stmt_range); + emit!(self, Instruction::Nop); } + self.set_source_range(stmt_range); self.unwind_fblock_stack(preserve_tos, false)?; + self.set_source_range(stmt_range); if let Some(constant) = folded_constant { self.emit_load_const(constant); } self.emit_return_value(); } None => { + self.set_source_range(stmt_range); + emit!(self, Instruction::Nop); // Unwind fblock stack with preserve_tos=false (no value to preserve) self.unwind_fblock_stack(false, false)?; + self.set_source_range(stmt_range); self.emit_return_const(ConstantData::None); } } + self.set_source_range(prev_source_range); let dead = self.new_block(); self.switch_to_block(dead); } @@ -4317,7 +4319,7 @@ impl Compiler { match body.last() { Some(ast::Stmt::Return(_)) => {} _ => { - self.emit_return_const(ConstantData::None); + self.emit_return_const_no_location(ConstantData::None); } } // Functions with no other constants should still have None in co_consts @@ -5586,7 +5588,12 @@ impl Compiler { return Ok(()); }; - emit!(self, PseudoInstruction::Jump { delta: end_block }); + self.set_no_location(); + emit!( + self, + PseudoInstruction::JumpNoInterrupt { delta: end_block } + ); + self.set_no_location(); self.switch_to_block(next_block); if let Some(test) = &clause.test { @@ -5623,6 +5630,7 @@ impl Compiler { self.compile_statements(body)?; self.ctx.loop_data = was_in_loop; emit!(self, PseudoInstruction::Jump { delta: while_block }); + self.set_no_location(); self.switch_to_block(else_block); self.pop_fblock(FBlockType::WhileLoop); @@ -5933,6 +5941,7 @@ impl Compiler { self.compile_statements(body)?; self.ctx.loop_data = was_in_loop; emit!(self, PseudoInstruction::Jump { delta: for_block }); + self.set_no_location(); self.switch_to_block(else_block); @@ -7539,69 +7548,21 @@ impl Compiler { // Compile expression for test, and jump to label if false let result = match &expression { ast::Expr::BoolOp(ast::ExprBoolOp { op, values, .. }) => { - match op { - ast::BoolOp::And => { - if condition { - // If all values are true. - let end_block = self.new_block(); - let (last_value, values) = values.split_last().unwrap(); - - // If any of the values is false, we can short-circuit. - for value in values { - self.compile_jump_if_inner(value, false, end_block, source_range)?; - } - - // It depends upon the last value now: will it be true? - self.compile_jump_if_inner( - last_value, - true, - target_block, - source_range, - )?; - self.switch_to_block(end_block); - } else { - // If any value is false, the whole condition is false. - for value in values { - self.compile_jump_if_inner( - value, - false, - target_block, - source_range, - )?; - } - } - } - ast::BoolOp::Or => { - if condition { - // If any of the values is true. - for value in values { - self.compile_jump_if_inner( - value, - true, - target_block, - source_range, - )?; - } - } else { - // If all of the values are false. - let end_block = self.new_block(); - let (last_value, values) = values.split_last().unwrap(); + let (last_value, prefix_values) = values.split_last().unwrap(); + let cond2 = matches!(op, ast::BoolOp::Or); + let next2 = if cond2 != condition { + self.new_block() + } else { + target_block + }; - // If any value is true, we can short-circuit: - for value in values { - self.compile_jump_if_inner(value, true, end_block, source_range)?; - } + for value in prefix_values { + self.compile_jump_if_inner(value, cond2, next2, source_range)?; + } + self.compile_jump_if_inner(last_value, condition, target_block, source_range)?; - // It all depends upon the last value now! - self.compile_jump_if_inner( - last_value, - false, - target_block, - source_range, - )?; - self.switch_to_block(end_block); - } - } + if next2 != target_block { + self.switch_to_block(next2); } Ok(()) } @@ -7610,6 +7571,23 @@ impl Compiler { operand, .. }) => self.compile_jump_if_inner(operand, !condition, target_block, source_range), + ast::Expr::If(ast::ExprIf { + test, body, orelse, .. + }) => { + let end = self.new_block(); + let next2 = self.new_block(); + self.compile_jump_if_inner(test, false, next2, source_range)?; + self.compile_jump_if_inner(body, condition, target_block, source_range)?; + self.set_no_location(); + emit!(self, PseudoInstruction::JumpNoInterrupt { delta: end }); + self.set_no_location(); + + self.switch_to_block(next2); + self.compile_jump_if_inner(orelse, condition, target_block, source_range)?; + + self.switch_to_block(end); + Ok(()) + } ast::Expr::Compare(ast::ExprCompare { left, ops, @@ -7692,54 +7670,84 @@ impl Compiler { /// Compile a boolean operation as an expression. /// This means, that the last value remains on the stack. fn compile_bool_op(&mut self, op: &ast::BoolOp, values: &[ast::Expr]) -> CompileResult<()> { - self.compile_bool_op_with_target(op, values, None) - } + fn flatten_same_boolop_values<'a>( + op: &ast::BoolOp, + value: &'a ast::Expr, + out: &mut Vec<&'a ast::Expr>, + ) { + if let ast::Expr::BoolOp(ast::ExprBoolOp { + op: inner_op, + values, + .. + }) = value + && inner_op == op + { + for value in values { + flatten_same_boolop_values(op, value, out); + } + } else { + out.push(value); + } + } + + let mut flattened = Vec::with_capacity(values.len()); + for value in values { + flatten_same_boolop_values(op, value, &mut flattened); + } - /// Compile a boolean operation as an expression, with an optional - /// short-circuit target override. When `short_circuit_target` is `Some`, - /// the short-circuit jumps go to that block instead of the default - /// `after_block`, enabling jump threading to avoid redundant `__bool__` calls. - fn compile_bool_op_with_target( - &mut self, - op: &ast::BoolOp, - values: &[ast::Expr], - short_circuit_target: Option, - ) -> CompileResult<()> { let after_block = self.new_block(); - let (last_value, values) = values.split_last().unwrap(); - let jump_target = short_circuit_target.unwrap_or(after_block); + let (last_value, prefix_values) = flattened.split_last().unwrap(); - for value in values { - // Optimization: when a non-last value is a BoolOp with the opposite - // operator, redirect its short-circuit exits to skip the outer's - // redundant __bool__ test (jump threading). - if short_circuit_target.is_none() - && let ast::Expr::BoolOp(ast::ExprBoolOp { + for value in prefix_values { + let continue_block = self.new_block(); + match value { + ast::Expr::BoolOp(ast::ExprBoolOp { op: inner_op, - values: inner_values, + values, .. - }) = value - && inner_op != op - { - let pop_block = self.new_block(); - self.compile_bool_op_with_target(inner_op, inner_values, Some(pop_block))?; - self.emit_short_circuit_test(op, after_block); - self.switch_to_block(pop_block); - emit!(self, Instruction::PopTop); - continue; + }) if inner_op != op => { + let (last_inner_value, inner_prefix_values) = values.split_last().unwrap(); + for inner_value in inner_prefix_values { + self.compile_expression(inner_value)?; + self.emit_short_circuit_test(inner_op, continue_block); + emit!(self, Instruction::PopTop); + } + self.compile_expression(last_inner_value)?; + } + _ => self.compile_expression(value)?, } - - self.compile_expression(value)?; - self.emit_short_circuit_test(op, jump_target); + self.emit_short_circuit_test(op, after_block); + self.switch_to_block(continue_block); emit!(self, Instruction::PopTop); } - // If all values did not qualify, take the value of the last value: self.compile_expression(last_value)?; self.switch_to_block(after_block); Ok(()) } + fn compile_bool_op_with_head_constant( + &mut self, + op: &ast::BoolOp, + head: ConstantData, + tail: &[ast::Expr], + ) -> CompileResult<()> { + self.emit_load_const(head); + self.mark_last_instruction_folded_from_nonliteral_expr(); + if tail.is_empty() { + return Ok(()); + } + + let after_block = self.new_block(); + for value in tail { + self.emit_short_circuit_test(op, after_block); + emit!(self, Instruction::PopTop); + self.compile_expression(value)?; + } + self.switch_to_block(after_block); + Ok(()) + } + /// Emit `Copy 1` + conditional jump for short-circuit evaluation. /// For `And`, emits `PopJumpIfFalse`; for `Or`, emits `PopJumpIfTrue`. fn emit_short_circuit_test(&mut self, op: &ast::BoolOp, target: BlockIdx) { @@ -8001,15 +8009,31 @@ impl Compiler { let range = expression.range(); self.set_source_range(range); + if let ast::Expr::Subscript(ast::ExprSubscript { + ctx: ast::ExprContext::Load, + .. + }) = expression + && let Some(constant) = self.try_fold_constant_expr(expression)? + { + self.emit_load_const(constant); + return Ok(()); + } + if !self.disable_const_boolop_folding && let ast::Expr::BoolOp(ast::ExprBoolOp { op, values, .. }) = expression { let mut simplified_prefix = 0usize; let mut last_constant = None; + let mut retained_head = None; for value in values { let Some(constant) = self.try_fold_constant_expr(value)? else { break; }; + if !Self::boolop_fast_fold_literal(value) { + retained_head = Some(constant); + simplified_prefix += 1; + break; + } let is_truthy = Self::constant_truthiness(&constant); last_constant = Some(constant); match op { @@ -8029,6 +8053,10 @@ impl Compiler { } } + if let Some(head) = retained_head { + self.compile_bool_op_with_head_constant(op, head, &values[simplified_prefix..])?; + return Ok(()); + } if simplified_prefix == values.len() { self.emit_load_const(last_constant.expect("missing folded boolop constant")); self.mark_last_instruction_folded_from_nonliteral_expr(); @@ -8458,7 +8486,12 @@ impl Compiler { // True case self.compile_expression(body)?; - emit!(self, PseudoInstruction::Jump { delta: after_block }); + self.set_no_location(); + emit!( + self, + PseudoInstruction::JumpNoInterrupt { delta: after_block } + ); + self.set_no_location(); // False case self.switch_to_block(else_block); @@ -9153,18 +9186,16 @@ impl Compiler { let is_inlined = self.is_inlined_comprehension_context(comprehension_type, &comp_table); if is_inlined && !has_an_async_gen && !element_contains_await { - // PEP 709: Inlined comprehension - compile inline without new scope - let was_in_inlined_comp = self.current_code_info().in_inlined_comp; - self.current_code_info().in_inlined_comp = true; - let result = self.compile_inlined_comprehension( + // PEP 709: Inlined comprehension - compile inline without new scope. + // CPython compiles the outermost iterable before entering the + // inlined-comprehension fast-hidden scope tweak. + return self.compile_inlined_comprehension( comp_table, init_collection, generators, compile_element, has_an_async_gen, ); - self.current_code_info().in_inlined_comp = was_in_inlined_comp; - return result; } // Non-inlined path: create a new code object (generator expressions, etc.) @@ -9373,6 +9404,22 @@ impl Compiler { compile_element: &dyn Fn(&mut Self) -> CompileResult<()>, has_async: bool, ) -> CompileResult<()> { + fn collect_bound_names(target: &ast::Expr, out: &mut Vec) { + match target { + ast::Expr::Name(ast::ExprName { id, .. }) => out.push(id.to_string()), + ast::Expr::Tuple(ast::ExprTuple { elts, .. }) + | ast::Expr::List(ast::ExprList { elts, .. }) => { + for elt in elts { + collect_bound_names(elt, out); + } + } + ast::Expr::Starred(ast::ExprStarred { value, .. }) => { + collect_bound_names(value, out); + } + _ => {} + } + } + // Compile the outermost iterator first. Its expression may reference // nested scopes (e.g. lambdas) whose sub_tables sit at the current // position in the parent's list. Those must be consumed before we @@ -9386,288 +9433,306 @@ impl Compiler { .expect("no current symbol table") .next_sub_table += 1; - // Splice the comprehension's children (e.g. nested inlined - // comprehensions) into the parent so the compiler can find them. - if !comp_table.sub_tables.is_empty() { - let current_table = self - .symbol_table_stack - .last_mut() - .expect("no current symbol table"); - let insert_pos = current_table.next_sub_table; - for (i, st) in comp_table.sub_tables.iter().enumerate() { - current_table.sub_tables.insert(insert_pos + i, st.clone()); - } - } - if has_async && generators[0].is_async { - emit!(self, Instruction::GetAIter); - } else { - emit!(self, Instruction::GetIter); - } - - // Collect local variables that need to be saved/restored. - // All DEF_LOCAL && !DEF_NONLOCAL names from the comp table, plus class block names. + let was_in_inlined_comp = self.current_code_info().in_inlined_comp; let in_class_block = { let ct = self.current_symbol_table(); - ct.typ == CompilerScope::Class && !self.current_code_info().in_inlined_comp + ct.typ == CompilerScope::Class && !was_in_inlined_comp }; - fn collect_bound_names(target: &ast::Expr, out: &mut Vec) { - match target { - ast::Expr::Name(ast::ExprName { id, .. }) => out.push(id.to_string()), - ast::Expr::Tuple(ast::ExprTuple { elts, .. }) - | ast::Expr::List(ast::ExprList { elts, .. }) => { - for elt in elts { - collect_bound_names(elt, out); - } - } - ast::Expr::Starred(ast::ExprStarred { value, .. }) => { - collect_bound_names(value, out); - } - _ => {} + self.current_code_info().in_inlined_comp = true; + + let mut temp_symbols: IndexMap = IndexMap::default(); + let mut changed_fast_hidden = Vec::new(); + + let result = (|| { + // Splice the comprehension's children (e.g. nested inlined + // comprehensions) into the parent so the compiler can find them. + if !comp_table.sub_tables.is_empty() { + let current_table = self + .symbol_table_stack + .last_mut() + .expect("no current symbol table"); + let insert_pos = current_table.next_sub_table; + for (i, st) in comp_table.sub_tables.iter().enumerate() { + current_table.sub_tables.insert(insert_pos + i, st.clone()); + } + } + if has_async && generators[0].is_async { + emit!(self, Instruction::GetAIter); + } else { + emit!(self, Instruction::GetIter); } - } - let mut source_order_bound_names = Vec::new(); - for generator in generators { - collect_bound_names(&generator.target, &mut source_order_bound_names); - } - let mut pushed_locals: Vec = Vec::new(); - for name in source_order_bound_names - .into_iter() - .chain(comp_table.symbols.keys().cloned()) - { - if pushed_locals.iter().any(|existing| existing == &name) { - continue; + + let mut source_order_bound_names = Vec::new(); + for generator in generators { + collect_bound_names(&generator.target, &mut source_order_bound_names); } - if let Some(sym) = comp_table.symbols.get(&name) { - if sym.flags.contains(SymbolFlags::PARAMETER) { + + let mut pushed_locals: Vec = Vec::new(); + for name in source_order_bound_names + .into_iter() + .chain(comp_table.symbols.keys().cloned()) + { + if pushed_locals.iter().any(|existing| existing == &name) { + continue; + } + if let Some(sym) = comp_table.symbols.get(&name) { + if sym.flags.contains(SymbolFlags::PARAMETER) { + continue; // skip .0 + } + // Walrus operator targets (ASSIGNED_IN_COMPREHENSION without ITER) + // are not local to the comprehension; they leak to the outer scope. + let is_walrus = sym.flags.contains(SymbolFlags::ASSIGNED_IN_COMPREHENSION) + && !sym.flags.contains(SymbolFlags::ITER); + let is_local = sym + .flags + .intersects(SymbolFlags::ASSIGNED | SymbolFlags::ITER) + && !sym.flags.contains(SymbolFlags::NONLOCAL) + && !is_walrus; + if is_local { + pushed_locals.push(name); + } + } + } + + // TweakInlinedComprehensionScopes: temporarily override parent + // symbols with comprehension scopes where they differ. For + // module/class scopes, also enable temporary fast locals for + // comprehension-bound names only. + for (name, comp_sym) in &comp_table.symbols { + if comp_sym.flags.contains(SymbolFlags::PARAMETER) { continue; // skip .0 } - // Walrus operator targets (ASSIGNED_IN_COMPREHENSION without ITER) - // are not local to the comprehension; they leak to the outer scope. - let is_walrus = sym.flags.contains(SymbolFlags::ASSIGNED_IN_COMPREHENSION) - && !sym.flags.contains(SymbolFlags::ITER); - let is_local = sym - .flags - .intersects(SymbolFlags::ASSIGNED | SymbolFlags::ITER) - && !sym.flags.contains(SymbolFlags::NONLOCAL) - && !is_walrus; - if is_local || in_class_block { - pushed_locals.push(name); + let comp_scope = comp_sym.scope; + + let current_table = self.symbol_table_stack.last().expect("no symbol table"); + if let Some(outer_sym) = current_table.symbols.get(name) { + let outer_scope = outer_sym.scope; + if (comp_scope != outer_scope + && comp_scope != SymbolScope::Free + && !(comp_scope == SymbolScope::Cell && outer_scope == SymbolScope::Free)) + || in_class_block + { + temp_symbols.insert(name.clone(), outer_sym.clone()); + let current_table = + self.symbol_table_stack.last_mut().expect("no symbol table"); + current_table.symbols.insert(name.clone(), comp_sym.clone()); + } + } + } + if !self.ctx.in_func() { + for name in &pushed_locals { + if self + .current_code_info() + .metadata + .fast_hidden + .get(name.as_str()) + .is_none_or(|&hidden| !hidden) + { + self.current_code_info() + .metadata + .fast_hidden + .insert(name.clone(), true); + changed_fast_hidden.push(name.clone()); + } } } - } - // TweakInlinedComprehensionScopes: temporarily override parent symbols - // with comp scopes where they differ. - let mut temp_symbols: IndexMap = IndexMap::default(); - for (name, comp_sym) in &comp_table.symbols { - if comp_sym.flags.contains(SymbolFlags::PARAMETER) { - continue; // skip .0 - } - let comp_scope = comp_sym.scope; - - let current_table = self.symbol_table_stack.last().expect("no symbol table"); - if let Some(outer_sym) = current_table.symbols.get(name) { - let outer_scope = outer_sym.scope; - if (comp_scope != outer_scope - && comp_scope != SymbolScope::Free - && !(comp_scope == SymbolScope::Cell && outer_scope == SymbolScope::Free)) - || in_class_block - { - temp_symbols.insert(name.clone(), outer_sym.clone()); - let current_table = - self.symbol_table_stack.last_mut().expect("no symbol table"); - current_table.symbols.insert(name.clone(), comp_sym.clone()); - } - } - } - - // Step 2: Save local variables that will be shadowed by the comprehension. - // For each variable, we push the fast local value via LoadFastAndClear. - // For merged CELL variables, LoadFastAndClear saves the cell object from - // the merged slot, and MAKE_CELL creates a new empty cell in-place. - // MAKE_CELL has no stack effect (operates only on fastlocals). - let mut total_stack_items: usize = 0; - for name in &pushed_locals { - let var_num = self.varname(name)?; - emit!(self, Instruction::LoadFastAndClear { var_num }); - total_stack_items += 1; - // If the comp symbol is CELL, emit MAKE_CELL to create fresh cell - if let Some(comp_sym) = comp_table.symbols.get(name) - && comp_sym.scope == SymbolScope::Cell - { - let i = if self - .current_symbol_table() - .symbols - .get(name) - .is_some_and(|s| s.scope == SymbolScope::Free) + // Step 2: Save local variables that will be shadowed by the comprehension. + // For each variable, we push the fast local value via LoadFastAndClear. + // For merged CELL variables, LoadFastAndClear saves the cell object from + // the merged slot, and MAKE_CELL creates a new empty cell in-place. + // MAKE_CELL has no stack effect (operates only on fastlocals). + let mut total_stack_items: usize = 0; + for name in &pushed_locals { + let var_num = self.varname(name)?; + emit!(self, Instruction::LoadFastAndClear { var_num }); + total_stack_items += 1; + // If the comp symbol is CELL, emit MAKE_CELL to create fresh cell + if let Some(comp_sym) = comp_table.symbols.get(name) + && comp_sym.scope == SymbolScope::Cell { - self.get_free_var_index(name)? - } else { - self.get_cell_var_index(name)? - }; - emit!(self, Instruction::MakeCell { i }); + let i = if self + .current_symbol_table() + .symbols + .get(name) + .is_some_and(|s| s.scope == SymbolScope::Free) + { + self.get_free_var_index(name)? + } else { + self.get_cell_var_index(name)? + }; + emit!(self, Instruction::MakeCell { i }); + } } - } - // Step 3: SWAP iterator to TOS (above saved locals + cell values) - if total_stack_items > 0 { - emit!( - self, - Instruction::Swap { - i: u32::try_from(total_stack_items + 1).unwrap() - } - ); - } + // Step 3: SWAP iterator to TOS (above saved locals + cell values) + if total_stack_items > 0 { + emit!( + self, + Instruction::Swap { + i: u32::try_from(total_stack_items + 1).unwrap() + } + ); + } - // Step 4: Create the collection (list/set/dict) - if let Some(init_collection) = init_collection { - self._emit(init_collection, OpArg::new(0), BlockIdx::NULL); - // SWAP to get iterator on top - emit!(self, Instruction::Swap { i: 2 }); - } + // Step 4: Create the collection (list/set/dict) + if let Some(init_collection) = init_collection { + self._emit(init_collection, OpArg::new(0), BlockIdx::NULL); + // SWAP to get iterator on top + emit!(self, Instruction::Swap { i: 2 }); + } - // Set up exception handler for cleanup on exception - let cleanup_block = self.new_block(); - let end_block = self.new_block(); + // Set up exception handler for cleanup on exception + let cleanup_block = self.new_block(); + let end_block = self.new_block(); - if !pushed_locals.is_empty() { - emit!( - self, - PseudoInstruction::SetupFinally { - delta: cleanup_block + if !pushed_locals.is_empty() { + emit!( + self, + PseudoInstruction::SetupFinally { + delta: cleanup_block + } + ); + self.push_fblock(FBlockType::TryExcept, cleanup_block, end_block)?; + } + + // Step 5: Compile the comprehension loop(s) + let mut loop_labels: Vec<(BlockIdx, BlockIdx, BlockIdx, bool, BlockIdx)> = vec![]; + for (i, generator) in generators.iter().enumerate() { + let loop_block = self.new_block(); + let if_cleanup_block = self.new_block(); + let after_block = self.new_block(); + + if i > 0 { + self.compile_for_iterable_expression(&generator.iter, generator.is_async)?; + if generator.is_async { + emit!(self, Instruction::GetAIter); + } else { + emit!(self, Instruction::GetIter); + } } - ); - self.push_fblock(FBlockType::TryExcept, cleanup_block, end_block)?; - } - // Step 5: Compile the comprehension loop(s) - let mut loop_labels: Vec<(BlockIdx, BlockIdx, BlockIdx, bool, BlockIdx)> = vec![]; - for (i, generator) in generators.iter().enumerate() { - let loop_block = self.new_block(); - let if_cleanup_block = self.new_block(); - let after_block = self.new_block(); + self.switch_to_block(loop_block); - if i > 0 { - self.compile_for_iterable_expression(&generator.iter, generator.is_async)?; + let mut end_async_for_target = BlockIdx::NULL; if generator.is_async { - emit!(self, Instruction::GetAIter); + emit!(self, PseudoInstruction::SetupFinally { delta: after_block }); + emit!(self, Instruction::GetANext); + self.push_fblock( + FBlockType::AsyncComprehensionGenerator, + loop_block, + after_block, + )?; + self.emit_load_const(ConstantData::None); + end_async_for_target = self.compile_yield_from_sequence(true)?; + emit!(self, PseudoInstruction::PopBlock); + self.pop_fblock(FBlockType::AsyncComprehensionGenerator); + self.compile_store(&generator.target)?; } else { - emit!(self, Instruction::GetIter); + emit!(self, Instruction::ForIter { delta: after_block }); + self.compile_store(&generator.target)?; } - } - - self.switch_to_block(loop_block); - let mut end_async_for_target = BlockIdx::NULL; - if generator.is_async { - emit!(self, PseudoInstruction::SetupFinally { delta: after_block }); - emit!(self, Instruction::GetANext); - self.push_fblock( - FBlockType::AsyncComprehensionGenerator, + loop_labels.push(( loop_block, + if_cleanup_block, after_block, - )?; - self.emit_load_const(ConstantData::None); - end_async_for_target = self.compile_yield_from_sequence(true)?; - emit!(self, PseudoInstruction::PopBlock); - self.pop_fblock(FBlockType::AsyncComprehensionGenerator); - self.compile_store(&generator.target)?; - } else { - emit!(self, Instruction::ForIter { delta: after_block }); - self.compile_store(&generator.target)?; - } - - loop_labels.push(( - loop_block, - if_cleanup_block, - after_block, - generator.is_async, - end_async_for_target, - )); + generator.is_async, + end_async_for_target, + )); - // CPython always lowers comprehension guards through codegen_jump_if - // and leaves constant-folding to later CFG optimization passes. - for if_condition in &generator.ifs { - self.compile_jump_if(if_condition, false, if_cleanup_block)?; + // CPython always lowers comprehension guards through codegen_jump_if + // and leaves constant-folding to later CFG optimization passes. + for if_condition in &generator.ifs { + self.compile_jump_if(if_condition, false, if_cleanup_block)?; + } } - } - // Step 6: Compile the element expression and append to collection - compile_element(self)?; + // Step 6: Compile the element expression and append to collection + compile_element(self)?; - // Step 7: Close all loops - for &(loop_block, if_cleanup_block, after_block, is_async, end_async_for_target) in - loop_labels.iter().rev() - { - emit!(self, PseudoInstruction::Jump { delta: loop_block }); + // Step 7: Close all loops + for &(loop_block, if_cleanup_block, after_block, is_async, end_async_for_target) in + loop_labels.iter().rev() + { + emit!(self, PseudoInstruction::Jump { delta: loop_block }); - self.switch_to_block(if_cleanup_block); - emit!(self, PseudoInstruction::Jump { delta: loop_block }); + self.switch_to_block(if_cleanup_block); + emit!(self, PseudoInstruction::Jump { delta: loop_block }); - self.switch_to_block(after_block); - if is_async { - self.emit_end_async_for(end_async_for_target); - } else { - emit!(self, Instruction::EndFor); - emit!(self, Instruction::PopIter); + self.switch_to_block(after_block); + if is_async { + self.emit_end_async_for(end_async_for_target); + } else { + emit!(self, Instruction::EndFor); + emit!(self, Instruction::PopIter); + } } - } - // Step 8: Clean up - restore saved locals (and cell values) - if total_stack_items > 0 { - emit!(self, PseudoInstruction::PopBlock); - self.pop_fblock(FBlockType::TryExcept); + // Step 8: Clean up - restore saved locals (and cell values) + if total_stack_items > 0 { + emit!(self, PseudoInstruction::PopBlock); + self.pop_fblock(FBlockType::TryExcept); - // Normal path: jump past cleanup - emit!(self, PseudoInstruction::Jump { delta: end_block }); + // Normal path: jump past cleanup + emit!(self, PseudoInstruction::Jump { delta: end_block }); - // Exception cleanup path - self.switch_to_block(cleanup_block); - // Stack: [saved_values..., collection, exception] - emit!(self, Instruction::Swap { i: 2 }); - emit!(self, Instruction::PopTop); // Pop incomplete collection + // Exception cleanup path + self.switch_to_block(cleanup_block); + // Stack: [saved_values..., collection, exception] + emit!(self, Instruction::Swap { i: 2 }); + emit!(self, Instruction::PopTop); // Pop incomplete collection - // Restore locals and cell values - emit!( - self, - Instruction::Swap { - i: u32::try_from(total_stack_items + 1).unwrap() + // Restore locals and cell values + emit!( + self, + Instruction::Swap { + i: u32::try_from(total_stack_items + 1).unwrap() + } + ); + for name in pushed_locals.iter().rev() { + let var_num = self.varname(name)?; + emit!(self, Instruction::StoreFast { var_num }); } - ); + // Re-raise the exception + emit!(self, Instruction::Reraise { depth: 0 }); + + // Normal end path + self.switch_to_block(end_block); + } + + // SWAP result to TOS (above saved values) + if total_stack_items > 0 { + emit!( + self, + Instruction::Swap { + i: u32::try_from(total_stack_items + 1).unwrap() + } + ); + } + + // Restore saved locals (StoreFast restores the saved cell object for merged cells) for name in pushed_locals.iter().rev() { let var_num = self.varname(name)?; emit!(self, Instruction::StoreFast { var_num }); } - // Re-raise the exception - emit!(self, Instruction::Reraise { depth: 0 }); - // Normal end path - self.switch_to_block(end_block); - } - - // SWAP result to TOS (above saved values) - if total_stack_items > 0 { - emit!( - self, - Instruction::Swap { - i: u32::try_from(total_stack_items + 1).unwrap() - } - ); - } - - // Restore saved locals (StoreFast restores the saved cell object for merged cells) - for name in pushed_locals.iter().rev() { - let var_num = self.varname(name)?; - emit!(self, Instruction::StoreFast { var_num }); - } + Ok(()) + })(); - // RevertInlinedComprehensionScopes: restore original symbols let current_table = self.symbol_table_stack.last_mut().expect("no symbol table"); for (name, original_sym) in temp_symbols { current_table.symbols.insert(name, original_sym); } + for name in changed_fast_hidden { + self.current_code_info() + .metadata + .fast_hidden + .insert(name, false); + } + self.current_code_info().in_inlined_comp = was_in_inlined_comp; - Ok(()) + result } fn compile_future_features(&mut self, features: &[ast::Alias]) -> Result<(), CodegenError> { @@ -10010,6 +10075,9 @@ impl Compiler { (ast::UnaryOp::Invert, ConstantData::Integer { value }) => { ConstantData::Integer { value: !value } } + (ast::UnaryOp::Not, value) => ConstantData::Boolean { + value: !Self::constant_truthiness(&value), + }, _ => return Ok(None), } } @@ -10094,6 +10162,13 @@ impl Compiler { emit!(self, Instruction::ReturnValue) } + fn emit_return_const_no_location(&mut self, constant: ConstantData) { + self.emit_load_const(constant); + self.set_no_location(); + emit!(self, Instruction::ReturnValue); + self.set_no_location(); + } + fn emit_end_async_for(&mut self, send_target: BlockIdx) { self._emit(Instruction::EndAsyncFor, OpArg::NULL, send_target); } @@ -10257,6 +10332,7 @@ impl Compiler { enum UnwindAction { With { is_async: bool, + range: TextRange, }, HandlerCleanup { name: Option, @@ -10276,10 +10352,16 @@ impl Compiler { for i in (loop_idx + 1..code.fblock.len()).rev() { match code.fblock[i].fb_type { FBlockType::With => { - unwind_actions.push(UnwindAction::With { is_async: false }); + unwind_actions.push(UnwindAction::With { + is_async: false, + range: code.fblock[i].fb_range, + }); } FBlockType::AsyncWith => { - unwind_actions.push(UnwindAction::With { is_async: true }); + unwind_actions.push(UnwindAction::With { + is_async: true, + range: code.fblock[i].fb_range, + }); } FBlockType::HandlerCleanup => { let name = match &code.fblock[i].fb_datum { @@ -10316,8 +10398,10 @@ impl Compiler { // Emit cleanup for each fblock for action in unwind_actions { match action { - UnwindAction::With { is_async } => { + UnwindAction::With { is_async, range } => { // Stack: [..., exit_func, self_exit] + let saved_range = self.current_source_range; + self.set_source_range(range); emit!(self, PseudoInstruction::PopBlock); self.emit_load_const(ConstantData::None); self.emit_load_const(ConstantData::None); @@ -10331,6 +10415,7 @@ impl Compiler { } emit!(self, Instruction::PopTop); + self.set_source_range(saved_range); } UnwindAction::HandlerCleanup { ref name } => { // codegen_unwind_fblock(HANDLER_CLEANUP) @@ -11299,6 +11384,23 @@ mod tests { compiler.exit_scope() } + fn scan_program_symbol_table(source: &str) -> SymbolTable { + let source_file = SourceFileBuilder::new("source_path", source).finish(); + let parsed = ruff_python_parser::parse( + source_file.source_text(), + ruff_python_parser::Mode::Module.into(), + ) + .unwrap(); + let ast = parsed.into_syntax(); + let ast = match ast { + ruff_python_ast::Mod::Module(stmts) => stmts, + _ => unreachable!(), + }; + SymbolTable::scan_program(&ast, source_file) + .map_err(|e| e.into_codegen_error("source_path".to_owned())) + .unwrap() + } + fn compile_exec_late_cfg_trace(source: &str) -> Vec<(String, String)> { let opts = CompileOpts::default(); let source_file = SourceFileBuilder::new("source_path", source).finish(); @@ -11350,30 +11452,38 @@ mod tests { }) .unwrap_or_else(|| panic!("missing function {function_name}")); + let name = &function.name; + let parameters = &function.parameters; + let body = &function.body; + let is_async = function.is_async; + let range = function.range(); + let mut compiler = Compiler::new(opts, source_file, "".to_owned()); compiler.future_annotations = symbol_table.future_annotations; compiler.symbol_table_stack.push(symbol_table); - compiler.set_source_range(function.range()); - compiler - .enter_function(function.name.as_str(), &function.parameters) - .unwrap(); + compiler.set_source_range(range); + compiler.enter_function(name.as_str(), parameters).unwrap(); compiler .current_code_info() .flags - .set(bytecode::CodeFlags::COROUTINE, false); + .set(bytecode::CodeFlags::COROUTINE, is_async); let prev_ctx = compiler.ctx; compiler.ctx = CompileContext { loop_data: None, in_class: prev_ctx.in_class, - func: FunctionContext::Function, - in_async_scope: false, + func: if is_async { + FunctionContext::AsyncFunction + } else { + FunctionContext::Function + }, + in_async_scope: is_async, }; compiler.set_qualname(); - compiler.compile_statements(&function.body).unwrap(); - match function.body.last() { + compiler.compile_statements(body).unwrap(); + match body.last() { Some(ast::Stmt::Return(_)) => {} - _ => compiler.emit_return_const(ConstantData::None), + _ => compiler.emit_return_const_no_location(ConstantData::None), } if compiler.current_code_info().metadata.consts.is_empty() { compiler.arg_constant(ConstantData::None); @@ -11489,6 +11599,51 @@ def f(self): } } + #[test] + fn test_import_originated_name_disables_method_call_optimization_even_with_local_import() { + let code = compile_exec( + "\ +import warnings + +def f(ch): + import warnings + warnings.warn( + '\"\\\\%c\" is an invalid escape sequence' % ch + if 0x20 <= ch < 0x7F + else '\"\\\\x%02x\" is an invalid escape sequence' % ch, + DeprecationWarning, + stacklevel=2, + ) +", + ); + let f = find_code(&code, "f").expect("missing f code"); + let ops: Vec<_> = f.instructions.iter().map(|unit| unit.op).collect(); + let warn_attr = ops + .iter() + .position(|op| matches!(op, Instruction::LoadAttr { .. })) + .expect("missing LOAD_ATTR for warnings.warn"); + let push_null = ops[warn_attr + 10..] + .iter() + .position(|op| matches!(op, Instruction::PushNull)) + .map(|idx| warn_attr + 10 + idx) + .expect("expected PUSH_NULL after plain LOAD_ATTR"); + + let load_attr = match f.instructions[warn_attr].op { + Instruction::LoadAttr { namei } => namei.get(OpArg::new(u32::from(u8::from( + f.instructions[warn_attr].arg, + )))), + _ => unreachable!(), + }; + assert!( + !load_attr.is_method(), + "import-originated names should use plain LOAD_ATTR" + ); + assert!( + matches!(ops[push_null + 1], Instruction::LoadSmallInt { .. }), + "expected warning message expression to start after PUSH_NULL, got ops={ops:?}" + ); + } + #[test] fn test_trace_constant_false_elif_chain() { let trace = compile_exec_late_cfg_trace( @@ -11783,6 +11938,44 @@ def outer(null): ); } + #[test] + fn test_nonliteral_constant_bool_op_preserves_short_circuit_shape() { + let code = compile_exec( + "\ +x = (\"a\"[0]) or 2 +", + ); + let ops: Vec<_> = code + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + assert!( + !code.instructions.iter().any(|unit| matches!( + unit.op, + Instruction::BinaryOp { op } + if op.get(OpArg::new(u32::from(u8::from(unit.arg)))) + == oparg::BinaryOperator::Subscr + )), + "constant subscript should fold before bool-op lowering, got ops={ops:?}" + ); + assert!( + ops.iter().any(|op| matches!(op, Instruction::Copy { .. })), + "folded non-literal BoolOp operand should keep COPY, got ops={ops:?}" + ); + assert!( + ops.iter().any(|op| matches!(op, Instruction::ToBool)), + "folded non-literal BoolOp operand should keep TO_BOOL, got ops={ops:?}" + ); + assert!( + ops.iter() + .any(|op| matches!(op, Instruction::PopJumpIfTrue { .. })), + "folded non-literal BoolOp operand should keep POP_JUMP_IF_TRUE, got ops={ops:?}" + ); + } + #[test] fn test_nested_double_async_with() { assert_dis_snapshot!(compile_exec( @@ -12076,22 +12269,39 @@ def f(obj): .filter(|op| !matches!(op, Instruction::Cache)) .collect(); - let cond_idx = ops - .iter() - .position(|op| matches!(op, Instruction::PopJumpIfNotNone { .. })) - .expect("missing POP_JUMP_IF_NOT_NONE"); - assert!( - matches!(ops.get(cond_idx + 1), Some(Instruction::NotTaken)), - "expected NOT_TAKEN after conditional jump, got {:?}; ops={ops:?}", - ops.get(cond_idx + 1) - ); - assert!( + let has_cpython_shape = ops.windows(7).any(|window| { matches!( - ops.get(cond_idx + 2), - Some(Instruction::JumpBackward { .. }) - ), - "expected loop backedge immediately after NOT_TAKEN, got {:?}; ops={ops:?}", - ops.get(cond_idx + 2) + window, + [ + Instruction::PopJumpIfNotNone { .. }, + Instruction::NotTaken, + Instruction::JumpBackward { .. }, + Instruction::LoadFastBorrow { .. } | Instruction::LoadFast { .. }, + Instruction::Swap { .. }, + Instruction::PopTop, + Instruction::ReturnValue, + ] + ) + }); + let has_conservative_shape = ops.windows(9).any(|window| { + matches!( + window, + [ + Instruction::PopJumpIfNone { .. }, + Instruction::NotTaken, + Instruction::LoadFastBorrow { .. } | Instruction::LoadFast { .. }, + Instruction::Swap { .. }, + Instruction::PopTop, + Instruction::ReturnValue, + Instruction::Nop, + Instruction::JumpBackward { .. }, + Instruction::EndFor, + ] + ) + }); + assert!( + has_cpython_shape || has_conservative_shape, + "expected loop return null-check to keep the backedge adjacent to the return cleanup, got ops={ops:?}" ); let end_for_idx = ops @@ -12187,6 +12397,80 @@ def outer(): ); } + #[test] + fn test_try_line_nop_is_preserved_before_setup_finally() { + let code = compile_exec( + "\ +def f(msg): + try: + fw = _wm.formatwarning + except AttributeError: + pass + else: + if fw is not _formatwarning_orig: + return fw(msg.message, msg.category, msg.filename, msg.lineno, msg.line) + return _wm._formatwarnmsg_impl(msg) +", + ); + 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(); + + assert!( + matches!( + ops.as_slice(), + [Instruction::Resume { .. }, Instruction::Nop, ..] + ), + "expected CPython try-line NOP before setup/fetch, got ops={ops:?}" + ); + } + + #[test] + fn test_try_else_return_keeps_nop_before_final_call_return() { + let code = compile_exec( + "\ +def f(msg): + try: + fw = _wm.formatwarning + except AttributeError: + pass + else: + if fw is not _formatwarning_orig: + return fw(msg.message, msg.category, msg.filename, msg.lineno, msg.line) + return _wm._formatwarnmsg_impl(msg) +", + ); + 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(); + + assert!( + ops.windows(7).any(|window| { + matches!( + window, + [ + Instruction::ReturnValue, + Instruction::Nop, + Instruction::LoadGlobal { .. }, + Instruction::LoadAttr { .. }, + Instruction::LoadFastBorrow { .. } | Instruction::LoadFast { .. }, + Instruction::Call { .. }, + Instruction::ReturnValue, + ] + ) + }), + "expected CPython-style NOP between conditional return and final call return, got ops={ops:?}" + ); + } + #[test] fn test_conditional_compare_uses_bool_compare_oparg() { let code = compile_exec( @@ -12950,8 +13234,171 @@ def f(w): } #[test] - fn test_genexpr_true_filter_omits_bool_scaffolding() { - let code = compile_exec( + #[ignore = "debug trace for loop bool-chain jump-back layout"] + fn test_debug_trace_loop_break_bool_chain_layout() { + let trace = compile_single_function_late_cfg_trace( + "\ +def f(filters, text, category, module, lineno, defaultaction): + for item in filters: + action, msg, cat, mod, ln = item + if ((msg is None or msg.match(text)) and + issubclass(category, cat) and + (mod is None or mod.match(module)) and + (ln == 0 or lineno == ln)): + break + else: + action = defaultaction + return action +", + "f", + ); + for (stage, dump) in trace { + eprintln!("=== {stage} ===\n{dump}"); + } + } + + #[test] + #[ignore = "debug trace for loop conditional body jump-back layout"] + fn test_debug_trace_loop_conditional_body_layout() { + let trace = compile_single_function_late_cfg_trace( + "\ +def f(new, old): + for replace in ['__module__', '__name__', '__qualname__', '__doc__']: + if hasattr(old, replace): + setattr(new, replace, getattr(old, replace)) + return new +", + "f", + ); + for (stage, dump) in trace { + eprintln!("=== {stage} ===\n{dump}"); + } + } + + #[test] + #[ignore = "debug trace for minimized utf7 encode nested-if layout"] + fn test_debug_trace_utf7_min_encode_layout() { + let trace = compile_single_function_late_cfg_trace( + "\ +def f(s, size, encodeSetO, encodeWhiteSpace): + inShift = True + base64bits = 0 + out = [] + for i, ch in enumerate(s): + if base64bits == 0: + if i + 1 < size: + ch2 = s[i + 1] + if E(ch2, encodeSetO, encodeWhiteSpace): + if B(ch2) or ch2 == '-': + out.append(b'-') + inShift = False + else: + out.append(b'-') + inShift = False + return out +", + "f", + ); + for (stage, dump) in trace { + eprintln!("=== {stage} ===\n{dump}"); + } + } + + #[test] + #[ignore = "debug trace for with-protected loop bool-chain layout"] + fn test_debug_trace_with_loop_break_bool_chain_layout() { + let trace = compile_single_function_late_cfg_trace( + "\ +def f(filters, text, category, module, lineno, defaultaction, _wm): + with _wm._lock: + for item in filters: + action, msg, cat, mod, ln = item + if ((msg is None or msg.match(text)) and + issubclass(category, cat) and + (mod is None or mod.match(module)) and + (ln == 0 or lineno == ln)): + break + else: + action = defaultaction + return action +", + "f", + ); + for (stage, dump) in trace { + eprintln!("=== {stage} ===\n{dump}"); + } + } + + #[test] + fn test_nested_boolop_same_or_prefixes_compile_without_extra_boolop_block() { + let code = compile_exec( + "\ +def f(c, encodeO, encodeWS): + return ( + (c > 127 or utf7_special[c] == 1) + or (encodeWS and (utf7_special[c] == 2)) + or (encodeO and (utf7_special[c] == 3)) + ) +", + ); + let f = find_code(&code, "f").expect("missing function code"); + let pop_jump_if_true_count = f + .instructions + .iter() + .filter(|unit| matches!(unit.op, Instruction::PopJumpIfTrue { .. })) + .count(); + + assert!( + pop_jump_if_true_count >= 3, + "expected nested boolop prefix path to compile short-circuit jumps, got ops={:?}", + f.instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + } + + #[test] + fn test_broad_exception_import_keeps_borrow_in_common_tail() { + let code = compile_exec( + "\ +def f(msg): + if msg.source is not None: + try: + import tracemalloc + except Exception: + suggest_tracemalloc = False + tb = None + suggest_tracemalloc = not tracemalloc.is_tracing() + tb = tracemalloc.get_object_traceback(msg.source) + if tb is not None: + for frame in tb: + pass + return 0 +", + ); + let f = find_code(&code, "f").expect("missing function code"); + let import_idx = f + .instructions + .iter() + .position(|unit| matches!(unit.op, Instruction::ImportName { .. })) + .expect("missing IMPORT_NAME"); + + assert!( + f.instructions[import_idx + 1..] + .iter() + .any(|unit| matches!(unit.op, Instruction::LoadFastBorrow { .. })), + "expected common tail after broad-exception import to keep LOAD_FAST_BORROW, got ops={:?}", + f.instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + } + + #[test] + fn test_genexpr_true_filter_omits_bool_scaffolding() { + let code = compile_exec( "\ def f(it): return (x for x in it if True) @@ -13340,7 +13787,10 @@ def f(names, cls): .filter(|unit| matches!(unit.op, Instruction::ReturnValue)) .count(); - assert_eq!(return_count, 1); + assert_eq!( + return_count, 2, + "expected CPython-style distinct return sites for normal and except paths" + ); } #[test] @@ -13385,6 +13835,58 @@ def f(p, s): ); } + #[test] + fn test_try_else_if_return_keeps_conditional_target_nop() { + let code = compile_exec( + "\ +def f(cond): + try: + x = cond + except E: + pass + else: + if x: + return 1 + return 2 +", + ); + 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 has_cpython_nop_target = ops.windows(5).any(|window| { + matches!( + window, + [ + Instruction::LoadSmallInt { .. } | Instruction::LoadConst { .. }, + Instruction::ReturnValue, + Instruction::Nop, + Instruction::LoadSmallInt { .. } | Instruction::LoadConst { .. }, + Instruction::ReturnValue, + ] + ) + }); + let has_direct_fallthrough = ops.windows(4).any(|window| { + matches!( + window, + [ + Instruction::LoadSmallInt { .. } | Instruction::LoadConst { .. }, + Instruction::ReturnValue, + Instruction::LoadSmallInt { .. } | Instruction::LoadConst { .. }, + Instruction::ReturnValue, + ] + ) + }); + assert!( + has_cpython_nop_target || has_direct_fallthrough, + "expected adjacent try-else return and final return targets, got ops={ops:?}" + ); + } + #[test] fn test_named_except_conditional_branch_duplicates_cleanup_return() { let code = compile_exec( @@ -13466,6 +13968,85 @@ def f(escaped_string, quote_types): ); } + #[test] + fn test_static_swap_triple_assign_keeps_store_fast_store_fast() { + let code = compile_exec( + "\ +def f(x, y, z): + a, b, a = x, y, z + return a +", + ); + 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(); + + assert!( + ops.windows(3).any(|window| { + matches!( + window, + [ + Instruction::Swap { .. }, + Instruction::StoreFastStoreFast { .. }, + Instruction::StoreFast { .. } + ] + ) + }), + "expected CPython-style SWAP/STORE_FAST_STORE_FAST/STORE_FAST sequence, got ops={ops:?}" + ); + } + + #[test] + fn test_constant_ifexp_stmt_in_loop_removes_empty_body() { + let code = compile_exec( + "\ +def f(x): + while x: + 0 if 1 else 0 +", + ); + 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(); + + assert!( + !ops.iter() + .any(|op| matches!(op, Instruction::LoadSmallInt { .. })), + "expected constant if-expression statement to compile away inside loop, got ops={ops:?}" + ); + } + + #[test] + fn test_ifexp_in_jump_context_skips_constant_true_arm_load() { + let code = compile_exec( + "\ +def f(): + a if (1 if b else c) else d +", + ); + 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(); + + assert!( + !ops.iter() + .any(|op| matches!(op, Instruction::LoadSmallInt { .. })), + "expected jump-context if-expression to avoid materializing constant truthy arm, got ops={ops:?}" + ); + } + #[test] fn test_with_suppress_tail_duplicates_final_return_none() { let code = compile_exec( @@ -13501,7 +14082,46 @@ def f(cm, cond): } #[test] - fn test_genexpr_compare_header_keeps_split_store_then_borrow_load() { + fn test_with_conditional_bare_return_keeps_return_line_nop_before_exit_cleanup() { + let code = compile_exec( + "\ +def f(cm, registry, altkey): + with cm: + if registry.get(altkey): + return + registry[altkey] = 1 +", + ); + 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(); + + assert!( + ops.windows(8).any(|window| { + matches!( + window, + [ + Instruction::Nop, + Instruction::LoadConst { .. }, + Instruction::LoadConst { .. }, + Instruction::LoadConst { .. }, + Instruction::Call { .. }, + Instruction::PopTop, + Instruction::LoadConst { .. }, + Instruction::ReturnValue, + ] + ) + }), + "expected CPython-style return-line NOP before with-exit cleanup return, got ops={ops:?}" + ); + } + + #[test] + fn test_genexpr_compare_header_uses_store_fast_load_fast_like_cpython() { let code = compile_exec( "\ def f(it): @@ -13517,23 +14137,17 @@ def f(it): .collect(); assert!( - !ops.iter() - .any(|op| matches!(op, Instruction::StoreFastLoadFast { .. })), - "expected compare header to keep split STORE_FAST/LOAD_FAST_BORROW, got ops={ops:?}" - ); - assert!( - ops.windows(4).any(|window| { + ops.windows(3).any(|window| { matches!( window, [ - Instruction::StoreFast { .. }, - Instruction::LoadFastBorrow { .. }, + Instruction::StoreFastLoadFast { .. }, Instruction::LoadConst { .. }, Instruction::CompareOp { .. }, ] ) }), - "expected split compare header sequence, got ops={ops:?}" + "expected CPython-style STORE_FAST_LOAD_FAST compare header, got ops={ops:?}" ); } @@ -13722,6 +14336,57 @@ values = [item for item in [r\"\\\\'a\\\\'\", r\"\\t3\", r\"\\\\\"[0]]]\n", ); } + #[test] + fn test_constant_subscript_folds_in_load_context() { + let cases = [ + ("value = (1, 2, 3)[0]\n", Some(BigInt::from(1)), None), + ("value = b\"abc\"[0]\n", Some(BigInt::from(97)), None), + ("value = \"abc\"[0]\n", None, Some("a")), + ]; + + for (source, expected_int, expected_str) in cases { + let code = compile_exec(source); + assert!( + !code.instructions.iter().any(|unit| matches!( + unit.op, + Instruction::BinaryOp { op } + if op.get(OpArg::new(u32::from(u8::from(unit.arg)))) + == oparg::BinaryOperator::Subscr + )), + "expected folded constant subscript for {source:?}, got instructions={:?}", + code.instructions + ); + + if let Some(expected_int) = expected_int.as_ref() { + let has_small_int = code.instructions.iter().any(|unit| { + matches!( + unit.op, + Instruction::LoadSmallInt { i } + if BigInt::from(i.get(OpArg::new(u32::from(u8::from(unit.arg))))) + == *expected_int + ) + }); + let has_const_int = code.constants.iter().any(|constant| { + matches!(constant, ConstantData::Integer { value } if value == expected_int) + }); + assert!( + has_small_int || has_const_int, + "missing folded integer constant {expected_int} for {source:?}, instructions={:?}", + code.instructions + ); + } + + if let Some(expected_str) = expected_str { + assert!( + code.constants.iter().any(|constant| { + matches!(constant, ConstantData::Str { value } if value.to_string() == expected_str) + }), + "missing folded string constant {expected_str:?} for {source:?}", + ); + } + } + } + #[test] fn test_list_of_constant_tuples_uses_list_extend() { let code = compile_exec( @@ -13739,27 +14404,105 @@ deprecated_cases = [('a', 'b'), ('c', 'd'), ('e', 'f'), ('g', 'h'), ('i', 'j')] } #[test] - fn test_constant_list_iterable_uses_tuple() { + fn test_large_list_of_unary_constants_uses_list_extend() { let code = compile_exec( "\ -def f(): - return {x: y for x, y in [(1, 2), ]} +values = [-1, not True, ~0, +True, 5] ", ); - let f = find_code(&code, "f").expect("missing function code"); assert!( - !f.instructions + code.instructions .iter() - .any(|unit| matches!(unit.op, Instruction::BuildList { .. })), - "constant list iterable should avoid BUILD_LIST before GET_ITER" + .any(|unit| matches!(unit.op, Instruction::ListExtend { .. })), + "expected unary-folded constants to participate in list folding, got instructions={:?}", + code.instructions ); - assert!(f.constants.iter().any(|constant| matches!( + assert!(code.constants.iter().any(|constant| matches!( constant, ConstantData::Tuple { elements } - if matches!( - elements.as_slice(), - [ConstantData::Tuple { elements: inner }] + if elements.len() == 5 + && matches!(&elements[0], ConstantData::Integer { value } if *value == BigInt::from(-1)) + && matches!(&elements[1], ConstantData::Boolean { value } if !value) + && matches!(&elements[2], ConstantData::Integer { value } if *value == BigInt::from(-1)) + && matches!(&elements[3], ConstantData::Integer { value } if *value == BigInt::from(1)) + && matches!(&elements[4], ConstantData::Integer { value } if *value == BigInt::from(5)) + ))); + } + + #[test] + fn test_large_constant_list_keeps_streaming_build() { + let source = format!( + "values = [{}]\n", + (0..31) + .map(|i| format!("'v{i}'")) + .collect::>() + .join(", ") + ); + let code = compile_exec(&source); + + assert!( + code.instructions + .iter() + .any(|unit| matches!(unit.op, Instruction::ListAppend { .. })), + "large constant lists should keep LIST_APPEND streaming form, got instructions={:?}", + code.instructions + ); + assert!( + !code + .instructions + .iter() + .any(|unit| matches!(unit.op, Instruction::ListExtend { .. })), + "large constant lists should not fold to LIST_EXTEND, got instructions={:?}", + code.instructions + ); + } + + #[test] + fn test_constant_tuple_binops_fold_like_cpython() { + let code = compile_exec("value = (1,) * 17 + ('spam',)\n"); + + assert!( + !code + .instructions + .iter() + .any(|unit| matches!(unit.op, Instruction::BinaryOp { .. })), + "tuple constant binops should fold away, got instructions={:?}", + code.instructions + ); + assert!(code.constants.iter().any(|constant| matches!( + constant, + ConstantData::Tuple { elements } + if elements.len() == 18 + && elements[..17] + .iter() + .all(|elt| matches!(elt, ConstantData::Integer { value } if *value == BigInt::from(1))) + && matches!(&elements[17], ConstantData::Str { value } if value.to_string() == "spam") + ))); + } + + #[test] + fn test_constant_list_iterable_uses_tuple() { + let code = compile_exec( + "\ +def f(): + return {x: y for x, y in [(1, 2), ]} +", + ); + let f = find_code(&code, "f").expect("missing function code"); + + assert!( + !f.instructions + .iter() + .any(|unit| matches!(unit.op, Instruction::BuildList { .. })), + "constant list iterable should avoid BUILD_LIST before GET_ITER" + ); + assert!(f.constants.iter().any(|constant| matches!( + constant, + ConstantData::Tuple { elements } + if matches!( + elements.as_slice(), + [ConstantData::Tuple { elements: inner }] if matches!( inner.as_slice(), [ @@ -13772,7 +14515,7 @@ def f(): } #[test] - fn test_constant_set_iterable_keeps_runtime_set_build() { + fn test_constant_set_iterable_uses_frozenset_const() { let code = compile_exec( "\ def f(): @@ -13782,10 +14525,10 @@ def f(): let f = find_code(&code, "f").expect("missing function code"); assert!( - f.instructions + !f.instructions .iter() .any(|unit| matches!(unit.op, Instruction::BuildSet { .. })), - "constant set iterable should keep BUILD_SET before GET_ITER" + "constant set iterable should avoid BUILD_SET before GET_ITER" ); assert!(f.constants.iter().any(|constant| matches!( constant, @@ -13801,6 +14544,89 @@ def f(): ))); } + #[test] + fn test_constant_list_membership_uses_tuple_const() { + let code = compile_exec( + "\ +f = lambda x: x in [1, 2, 3] +", + ); + let lambda = find_code(&code, "").expect("missing lambda code"); + + assert!( + !lambda + .instructions + .iter() + .any(|unit| matches!(unit.op, Instruction::BuildList { .. })), + "constant list membership should avoid BUILD_LIST before CONTAINS_OP" + ); + assert!(lambda.constants.iter().any(|constant| matches!( + constant, + ConstantData::Tuple { elements } + if matches!( + elements.as_slice(), + [ + ConstantData::Integer { .. }, + ConstantData::Integer { .. }, + ConstantData::Integer { .. } + ] + ) + ))); + } + + #[test] + fn test_small_constant_set_membership_uses_frozenset_const() { + let code = compile_exec( + "\ +f = lambda x: x in {0} +", + ); + let lambda = find_code(&code, "").expect("missing lambda code"); + + assert!( + !lambda + .instructions + .iter() + .any(|unit| matches!(unit.op, Instruction::BuildSet { .. })), + "constant set membership should avoid BUILD_SET before CONTAINS_OP" + ); + assert!(lambda.constants.iter().any(|constant| matches!( + constant, + ConstantData::Frozenset { elements } + if matches!(elements.as_slice(), [ConstantData::Integer { value }] if *value == BigInt::from(0)) + ))); + } + + #[test] + fn test_nonconstant_list_membership_uses_tuple() { + let code = compile_exec( + "\ +def f(a, b, c, x): + return x in [a, b, c] +", + ); + let f = find_code(&code, "f").expect("missing f code"); + let ops: Vec<_> = f + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + assert!( + ops.windows(2).any(|window| { + matches!( + window, + [ + Instruction::BuildTuple { .. }, + Instruction::ContainsOp { .. } + ] + ) + }), + "expected BUILD_TUPLE before CONTAINS_OP for non-constant list membership, got ops={ops:?}" + ); + } + #[test] fn test_starred_tuple_iterable_drops_list_to_tuple_before_get_iter() { let code = compile_exec( @@ -13825,15 +14651,654 @@ def f(a, b, c): } #[test] - fn test_comprehension_single_list_iterable_uses_tuple() { + fn test_comprehension_single_list_iterable_uses_tuple() { + let code = compile_exec( + "\ +def g(): + [x for x in [(yield 1)]] +", + ); + let g = find_code(&code, "g").expect("missing g code"); + let ops: Vec<_> = g + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + assert!( + ops.windows(2).any(|window| { + matches!( + window, + [Instruction::BuildTuple { .. }, Instruction::GetIter] + ) + }), + "expected BUILD_TUPLE before GET_ITER for single-item list iterable in comprehension, got ops={ops:?}" + ); + } + + #[test] + fn test_nested_comprehension_list_iterable_uses_tuple() { + let code = compile_exec( + "\ +def f(): + return [[y for y in [x, x + 1]] for x in [1, 3, 5]] +", + ); + let f = find_code(&code, "f").expect("missing f code"); + let ops: Vec<_> = f + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + assert!( + ops.windows(2).any(|window| { + matches!( + window, + [Instruction::BuildTuple { .. }, Instruction::GetIter] + ) + }), + "expected BUILD_TUPLE before GET_ITER for nested list iterable in comprehension, got ops={ops:?}" + ); + } + + #[test] + fn test_constant_comprehension_iterable_with_unary_int_uses_tuple_const() { + let code = compile_exec( + "\ +l = lambda : [2 < x for x in [-1, 3, 0]] +", + ); + let lambda = find_code(&code, "").expect("missing lambda code"); + + assert!( + lambda.constants.iter().any(|constant| matches!( + constant, + ConstantData::Tuple { elements } + if matches!( + elements.as_slice(), + [ + ConstantData::Integer { .. }, + ConstantData::Integer { .. }, + ConstantData::Integer { .. } + ] + ) + )), + "expected folded tuple constant for comprehension iterable" + ); + } + + #[test] + fn test_module_scope_listcomp_is_inlined() { + let code = compile_exec("values = [i for i in range(3)]\n"); + + assert!( + find_code(&code, "").is_none(), + "module-scope list comprehension should be inlined" + ); + assert!( + code.instructions + .iter() + .any(|unit| matches!(unit.op, Instruction::LoadFastAndClear { .. })), + "inlined module-scope list comprehension should use LOAD_FAST_AND_CLEAR, got instructions={:?}", + code.instructions + ); + } + + #[test] + fn test_module_scope_dictcomp_is_inlined() { + let code = compile_exec("mapping = {i: i for i in range(3)}\n"); + + assert!( + find_code(&code, "").is_none(), + "module-scope dict comprehension should be inlined" + ); + assert!( + code.instructions + .iter() + .any(|unit| matches!(unit.op, Instruction::LoadFastAndClear { .. })), + "inlined module-scope dict comprehension should use LOAD_FAST_AND_CLEAR, got instructions={:?}", + code.instructions + ); + } + + #[test] + fn test_nested_module_scope_dictcomp_symbols_are_local() { + let symbol_table = scan_program_symbol_table( + "\ +deoptmap = { + specialized: base + for base, family in _specializations.items() + for specialized in family +} +", + ); + + for name in ["base", "family", "specialized"] { + let symbol = symbol_table + .lookup(name) + .unwrap_or_else(|| panic!("missing module symbol {name}")); + assert_eq!( + symbol.scope, + SymbolScope::Local, + "expected module-scope inlined comprehension symbol {name} to be Local, got {symbol:?}" + ); + } + + let comp = symbol_table + .sub_tables + .first() + .expect("missing comprehension symbol table"); + assert!(comp.comp_inlined, "expected comprehension to be inlined"); + for name in ["base", "family", "specialized"] { + let symbol = comp + .lookup(name) + .unwrap_or_else(|| panic!("missing comprehension symbol {name}")); + assert_eq!( + symbol.scope, + SymbolScope::Local, + "expected comprehension symbol {name} to be Local, got {symbol:?}" + ); + } + } + + #[test] + fn test_nested_module_scope_dictcomp_uses_fast_locals() { + let code = compile_exec( + "\ +deoptmap = { + specialized: base + for base, family in _specializations.items() + for specialized in family +} +", + ); + let ops: Vec<_> = code + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + assert!( + ops.iter() + .any(|op| matches!(op, Instruction::StoreFastStoreFast { .. })), + "expected outer target unpack to use STORE_FAST_STORE_FAST, got ops={ops:?}" + ); + assert!( + ops.iter().any(|op| matches!( + op, + Instruction::StoreFastLoadFast { .. } + | Instruction::LoadFastBorrowLoadFastBorrow { .. } + )), + "expected inner target/store-use path to use fast locals, got ops={ops:?}" + ); + assert!( + ops.iter() + .filter(|op| matches!(op, Instruction::LoadName { .. })) + .count() + <= 1, + "unexpected extra LOAD_NAME ops in nested inlined comprehension, got ops={ops:?}" + ); + assert!( + ops.iter() + .filter(|op| matches!(op, Instruction::StoreName { .. })) + .count() + <= 1, + "unexpected extra STORE_NAME ops in nested inlined comprehension, got ops={ops:?}" + ); + } + + #[test] + fn test_module_scope_inlined_comprehension_keeps_outer_iter_as_name_lookup() { + let code = compile_exec( + "\ +path_separators = ['/'] +_pathseps_with_colon = {f':{s}' for s in path_separators} +", + ); + let ops: Vec<_> = code + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + let load_name_path = ops + .windows(2) + .any(|window| matches!(window, [Instruction::LoadName { .. }, Instruction::GetIter])); + assert!( + load_name_path, + "expected outer iterable to stay a NAME lookup before GET_ITER, got ops={ops:?}" + ); + assert!( + !ops.windows(2).any(|window| matches!( + window, + [Instruction::LoadFast { .. }, Instruction::GetIter] + | [Instruction::LoadFastCheck { .. }, Instruction::GetIter] + )), + "module local outer iterable should not become a fast local, got ops={ops:?}" + ); + assert!( + ops.iter().any(|op| matches!( + op, + Instruction::StoreFastLoadFast { .. } | Instruction::StoreFast { .. } + )), + "comprehension target should still use fast locals, got ops={ops:?}" + ); + } + + #[test] + fn test_or_condition_in_jump_context_uses_shared_true_fallthrough() { + let code = compile_exec( + "\ +def f(lines): + for line in lines: + if line.startswith('--') or not line.strip(): + continue + return line +", + ); + let f = find_code(&code, "f").expect("missing f code"); + let ops: Vec<_> = f + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + let first_pop_jump = ops + .iter() + .find(|op| { + matches!( + op, + Instruction::PopJumpIfTrue { .. } | Instruction::PopJumpIfFalse { .. } + ) + }) + .copied() + .expect("missing conditional jump"); + assert!( + matches!(first_pop_jump, Instruction::PopJumpIfTrue { .. }), + "expected first OR branch to jump on true into shared fallthrough, got ops={ops:?}" + ); + } + + #[test] + fn test_loop_break_bool_chain_reorders_false_path_to_jump_back() { + let code = compile_exec( + "\ +def f(filters, text, category, module, lineno, defaultaction): + for item in filters: + action, msg, cat, mod, ln = item + if ((msg is None or msg.match(text)) and + issubclass(category, cat) and + (mod is None or mod.match(module)) and + (ln == 0 or lineno == ln)): + break + else: + action = defaultaction + return action +", + ); + let f = find_code(&code, "f").expect("missing f code"); + let ops: Vec<_> = f + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + assert!( + ops.windows(5).any(|window| { + matches!( + window, + [ + Instruction::ToBool, + Instruction::PopJumpIfTrue { .. }, + Instruction::NotTaken, + Instruction::JumpBackward { .. } + | Instruction::JumpBackwardNoInterrupt { .. }, + Instruction::LoadGlobal { .. }, + ] + ) + }), + "expected CPython-style false path to fall through into loop jump-back, got ops={ops:?}" + ); + } + + #[test] + fn test_loop_conditional_body_keeps_duplicate_jump_back_paths() { + let code = compile_exec( + "\ +def f(new, old): + for replace in ['__module__', '__name__', '__qualname__', '__doc__']: + if hasattr(old, replace): + setattr(new, replace, getattr(old, replace)) + return new +", + ); + let f = find_code(&code, "f").expect("missing f code"); + let ops: Vec<_> = f + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + let jump_back_count = ops + .iter() + .filter(|op| { + matches!( + op, + Instruction::JumpBackward { .. } | Instruction::JumpBackwardNoInterrupt { .. } + ) + }) + .count(); + assert!( + jump_back_count >= 2, + "expected separate false-path and body jump-back blocks, got ops={ops:?}" + ); + assert!( + ops.windows(5).any(|window| { + matches!( + window, + [ + Instruction::ToBool, + Instruction::PopJumpIfTrue { .. }, + Instruction::NotTaken, + Instruction::JumpBackward { .. } + | Instruction::JumpBackwardNoInterrupt { .. }, + Instruction::LoadGlobal { .. }, + ] + ) + }), + "expected false path to jump back before body, got ops={ops:?}" + ); + } + + #[test] + fn test_loop_if_pass_uses_line_bearing_jump_back_instead_of_nop() { + let code = compile_exec( + "\ +def f(x, y): + for i in x: + if y: + pass +", + ); + let f = find_code(&code, "f").expect("missing f code"); + let ops: Vec<_> = f + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + assert!( + ops.windows(5).any(|window| { + matches!( + window, + [ + Instruction::ToBool, + Instruction::PopJumpIfTrue { .. }, + Instruction::NotTaken, + Instruction::JumpBackward { .. } + | Instruction::JumpBackwardNoInterrupt { .. }, + Instruction::JumpBackward { .. } + | Instruction::JumpBackwardNoInterrupt { .. }, + ] + ) + }), + "expected CPython-style synthetic false-path jump-back plus body jump-back, got ops={ops:?}" + ); + assert!( + !ops.iter().any(|op| matches!(op, Instruction::Nop)), + "expected pass body line to attach to loop backedge instead of leaving a NOP, got ops={ops:?}" + ); + } + + #[test] + fn test_nested_if_shared_jump_back_target_is_duplicated() { + let code = compile_exec( + "\ +def f(s, size, encodeSetO, encodeWhiteSpace): + inShift = True + base64bits = 0 + out = [] + for i, ch in enumerate(s): + if base64bits == 0: + if i + 1 < size: + ch2 = s[i + 1] + if E(ch2, encodeSetO, encodeWhiteSpace): + if B(ch2) or ch2 == '-': + out.append(b'-') + inShift = False + else: + out.append(b'-') + inShift = False + return out +", + ); + let f = find_code(&code, "f").expect("missing f code"); + let ops: Vec<_> = f + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + assert!( + ops.windows(6).any(|window| { + matches!( + window, + [ + Instruction::PopTop, + Instruction::LoadConst { .. }, + Instruction::StoreFast { .. }, + Instruction::JumpBackward { .. } + | Instruction::JumpBackwardNoInterrupt { .. }, + Instruction::JumpBackward { .. } + | Instruction::JumpBackwardNoInterrupt { .. }, + Instruction::LoadFast { .. } | Instruction::LoadFastBorrow { .. }, + ] + ) + }), + "expected separate nested-if and outer-if jump-back tails, got ops={ops:?}" + ); + } + + #[test] + fn test_protected_loop_conditional_keeps_forward_body_entry() { + let code = compile_exec( + "\ +def outer(it, C1): + def f(): + for x in it: + try: + if C1: + yield 2 + except OSError: + pass + return f +", + ); + let outer = find_code(&code, "outer").expect("missing outer code"); + let f = find_code(outer, "f").expect("missing f code"); + let ops: Vec<_> = f + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + assert!( + ops.windows(7).any(|window| { + matches!( + window, + [ + Instruction::ToBool, + Instruction::PopJumpIfFalse { .. }, + Instruction::NotTaken, + Instruction::LoadSmallInt { .. }, + Instruction::YieldValue { .. }, + Instruction::Resume { .. }, + Instruction::PopTop, + ] + ) + }), + "expected protected conditional to keep CPython-style forward body entry, got ops={ops:?}" + ); + } + + #[test] + fn test_nested_except_false_path_duplicates_pop_except_jump_back_tail() { + let code = compile_exec( + "\ +def f(it, C3): + for x in it: + try: + X = 3 + except OSError: + try: + if C3: + X = 4 + except OSError: + pass + return 42 +", + ); + let f = find_code(&code, "f").expect("missing f code"); + let ops: Vec<_> = f + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + assert!( + ops.windows(6).any(|window| { + matches!( + window, + [ + Instruction::LoadSmallInt { .. }, + Instruction::StoreFast { .. }, + Instruction::PopExcept, + Instruction::JumpBackward { .. } + | Instruction::JumpBackwardNoInterrupt { .. }, + Instruction::PopExcept, + Instruction::JumpBackward { .. } + | Instruction::JumpBackwardNoInterrupt { .. }, + ] + ) + }), + "expected CPython-style duplicated false-path exit tail, got ops={ops:?}" + ); + } + + #[test] + fn test_more_nested_except_false_paths_duplicate_all_jump_back_tails() { + let code = compile_exec( + "\ +def f(it, C3, C4): + for x in it: + try: + X = 3 + except OSError: + try: + if C3: + if C4: + X = 4 + except OSError: + try: + if C3: + if C4: + X = 5 + except OSError: + pass + return 42 +", + ); + let f = find_code(&code, "f").expect("missing f code"); + let ops: Vec<_> = f + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + assert!( + ops.windows(8).any(|window| { + matches!( + window, + [ + Instruction::LoadSmallInt { .. }, + Instruction::StoreFast { .. }, + Instruction::PopExcept, + Instruction::JumpBackward { .. } + | Instruction::JumpBackwardNoInterrupt { .. }, + Instruction::PopExcept, + Instruction::JumpBackward { .. } + | Instruction::JumpBackwardNoInterrupt { .. }, + Instruction::PopExcept, + Instruction::JumpBackward { .. } + | Instruction::JumpBackwardNoInterrupt { .. }, + ] + ) + }), + "expected CPython-style duplicated nested false-path exit tails, got ops={ops:?}" + ); + } + + #[test] + fn test_no_wraparound_jump_keeps_forward_hop_before_loop_backedge() { + let code = compile_exec( + "\ +def while_not_chained(a, b, c): + while not (a < b < c): + pass +", + ); + let f = find_code(&code, "while_not_chained").expect("missing while_not_chained code"); + let ops: Vec<_> = f + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + assert!( + ops.windows(5).any(|window| { + matches!( + window, + [ + Instruction::PopJumpIfTrue { .. }, + Instruction::NotTaken, + Instruction::JumpForward { .. }, + Instruction::PopTop, + Instruction::JumpBackward { .. } + | Instruction::JumpBackwardNoInterrupt { .. }, + ] + ) + }), + "expected CPython-style no-wraparound forward hop before the loop backedge, got ops={ops:?}" + ); + } + + #[test] + fn test_and_is_not_none_loop_guard_uses_direct_jump_back_false_path() { let code = compile_exec( "\ -def g(): - [x for x in [(yield 1)]] +def f(code): + last_line = -2 + for _, _, line in code.co_lines(): + if line is not None and line != last_line: + last_line = line ", ); - let g = find_code(&code, "g").expect("missing g code"); - let ops: Vec<_> = g + let f = find_code(&code, "f").expect("missing f code"); + let ops: Vec<_> = f .instructions .iter() .map(|unit| unit.op) @@ -13841,22 +15306,34 @@ def g(): .collect(); assert!( - ops.windows(2).any(|window| { + ops.windows(6).any(|window| { matches!( window, - [Instruction::BuildTuple { .. }, Instruction::GetIter] + [ + Instruction::LoadFastBorrow { .. } | Instruction::LoadFast { .. }, + Instruction::PopJumpIfNotNone { .. }, + Instruction::NotTaken, + Instruction::JumpBackward { .. } + | Instruction::JumpBackwardNoInterrupt { .. }, + Instruction::LoadFastBorrowLoadFastBorrow { .. } + | Instruction::LoadFastLoadFast { .. }, + Instruction::CompareOp { .. }, + ] ) }), - "expected BUILD_TUPLE before GET_ITER for single-item list iterable in comprehension, got ops={ops:?}" + "expected CPython-style direct jump-back false path for 'is not None and ...', got ops={ops:?}" ); } #[test] - fn test_nested_comprehension_list_iterable_uses_tuple() { + fn test_continue_inside_with_keeps_line_marker_nop_before_exit_cleanup() { let code = compile_exec( "\ -def f(): - return [[y for y in [x, x + 1]] for x in [1, 3, 5]] +def f(it): + for func in it: + with cm(): + if cond(): + continue ", ); let f = find_code(&code, "f").expect("missing f code"); @@ -13868,39 +15345,88 @@ def f(): .collect(); assert!( - ops.windows(2).any(|window| { + ops.windows(9).any(|window| { matches!( window, - [Instruction::BuildTuple { .. }, Instruction::GetIter] + [ + Instruction::PopJumpIfFalse { .. }, + Instruction::NotTaken, + Instruction::Nop, + Instruction::LoadConst { .. }, + Instruction::LoadConst { .. }, + Instruction::LoadConst { .. }, + Instruction::Call { .. }, + Instruction::PopTop, + Instruction::JumpBackward { .. } + | Instruction::JumpBackwardNoInterrupt { .. }, + ] ) }), - "expected BUILD_TUPLE before GET_ITER for nested list iterable in comprehension, got ops={ops:?}" + "expected CPython-style line-marker NOP before with-exit cleanup on continue, got ops={ops:?}" ); } #[test] - fn test_constant_comprehension_iterable_with_unary_int_uses_tuple_const() { + fn test_try_loop_elif_places_return_before_orelse_tail() { let code = compile_exec( "\ -l = lambda : [2 < x for x in [-1, 3, 0]] +def f(source, suggest, tb, s): + if source is not None: + try: + tb = tb + except Exception: + suggest = False + tb = None + if tb is not None: + for frame in tb: + s += frame + elif suggest: + s += 'x' + return s ", ); - let lambda = find_code(&code, "").expect("missing lambda code"); + let f = find_code(&code, "f").expect("missing f code"); + let ops: Vec<_> = f + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + let has_direct_return = ops.windows(8).any(|window| { + matches!( + window, + [ + Instruction::EndFor, + Instruction::PopIter, + Instruction::LoadFastBorrow { .. } | Instruction::LoadFast { .. }, + Instruction::ReturnValue, + Instruction::LoadFastBorrow { .. } | Instruction::LoadFast { .. }, + Instruction::ToBool, + Instruction::PopJumpIfFalse { .. }, + Instruction::NotTaken, + ] + ) + }); + let has_nop_anchored_return = ops.windows(9).any(|window| { + matches!( + window, + [ + Instruction::EndFor, + Instruction::PopIter, + Instruction::Nop, + Instruction::LoadFastBorrow { .. } | Instruction::LoadFast { .. }, + Instruction::ReturnValue, + Instruction::LoadFastBorrow { .. } | Instruction::LoadFast { .. }, + Instruction::ToBool, + Instruction::PopJumpIfFalse { .. }, + Instruction::NotTaken, + ] + ) + }); assert!( - lambda.constants.iter().any(|constant| matches!( - constant, - ConstantData::Tuple { elements } - if matches!( - elements.as_slice(), - [ - ConstantData::Integer { .. }, - ConstantData::Integer { .. }, - ConstantData::Integer { .. } - ] - ) - )), - "expected folded tuple constant for comprehension iterable" + has_direct_return || has_nop_anchored_return, + "expected CPython-style duplicated return between loop exit and elif tail, got ops={ops:?}" ); } @@ -14177,4 +15703,385 @@ else: .collect::>() ); } + + #[test] + fn test_nested_try_else_multi_resume_join_keeps_strong_load_fast_tail() { + let code = compile_exec( + "\ +def f(msg): + s = '' + try: + import a + except Exception: + suggest = False + tb = None + else: + try: + suggest = not t() + tb = g(msg) + except Exception: + suggest = False + tb = None + if tb is not None: + for frame in tb: + s += frame + elif suggest: + s += 'y' + return s +", + ); + let f = find_code(&code, "f").expect("missing f code"); + let ops: Vec<_> = f + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + let tail_start = ops + .iter() + .position(|op| matches!(op, Instruction::PopJumpIfNone { .. })) + .expect("missing tail POP_JUMP_IF_NONE") + .saturating_sub(1); + let handler_start = ops + .iter() + .position(|op| matches!(op, Instruction::PushExcInfo)) + .expect("missing handler entry"); + let tail = &ops[tail_start..handler_start]; + + assert!( + !tail.iter().any(|op| { + matches!( + op, + Instruction::LoadFastBorrow { .. } + | Instruction::LoadFastBorrowLoadFastBorrow { .. } + ) + }), + "expected nested try/except else-resume tail to keep strong LOAD_FAST ops, got tail={tail:?}" + ); + + assert!( + tail.iter() + .any(|op| matches!(op, Instruction::LoadFastLoadFast { .. })), + "expected loop body to keep LOAD_FAST_LOAD_FAST in the resume tail, got tail={tail:?}" + ); + } + + #[test] + fn test_protected_conditional_tail_keeps_strong_load_fast() { + let code = compile_exec( + "\ +def f(m, klass, category, warning_base): + try: + cat = getattr(m, klass) + except AttributeError: + raise ValueError(category) + if not issubclass(cat, warning_base): + raise TypeError(category) + return cat +", + ); + let f = find_code(&code, "f").expect("missing f code"); + let ops: Vec<_> = f + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + let tail_start = ops + .iter() + .position(|op| matches!(op, Instruction::StoreFast { .. })) + .expect("missing STORE_FAST cat"); + let handler_start = ops + .iter() + .position(|op| matches!(op, Instruction::PushExcInfo)) + .expect("missing handler entry"); + let tail = &ops[tail_start + 1..handler_start]; + + assert!( + !tail.iter().any(|op| { + matches!( + op, + Instruction::LoadFastBorrow { .. } + | Instruction::LoadFastBorrowLoadFastBorrow { .. } + ) + }), + "expected protected conditional tail to keep strong LOAD_FAST ops, got tail={tail:?}" + ); + + assert!( + tail.iter() + .any(|op| matches!(op, Instruction::LoadFastLoadFast { .. })), + "expected protected tail to keep LOAD_FAST_LOAD_FAST for issubclass args, got tail={tail:?}" + ); + } + + #[test] + fn test_protected_import_tail_keeps_strong_load_fast() { + let code = compile_exec( + "\ +def f(s, size, pos, errors): + message = 'x' + look = pos + try: + import unicodedata + except ImportError: + return None + if look < size and chr(s[look]) == '{': + while look < size and chr(s[look]) != '}': + look += 1 + if look > pos + 1 and look < size and chr(s[look]) == '}': + message = 'y' + return message +", + ); + let f = find_code(&code, "f").expect("missing f code"); + let ops: Vec<_> = f + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + let import_idx = ops + .iter() + .position(|op| matches!(op, Instruction::ImportName { .. })) + .expect("missing IMPORT_NAME"); + let protected_tail = &ops[import_idx + 1..]; + + assert!( + !protected_tail.iter().any(|op| { + matches!( + op, + Instruction::LoadFastBorrow { .. } + | Instruction::LoadFastBorrowLoadFastBorrow { .. } + ) + }), + "expected protected import tail to keep strong LOAD_FAST ops, got tail={protected_tail:?}" + ); + + assert!( + protected_tail + .iter() + .any(|op| matches!(op, Instruction::LoadFastLoadFast { .. })), + "expected protected import tail to keep LOAD_FAST_LOAD_FAST ops, got tail={protected_tail:?}" + ); + } + + #[test] + fn test_handler_resume_join_keeps_strong_load_fast_common_tail() { + let code = compile_exec( + "\ +def f(p, errors, s, pos, look, final, escape_start, st): + try: + chr_codec = unicodedata.lookup('%s' % st) + except LookupError as e: + x = unicode_call_errorhandler( + errors, 'unicodeescape', 'unknown Unicode character name', s, pos - 1, look + 1 + ) + else: + x = chr_codec, look + 1 + p.append(x[0]) + pos = x[1] + if not final: + pos = escape_start + return p, pos + return unicode_call_errorhandler( + errors, 'unicodeescape', 'unknown Unicode character name', s, pos - 1, look + 1 + ) +", + ); + let f = find_code(&code, "f").expect("missing f code"); + let append_idx = f + .instructions + .iter() + .position(|unit| match unit.op { + Instruction::LoadAttr { namei } => { + let load_attr = namei.get(OpArg::new(u32::from(u8::from(unit.arg)))); + f.names[usize::try_from(load_attr.name_idx()).unwrap()].as_str() == "append" + } + _ => false, + }) + .expect("missing append tail"); + let tail: Vec<_> = f.instructions[append_idx.saturating_sub(1)..] + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + assert!( + !tail.iter().any(|op| { + matches!( + op, + Instruction::LoadFastBorrow { .. } + | Instruction::LoadFastBorrowLoadFastBorrow { .. } + ) + }), + "expected handler resume common tail to keep strong LOAD_FAST ops, got tail={tail:?}" + ); + + assert!( + tail.iter() + .any(|op| matches!(op, Instruction::LoadFastLoadFast { .. })), + "expected handler resume common tail to keep LOAD_FAST_LOAD_FAST ops, got tail={tail:?}" + ); + } + + #[test] + fn test_named_except_cleanup_loop_header_keeps_borrow_in_for_loop() { + let code = compile_exec( + "\ +def f(args): + for arg in args: + try: + _wm._setoption(arg) + except _wm._OptionError as msg: + print('Invalid -W option ignored:', msg, file=sys.stderr) +", + ); + let f = find_code(&code, "f").expect("missing f code"); + let attr_idx = f + .instructions + .iter() + .position(|unit| match unit.op { + Instruction::LoadAttr { namei } => { + let load_attr = namei.get(OpArg::new(u32::from(u8::from(unit.arg)))); + f.names[usize::try_from(load_attr.name_idx()).unwrap()].as_str() == "_setoption" + } + _ => false, + }) + .expect("missing _setoption attr load"); + let window: Vec<_> = f.instructions[attr_idx + 1..] + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .take(3) + .collect(); + assert!( + matches!( + window.as_slice(), + [ + Instruction::LoadFastBorrow { .. }, + Instruction::Call { .. }, + Instruction::PopTop + ] + ), + "expected loop body call to keep borrowed arg load after named-except cleanup, got window={window:?}" + ); + } + + #[test] + fn test_named_except_cleanup_deopts_same_guard_fallbacks_not_outer_tail() { + let code = compile_exec( + r#" +def f(s, size, errors, final): + found_invalid_escape = False + p = [] + pos = 0 + while pos < size: + ch = chr(s[pos]) + pos += 1 + if ch == "N": + message = "malformed \\N character escape" + look = pos + try: + import unicodedata + except ImportError: + message = "\\N escapes not supported (can't load unicodedata module)" + unicode_call_errorhandler( + errors, "unicodeescape", message, s, pos - 1, size + ) + continue + if look < size and chr(s[look]) == "{": + while look < size and chr(s[look]) != "}": + look += 1 + if look > pos + 1 and look < size and chr(s[look]) == "}": + message = "unknown Unicode character name" + st = s[pos + 1 : look] + try: + chr_codec = unicodedata.lookup("%s" % st) + except LookupError as e: + x = unicode_call_errorhandler( + errors, "unicodeescape", message, s, pos - 1, look + 1 + ) + else: + x = chr_codec, look + 1 + p.append(x[0]) + pos = x[1] + else: + if not final: + pos = 0 + break + x = unicode_call_errorhandler( + errors, "unicodeescape", message, s, pos - 1, look + 1 + ) + p.append(x[0]) + pos = x[1] + else: + if not final: + pos = 0 + break + x = unicode_call_errorhandler( + errors, "unicodeescape", message, s, pos - 1, look + 1 + ) + p.append(x[0]) + pos = x[1] + else: + if not found_invalid_escape: + found_invalid_escape = True + warnings.warn( + "invalid escape sequence '\\%c'" % ch, DeprecationWarning, 2 + ) + p.append("\\") + p.append(ch) + return p, pos +"#, + ); + let f = find_code(&code, "f").expect("missing f code"); + + let mut saw_strong_final = false; + let mut saw_borrow_p_after_warn = false; + let mut saw_borrow_ch_after_warn = false; + let mut after_warn_attr = false; + + for unit in f.instructions.iter() { + match unit.op { + Instruction::LoadAttr { namei } => { + let load_attr = namei.get(OpArg::new(u32::from(u8::from(unit.arg)))); + let name = f.names[usize::try_from(load_attr.name_idx()).unwrap()].as_str(); + if name == "warn" { + after_warn_attr = true; + } + } + Instruction::LoadFast { var_num } => { + let idx = usize::from(var_num.get(OpArg::new(u32::from(u8::from(unit.arg))))); + let name = f.varnames[idx].as_str(); + if name == "final" { + saw_strong_final = true; + } + } + Instruction::LoadFastBorrow { var_num } => { + let idx = usize::from(var_num.get(OpArg::new(u32::from(u8::from(unit.arg))))); + let name = f.varnames[idx].as_str(); + if after_warn_attr && name == "p" { + saw_borrow_p_after_warn = true; + } + if after_warn_attr && name == "ch" { + saw_borrow_ch_after_warn = true; + } + } + _ => {} + } + } + + assert!( + saw_strong_final, + "expected named-except fallback guards to deopt final to strong LOAD_FAST" + ); + assert!( + saw_borrow_p_after_warn && saw_borrow_ch_after_warn, + "expected outer invalid-escape tail to keep borrowed p/ch loads" + ); + } } diff --git a/crates/codegen/src/ir.rs b/crates/codegen/src/ir.rs index 7d09c2cef83..73c8028d74b 100644 --- a/crates/codegen/src/ir.rs +++ b/crates/codegen/src/ir.rs @@ -28,7 +28,10 @@ struct LineTableLocation { } const MAX_INT_SIZE_BITS: u64 = 128; +const MAX_COLLECTION_SIZE: usize = 256; +const MAX_TOTAL_ITEMS: isize = 1024; const MIN_CONST_SEQUENCE_SIZE: usize = 3; +const STACK_USE_GUIDELINE: usize = 30; /// Metadata for a code unit // = _PyCompile_CodeUnitMetadata @@ -136,6 +139,25 @@ fn set_to_nop(info: &mut InstructionInfo) { info.cache_entries = 0; } +fn is_named_except_cleanup_normal_exit_block(block: &Block) -> bool { + let len = block.instructions.len(); + if len < 5 { + return false; + } + let tail = &block.instructions[len - 5..]; + matches!(tail[0].instr.real(), Some(Instruction::PopExcept)) + && matches!(tail[1].instr.real(), Some(Instruction::LoadConst { .. })) + && matches!( + tail[2].instr.real(), + Some(Instruction::StoreName { .. } | Instruction::StoreFast { .. }) + ) + && matches!( + tail[3].instr.real(), + Some(Instruction::DeleteName { .. } | Instruction::DeleteFast { .. }) + ) + && tail[4].instr.is_unconditional_jump() +} + // spell-checker:ignore petgraph // TODO: look into using petgraph for handling blocks and stuff? it's heavier than this, but it // might enable more analysis/optimizations @@ -215,7 +237,7 @@ impl CodeInfo { self.fold_tuple_constants(); self.fold_list_constants(); self.fold_set_constants(); - self.fold_const_iterable_for_iter(); + self.optimize_lists_and_sets(); self.convert_to_load_small_int(); self.remove_unused_consts(); @@ -230,12 +252,27 @@ impl CodeInfo { self.apply_static_swaps(); // Peephole optimizer handles constant and compare folding. self.peephole_optimize(); + self.fold_tuple_constants(); + self.fold_list_constants(); + self.fold_set_constants(); + self.optimize_lists_and_sets(); + self.convert_to_load_small_int(); + self.remove_unused_consts(); + self.dce(); // Phase 1: _PyCfg_OptimizeCodeUnit (flowgraph.c) // Split blocks so each block has at most one branch as its last instruction split_blocks_at_jumps(&mut self.blocks); mark_except_handlers(&mut self.blocks); label_exception_targets(&mut self.blocks); + // CPython's CFG builder does not leave empty unconditional-jump targets + // in front of small exit blocks. Redirect only unconditional jumps + // here so inline_small_or_no_lineno_blocks() can see direct exit + // targets without erasing conditional target NOP anchors. + redirect_empty_unconditional_jump_targets(&mut self.blocks); + // CPython optimize_cfg starts by inlining tiny exit/no-lineno blocks + // before unreachable elimination and later jump cleanup. + inline_small_or_no_lineno_blocks(&mut self.blocks); // optimize_cfg: jump threading (before push_cold_blocks_to_end) jump_threading(&mut self.blocks); self.eliminate_unreachable_blocks(); @@ -245,19 +282,24 @@ impl CodeInfo { // later jump normalization / block reordering can create adjacencies // that never exist at this stage in flowgraph.c. self.insert_superinstructions(); + // CPython resolves line numbers once before cold-block extraction and + // again after reordering blocks. + resolve_line_numbers(&mut self.blocks); push_cold_blocks_to_end(&mut self.blocks); + reorder_conditional_chain_and_jump_back_blocks(&mut self.blocks); // Phase 2: _PyCfg_OptimizedCfgToInstructionSequence (flowgraph.c) normalize_jumps(&mut self.blocks); reorder_conditional_exit_and_jump_blocks(&mut self.blocks); reorder_conditional_jump_and_exit_blocks(&mut self.blocks); reorder_jump_over_exception_cleanup_blocks(&mut self.blocks); - inline_small_or_no_lineno_blocks(&mut self.blocks); self.dce(); // re-run within-block DCE after normalize_jumps creates new instructions self.eliminate_unreachable_blocks(); resolve_line_numbers(&mut self.blocks); + materialize_empty_conditional_exit_targets(&mut self.blocks); redirect_empty_block_targets(&mut self.blocks); duplicate_end_returns(&mut self.blocks, &self.metadata); + duplicate_shared_jump_back_targets(&mut self.blocks); self.dce(); // truncate after terminal in blocks that got return duplicated self.eliminate_unreachable_blocks(); // remove now-unreachable last block self.remove_redundant_const_pop_top_pairs(); @@ -266,8 +308,6 @@ impl CodeInfo { // once more so loop backedges stay direct instead of becoming // JUMP_FORWARD -> JUMP_BACKWARD chains. jump_threading_unconditional(&mut self.blocks); - reorder_conditional_exit_and_jump_blocks(&mut self.blocks); - reorder_conditional_jump_and_exit_blocks(&mut self.blocks); reorder_jump_over_exception_cleanup_blocks(&mut self.blocks); self.eliminate_unreachable_blocks(); remove_redundant_nops_and_jumps(&mut self.blocks); @@ -275,10 +315,6 @@ impl CodeInfo { inline_pop_except_return_blocks(&mut self.blocks); duplicate_named_except_cleanup_returns(&mut self.blocks, &self.metadata); 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. - self.insert_superinstructions(); let cellfixedoffsets = build_cellfixedoffsets( &self.metadata.varnames, &self.metadata.cellvars, @@ -299,12 +335,15 @@ impl CodeInfo { self.compute_load_fast_start_depths(); // optimize_load_fast: after normalize_jumps self.optimize_load_fast_borrow(); + self.deoptimize_borrow_after_multi_handler_resume_join(); + self.deoptimize_borrow_after_named_except_cleanup_join(); + self.deoptimize_borrow_in_protected_conditional_tail(); + self.deoptimize_borrow_after_protected_import(); self.deoptimize_borrow_after_push_exc_info(); self.deoptimize_borrow_for_handler_return_paths(); self.deoptimize_borrow_for_match_keys_attr(); 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(); @@ -378,7 +417,9 @@ impl CodeInfo { if lineno < 0 || prev_lineno == lineno { remove = true; } else if src < src_instructions.len() - 1 { - if src_instructions[src + 1].folded_from_nonliteral_expr { + if src_instructions[src + 1].instr.is_block_push() { + remove = false; + } else if src_instructions[src + 1].folded_from_nonliteral_expr { remove = true; } else { let next_lineno = src_instructions[src + 1] @@ -583,6 +624,11 @@ impl CodeInfo { first_line_number.get() as i32, opts.debug_ranges, ); + let locations = rustpython_compiler_core::marshal::linetable_to_locations( + &linetable, + first_line_number.get() as i32, + instructions.len(), + ); // Generate exception table before moving source_path let exceptiontable = generate_exception_table(&blocks, &block_to_index); @@ -634,7 +680,7 @@ impl CodeInfo { max_stackdepth, instructions: CodeUnits::from(instructions), - locations: locations.into_boxed_slice(), + locations, constants: constants.into_iter().collect(), names: name_cache.into_iter().collect(), varnames: varname_cache.into_iter().collect(), @@ -855,14 +901,13 @@ impl CodeInfo { && let Some(folded_const) = Self::eval_unary_constant(&operand, op, intrinsic) { let (const_idx, _) = self.metadata.consts.insert_full(folded_const); - let folded_from_nonliteral_expr = true; set_to_nop(&mut block.instructions[operand_index]); block.instructions[i].instr = Instruction::LoadConst { consti: Arg::marker(), } .into(); block.instructions[i].arg = OpArg::new(const_idx as u32); - block.instructions[i].folded_from_nonliteral_expr = folded_from_nonliteral_expr; + block.instructions[i].folded_from_nonliteral_expr = false; i = i.saturating_sub(1); } else { i += 1; @@ -892,6 +937,45 @@ impl CodeInfo { Some(indices) } + fn get_const_sequence( + metadata: &CodeUnitMetadata, + block: &Block, + build_index: usize, + size: usize, + ) -> Option<(Vec, Vec)> { + if size == 0 { + return Some((Vec::new(), Vec::new())); + } + + let operand_indices = build_index + .checked_sub(1) + .and_then(|start| Self::get_const_loading_instr_indices(block, start, size))?; + let mut elements = Vec::with_capacity(size); + + for &j in &operand_indices { + let load_instr = &block.instructions[j]; + if load_instr.folded_from_nonliteral_expr { + return None; + } + elements.push(Self::get_const_value_from(metadata, load_instr)?); + } + + Some((operand_indices, elements)) + } + + fn get_non_nop_instr_indices(block: &Block, start: usize, count: usize) -> Option> { + let mut indices = Vec::with_capacity(count); + for idx in start..block.instructions.len() { + if !matches!(block.instructions[idx].instr.real(), Some(Instruction::Nop)) { + indices.push(idx); + if indices.len() == count { + return Some(indices); + } + } + } + None + } + /// Constant folding: fold LOAD_CONST/LOAD_SMALL_INT + LOAD_CONST/LOAD_SMALL_INT + BINARY_OP /// into a single LOAD_CONST when the result is computable at compile time. /// = fold_binops_on_constants in CPython flowgraph.c @@ -991,6 +1075,19 @@ impl CodeInfo { } } + fn const_folding_check_complexity(obj: &ConstantData, mut limit: isize) -> Option { + if let ConstantData::Tuple { elements } = obj { + limit -= isize::try_from(elements.len()).ok()?; + if limit < 0 { + return None; + } + for element in elements { + limit = Self::const_folding_check_complexity(element, limit)?; + } + } + Some(limit) + } + fn eval_binop( left: &ConstantData, right: &ConstantData, @@ -1198,6 +1295,55 @@ impl CodeInfo { value: result.into(), }) } + (ConstantData::Tuple { elements: l }, ConstantData::Tuple { elements: r }) + if matches!(op, BinOp::Add) => + { + let mut result = l.clone(); + result.extend(r.iter().cloned()); + Some(ConstantData::Tuple { elements: result }) + } + (ConstantData::Tuple { elements }, ConstantData::Integer { value: n }) + if matches!(op, BinOp::Multiply) => + { + let n = n.to_usize()?; + if n != 0 && !elements.is_empty() { + if n > MAX_COLLECTION_SIZE / elements.len() { + return None; + } + Self::const_folding_check_complexity( + &ConstantData::Tuple { + elements: elements.clone(), + }, + MAX_TOTAL_ITEMS / isize::try_from(n).ok()?, + )?; + } + let mut result = Vec::with_capacity(elements.len() * n); + for _ in 0..n { + result.extend(elements.iter().cloned()); + } + Some(ConstantData::Tuple { elements: result }) + } + (ConstantData::Integer { value: n }, ConstantData::Tuple { elements }) + if matches!(op, BinOp::Multiply) => + { + let n = n.to_usize()?; + if n != 0 && !elements.is_empty() { + if n > MAX_COLLECTION_SIZE / elements.len() { + return None; + } + Self::const_folding_check_complexity( + &ConstantData::Tuple { + elements: elements.clone(), + }, + MAX_TOTAL_ITEMS / isize::try_from(n).ok()?, + )?; + } + let mut result = Vec::with_capacity(elements.len() * n); + for _ in 0..n { + result.extend(elements.iter().cloned()); + } + Some(ConstantData::Tuple { elements: result }) + } (ConstantData::Integer { value: n }, ConstantData::Str { value: s }) if matches!(op, BinOp::Multiply) => { @@ -1374,52 +1520,18 @@ impl CodeInfo { }; let list_size = u32::from(instr.arg) as usize; - if list_size == 0 { + if list_size == 0 || list_size > STACK_USE_GUIDELINE { i += 1; continue; } - let Some(operand_indices) = i.checked_sub(1).and_then(|start| { - Self::get_const_loading_instr_indices(block, start, list_size) - }) else { + let Some((operand_indices, elements)) = + Self::get_const_sequence(&self.metadata, block, i, list_size) + else { i += 1; continue; }; - let mut elements = Vec::with_capacity(list_size); - let mut all_const = true; - - for &j in &operand_indices { - let load_instr = &block.instructions[j]; - if load_instr.folded_from_nonliteral_expr { - all_const = false; - break; - } - match load_instr.instr.real() { - Some(Instruction::LoadConst { .. }) => { - let const_idx = u32::from(load_instr.arg) as usize; - if let Some(constant) = - self.metadata.consts.get_index(const_idx).cloned() - { - elements.push(constant); - } else { - all_const = false; - break; - } - } - Some(Instruction::LoadSmallInt { .. }) => { - let value = u32::from(load_instr.arg) as i32; - elements.push(ConstantData::Integer { - value: BigInt::from(value), - }); - } - _ => { - all_const = false; - break; - } - } - } - - if !all_const || list_size < MIN_CONST_SEQUENCE_SIZE { + if list_size < MIN_CONST_SEQUENCE_SIZE { i += 1; continue; } @@ -1467,10 +1579,14 @@ impl CodeInfo { } } - /// Convert constant list construction before GET_ITER to just LOAD_CONST tuple. - /// BUILD_LIST 0 + LOAD_CONST (tuple) + LIST_EXTEND 1 + GET_ITER - /// → LOAD_CONST (tuple) + GET_ITER - fn fold_const_iterable_for_iter(&mut self) { + /// Port of CPython's flowgraph.c optimize_lists_and_sets(). + /// + /// For GET_ITER / CONTAINS_OP users: + /// - Constant BUILD_LIST/BUILD_SET becomes LOAD_CONST tuple/frozenset. + /// - Non-constant BUILD_LIST becomes BUILD_TUPLE. + /// - Previously folded BUILD_LIST 0 + LOAD_CONST + LIST_EXTEND and + /// BUILD_SET 0 + LOAD_CONST + SET_UPDATE collapse back to LOAD_CONST. + fn optimize_lists_and_sets(&mut self) { for block in &mut self.blocks { let mut i = 0; while i + 1 < block.instructions.len() { @@ -1490,95 +1606,140 @@ impl CodeInfo { continue; } - let is_build = matches!( - block.instructions[i].instr.real(), - Some(Instruction::BuildList { .. }) - ) && u32::from(block.instructions[i].arg) == 0; + if let Some(non_nop4) = Self::get_non_nop_instr_indices(block, i, 4) { + let is_build_list = non_nop4[0] == i + && matches!( + block.instructions[non_nop4[0]].instr.real(), + Some(Instruction::BuildList { .. }) + ) + && u32::from(block.instructions[non_nop4[0]].arg) == 0; + let is_const = matches!( + block.instructions[non_nop4[1]].instr.real(), + Some(Instruction::LoadConst { .. }) + ); + let is_list_extend = matches!( + block.instructions[non_nop4[2]].instr.real(), + Some(Instruction::ListExtend { .. }) + ) && u32::from(block.instructions[non_nop4[2]].arg) == 1; + let uses_iter_or_contains = matches!( + block.instructions[non_nop4[3]].instr.real(), + Some(Instruction::GetIter | Instruction::ContainsOp { .. }) + ); - let is_const = matches!( - block - .instructions - .get(i + 1) - .and_then(|instr| instr.instr.real()), - Some(Instruction::LoadConst { .. }) - ); + if is_build_list && is_const && is_list_extend && uses_iter_or_contains { + let loc = block.instructions[i].location; + set_to_nop(&mut block.instructions[i]); + block.instructions[i].location = loc; + set_to_nop(&mut block.instructions[non_nop4[2]]); + block.instructions[non_nop4[2]].location = loc; + i += 1; + continue; + } - let is_extend = matches!( - block - .instructions - .get(i + 2) - .and_then(|instr| instr.instr.real()), - Some(Instruction::ListExtend { .. }) - ) && block - .instructions - .get(i + 2) - .is_some_and(|instr| u32::from(instr.arg) == 1); + let is_build_set = non_nop4[0] == i + && matches!( + block.instructions[non_nop4[0]].instr.real(), + Some(Instruction::BuildSet { .. }) + ) + && u32::from(block.instructions[non_nop4[0]].arg) == 0; + let is_set_update = matches!( + block.instructions[non_nop4[2]].instr.real(), + Some(Instruction::SetUpdate { .. }) + ) && u32::from(block.instructions[non_nop4[2]].arg) == 1; + + if is_build_set && is_const && is_set_update && uses_iter_or_contains { + let loc = block.instructions[i].location; + set_to_nop(&mut block.instructions[i]); + block.instructions[i].location = loc; + set_to_nop(&mut block.instructions[non_nop4[2]]); + block.instructions[non_nop4[2]].location = loc; + i += 1; + continue; + } + } - let is_iter = matches!( - block - .instructions - .get(i + 3) - .and_then(|instr| instr.instr.real()), - Some(Instruction::GetIter) - ); + let Some(non_nop2) = Self::get_non_nop_instr_indices(block, i, 2) else { + i += 1; + continue; + }; + let uses_iter_or_contains = non_nop2[0] == i + && matches!( + block.instructions[non_nop2[1]].instr.real(), + Some(Instruction::GetIter | Instruction::ContainsOp { .. }) + ); + if !uses_iter_or_contains { + i += 1; + continue; + } - if is_build && is_const && is_extend && is_iter { - // Replace: BUILD_X 0 → NOP, keep LOAD_CONST, LIST_EXTEND → NOP - let loc = block.instructions[i].location; - set_to_nop(&mut block.instructions[i]); - block.instructions[i].location = loc; - set_to_nop(&mut block.instructions[i + 2]); - block.instructions[i + 2].location = loc; - i += 4; - } else if matches!( + if matches!( block.instructions[i].instr.real(), Some(Instruction::BuildList { .. }) - ) && matches!( - block.instructions[i + 1].instr.real(), - Some(Instruction::GetIter) ) { let seq_size = u32::from(block.instructions[i].arg) as usize; - - if seq_size != 0 { - let Some(operand_indices) = i.checked_sub(1).and_then(|start| { - Self::get_const_loading_instr_indices(block, start, seq_size) - }) else { - i += 2; - continue; - }; - let mut elements = Vec::with_capacity(seq_size); - let mut all_const = true; + if seq_size > STACK_USE_GUIDELINE { + i += 2; + continue; + } + if let Some((operand_indices, elements)) = + Self::get_const_sequence(&self.metadata, block, i, seq_size) + { + let const_data = ConstantData::Tuple { elements }; + let (const_idx, _) = self.metadata.consts.insert_full(const_data); + let folded_loc = block.instructions[i].location; + let end_loc = block.instructions[i].end_location; + let eh = block.instructions[i].except_handler; for &j in &operand_indices { - match Self::get_const_value_from(&self.metadata, &block.instructions[j]) - { - Some(constant) => elements.push(constant), - None => { - all_const = false; - break; - } - } + set_to_nop(&mut block.instructions[j]); + block.instructions[j].location = folded_loc; + block.instructions[j].end_location = end_loc; } - if all_const { - let const_data = ConstantData::Tuple { elements }; - let (const_idx, _) = self.metadata.consts.insert_full(const_data); - let folded_loc = block.instructions[i].location; - - for &j in &operand_indices { - set_to_nop(&mut block.instructions[j]); - block.instructions[j].location = folded_loc; - } - - block.instructions[i].instr = Opcode::LoadConst.into(); - block.instructions[i].arg = OpArg::new(const_idx as u32); - i += 2; - continue; - } + block.instructions[i].instr = Opcode::LoadConst.into(); + block.instructions[i].arg = OpArg::new(const_idx as u32); + block.instructions[i].location = folded_loc; + block.instructions[i].end_location = end_loc; + block.instructions[i].except_handler = eh; + i += 2; + continue; } block.instructions[i].instr = Opcode::BuildTuple.into(); i += 2; + } else if matches!( + block.instructions[i].instr.real(), + Some(Instruction::BuildSet { .. }) + ) { + let seq_size = u32::from(block.instructions[i].arg) as usize; + if seq_size > STACK_USE_GUIDELINE { + i += 2; + continue; + } + let Some((operand_indices, elements)) = + Self::get_const_sequence(&self.metadata, block, i, seq_size) + else { + i += 2; + continue; + }; + let const_data = ConstantData::Frozenset { elements }; + let (const_idx, _) = self.metadata.consts.insert_full(const_data); + let folded_loc = block.instructions[i].location; + let end_loc = block.instructions[i].end_location; + let eh = block.instructions[i].except_handler; + + for &j in &operand_indices { + set_to_nop(&mut block.instructions[j]); + block.instructions[j].location = folded_loc; + block.instructions[j].end_location = end_loc; + } + + block.instructions[i].instr = Opcode::LoadConst.into(); + block.instructions[i].arg = OpArg::new(const_idx as u32); + block.instructions[i].location = folded_loc; + block.instructions[i].end_location = end_loc; + block.instructions[i].except_handler = eh; + i += 2; } else { i += 1; } @@ -1599,56 +1760,17 @@ impl CodeInfo { }; let set_size = u32::from(instr.arg) as usize; - if set_size < 3 { + if !(3..=STACK_USE_GUIDELINE).contains(&set_size) { i += 1; continue; } - let Some(operand_indices) = i.checked_sub(1).and_then(|start| { - Self::get_const_loading_instr_indices(block, start, set_size) - }) else { + let Some((operand_indices, elements)) = + Self::get_const_sequence(&self.metadata, block, i, set_size) + else { i += 1; continue; }; - let mut elements = Vec::with_capacity(set_size); - let mut all_const = true; - - for &j in &operand_indices { - let load_instr = &block.instructions[j]; - if load_instr.folded_from_nonliteral_expr { - all_const = false; - break; - } - match load_instr.instr.real() { - Some(Instruction::LoadConst { .. }) => { - let const_idx = u32::from(load_instr.arg) as usize; - if let Some(constant) = - self.metadata.consts.get_index(const_idx).cloned() - { - elements.push(constant); - } else { - all_const = false; - break; - } - } - Some(Instruction::LoadSmallInt { .. }) => { - let value = u32::from(load_instr.arg) as i32; - elements.push(ConstantData::Integer { - value: BigInt::from(value), - }); - } - _ => { - all_const = false; - break; - } - } - } - - if !all_const { - i += 1; - continue; - } - let const_data = ConstantData::Frozenset { elements }; let (const_idx, _) = self.metadata.consts.insert_full(const_data); @@ -1978,6 +2100,24 @@ impl CodeInfo { } i = run_end.max(i + 1); } + + // General same-line duplicate STORE_FAST elimination from + // flowgraph.c optimize_basic_block(). This is required for + // apply_static_swaps() patterns such as `a, a = x, y`. + for i in 0..instructions.len().saturating_sub(1) { + let lhs = &instructions[i]; + let rhs = &instructions[i + 1]; + if !matches!(lhs.instr.real(), Some(Instruction::StoreFast { .. })) + || !matches!(rhs.instr.real(), Some(Instruction::StoreFast { .. })) + || u32::from(lhs.arg) != u32::from(rhs.arg) + || instruction_lineno(lhs) != instruction_lineno(rhs) + { + continue; + } + instructions[i].instr = Instruction::PopTop.into(); + instructions[i].arg = OpArg::NULL; + instructions[i].target = BlockIdx::NULL; + } } } @@ -2365,43 +2505,20 @@ impl CodeInfo { } } - /// insert_superinstructions (flowgraph.c): combine a narrow subset of - /// STORE_FAST + LOAD_FAST patterns that CPython uses in comprehension loop - /// headers. Keeping this scoped avoids reintroducing earlier mismatches in - /// non-loop code while we continue aligning the surrounding borrow rules. + /// insert_superinstructions (flowgraph.c): combine adjacent same-line + /// LOAD_FAST / STORE_FAST pairs before later flowgraph passes change + /// block layout. fn insert_superinstructions(&mut self) { for block in &mut self.blocks { let mut i = 0; while i + 1 < block.instructions.len() { let curr = &block.instructions[i]; - let line = curr.location.line; - - let mut j = i + 1; - while j < block.instructions.len() - && matches!(block.instructions[j].instr.real(), Some(Instruction::Nop)) - && block.instructions[j].location.line == line - { - j += 1; - } - if j >= block.instructions.len() { - i += 1; - continue; - } - - let next = &block.instructions[j]; - if next.location.line != line { + let next = &block.instructions[i + 1]; + if instruction_lineno(curr) != instruction_lineno(next) { i += 1; continue; } - let prev_real = block.instructions[..i] - .iter() - .rev() - .find_map(|info| info.instr.real()); - let next_real = block.instructions[(j + 1)..] - .iter() - .find_map(|info| info.instr.real()); - match (curr.instr.real(), next.instr.real()) { (Some(Instruction::LoadFast { .. }), Some(Instruction::LoadFast { .. })) => { let idx1 = u32::from(curr.arg); @@ -2416,19 +2533,9 @@ impl CodeInfo { } .into(); block.instructions[i].arg = OpArg::new(packed); - block.instructions.drain(i + 1..=j); + block.instructions.remove(i + 1); } - ( - Some(Instruction::StoreFast { .. }), - Some(Instruction::LoadFast { .. } | Instruction::LoadFastBorrow { .. }), - ) => { - if self.flags.contains(CodeFlags::GENERATOR) - && matches!(prev_real, Some(Instruction::ForIter { .. })) - && !matches!(next_real, Some(Instruction::ToBool)) - { - i += 1; - continue; - } + (Some(Instruction::StoreFast { .. }), Some(Instruction::LoadFast { .. })) => { let store_idx = u32::from(curr.arg); let load_idx = u32::from(next.arg); if store_idx >= 16 || load_idx >= 16 { @@ -2441,7 +2548,7 @@ impl CodeInfo { } .into(); block.instructions[i].arg = OpArg::new(packed); - block.instructions.drain(i + 1..=j); + block.instructions.remove(i + 1); } (Some(Instruction::StoreFast { .. }), Some(Instruction::StoreFast { .. })) => { let idx1 = u32::from(curr.arg); @@ -2456,7 +2563,7 @@ impl CodeInfo { } .into(); block.instructions[i].arg = OpArg::new(packed); - block.instructions.drain(i + 1..=j); + block.instructions.remove(i + 1); } _ => i += 1, } @@ -2901,26 +3008,27 @@ impl CodeInfo { } } - fn deoptimize_borrow_after_push_exc_info(&mut self) { - for block in &mut self.blocks { - let mut in_exception_state = false; + fn deoptimize_borrow_after_multi_handler_resume_join(&mut self) { + fn second_last_real_instr(block: &Block) -> Option { + let mut reals = block + .instructions + .iter() + .rev() + .filter_map(|info| info.instr.real()); + let _last = reals.next()?; + reals.next() + } + + fn deoptimize_block_borrows(block: &mut Block) { for info in &mut block.instructions { match info.instr.real() { - Some(Instruction::PushExcInfo) => { - in_exception_state = true; - } - Some(Instruction::PopExcept) | Some(Instruction::Reraise { .. }) => { - in_exception_state = false; - } - Some(Instruction::LoadFastBorrow { .. }) if in_exception_state => { + Some(Instruction::LoadFastBorrow { .. }) => { info.instr = Instruction::LoadFast { var_num: Arg::marker(), } .into(); } - Some(Instruction::LoadFastBorrowLoadFastBorrow { .. }) - if in_exception_state => - { + Some(Instruction::LoadFastBorrowLoadFastBorrow { .. }) => { info.instr = Instruction::LoadFastLoadFast { var_nums: Arg::marker(), } @@ -2930,9 +3038,554 @@ impl CodeInfo { } } } - } - fn deoptimize_borrow_for_match_keys_attr(&mut self) { + let mut handler_resume_predecessors = vec![0usize; self.blocks.len()]; + let mut is_handler_resume_block = vec![false; self.blocks.len()]; + let mut predecessors = vec![Vec::new(); self.blocks.len()]; + for (block_idx, block) in self.blocks.iter().enumerate() { + let Some(last_info) = block.instructions.last() else { + continue; + }; + if last_info.target == BlockIdx::NULL || !last_info.instr.is_unconditional_jump() { + continue; + } + if !matches!(second_last_real_instr(block), Some(Instruction::PopExcept)) { + continue; + } + is_handler_resume_block[block_idx] = true; + handler_resume_predecessors[last_info.target.idx()] += 1; + } + for (pred_idx, block) in self.blocks.iter().enumerate() { + if block.next != BlockIdx::NULL { + predecessors[block.next.idx()].push(BlockIdx::new(pred_idx as u32)); + } + for info in &block.instructions { + if info.target != BlockIdx::NULL { + predecessors[info.target.idx()].push(BlockIdx::new(pred_idx as u32)); + } + } + } + + let mut visited = vec![false; self.blocks.len()]; + for (idx, &count) in handler_resume_predecessors.iter().enumerate() { + if count < 2 { + continue; + } + let seed = BlockIdx::new(idx as u32); + let mut segment = Vec::new(); + let mut cursor = seed; + while cursor != BlockIdx::NULL { + if block_is_exceptional(&self.blocks[cursor.idx()]) { + break; + } + segment.push(cursor); + cursor = self.blocks[cursor.idx()].next; + } + + let mut in_segment = vec![false; self.blocks.len()]; + for block_idx in &segment { + in_segment[block_idx.idx()] = true; + } + + for block_idx in segment { + if visited[block_idx.idx()] { + continue; + } + if block_idx != seed + && predecessors[block_idx.idx()] + .iter() + .any(|pred| !in_segment[pred.idx()] && !is_handler_resume_block[pred.idx()]) + { + continue; + } + visited[block_idx.idx()] = true; + deoptimize_block_borrows(&mut self.blocks[block_idx.idx()]); + } + } + } + + fn deoptimize_borrow_after_named_except_cleanup_join(&mut self) { + fn first_real_instr(block: &Block) -> Option { + block.instructions.iter().find_map(|info| info.instr.real()) + } + + fn leading_bool_guard_local(block: &Block) -> Option { + let infos: Vec<_> = block + .instructions + .iter() + .filter(|info| info.instr.real().is_some()) + .take(3) + .collect(); + if infos.len() < 3 { + return None; + } + let load_local = match infos[0].instr.real() { + Some(Instruction::LoadFast { var_num }) => usize::from(var_num.get(infos[0].arg)), + Some(Instruction::LoadFastBorrow { var_num }) => { + usize::from(var_num.get(infos[0].arg)) + } + _ => return None, + }; + if !matches!(infos[1].instr.real(), Some(Instruction::ToBool)) { + return None; + } + if !matches!( + infos[2].instr.real(), + Some( + Instruction::PopJumpIfFalse { .. } + | Instruction::PopJumpIfTrue { .. } + | Instruction::PopJumpIfNone { .. } + | Instruction::PopJumpIfNotNone { .. } + ) + ) { + return None; + } + Some(load_local) + } + + fn deoptimize_block_borrows(block: &mut Block) { + for info in &mut block.instructions { + match info.instr.real() { + Some(Instruction::LoadFastBorrow { .. }) => { + info.instr = Instruction::LoadFast { + var_num: Arg::marker(), + } + .into(); + } + Some(Instruction::LoadFastBorrowLoadFastBorrow { .. }) => { + info.instr = Instruction::LoadFastLoadFast { + var_nums: Arg::marker(), + } + .into(); + } + _ => {} + } + } + } + + let mut named_cleanup_predecessors = vec![0usize; self.blocks.len()]; + let mut is_named_cleanup_resume_block = vec![false; self.blocks.len()]; + let mut predecessors = vec![Vec::new(); self.blocks.len()]; + + for (block_idx, block) in self.blocks.iter().enumerate() { + let Some(last_info) = block.instructions.last() else { + continue; + }; + if last_info.target == BlockIdx::NULL || !last_info.instr.is_unconditional_jump() { + continue; + } + if !is_named_except_cleanup_normal_exit_block(block) { + continue; + } + if matches!( + first_real_instr(&self.blocks[last_info.target.idx()]), + Some(Instruction::ForIter { .. }) + ) { + continue; + } + is_named_cleanup_resume_block[block_idx] = true; + named_cleanup_predecessors[last_info.target.idx()] += 1; + } + for (pred_idx, block) in self.blocks.iter().enumerate() { + if block.next != BlockIdx::NULL { + predecessors[block.next.idx()].push(BlockIdx::new(pred_idx as u32)); + } + for info in &block.instructions { + if info.target != BlockIdx::NULL { + predecessors[info.target.idx()].push(BlockIdx::new(pred_idx as u32)); + } + } + } + + let mut visited = vec![false; self.blocks.len()]; + for (idx, &count) in named_cleanup_predecessors.iter().enumerate() { + if count == 0 { + continue; + } + let seed = BlockIdx::new(idx as u32); + let mut segment = Vec::new(); + let mut cursor = seed; + let mut fallback_guard_local = None; + while cursor != BlockIdx::NULL { + let block = &self.blocks[cursor.idx()]; + if block_is_exceptional(block) { + break; + } + if cursor != seed + && let Some(local) = leading_bool_guard_local(block) + { + match fallback_guard_local { + None => fallback_guard_local = Some(local), + Some(expected) if expected != local => break, + Some(_) => {} + } + } + segment.push(cursor); + cursor = block.next; + } + + let mut in_segment = vec![false; self.blocks.len()]; + for block_idx in &segment { + in_segment[block_idx.idx()] = true; + } + + for block_idx in segment { + if visited[block_idx.idx()] { + continue; + } + let is_same_guard_fallback = fallback_guard_local.is_some_and(|local| { + leading_bool_guard_local(&self.blocks[block_idx.idx()]) == Some(local) + }); + if block_idx != seed + && !is_same_guard_fallback + && predecessors[block_idx.idx()].iter().any(|pred| { + !in_segment[pred.idx()] && !is_named_cleanup_resume_block[pred.idx()] + }) + { + continue; + } + visited[block_idx.idx()] = true; + deoptimize_block_borrows(&mut self.blocks[block_idx.idx()]); + } + } + } + + fn deoptimize_borrow_in_protected_conditional_tail(&mut self) { + fn second_last_real_instr(block: &Block) -> Option { + let mut reals = block + .instructions + .iter() + .rev() + .filter_map(|info| info.instr.real()); + let _last = reals.next()?; + reals.next() + } + + fn deoptimize_block_borrows(block: &mut Block) { + for info in &mut block.instructions { + match info.instr.real() { + Some(Instruction::LoadFastBorrow { .. }) => { + info.instr = Instruction::LoadFast { + var_num: Arg::marker(), + } + .into(); + } + Some(Instruction::LoadFastBorrowLoadFastBorrow { .. }) => { + info.instr = Instruction::LoadFastLoadFast { + var_nums: Arg::marker(), + } + .into(); + } + _ => {} + } + } + } + + let mut predecessors = vec![Vec::new(); self.blocks.len()]; + let mut is_handler_resume_block = vec![false; self.blocks.len()]; + for (pred_idx, block) in self.blocks.iter().enumerate() { + if matches!(second_last_real_instr(block), Some(Instruction::PopExcept)) + && block.instructions.last().is_some_and(|info| { + info.target != BlockIdx::NULL && info.instr.is_unconditional_jump() + }) + { + is_handler_resume_block[pred_idx] = true; + } + if block.next != BlockIdx::NULL { + predecessors[block.next.idx()].push(BlockIdx::new(pred_idx as u32)); + } + for info in &block.instructions { + if info.target != BlockIdx::NULL { + predecessors[info.target.idx()].push(BlockIdx::new(pred_idx as u32)); + } + } + } + + let seeds: Vec<_> = + self.blocks + .iter() + .enumerate() + .filter_map(|(idx, block)| { + let prev_protected = predecessors[idx].iter().any(|pred| { + self.blocks[pred.idx()] + .instructions + .iter() + .any(|info| info.except_handler.is_some()) + }); + (!block_is_exceptional(block) + && trailing_conditional_jump_index(block).is_some() + && prev_protected + && block.instructions.iter().any(|info| { + matches!(info.instr.real(), Some(Instruction::Call { .. })) + })) + .then_some(BlockIdx::new(idx as u32)) + }) + .collect(); + + let mut visited = vec![false; self.blocks.len()]; + for seed in seeds { + let mut segment = Vec::new(); + let mut cursor = seed; + while cursor != BlockIdx::NULL { + if block_is_exceptional(&self.blocks[cursor.idx()]) { + break; + } + segment.push(cursor); + cursor = self.blocks[cursor.idx()].next; + } + + let segment_ops: Vec<_> = segment + .iter() + .flat_map(|block_idx| { + self.blocks[block_idx.idx()] + .instructions + .iter() + .filter_map(|info| info.instr.real()) + }) + .collect(); + let call_count = segment_ops + .iter() + .filter(|instr| matches!(instr, Instruction::Call { .. })) + .count(); + let raise_count = segment_ops + .iter() + .filter(|instr| matches!(instr, Instruction::RaiseVarargs { .. })) + .count(); + let return_count = segment_ops + .iter() + .filter(|instr| matches!(instr, Instruction::ReturnValue)) + .count(); + let conditional_count = segment_ops + .iter() + .filter(|instr| { + matches!( + instr, + Instruction::PopJumpIfFalse { .. } + | Instruction::PopJumpIfTrue { .. } + | Instruction::PopJumpIfNone { .. } + | Instruction::PopJumpIfNotNone { .. } + ) + }) + .count(); + let has_complex_tail = segment_ops.iter().any(|instr| { + matches!( + instr, + Instruction::StoreFast { .. } + | Instruction::StoreFastLoadFast { .. } + | Instruction::StoreFastStoreFast { .. } + | Instruction::ForIter { .. } + | Instruction::JumpBackward { .. } + | Instruction::JumpBackwardNoInterrupt { .. } + | Instruction::EndFor + | Instruction::PopIter + | Instruction::LoadFastAndClear { .. } + | Instruction::LoadFastCheck { .. } + | Instruction::ListAppend { .. } + | Instruction::MapAdd { .. } + | Instruction::SetAdd { .. } + ) + }); + if has_complex_tail + || call_count != 2 + || raise_count != 1 + || return_count != 1 + || conditional_count != 1 + { + continue; + } + + let mut in_segment = vec![false; self.blocks.len()]; + for block_idx in &segment { + in_segment[block_idx.idx()] = true; + } + + for block_idx in segment { + if visited[block_idx.idx()] { + continue; + } + if block_idx != seed + && predecessors[block_idx.idx()] + .iter() + .any(|pred| !in_segment[pred.idx()] && !is_handler_resume_block[pred.idx()]) + { + continue; + } + visited[block_idx.idx()] = true; + deoptimize_block_borrows(&mut self.blocks[block_idx.idx()]); + } + } + } + + fn deoptimize_borrow_after_push_exc_info(&mut self) { + for block in &mut self.blocks { + let mut in_exception_state = false; + for info in &mut block.instructions { + match info.instr.real() { + Some(Instruction::PushExcInfo) => { + in_exception_state = true; + } + Some(Instruction::PopExcept) | Some(Instruction::Reraise { .. }) => { + in_exception_state = false; + } + Some(Instruction::LoadFastBorrow { .. }) if in_exception_state => { + info.instr = Instruction::LoadFast { + var_num: Arg::marker(), + } + .into(); + } + Some(Instruction::LoadFastBorrowLoadFastBorrow { .. }) + if in_exception_state => + { + info.instr = Instruction::LoadFastLoadFast { + var_nums: Arg::marker(), + } + .into(); + } + _ => {} + } + } + } + } + + fn deoptimize_borrow_after_protected_import(&mut self) { + fn deoptimize_block_borrows(block: &mut Block, after_import_only: bool) { + let mut after_import = !after_import_only; + for info in &mut block.instructions { + if matches!(info.instr.real(), Some(Instruction::ImportName { .. })) { + after_import = true; + continue; + } + if !after_import { + continue; + } + match info.instr.real() { + Some(Instruction::LoadFastBorrow { .. }) => { + info.instr = Instruction::LoadFast { + var_num: Arg::marker(), + } + .into(); + } + Some(Instruction::LoadFastBorrowLoadFastBorrow { .. }) => { + info.instr = Instruction::LoadFastLoadFast { + var_nums: Arg::marker(), + } + .into(); + } + _ => {} + } + } + } + + fn is_handler_resume_predecessor(block: &Block, target: BlockIdx) -> bool { + let has_pop_except = block + .instructions + .iter() + .any(|info| matches!(info.instr.real(), Some(Instruction::PopExcept))); + let jumps_to_target = block.instructions.iter().any(|info| { + info.target == target + && matches!( + info.instr.real(), + Some( + Instruction::JumpForward { .. } + | Instruction::JumpBackward { .. } + | Instruction::JumpBackwardNoInterrupt { .. } + ) + ) + }); + has_pop_except && jumps_to_target + } + + let mut predecessors = vec![Vec::new(); self.blocks.len()]; + for (pred_idx, block) in self.blocks.iter().enumerate() { + if block.next != BlockIdx::NULL { + predecessors[block.next.idx()].push(BlockIdx::new(pred_idx as u32)); + } + for info in &block.instructions { + if info.target != BlockIdx::NULL { + predecessors[info.target.idx()].push(BlockIdx::new(pred_idx as u32)); + } + } + } + + let seeds: Vec<_> = self + .blocks + .iter() + .enumerate() + .filter_map(|(idx, block)| { + (!block_is_exceptional(block) + && block + .instructions + .iter() + .any(|info| info.except_handler.is_some()) + && block.instructions.iter().any(|info| { + matches!(info.instr.real(), Some(Instruction::ImportName { .. })) + })) + .then_some(BlockIdx::new(idx as u32)) + }) + .collect(); + + let mut visited = vec![false; self.blocks.len()]; + for seed in seeds { + let mut seed_handler_chain = vec![false; self.blocks.len()]; + let seed_handler_blocks: Vec<_> = self.blocks[seed.idx()] + .instructions + .iter() + .filter_map(|info| info.except_handler.map(|handler| handler.handler_block)) + .collect(); + for handler_block in seed_handler_blocks { + let mut cursor = handler_block; + while cursor != BlockIdx::NULL && !seed_handler_chain[cursor.idx()] { + seed_handler_chain[cursor.idx()] = true; + cursor = self.blocks[cursor.idx()].next; + } + } + + let mut in_segment = vec![false; self.blocks.len()]; + in_segment[seed.idx()] = true; + let mut segment = vec![seed]; + let mut cursor = self.blocks[seed.idx()].next; + while cursor != BlockIdx::NULL && !block_is_exceptional(&self.blocks[cursor.idx()]) { + if predecessors[cursor.idx()].iter().any(|pred| { + !in_segment[pred.idx()] + && seed_handler_chain[pred.idx()] + && is_handler_resume_predecessor(&self.blocks[pred.idx()], cursor) + }) { + break; + } + in_segment[cursor.idx()] = true; + segment.push(cursor); + cursor = self.blocks[cursor.idx()].next; + } + + let has_backward_jump = segment.iter().any(|block_idx| { + self.blocks[block_idx.idx()] + .instructions + .iter() + .any(|info| { + matches!( + info.instr.real(), + Some( + Instruction::JumpBackward { .. } + | Instruction::JumpBackwardNoInterrupt { .. } + ) + ) + }) + }); + if !has_backward_jump { + continue; + } + + for (i, block_idx) in segment.into_iter().enumerate() { + if visited[block_idx.idx()] { + continue; + } + visited[block_idx.idx()] = true; + deoptimize_block_borrows(&mut self.blocks[block_idx.idx()], i == 0); + } + } + } + + fn deoptimize_borrow_for_match_keys_attr(&mut self) { let Some(key_name_idx) = self.metadata.names.get_index_of("KEY") else { return; }; @@ -3076,10 +3729,7 @@ impl CodeInfo { Some(Instruction::StoreFastStoreFast { .. }) ) && (new_instructions.last().is_some_and( |prev: &InstructionInfo| { - matches!( - prev.instr.real(), - Some(Instruction::PopIter) | Some(Instruction::Swap { .. }) - ) + matches!(prev.instr.real(), Some(Instruction::PopIter)) }, ) || (i == 0 && starts_after_cleanup[block_idx]) || in_restore_prefix); @@ -3469,7 +4119,7 @@ impl CodeInfo { self.fold_tuple_constants(); self.fold_list_constants(); self.fold_set_constants(); - self.fold_const_iterable_for_iter(); + self.optimize_lists_and_sets(); self.convert_to_load_small_int(); self.remove_unused_consts(); self.dce(); @@ -3481,6 +4131,13 @@ impl CodeInfo { "after_peephole_optimize".to_owned(), self.debug_block_dump(), )); + self.fold_tuple_constants(); + self.fold_list_constants(); + self.fold_set_constants(); + self.optimize_lists_and_sets(); + self.convert_to_load_small_int(); + self.remove_unused_consts(); + self.dce(); split_blocks_at_jumps(&mut self.blocks); trace.push(( "after_split_blocks_at_jumps".to_owned(), @@ -3496,9 +4153,24 @@ impl CodeInfo { "after_early_remove_nops".to_owned(), self.debug_block_dump(), )); + inline_small_or_no_lineno_blocks(&mut self.blocks); + trace.push(( + "after_inline_small_or_no_lineno_blocks".to_owned(), + self.debug_block_dump(), + )); self.add_checks_for_loads_of_uninitialized_variables(); self.insert_superinstructions(); + resolve_line_numbers(&mut self.blocks); + trace.push(( + "after_first_resolve_line_numbers".to_owned(), + self.debug_block_dump(), + )); push_cold_blocks_to_end(&mut self.blocks); + trace.push(( + "after_push_cold_before_chain_reorder".to_owned(), + self.debug_block_dump(), + )); + reorder_conditional_chain_and_jump_back_blocks(&mut self.blocks); trace.push(( "after_push_cold_blocks_to_end".to_owned(), @@ -3507,18 +4179,11 @@ impl CodeInfo { normalize_jumps(&mut self.blocks); trace.push(("after_normalize_jumps".to_owned(), self.debug_block_dump())); - reorder_conditional_exit_and_jump_blocks(&mut self.blocks); reorder_conditional_jump_and_exit_blocks(&mut self.blocks); reorder_jump_over_exception_cleanup_blocks(&mut self.blocks); trace.push(("after_reorder".to_owned(), self.debug_block_dump())); - inline_small_or_no_lineno_blocks(&mut self.blocks); - trace.push(( - "after_inline_small_or_no_lineno_blocks".to_owned(), - self.debug_block_dump(), - )); - self.dce(); self.eliminate_unreachable_blocks(); trace.push(("after_dce_unreachable".to_owned(), self.debug_block_dump())); @@ -3529,6 +4194,11 @@ impl CodeInfo { self.debug_block_dump(), )); + materialize_empty_conditional_exit_targets(&mut self.blocks); + trace.push(( + "after_materialize_empty_conditional_exit_targets".to_owned(), + self.debug_block_dump(), + )); redirect_empty_block_targets(&mut self.blocks); trace.push(( "after_redirect_empty_block_targets".to_owned(), @@ -3536,6 +4206,7 @@ impl CodeInfo { )); duplicate_end_returns(&mut self.blocks, &self.metadata); + duplicate_shared_jump_back_targets(&mut self.blocks); trace.push(( "after_duplicate_end_returns".to_owned(), self.debug_block_dump(), @@ -3573,6 +4244,10 @@ impl CodeInfo { self.debug_block_dump(), )); self.optimize_load_fast_borrow(); + self.deoptimize_borrow_after_multi_handler_resume_join(); + self.deoptimize_borrow_after_named_except_cleanup_join(); + self.deoptimize_borrow_in_protected_conditional_tail(); + self.deoptimize_borrow_after_protected_import(); trace.push(( "after_optimize_load_fast_borrow".to_owned(), self.debug_block_dump(), @@ -4080,6 +4755,7 @@ fn push_cold_blocks_to_end(blocks: &mut Vec) { last = blocks[last.idx()].next; } blocks[last.idx()].next = cold_head; + remove_redundant_nops_and_jumps(blocks); } } @@ -4195,6 +4871,14 @@ fn jump_threading_impl(blocks: &mut [Block], include_conditional: bool) { let mut changed = true; while changed { changed = false; + let mut block_order = vec![u32::MAX; blocks.len()]; + let mut cursor = BlockIdx(0); + let mut pos = 0u32; + while cursor != BlockIdx::NULL { + block_order[cursor.idx()] = pos; + pos += 1; + cursor = blocks[cursor.idx()].next; + } for bi in 0..blocks.len() { let last_idx = match blocks[bi].instructions.len().checked_sub(1) { Some(i) => i, @@ -4225,26 +4909,62 @@ fn jump_threading_impl(blocks: &mut [Block], include_conditional: bool) { if target == BlockIdx::NULL { continue; } - // Thread through blocks that are only leading NOPs followed by an - // unconditional jump so late line anchors do not leave - // JUMP_FORWARD -> NOP -> JUMP_BACKWARD chains behind. - let target_jump = blocks[target.idx()] - .instructions - .iter() - .find(|info| !matches!(info.instr.real(), Some(Instruction::Nop))) - .copied(); + if include_conditional && is_conditional_jump(&ins.instr) { + let source_pos = block_order[bi]; + let target_pos = block_order.get(target.idx()).copied().unwrap_or(u32::MAX); + if target_pos <= source_pos { + continue; + } + } + // Match CPython's early flowgraph jump threading: inspect the + // target block's first instruction only. A later unconditional-only + // cleanup pass may thread through line-anchor NOPs introduced after + // jump normalization. + let target_jump = if include_conditional { + blocks[target.idx()].instructions.first().copied() + } else { + blocks[target.idx()] + .instructions + .iter() + .find(|info| !matches!(info.instr.real(), Some(Instruction::Nop))) + .copied() + }; if let Some(target_ins) = target_jump && target_ins.instr.is_unconditional_jump() && target_ins.target != BlockIdx::NULL && target_ins.target != target { + let source_pos = block_order[bi]; + let target_pos = block_order.get(target.idx()).copied().unwrap_or(u32::MAX); + let final_target = target_ins.target; + let final_target_pos = block_order + .get(final_target.idx()) + .copied() + .unwrap_or(u32::MAX); + if !include_conditional && source_pos < target_pos && final_target_pos < target_pos + { + // Keep the forward hop when threading would turn it into a + // backward edge. CPython preserves this shape for chained + // compare loop exits to avoid wraparound-style jumps. + continue; + } let conditional = is_conditional_jump(&ins.instr); + if conditional + && !matches!( + jump_thread_kind(target_ins.instr), + Some(JumpThreadKind::Plain) + ) + { + continue; + } let Some(threaded_instr) = threaded_jump_instr(ins.instr, target_ins.instr, conditional) else { continue; }; - let final_target = target_ins.target; + if conditional && final_target_pos <= source_pos { + continue; + } if ins.target == final_target { continue; } @@ -4275,6 +4995,13 @@ fn is_conditional_jump(instr: &AnyInstruction) -> bool { ) } +fn is_false_path_conditional_jump(instr: &AnyInstruction) -> bool { + matches!( + instr.real().map(Into::into), + Some(Opcode::PopJumpIfFalse | Opcode::PopJumpIfNone | Opcode::PopJumpIfNotNone) + ) +} + /// Invert a conditional jump opcode. fn reversed_conditional(instr: &AnyInstruction) -> Option { Some(match AnyOpcode::from(*instr).real()? { @@ -4286,7 +5013,7 @@ fn reversed_conditional(instr: &AnyInstruction) -> Option { }) } -/// flowgraph.c normalize_jumps + remove_redundant_jumps +/// flowgraph.c normalize_jumps fn normalize_jumps(blocks: &mut Vec) { let mut visit_order = Vec::new(); let mut visited = vec![false; blocks.len()]; @@ -4303,20 +5030,6 @@ fn normalize_jumps(blocks: &mut Vec) { let idx = block_idx.idx(); visited[idx] = true; - // Remove redundant unconditional jump to next block - let next = blocks[idx].next; - if next != BlockIdx::NULL { - let last = blocks[idx].instructions.last(); - let is_jump_to_next = last.is_some_and(|ins| { - ins.instr.is_unconditional_jump() - && ins.target != BlockIdx::NULL - && ins.target == next - }); - if is_jump_to_next && let Some(last_instr) = blocks[idx].instructions.last_mut() { - set_to_nop(last_instr); - } - } - // Normalize conditional jumps: forward gets NOT_TAKEN, backward gets inverted let last = blocks[idx].instructions.last(); if let Some(last_ins) = last @@ -4459,24 +5172,6 @@ fn inline_small_or_no_lineno_blocks(blocks: &mut [Block]) { .iter() .all(|ins| !instruction_has_lineno(ins)) }; - let current_is_named_except_cleanup_normal_exit = |block: &Block| { - let len = block.instructions.len(); - if len < 5 { - return false; - } - let tail = &block.instructions[len - 5..]; - matches!(tail[0].instr.real(), Some(Instruction::PopExcept)) - && matches!(tail[1].instr.real(), Some(Instruction::LoadConst { .. })) - && matches!( - tail[2].instr.real(), - Some(Instruction::StoreName { .. } | Instruction::StoreFast { .. }) - ) - && matches!( - tail[3].instr.real(), - Some(Instruction::DeleteName { .. } | Instruction::DeleteFast { .. }) - ) - && tail[4].instr.is_unconditional_jump() - }; let target_pushes_handler = |block: &Block| { block .instructions @@ -4500,7 +5195,7 @@ fn inline_small_or_no_lineno_blocks(blocks: &mut [Block]) { let target = last.target; if block_is_exceptional(&blocks[current.idx()]) || block_is_exceptional(&blocks[target.idx()]) - || (current_is_named_except_cleanup_normal_exit(&blocks[current.idx()]) + || (is_named_except_cleanup_normal_exit_block(&blocks[current.idx()]) && target_pushes_handler(&blocks[target.idx()])) { current = next; @@ -4511,16 +5206,20 @@ fn inline_small_or_no_lineno_blocks(blocks: &mut [Block]) { let no_lineno_no_fallthrough = block_has_no_lineno(&blocks[target.idx()]) && !block_has_fallthrough(&blocks[target.idx()]); if small_exit_block || no_lineno_no_fallthrough { - let removed_jump_location = last.location; - let removed_jump_end_location = last.end_location; - if let Some(last_instr) = blocks[current.idx()].instructions.last_mut() { - set_to_nop(last_instr); - } - let mut appended = blocks[target.idx()].instructions.clone(); - if let Some(first) = appended.first_mut() { - overwrite_location(first, removed_jump_location, removed_jump_end_location); + let removed_jump_had_lineno = blocks[current.idx()] + .instructions + .last() + .is_some_and(instruction_has_lineno); + if removed_jump_had_lineno { + if let Some(last_instr) = blocks[current.idx()].instructions.last_mut() { + set_to_nop(last_instr); + } + } else { + let _ = blocks[current.idx()].instructions.pop(); } - blocks[current.idx()].instructions.extend(appended); + blocks[current.idx()] + .instructions + .extend(blocks[target.idx()].instructions.clone()); changes = true; } @@ -4557,7 +5256,9 @@ fn remove_redundant_nops_in_blocks(blocks: &mut [Block]) -> usize { if lineno < 0 || prev_lineno == lineno { remove = true; } else if src < src_instructions.len() - 1 { - if src_instructions[src + 1].instr.is_unconditional_jump() { + if src_instructions[src + 1].instr.is_block_push() { + remove = false; + } else if src_instructions[src + 1].instr.is_unconditional_jump() { src_instructions[src + 1].lineno_override = Some(lineno); remove = true; } else if src_instructions[src + 1].folded_from_nonliteral_expr { @@ -4674,6 +5375,99 @@ fn redirect_empty_block_targets(blocks: &mut [Block]) { } } +fn redirect_empty_unconditional_jump_targets(blocks: &mut [Block]) { + let redirected_targets: Vec> = blocks + .iter() + .map(|block| { + block + .instructions + .iter() + .map(|instr| { + if instr.target == BlockIdx::NULL || !instr.instr.is_unconditional_jump() { + instr.target + } else { + next_nonempty_block(blocks, instr.target) + } + }) + .collect() + }) + .collect(); + + for (block, block_targets) in blocks.iter_mut().zip(redirected_targets) { + for (instr, target) in block.instructions.iter_mut().zip(block_targets) { + if target != BlockIdx::NULL { + instr.target = target; + } + } + } +} + +fn materialize_empty_conditional_exit_targets(blocks: &mut [Block]) { + let mut jump_back_inserts = Vec::new(); + let mut inserts = Vec::new(); + for (block_idx, block) in blocks.iter().enumerate() { + let Some(last) = block.instructions.last() else { + continue; + }; + if !is_conditional_jump(&last.instr) || last.target == BlockIdx::NULL { + continue; + } + let target = last.target; + if !blocks[target.idx()].instructions.is_empty() { + continue; + } + let next = next_nonempty_block(blocks, blocks[target.idx()].next); + if next != BlockIdx::NULL + && is_jump_only_block(&blocks[next.idx()]) + && block_has_no_lineno(&blocks[next.idx()]) + && comes_before( + blocks, + next_nonempty_block(blocks, blocks[next.idx()].instructions[0].target), + next, + ) + { + jump_back_inserts.push((BlockIdx(block_idx as u32), target, next)); + continue; + } + if next == BlockIdx::NULL || !is_scope_exit_block(&blocks[next.idx()]) { + continue; + } + inserts.push((BlockIdx(block_idx as u32), target)); + } + + for (source, target, next) in jump_back_inserts { + if !blocks[target.idx()].instructions.is_empty() { + continue; + } + let Some(last) = blocks[source.idx()].instructions.last().copied() else { + continue; + }; + let mut cloned = blocks[next.idx()].instructions[0]; + overwrite_location(&mut cloned, last.location, last.end_location); + blocks[target.idx()].instructions.push(cloned); + } + + for (source, target) in inserts { + if !blocks[target.idx()].instructions.is_empty() { + continue; + } + let Some(last) = blocks[source.idx()].instructions.last().copied() else { + continue; + }; + blocks[target.idx()].instructions.push(InstructionInfo { + instr: Instruction::Nop.into(), + arg: OpArg::NULL, + target: BlockIdx::NULL, + location: last.location, + end_location: last.end_location, + except_handler: None, + folded_from_nonliteral_expr: false, + lineno_override: None, + cache_entries: 0, + }); + } +} + fn merge_unsafe_mask(slot: &mut Option>, incoming: &[bool]) -> bool { match slot { Some(existing) => { @@ -4754,6 +5548,10 @@ fn instruction_has_lineno(instr: &InstructionInfo) -> bool { instruction_lineno(instr) > 0 } +fn propagation_location(instr: &InstructionInfo) -> Option<(SourceLocation, SourceLocation)> { + instruction_has_lineno(instr).then_some((instr.location, instr.end_location)) +} + fn block_has_fallthrough(block: &Block) -> bool { block .instructions @@ -4769,10 +5567,43 @@ fn is_exit_without_lineno(block: &Block) -> bool { let Some(first) = block.instructions.first() else { return false; }; - let Some(last) = block.instructions.last() else { + if instruction_has_lineno(first) || !block_has_no_lineno(block) { + return false; + } + + if block + .instructions + .last() + .is_some_and(|last| last.instr.is_scope_exit()) + { + return true; + } + + // CPython duplicates no-lineno exit blocks before propagating locations. + // RustPython's late CFG can inline the following synthetic jump-back block + // into that exit block first, collapsing `POP_EXCEPT; JUMP_BACKWARD` into a + // single block. Treat that merged tail as exit-like so resolve_line_numbers() + // can still duplicate it per predecessor and recover CPython's structure. + let Some((last, prefix)) = block.instructions.split_last() else { return false; }; - !instruction_has_lineno(first) && last.instr.is_scope_exit() + last.instr.is_unconditional_jump() + && prefix.iter().all(|info| { + matches!( + info.instr.real(), + Some(Instruction::PopExcept) | Some(Instruction::Nop) + ) + }) + && prefix + .iter() + .any(|info| matches!(info.instr.real(), Some(Instruction::PopExcept))) +} + +fn block_has_no_lineno(block: &Block) -> bool { + block + .instructions + .iter() + .all(|ins| !instruction_has_lineno(ins)) } fn is_jump_only_block(block: &Block) -> bool { @@ -4782,6 +5613,23 @@ fn is_jump_only_block(block: &Block) -> bool { instr.instr.is_unconditional_jump() && instr.target != BlockIdx::NULL } +fn is_pop_top_jump_block(block: &Block) -> bool { + let mut real_instrs = block + .instructions + .iter() + .filter(|info| !matches!(info.instr.real(), Some(Instruction::Nop))); + let Some(first) = real_instrs.next() else { + return false; + }; + let Some(second) = real_instrs.next() else { + return false; + }; + real_instrs.next().is_none() + && matches!(first.instr.real(), Some(Instruction::PopTop)) + && second.instr.is_unconditional_jump() + && second.target != BlockIdx::NULL +} + fn is_scope_exit_block(block: &Block) -> bool { block .instructions @@ -4789,6 +5637,19 @@ fn is_scope_exit_block(block: &Block) -> bool { .is_some_and(|instr| instr.instr.is_scope_exit()) } +fn is_loop_cleanup_block(block: &Block) -> bool { + block + .instructions + .iter() + .find_map(|info| info.instr.real()) + .is_some_and(|instr| { + matches!( + instr, + Instruction::EndFor | Instruction::EndAsyncFor | Instruction::PopIter + ) + }) +} + fn is_exception_cleanup_block(block: &Block) -> bool { block .instructions @@ -4800,6 +5661,29 @@ fn is_exception_cleanup_block(block: &Block) -> bool { .is_some_and(|instr| matches!(instr.instr.real(), Some(Instruction::Reraise { .. }))) } +fn block_is_protected(block: &Block) -> bool { + block + .instructions + .iter() + .any(|info| info.except_handler.is_some()) +} + +fn block_contains_suspension_point(block: &Block) -> bool { + block + .instructions + .iter() + .filter_map(|info| info.instr.real()) + .any(|instr| { + matches!( + instr, + Instruction::YieldValue { .. } + | Instruction::GetAwaitable { .. } + | Instruction::GetANext + | Instruction::EndAsyncFor + ) + }) +} + fn block_is_exceptional(block: &Block) -> bool { block.except_handler || block.preserve_lasti || is_exception_cleanup_block(block) } @@ -4882,7 +5766,9 @@ fn reorder_conditional_exit_and_jump_blocks(blocks: &mut [Block]) { let mut jump_block = BlockIdx::NULL; cursor = jump_start; while cursor != BlockIdx::NULL { - if block_is_exceptional(&blocks[cursor.idx()]) { + if block_is_exceptional(&blocks[cursor.idx()]) + || block_is_protected(&blocks[cursor.idx()]) + { jump_block = BlockIdx::NULL; break; } @@ -4977,6 +5863,9 @@ fn reorder_conditional_jump_and_exit_blocks(blocks: &mut [Block]) { break BlockIdx::NULL; } if block_is_exceptional(&blocks[cursor.idx()]) { + if exit_block != BlockIdx::NULL { + break cursor; + } exit_block = BlockIdx::NULL; break BlockIdx::NULL; } @@ -5010,6 +5899,161 @@ fn reorder_conditional_jump_and_exit_blocks(blocks: &mut [Block]) { } } +fn reorder_conditional_chain_and_jump_back_blocks(blocks: &mut Vec) { + let target_comes_before = |target: BlockIdx, block: BlockIdx, blocks: &[Block]| -> bool { + let mut current = BlockIdx(0); + while current != BlockIdx::NULL { + if current == target { + return true; + } + if current == block { + return false; + } + current = blocks[current.idx()].next; + } + false + }; + + let mut current = BlockIdx(0); + while current != BlockIdx::NULL { + let idx = current.idx(); + let next = blocks[idx].next; + let Some(cond_idx) = trailing_conditional_jump_index(&blocks[idx]) else { + current = next; + continue; + }; + let last = blocks[idx].instructions[cond_idx]; + + let Some(reversed) = reversed_conditional(&last.instr) else { + current = next; + continue; + }; + if !is_false_path_conditional_jump(&last.instr) { + current = next; + continue; + } + + let chain_start = next; + let jump_start = last.target; + if chain_start == BlockIdx::NULL + || jump_start == BlockIdx::NULL + || chain_start == jump_start + { + current = next; + continue; + } + if block_is_protected(&blocks[idx]) && block_contains_suspension_point(&blocks[idx]) { + current = next; + continue; + } + if let Some(chain_cond_idx) = trailing_conditional_jump_index(&blocks[chain_start.idx()]) { + let chain_cond = blocks[chain_start.idx()].instructions[chain_cond_idx]; + if matches!( + chain_cond.instr.real().map(Into::into), + Some(Opcode::PopJumpIfTrue) + ) { + let chain_true_target = next_nonempty_block(blocks, chain_cond.target); + if chain_true_target != BlockIdx::NULL + && !is_scope_exit_block(&blocks[chain_true_target.idx()]) + && !is_jump_only_block(&blocks[chain_true_target.idx()]) + && !is_pop_top_jump_block(&blocks[chain_true_target.idx()]) + { + current = next; + continue; + } + } + } + + let mut chain_end = BlockIdx::NULL; + let mut saw_nonempty = false; + let mut nonempty_blocks = 0usize; + let mut real_instr_count = 0usize; + let mut cursor = chain_start; + let mut chain_valid = true; + while cursor != BlockIdx::NULL && cursor != jump_start { + if block_is_exceptional(&blocks[cursor.idx()]) + || block_is_protected(&blocks[cursor.idx()]) + && block_contains_suspension_point(&blocks[cursor.idx()]) + { + chain_valid = false; + break; + } + if !blocks[cursor.idx()].instructions.is_empty() { + saw_nonempty = true; + nonempty_blocks += 1; + real_instr_count += blocks[cursor.idx()] + .instructions + .iter() + .filter(|info| info.instr.real().is_some()) + .count(); + } + chain_end = cursor; + cursor = blocks[cursor.idx()].next; + } + if !chain_valid + || !saw_nonempty + || chain_end == BlockIdx::NULL + || cursor != jump_start + || nonempty_blocks > 8 + || real_instr_count > 80 + { + current = next; + continue; + } + + let mut jump_end = BlockIdx::NULL; + let mut jump_block = BlockIdx::NULL; + cursor = jump_start; + while cursor != BlockIdx::NULL { + if block_is_exceptional(&blocks[cursor.idx()]) { + jump_block = BlockIdx::NULL; + break; + } + jump_end = cursor; + if blocks[cursor.idx()].instructions.is_empty() { + cursor = blocks[cursor.idx()].next; + continue; + } + if !is_jump_only_block(&blocks[cursor.idx()]) + || !target_comes_before(blocks[cursor.idx()].instructions[0].target, cursor, blocks) + { + jump_block = BlockIdx::NULL; + } else { + jump_block = cursor; + } + break; + } + if jump_block == BlockIdx::NULL || jump_end == BlockIdx::NULL { + current = next; + continue; + } + + let after_jump = next_nonempty_block(blocks, blocks[jump_block.idx()].next); + if nonempty_blocks == 1 + && after_jump != BlockIdx::NULL + && !blocks[after_jump.idx()].cold + && !block_is_exceptional(&blocks[after_jump.idx()]) + && !is_scope_exit_block(&blocks[after_jump.idx()]) + && !is_loop_cleanup_block(&blocks[after_jump.idx()]) + { + current = next; + continue; + } + + let mut cloned_jump = blocks[jump_block.idx()].clone(); + cloned_jump.next = chain_start; + cloned_jump.start_depth = None; + let cloned_idx = BlockIdx::new(blocks.len() as u32); + blocks.push(cloned_jump); + blocks[idx].next = cloned_idx; + let cond_mut = &mut blocks[idx].instructions[cond_idx]; + cond_mut.instr = reversed; + cond_mut.target = chain_start; + + current = next; + } +} + #[allow(dead_code)] fn reorder_jump_over_exception_cleanup_blocks(blocks: &mut [Block]) { let mut current = BlockIdx(0); @@ -5122,18 +6166,46 @@ fn overwrite_location( instr.lineno_override = None; } -fn propagate_locations_in_block( - block: &mut Block, - location: SourceLocation, - end_location: SourceLocation, -) { - let mut prev_location = location; - let mut prev_end_location = end_location; - for instr in &mut block.instructions { - maybe_propagate_location(instr, prev_location, prev_end_location); - prev_location = instr.location; - prev_end_location = instr.end_location; +fn compute_reachable_blocks(blocks: &[Block]) -> Vec { + let mut reachable = vec![false; blocks.len()]; + if blocks.is_empty() { + return reachable; + } + + reachable[0] = true; + let mut changed = true; + while changed { + changed = false; + for i in 0..blocks.len() { + if !reachable[i] { + continue; + } + for ins in &blocks[i].instructions { + if ins.target != BlockIdx::NULL && !reachable[ins.target.idx()] { + reachable[ins.target.idx()] = true; + changed = true; + } + if let Some(eh) = &ins.except_handler + && !reachable[eh.handler_block.idx()] + { + reachable[eh.handler_block.idx()] = true; + changed = true; + } + } + let next = blocks[i].next; + if next != BlockIdx::NULL + && !reachable[next.idx()] + && !blocks[i].instructions.last().is_some_and(|ins| { + ins.instr.is_scope_exit() || ins.instr.is_unconditional_jump() + }) + { + reachable[next.idx()] = true; + changed = true; + } + } } + + reachable } fn compute_predecessors(blocks: &[Block]) -> Vec { @@ -5142,20 +6214,26 @@ fn compute_predecessors(blocks: &[Block]) -> Vec { return predecessors; } + let reachable = compute_reachable_blocks(blocks); predecessors[0] = 1; let mut current = BlockIdx(0); while current != BlockIdx::NULL { + if !reachable[current.idx()] { + current = blocks[current.idx()].next; + continue; + } + let block = &blocks[current.idx()]; if block_has_fallthrough(block) { let next = next_nonempty_block(blocks, block.next); - if next != BlockIdx::NULL { + if next != BlockIdx::NULL && reachable[next.idx()] { predecessors[next.idx()] += 1; } } for ins in &block.instructions { if ins.target != BlockIdx::NULL { let target = next_nonempty_block(blocks, ins.target); - if target != BlockIdx::NULL { + if target != BlockIdx::NULL && reachable[target.idx()] { predecessors[target.idx()] += 1; } } @@ -5187,25 +6265,26 @@ fn duplicate_exits_without_lineno(blocks: &mut Vec, predecessors: &mut Ve continue; } - // Copy the exit block and splice it into the linked list after current + // Copy the exit block and splice it into the linked list after the + // original target block, matching CPython's copy_basicblock() layout. let new_idx = BlockIdx(blocks.len() as u32); let mut new_block = blocks[target.idx()].clone(); - let jump_loc = last.location; - let jump_end_loc = last.end_location; - propagate_locations_in_block(&mut new_block, jump_loc, jump_end_loc); - let old_next = blocks[current.idx()].next; + if let Some(first) = new_block.instructions.first_mut() + && let Some((location, end_location)) = propagation_location(last) + { + overwrite_location(first, location, end_location); + } + let old_next = blocks[target.idx()].next; new_block.next = old_next; blocks.push(new_block); - blocks[current.idx()].next = new_idx; + blocks[target.idx()].next = new_idx; // Update the jump target let last_mut = blocks[current.idx()].instructions.last_mut().unwrap(); last_mut.target = new_idx; predecessors[target.idx()] -= 1; predecessors.push(1); - - // Skip past the newly inserted block - current = old_next; + current = blocks[current.idx()].next; } current = BlockIdx(0); @@ -5218,59 +6297,16 @@ fn duplicate_exits_without_lineno(blocks: &mut Vec, predecessors: &mut Ve if target != BlockIdx::NULL && predecessors[target.idx()] == 1 && is_exit_without_lineno(&blocks[target.idx()]) + && let Some((location, end_location)) = propagation_location(last) + && let Some(first) = blocks[target.idx()].instructions.first_mut() { - let last_location = last.location; - let last_end_location = last.end_location; - propagate_locations_in_block( - &mut blocks[target.idx()], - last_location, - last_end_location, - ); + maybe_propagate_location(first, location, end_location); } } current = blocks[current.idx()].next; } } -fn duplicate_jump_targets_without_lineno(blocks: &mut Vec, predecessors: &mut Vec) { - let mut current = BlockIdx(0); - while current != BlockIdx::NULL { - let block = &blocks[current.idx()]; - let last = match block.instructions.last() { - Some(ins) if ins.instr.is_unconditional_jump() && ins.target != BlockIdx::NULL => *ins, - _ => { - current = blocks[current.idx()].next; - continue; - } - }; - - let target = next_nonempty_block(blocks, last.target); - if target == BlockIdx::NULL || !is_jump_only_block(&blocks[target.idx()]) { - current = blocks[current.idx()].next; - continue; - } - if predecessors[target.idx()] <= 1 { - current = blocks[current.idx()].next; - continue; - } - - let new_idx = BlockIdx(blocks.len() as u32); - let mut new_block = blocks[target.idx()].clone(); - propagate_locations_in_block(&mut new_block, last.location, last.end_location); - let old_next = blocks[current.idx()].next; - new_block.next = old_next; - blocks.push(new_block); - blocks[current.idx()].next = new_idx; - - let last_mut = blocks[current.idx()].instructions.last_mut().unwrap(); - last_mut.target = new_idx; - predecessors[target.idx()] -= 1; - predecessors.push(1); - - current = old_next; - } -} - fn propagate_line_numbers(blocks: &mut [Block], predecessors: &[u32]) { let mut current = BlockIdx(0); while current != BlockIdx::NULL { @@ -5288,29 +6324,35 @@ fn propagate_line_numbers(blocks: &mut [Block], predecessors: &[u32]) { if let Some((location, end_location)) = prev_location { maybe_propagate_location(instr, location, end_location); } - prev_location = Some((instr.location, instr.end_location)); + prev_location = propagation_location(instr); } } if has_fallthrough { let target = next_nonempty_block(blocks, next_block); - if target != BlockIdx::NULL && predecessors[target.idx()] == 1 { - propagate_locations_in_block( - &mut blocks[target.idx()], - last.location, - last.end_location, - ); + if target != BlockIdx::NULL + && predecessors[target.idx()] == 1 + && let Some((location, end_location)) = propagation_location(&last) + && let Some(first) = blocks[target.idx()].instructions.first_mut() + { + maybe_propagate_location(first, location, end_location); } } if is_jump_instruction(&last) { - let target = next_nonempty_block(blocks, last.target); - if target != BlockIdx::NULL && predecessors[target.idx()] == 1 { - propagate_locations_in_block( - &mut blocks[target.idx()], - last.location, - last.end_location, - ); + let mut target = next_nonempty_block(blocks, last.target); + while target != BlockIdx::NULL + && blocks[target.idx()].instructions.is_empty() + && predecessors[target.idx()] == 1 + { + target = blocks[target.idx()].next; + } + if target != BlockIdx::NULL + && predecessors[target.idx()] == 1 + && let Some((location, end_location)) = propagation_location(&last) + && let Some(first) = blocks[target.idx()].instructions.first_mut() + { + maybe_propagate_location(first, location, end_location); } } } @@ -5321,7 +6363,6 @@ fn propagate_line_numbers(blocks: &mut [Block], predecessors: &[u32]) { fn resolve_line_numbers(blocks: &mut Vec) { let mut predecessors = compute_predecessors(blocks); duplicate_exits_without_lineno(blocks, &mut predecessors); - duplicate_jump_targets_without_lineno(blocks, &mut predecessors); propagate_line_numbers(blocks, &predecessors); } @@ -5339,6 +6380,82 @@ fn find_layout_predecessor(blocks: &[Block], target: BlockIdx) -> BlockIdx { BlockIdx::NULL } +fn comes_before(blocks: &[Block], first: BlockIdx, second: BlockIdx) -> bool { + let mut current = BlockIdx(0); + while current != BlockIdx::NULL { + if current == first { + return true; + } + if current == second { + return false; + } + current = blocks[current.idx()].next; + } + false +} + +fn duplicate_shared_jump_back_targets(blocks: &mut Vec) { + let predecessors = compute_predecessors(blocks); + let mut clones = Vec::new(); + + for target in 0..blocks.len() { + let target = BlockIdx(target as u32); + if !is_jump_only_block(&blocks[target.idx()]) || !block_has_no_lineno(&blocks[target.idx()]) + { + continue; + } + + let jump_target = next_nonempty_block(blocks, blocks[target.idx()].instructions[0].target); + if jump_target == BlockIdx::NULL || !comes_before(blocks, jump_target, target) { + continue; + } + + let layout_pred = find_layout_predecessor(blocks, target); + if layout_pred == BlockIdx::NULL + || !block_has_fallthrough(&blocks[layout_pred.idx()]) + || next_nonempty_block(blocks, blocks[layout_pred.idx()].next) != target + || predecessors[target.idx()] < 2 + { + continue; + } + + for block_idx in 0..blocks.len() { + let block_idx = BlockIdx(block_idx as u32); + if block_idx == target || block_idx == layout_pred { + continue; + } + + let Some(instr_idx) = trailing_conditional_jump_index(&blocks[block_idx.idx()]) else { + continue; + }; + if next_nonempty_block( + blocks, + blocks[block_idx.idx()].instructions[instr_idx].target, + ) != target + { + continue; + } + + clones.push((target, block_idx, instr_idx)); + } + } + + for (target, block_idx, instr_idx) in clones.into_iter().rev() { + let jump = blocks[block_idx.idx()].instructions[instr_idx]; + let mut cloned = blocks[target.idx()].clone(); + if let Some(first) = cloned.instructions.first_mut() { + overwrite_location(first, jump.location, jump.end_location); + } + + let new_idx = BlockIdx(blocks.len() as u32); + let old_next = blocks[target.idx()].next; + cloned.next = old_next; + blocks.push(cloned); + blocks[target.idx()].next = new_idx; + blocks[block_idx.idx()].instructions[instr_idx].target = new_idx; + } +} + /// Duplicate `LOAD_CONST None + RETURN_VALUE` for blocks that fall through /// to the final return block. fn duplicate_end_returns(blocks: &mut Vec, metadata: &CodeUnitMetadata) { diff --git a/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__nested_bool_op.snap b/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__nested_bool_op.snap index bc9268b45bd..3ad96c56454 100644 --- a/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__nested_bool_op.snap +++ b/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__nested_bool_op.snap @@ -1,6 +1,6 @@ --- source: crates/codegen/src/compile.rs -assertion_line: 9688 +assertion_line: 11769 expression: "compile_exec(\"\\\nx = Test() and False or False\n\")" --- 1 0 RESUME (0) diff --git a/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__nested_double_async_with.snap b/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__nested_double_async_with.snap index fccc7b7c336..27ae2ae18bc 100644 --- a/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__nested_double_async_with.snap +++ b/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__nested_double_async_with.snap @@ -1,6 +1,6 @@ --- source: crates/codegen/src/compile.rs -assertion_line: 11626 +assertion_line: 11847 expression: "compile_exec(\"\\\nasync def test():\n for stop_exc in (StopIteration('spam'), StopAsyncIteration('ham')):\n with self.subTest(type=type(stop_exc)):\n try:\n async with egg():\n raise stop_exc\n except Exception as ex:\n self.assertIs(ex, stop_exc)\n else:\n self.fail(f'{stop_exc} was suppressed')\n\")" --- 1 0 RESUME (0) @@ -213,7 +213,8 @@ expression: "compile_exec(\"\\\nasync def test():\n for stop_exc in (StopIter 196 STORE_FAST (1, ex) 197 DELETE_FAST (1, ex) 198 RERAISE (1) - 199 RERAISE (0) + + 7 199 RERAISE (0) 200 COPY (3) 201 POP_EXCEPT 202 RERAISE (1) @@ -248,8 +249,7 @@ expression: "compile_exec(\"\\\nasync def test():\n for stop_exc in (StopIter 230 COPY (3) 231 POP_EXCEPT 232 RERAISE (1) - - 2 233 CALL_INTRINSIC_1 (StopIterationError) + 233 CALL_INTRINSIC_1 (StopIterationError) 234 RERAISE (1) 2 MAKE_FUNCTION diff --git a/crates/codegen/src/symboltable.rs b/crates/codegen/src/symboltable.rs index e87aa3d4aea..4ab2ba65548 100644 --- a/crates/codegen/src/symboltable.rs +++ b/crates/codegen/src/symboltable.rs @@ -385,15 +385,11 @@ fn inline_comprehension( } } } else { - // Name doesn't exist in parent, copy from comprehension. - // Reset scope to Unknown so analyze_symbol will resolve it - // in the parent's context. + // Name doesn't exist in parent, copy the comprehension binding. + // This matches CPython's inline_comprehension(): newly introduced + // comprehension locals stay locals in the parent scope. let mut symbol = sub_symbol.clone(); - symbol.scope = if sub_symbol.is_bound() { - SymbolScope::Unknown - } else { - scope - }; + symbol.scope = scope; parent_symbols.insert(name.clone(), symbol); } } @@ -2211,24 +2207,13 @@ impl SymbolTableBuilder { self.tables.last_mut().unwrap().is_generator = is_generator; // PEP 709: Mark non-generator comprehensions for inlining. - // Only in function-like scopes for now. Module/class scope inlining - // needs more work (Cell name resolution, __class__ handling). - // Also excluded: generator expressions, async comprehensions, + // Excluded: generator expressions, async comprehensions, // and annotation scopes nested in classes (can_see_class_scope). let element_has_await = expr_contains_await(elt1) || elt2.is_some_and(expr_contains_await); if !is_generator && !has_async_gen && !element_has_await { let parent = self.tables.iter().rev().nth(1); let parent_can_see_class = parent.is_some_and(|t| t.can_see_class_scope); - let parent_is_func = parent.is_some_and(|t| { - matches!( - t.typ, - CompilerScope::Function - | CompilerScope::AsyncFunction - | CompilerScope::Lambda - | CompilerScope::Comprehension - ) - }); - if !parent_can_see_class && parent_is_func { + if !parent_can_see_class { self.tables.last_mut().unwrap().comp_inlined = true; } } diff --git a/crates/compiler-core/src/marshal.rs b/crates/compiler-core/src/marshal.rs index 5594d0f4c45..56d9e39c03a 100644 --- a/crates/compiler-core/src/marshal.rs +++ b/crates/compiler-core/src/marshal.rs @@ -1211,20 +1211,24 @@ pub fn linetable_to_locations( }; line += line_delta; + let mk = |l: i32| { + if l > 0 { + OneIndexed::new(l as usize).unwrap_or(OneIndexed::MIN) + } else { + OneIndexed::MIN + } + }; for _ in 0..length { if locations.len() >= num_instructions { break; } if kind == PyCodeLocationInfoKind::None { - locations.push(default_loc()); - } else { - let mk = |l: i32| { - if l > 0 { - OneIndexed::new(l as usize).unwrap_or(OneIndexed::MIN) - } else { - OneIndexed::MIN - } + let loc = SourceLocation { + line: mk(line), + character_offset: OneIndexed::from_zero_indexed(0), }; + locations.push((loc, loc)); + } else { locations.push(( SourceLocation { line: mk(line), diff --git a/scripts/compare_bytecode.py b/scripts/compare_bytecode.py index f7b5ed916ed..48da77688d8 100644 --- a/scripts/compare_bytecode.py +++ b/scripts/compare_bytecode.py @@ -72,11 +72,15 @@ def _start_one(interpreter, targets, base_dir): files_file = tempfile.NamedTemporaryFile( mode="w", encoding="utf-8", delete=False, dir=PROJECT_ROOT ) + output_file = tempfile.NamedTemporaryFile( + mode="w", encoding="utf-8", delete=False, dir=PROJECT_ROOT + ) try: for _, path in targets: files_file.write(path) files_file.write("\n") files_file.close() + output_file.close() cmd = [ interpreter, DIS_DUMP, @@ -84,6 +88,8 @@ def _start_one(interpreter, targets, base_dir): base_dir, "--files-from", files_file.name, + "--output", + output_file.name, "--progress", "10", ] @@ -94,44 +100,128 @@ def _start_one(interpreter, targets, base_dir): env=env, cwd=PROJECT_ROOT, ) - return proc, files_file.name + return { + "proc": proc, + "files_file": files_file.name, + "output_file": output_file.name, + "targets": targets, + "interpreter": interpreter, + "base_dir": base_dir, + } except Exception: os.unlink(files_file.name) + os.unlink(output_file.name) raise -def _finish_one(proc, files_file): - """Wait for a single dis_dump.py process and return parsed JSON.""" +def _load_dump_output(output_file): + try: + with open(output_file, encoding="utf-8") as f: + content = f.read().strip() + except OSError as e: + print(" Failed to read dump output: %s" % e, file=sys.stderr) + return None + if not content: + return {} + try: + return json.loads(content) + except json.JSONDecodeError as e: + print(" JSON parse error: %s" % e, file=sys.stderr) + return None + + +def _run_sync_dump(interpreter, targets, base_dir, timeout=600): + job = _start_one(interpreter, targets, base_dir) + proc = job["proc"] try: stdout = proc.communicate(timeout=600)[0] except subprocess.TimeoutExpired: proc.kill() proc.communicate() print(" Timeout (600s)", file=sys.stderr) - os.unlink(files_file) - return {} + stdout = b"" + timed_out = True + finally: + timed_out = locals().get("timed_out", False) + + try: + data = _load_dump_output(job["output_file"]) finally: - if os.path.exists(files_file): - os.unlink(files_file) + for path in (job["files_file"], job["output_file"]): + if os.path.exists(path): + os.unlink(path) + if timed_out: + return {} if proc.returncode != 0: print(" Warning: exited with code %d" % proc.returncode, file=sys.stderr) + stray = stdout.decode(errors="replace").strip() + if stray: + print(" Warning: unexpected stdout from dump helper", file=sys.stderr) + return data - content = stdout.decode(errors="replace").strip() - if not content: - return {} + +def _rerun_missing_targets(interpreter, targets, base_dir): + recovered = {} + for target in targets: + data = _run_sync_dump(interpreter, [target], base_dir) + if data: + recovered.update(data) + return recovered + + +def _finish_one(job): + """Wait for a single dis_dump.py process and return parsed JSON.""" + proc = job["proc"] + expected = {relpath for relpath, _ in job["targets"]} try: - return json.loads(content) - except json.JSONDecodeError as e: - print(" JSON parse error: %s" % e, file=sys.stderr) - return {} + stdout = proc.communicate(timeout=600)[0] + except subprocess.TimeoutExpired: + proc.kill() + proc.communicate() + print( + " Timeout (600s), retrying %d file(s) serially" + % len(job["targets"]), + file=sys.stderr, + ) + data = None + else: + data = _load_dump_output(job["output_file"]) + finally: + for path in (job["files_file"], job["output_file"]): + if os.path.exists(path): + os.unlink(path) + + if proc.returncode != 0: + print(" Warning: exited with code %d" % proc.returncode, file=sys.stderr) + + stray = stdout.decode(errors="replace").strip() if "stdout" in locals() else "" + if stray: + print(" Warning: unexpected stdout from dump helper", file=sys.stderr) + + if data is None: + return _rerun_missing_targets( + job["interpreter"], job["targets"], job["base_dir"] + ) + + missing = [ + target for target in job["targets"] if target[0] not in data and target[0] in expected + ] + if missing: + print( + " Re-running %d missing file(s) serially" % len(missing), + file=sys.stderr, + ) + data.update( + _rerun_missing_targets(job["interpreter"], missing, job["base_dir"]) + ) + return data def start_dump(interpreter, targets, base_dir, num_workers=1): """Start dis_dump.py under the given interpreter, split across workers.""" if num_workers <= 1 or len(targets) <= num_workers: - proc, ff = _start_one(interpreter, targets, base_dir) - return [(proc, ff)] + return [_start_one(interpreter, targets, base_dir)] chunks = [[] for _ in range(num_workers)] for i, t in enumerate(targets): @@ -143,8 +233,8 @@ def start_dump(interpreter, targets, base_dir, num_workers=1): def finish_dump(procs): """Wait for all dis_dump.py processes and merge results.""" merged = {} - for proc, files_file in procs: - merged.update(_finish_one(proc, files_file)) + for job in procs: + merged.update(_finish_one(job)) return merged diff --git a/scripts/dis_dump.py b/scripts/dis_dump.py index e8b9c1bf5f8..23dfd1e0cda 100755 --- a/scripts/dis_dump.py +++ b/scripts/dis_dump.py @@ -8,6 +8,7 @@ Usage: python dis_dump.py Lib/ python dis_dump.py --base-dir Lib path/to/file.py + python dis_dump.py --base-dir Lib --output dump.json path/to/file.py """ import argparse @@ -352,6 +353,11 @@ def main(): default=0, help="Print a dot to stderr every N files processed", ) + parser.add_argument( + "--output", + default=None, + help="Write JSON output to this file instead of stdout", + ) args = parser.parse_args() targets = list(args.targets) @@ -386,7 +392,16 @@ def main(): sys.stderr.write(".") sys.stderr.flush() - json.dump(results, sys.stdout, ensure_ascii=False, separators=(",", ":")) + output = ( + open(args.output, "w", encoding="utf-8") + if args.output + else sys.stdout + ) + try: + json.dump(results, output, ensure_ascii=False, separators=(",", ":")) + finally: + if args.output: + output.close() if __name__ == "__main__": From f4674b6d9de0f22d9dd6bdfd37466ffaa2158347 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Sun, 19 Apr 2026 23:23:16 +0900 Subject: [PATCH 2/3] Bytecode parity - slice augassign, async comp inline - Augmented assignment with two-part slices uses BINARY_SLICE/STORE_SLICE - Inline async comprehensions (remove await/async guards) - Inlined comprehension cleanup jump uses JumpNoInterrupt - Class firstlineno uses first decorator line when decorated - Recurse into nested functions for static attribute collection - Fold unary positive complex constants (+0.0j) - Add deoptimize_borrow_for_folded_nonliteral_exprs pass - Add inline_single_predecessor_artificial_expr_exit_blocks pass - Skip shared artificial expr-stmt exit blocks in small-block inlining - Mark folded boolop tail as folded_from_nonliteral_expr --- Lib/test/test_code.py | 2 - Lib/test/test_compile.py | 4 - Lib/test/test_inspect/test_inspect.py | 1 - Lib/test/test_peepholer.py | 3 - crates/codegen/src/compile.rs | 506 ++++++++++++++++++++++---- crates/codegen/src/ir.rs | 314 +++++++++++++++- crates/codegen/src/symboltable.rs | 92 +++-- crates/vm/src/frame.rs | 42 ++- scripts/compare_bytecode.py | 7 +- scripts/dis_dump.py | 6 +- 10 files changed, 838 insertions(+), 139 deletions(-) diff --git a/Lib/test/test_code.py b/Lib/test/test_code.py index b4b15e29f26..6602c24353f 100644 --- a/Lib/test/test_code.py +++ b/Lib/test/test_code.py @@ -508,8 +508,6 @@ def foo(): with self.assertRaisesRegex(SystemError, msg): foo() - # TODO: RUSTPYTHON - @unittest.expectedFailure # @requires_debug_ranges() def test_co_positions_artificial_instructions(self): import dis diff --git a/Lib/test/test_compile.py b/Lib/test/test_compile.py index 94a9ae899b0..118ace15eaf 100644 --- a/Lib/test/test_compile.py +++ b/Lib/test/test_compile.py @@ -999,7 +999,6 @@ class C: dis.dis(code) self.assertNotIn('NOP', output.getvalue()) - @unittest.expectedFailure # TODO: RUSTPYTHON; AssertionError: unable to find constant -0.0 in (0.0,) def test_dont_merge_constants(self): # Issue #25843: compile() must not merge constants which are equal # but have a different type. @@ -1193,7 +1192,6 @@ def call(): line1 = call.__code__.co_firstlineno + 1 assert line1 not in [line for (_, _, line) in call.__code__.co_lines()] - @unittest.expectedFailure # TODO: RUSTPYTHON def test_lineno_after_implicit_return(self): TRUE = True # Don't use constant True or False, as compiler will remove test @@ -2495,7 +2493,6 @@ def f(self): self.assertIsInstance(C.__static_attributes__, tuple) self.assertEqual(sorted(C.__static_attributes__), ['a', 'b']) - @unittest.expectedFailure # TODO: RUSTPYTHON; AttributeError: type object 'C' has no attribute '__static_attributes__' def test_nested_function(self): class C: def f(self): @@ -2579,7 +2576,6 @@ def test_binop(self): def test_list(self): self.check_stack_size("[" + "x, " * self.N + "x]") - @unittest.expectedFailure # TODO: RUSTPYTHON; AssertionError: 101 not less than or equal to 6 def test_tuple(self): self.check_stack_size("(" + "x, " * self.N + "x)") diff --git a/Lib/test/test_inspect/test_inspect.py b/Lib/test/test_inspect/test_inspect.py index ea32d11b637..d479f21f558 100644 --- a/Lib/test/test_inspect/test_inspect.py +++ b/Lib/test/test_inspect/test_inspect.py @@ -1161,7 +1161,6 @@ def test_nested_class_definition(self): self.assertSourceEqual(mod2.cls183, 183, 188) self.assertSourceEqual(mod2.cls183.cls185, 185, 188) - @unittest.expectedFailure # TODO: RUSTPYTHON; pass def test_class_decorator(self): self.assertSourceEqual(mod2.cls196, 194, 201) self.assertSourceEqual(mod2.cls196.cls200, 198, 201) diff --git a/Lib/test/test_peepholer.py b/Lib/test/test_peepholer.py index 53ff218c4e1..8e76500f26f 100644 --- a/Lib/test/test_peepholer.py +++ b/Lib/test/test_peepholer.py @@ -199,7 +199,6 @@ def crater(): ],) self.check_lnotab(crater) - @unittest.expectedFailure # TODO: RUSTPYTHON def test_constant_folding_lists_of_constants(self): for line, elem in ( # in/not in constants with BUILD_LIST should be folded to a tuple: @@ -214,7 +213,6 @@ def test_constant_folding_lists_of_constants(self): self.assertNotInBytecode(code, 'BUILD_LIST') self.check_lnotab(code) - @unittest.expectedFailure # TODO: RUSTPYTHON def test_constant_folding_sets_of_constants(self): for line, elem in ( # in/not in constants with BUILD_SET should be folded to a frozenset: @@ -634,7 +632,6 @@ def g()->1+1: self.assertNotInBytecode(f, 'BINARY_OP') self.check_lnotab(f) - @unittest.expectedFailure # TODO: RUSTPYTHON; no BUILD_LIST to BUILD_TUPLE optimization def test_in_literal_list(self): def containtest(): return x in [a, b] diff --git a/crates/codegen/src/compile.rs b/crates/codegen/src/compile.rs index 34ea9aead72..2a3dc9e71db 100644 --- a/crates/codegen/src/compile.rs +++ b/crates/codegen/src/compile.rs @@ -504,6 +504,7 @@ impl Compiler { cellvars: IndexSet::default(), freevars: IndexSet::default(), fast_hidden: IndexMap::default(), + fast_hidden_final: IndexSet::default(), argcount: 0, posonlyargcount: 0, kwonlyargcount: 0, @@ -1323,6 +1324,7 @@ impl Compiler { cellvars: cellvar_cache, freevars: freevar_cache, fast_hidden: IndexMap::default(), + fast_hidden_final: IndexSet::default(), argcount: arg_count, posonlyargcount: posonlyarg_count, kwonlyargcount: kwonlyarg_count, @@ -2222,6 +2224,7 @@ impl Compiler { let symbol_scope = { let current_table = self.current_symbol_table(); if current_table.typ == CompilerScope::Class + && !self.current_code_info().in_inlined_comp && ((usage == NameUsage::Load && (name == "__class__" || name == "__classdict__" @@ -2354,7 +2357,15 @@ impl Compiler { NameOp::Name => { let namei = self.get_global_name_index(&name); match usage { - NameUsage::Load => emit!(self, Instruction::LoadName { namei }), + NameUsage::Load => { + if self.current_symbol_table().typ == CompilerScope::Class + && self.current_code_info().in_inlined_comp + { + self.emit_load_global(namei, false); + } else { + emit!(self, Instruction::LoadName { namei }); + } + } NameUsage::Store => emit!(self, Instruction::StoreName { namei }), NameUsage::Delete => emit!(self, Instruction::DeleteName { namei }), }; @@ -5006,39 +5017,42 @@ impl Compiler { fn collect_static_attributes(body: &[ast::Stmt], attrs: Option<&mut IndexSet>) { let Some(attrs) = attrs else { return }; for stmt in body { - let f = match stmt { - ast::Stmt::FunctionDef(f) => f, - _ => continue, - }; - // Skip @staticmethod and @classmethod decorated functions - let has_special_decorator = f.decorator_list.iter().any(|d| { - matches!(&d.expression, ast::Expr::Name(n) - if n.id.as_str() == "staticmethod" || n.id.as_str() == "classmethod") - }); - if has_special_decorator { - continue; - } - // Skip implicit classmethods (__init_subclass__, __class_getitem__) - let fname = f.name.as_str(); - if fname == "__init_subclass__" || fname == "__class_getitem__" { - continue; - } - // For __new__, scan for "self" (not the first param "cls") - if fname == "__new__" { - Self::scan_store_attrs(&f.body, "self", attrs); - continue; - } - let first_param = f - .parameters - .posonlyargs - .first() - .or(f.parameters.args.first()) - .map(|p| &p.parameter.name); - let Some(self_name) = first_param else { + let ast::Stmt::FunctionDef(f) = stmt else { continue; }; - Self::scan_store_attrs(&f.body, self_name.as_str(), attrs); + Self::scan_function_store_attrs(f, attrs); + } + } + + fn scan_function_store_attrs(f: &ast::StmtFunctionDef, attrs: &mut IndexSet) { + // Skip @staticmethod and @classmethod decorated functions + let has_special_decorator = f.decorator_list.iter().any(|d| { + matches!(&d.expression, ast::Expr::Name(n) + if n.id.as_str() == "staticmethod" || n.id.as_str() == "classmethod") + }); + if has_special_decorator { + return; + } + // Skip implicit classmethods (__init_subclass__, __class_getitem__) + let fname = f.name.as_str(); + if fname == "__init_subclass__" || fname == "__class_getitem__" { + return; + } + // For __new__, scan for "self" (not the first param "cls") + if fname == "__new__" { + Self::scan_store_attrs(&f.body, "self", attrs); + return; } + let first_param = f + .parameters + .posonlyargs + .first() + .or(f.parameters.args.first()) + .map(|p| &p.parameter.name); + let Some(self_name) = first_param else { + return; + }; + Self::scan_store_attrs(&f.body, self_name.as_str(), attrs); } /// Extract self.attr patterns from an assignment target expression. @@ -5120,6 +5134,9 @@ impl Compiler { Self::scan_store_attrs(&case.body, name, attrs); } } + ast::Stmt::FunctionDef(f) => { + Self::scan_function_store_attrs(f, attrs); + } _ => {} } } @@ -5364,7 +5381,16 @@ impl Compiler { self.prepare_decorators(decorator_list)?; let is_generic = type_params.is_some(); - let firstlineno = self.get_source_line_number().get().to_u32(); + let firstlineno = decorator_list + .first() + .map(|decorator| { + self.source_file + .to_source_code() + .line_index(decorator.expression.range().start()) + .get() + .to_u32() + }) + .unwrap_or_else(|| self.get_source_line_number().get().to_u32()); // Save context before entering any scopes let saved_ctx = self.ctx; @@ -5588,7 +5614,6 @@ impl Compiler { return Ok(()); }; - self.set_no_location(); emit!( self, PseudoInstruction::JumpNoInterrupt { delta: end_block } @@ -7439,7 +7464,7 @@ impl Compiler { ) -> CompileResult<()> { enum AugAssignKind<'a> { Name { id: &'a str }, - Subscript, + Subscript { use_slice_opt: bool }, Attr { idx: bytecode::NameIdx }, } @@ -7455,19 +7480,31 @@ impl Compiler { ctx: _, .. }) => { - // For augmented assignment, we need to load the value first - // But we can't use compile_subscript directly because we need DUP_TOP2 + let use_slice_opt = slice.should_use_slice_optimization(); self.compile_expression(value)?; - self.compile_expression(slice)?; - emit!(self, Instruction::Copy { i: 2 }); - emit!(self, Instruction::Copy { i: 2 }); - emit!( - self, - Instruction::BinaryOp { - op: BinaryOperator::Subscr - } - ); - AugAssignKind::Subscript + if use_slice_opt { + let ast::Expr::Slice(slice_expr) = slice.as_ref() else { + unreachable!( + "should_use_slice_optimization should only return true for ast::Expr::Slice" + ); + }; + self.compile_slice_two_parts(slice_expr)?; + emit!(self, Instruction::Copy { i: 3 }); + emit!(self, Instruction::Copy { i: 3 }); + emit!(self, Instruction::Copy { i: 3 }); + emit!(self, Instruction::BinarySlice); + } else { + self.compile_expression(slice)?; + emit!(self, Instruction::Copy { i: 2 }); + emit!(self, Instruction::Copy { i: 2 }); + emit!( + self, + Instruction::BinaryOp { + op: BinaryOperator::Subscr + } + ); + } + AugAssignKind::Subscript { use_slice_opt } } ast::Expr::Attribute(ast::ExprAttribute { value, attr, .. }) => { let attr = attr.as_str(); @@ -7490,11 +7527,19 @@ impl Compiler { // stack: RESULT self.compile_name(id, NameUsage::Store)?; } - AugAssignKind::Subscript => { - // stack: CONTAINER SLICE RESULT - emit!(self, Instruction::Swap { i: 3 }); - emit!(self, Instruction::Swap { i: 2 }); - emit!(self, Instruction::StoreSubscr); + AugAssignKind::Subscript { use_slice_opt } => { + if use_slice_opt { + // stack: CONTAINER START STOP RESULT + emit!(self, Instruction::Swap { i: 4 }); + emit!(self, Instruction::Swap { i: 3 }); + emit!(self, Instruction::Swap { i: 2 }); + emit!(self, Instruction::StoreSlice); + } else { + // stack: CONTAINER SLICE RESULT + emit!(self, Instruction::Swap { i: 3 }); + emit!(self, Instruction::Swap { i: 2 }); + emit!(self, Instruction::StoreSubscr); + } } AugAssignKind::Attr { idx } => { // stack: CONTAINER RESULT @@ -7578,7 +7623,6 @@ impl Compiler { let next2 = self.new_block(); self.compile_jump_if_inner(test, false, next2, source_range)?; self.compile_jump_if_inner(body, condition, target_block, source_range)?; - self.set_no_location(); emit!(self, PseudoInstruction::JumpNoInterrupt { delta: end }); self.set_no_location(); @@ -8069,6 +8113,7 @@ impl Compiler { } else { self.compile_bool_op(op, tail)?; } + self.mark_last_instruction_folded_from_nonliteral_expr(); return Ok(()); } } @@ -8486,7 +8531,6 @@ impl Compiler { // True case self.compile_expression(body)?; - self.set_no_location(); emit!( self, PseudoInstruction::JumpNoInterrupt { delta: after_block } @@ -8514,6 +8558,7 @@ impl Compiler { let name = self.mangle(id.as_str()); let info = self.code_stack.last_mut().unwrap(); info.metadata.fast_hidden.insert(name.to_string(), false); + info.metadata.fast_hidden_final.swap_remove(name.as_ref()); } self.compile_expression(value)?; emit!(self, Instruction::Copy { i: 1 }); @@ -9185,10 +9230,12 @@ impl Compiler { let is_inlined = self.is_inlined_comprehension_context(comprehension_type, &comp_table); - if is_inlined && !has_an_async_gen && !element_contains_await { - // PEP 709: Inlined comprehension - compile inline without new scope. - // CPython compiles the outermost iterable before entering the - // inlined-comprehension fast-hidden scope tweak. + if is_inlined { + // CPython inlines every non-generator comprehension that the + // symtable marked as comp_inlined, including async variants. + // codegen_comprehension() only branches on ste_comp_inlined here + // and relies on the inlined path itself to handle GET_AITER / + // async-comprehension cleanup. return self.compile_inlined_comprehension( comp_table, init_collection, @@ -9532,6 +9579,10 @@ impl Compiler { .metadata .fast_hidden .insert(name.clone(), true); + self.current_code_info() + .metadata + .fast_hidden_final + .insert(name.clone()); changed_fast_hidden.push(name.clone()); } } @@ -9674,8 +9725,14 @@ impl Compiler { emit!(self, PseudoInstruction::PopBlock); self.pop_fblock(FBlockType::TryExcept); - // Normal path: jump past cleanup - emit!(self, PseudoInstruction::Jump { delta: end_block }); + // Match CPython codegen_pop_inlined_comprehension_locals(): + // the synthetic jump that skips the exception cleanup uses + // JUMP_NO_INTERRUPT, which becomes JUMP_BACKWARD_NO_INTERRUPT + // when the cleanup tail sits above the final restore block. + emit!( + self, + PseudoInstruction::JumpNoInterrupt { delta: end_block } + ); // Exception cleanup path self.switch_to_block(cleanup_block); @@ -11356,6 +11413,18 @@ mod tests { compile_exec_with_options(source, opts) } + fn compile_single(source: &str) -> CodeObject { + let opts = CompileOpts::default(); + let source_file = SourceFileBuilder::new("source_path", source).finish(); + let parsed = ruff_python_parser::parse( + source_file.source_text(), + ruff_python_parser::Mode::Module.into(), + ) + .unwrap() + .into_syntax(); + compile_top(parsed, source_file, Mode::Single, opts).unwrap() + } + fn compile_exec_optimized(source: &str) -> CodeObject { let opts = CompileOpts { optimize: 1, @@ -11976,6 +12045,79 @@ x = (\"a\"[0]) or 2 ); } + #[test] + fn test_unary_positive_complex_constant_folds_to_load_const() { + let code = compile_exec( + "\ +x = +0.0j +", + ); + let ops: Vec<_> = code + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + assert!( + !ops.iter() + .any(|op| matches!(op, Instruction::CallIntrinsic1 { .. })), + "unary positive complex constant should not leave CALL_INTRINSIC_1, got ops={ops:?}" + ); + assert!( + matches!( + ops.as_slice(), + [ + Instruction::Resume { .. }, + Instruction::LoadConst { .. }, + Instruction::StoreName { .. }, + Instruction::LoadConst { .. }, + Instruction::ReturnValue + ] + ), + "expected module assignment to fold +0.0j into LOAD_CONST, got ops={ops:?}" + ); + } + + #[test] + fn test_folded_nonliteral_bool_op_tail_keeps_plain_load_fast() { + let code = compile_exec( + "\ +def and_true(x): + return True and x + +def or_false(x): + return False or x +", + ); + + for name in ["and_true", "or_false"] { + let function = find_code(&code, name).unwrap_or_else(|| panic!("missing {name} code")); + let ops: Vec<_> = function + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + assert!( + ops.iter() + .any(|op| matches!(op, Instruction::LoadFast { .. })), + "expected folded bool-op tail to keep LOAD_FAST in {name}, got ops={ops:?}" + ); + assert!( + !ops.iter().any(|op| { + matches!( + op, + Instruction::LoadFastBorrow { .. } + | Instruction::LoadFastBorrowLoadFastBorrow { .. } + ) + }), + "folded bool-op tail should not introduce borrow loads in {name}, got ops={ops:?}" + ); + } + } + #[test] fn test_nested_double_async_with() { assert_dis_snapshot!(compile_exec( @@ -12249,6 +12391,67 @@ def f(kwonlyargs, kw_only_defaults, arg2value): ); } + #[test] + fn test_augassign_two_part_slice_uses_slice_opcodes() { + let code = compile_exec( + "\ +def aug(x, a, b, y): + x[a:b] += y +", + ); + let aug = find_code(&code, "aug").expect("missing aug code"); + let ops: Vec<_> = aug + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + assert_eq!( + ops.iter() + .filter(|op| matches!(op, Instruction::BinarySlice)) + .count(), + 1, + "expected one BINARY_SLICE in augassign slice path, got ops={ops:?}" + ); + assert_eq!( + ops.iter() + .filter(|op| matches!(op, Instruction::StoreSlice)) + .count(), + 1, + "expected one STORE_SLICE in augassign slice path, got ops={ops:?}" + ); + assert!( + !ops.iter().any(|op| { + matches!( + op, + Instruction::BuildSlice { .. } | Instruction::StoreSubscr + ) + }), + "two-part augassign slice should avoid BUILD_SLICE/STORE_SUBSCR, got ops={ops:?}" + ); + assert!( + ops.windows(10).any(|window| { + matches!( + window, + [ + Instruction::Copy { .. }, + Instruction::Copy { .. }, + Instruction::Copy { .. }, + Instruction::BinarySlice, + Instruction::LoadFastBorrow { .. } | Instruction::LoadFast { .. }, + Instruction::BinaryOp { .. }, + Instruction::Swap { .. }, + Instruction::Swap { .. }, + Instruction::Swap { .. }, + Instruction::StoreSlice, + ] + ) + }), + "expected CPython-style augassign slice window, got ops={ops:?}" + ); + } + #[test] fn test_loop_return_reorders_backedge_before_exit_cleanup() { let code = compile_exec( @@ -13476,6 +13679,96 @@ class C: ); } + #[test] + fn test_nested_function_static_attributes_are_collected() { + let code = compile_exec( + "\ +class C: + def f(self): + self.x = 1 + self.y = 2 + self.x = 3 + + def g(self, obj): + self.y = 4 + self.z = 5 + + def h(self, a): + self.u = 6 + self.v = 7 + + obj.self = 8 +", + ); + let class_code = find_code(&code, "C").expect("missing class code"); + + assert!( + class_code.constants.iter().any(|constant| matches!( + constant, + ConstantData::Tuple { elements } + if elements + == &[ + ConstantData::Str { value: "u".into() }, + ConstantData::Str { value: "v".into() }, + ConstantData::Str { value: "x".into() }, + ConstantData::Str { value: "y".into() }, + ConstantData::Str { value: "z".into() }, + ] + )), + "expected nested function static attributes in class consts" + ); + } + + #[test] + fn test_decorated_class_uses_first_decorator_for_firstlineno() { + let code = compile_exec( + "\ +@dec1 +@dec2 +class C: + pass +", + ); + let class_code = find_code(&code, "C").expect("missing class code"); + let store_firstlineno = class_code + .instructions + .iter() + .position(|unit| { + matches!( + unit.op, + Instruction::StoreName { namei } + if class_code.names + [namei.get(OpArg::new(u32::from(u8::from(unit.arg)))) as usize] + .as_str() + == "__firstlineno__" + ) + }) + .expect("missing STORE_NAME __firstlineno__"); + let load_firstlineno = class_code + .instructions + .get(store_firstlineno.saturating_sub(1)) + .expect("missing LOAD_CONST for __firstlineno__"); + + let expected = ConstantData::Integer { + value: BigInt::from(1), + }; + assert!( + matches!( + load_firstlineno.op, + Instruction::LoadSmallInt { .. } | Instruction::LoadConst { .. } + ), + "expected LOAD_SMALL_INT/LOAD_CONST before __firstlineno__, got {:?}", + load_firstlineno.op + ); + if let Instruction::LoadConst { consti } = load_firstlineno.op { + let value = &class_code.constants + [consti.get(OpArg::new(u32::from(u8::from(load_firstlineno.arg))))]; + assert_eq!(value, &expected); + } else { + assert_eq!(u32::from(u8::from(load_firstlineno.arg)), 1); + } + } + #[test] fn test_future_annotations_class_keeps_conditional_annotations_cell() { let code = compile_exec( @@ -14001,7 +14294,7 @@ def f(x, y, z): } #[test] - fn test_constant_ifexp_stmt_in_loop_removes_empty_body() { + fn test_constant_if_expression_stmt_in_loop_removes_empty_body() { let code = compile_exec( "\ def f(x): @@ -14025,7 +14318,7 @@ def f(x): } #[test] - fn test_ifexp_in_jump_context_skips_constant_true_arm_load() { + fn test_if_expression_in_jump_context_skips_constant_true_arm_load() { let code = compile_exec( "\ def f(): @@ -14764,6 +15057,41 @@ l = lambda : [2 < x for x in [-1, 3, 0]] ); } + #[test] + fn test_async_dictcomp_in_async_function_is_inlined() { + let code = compile_exec( + "\ +async def f(items): + return {item: item async for item in items} +", + ); + let f = find_code(&code, "f").expect("missing f code"); + let ops: Vec<_> = f + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + assert!( + find_code(&code, "").is_none(), + "async dict comprehension should be inlined" + ); + assert!( + ops.iter().any(|op| matches!(op, Instruction::GetAIter)), + "inlined async dict comprehension should keep GET_AITER in outer code, got ops={ops:?}" + ); + assert!( + ops.iter() + .any(|op| matches!(op, Instruction::LoadFastAndClear { .. })), + "inlined async dict comprehension should use LOAD_FAST_AND_CLEAR, got ops={ops:?}" + ); + assert!( + !ops.iter().any(|op| matches!(op, Instruction::MakeFunction)), + "inlined async dict comprehension should not materialize MAKE_FUNCTION, got ops={ops:?}" + ); + } + #[test] fn test_nested_module_scope_dictcomp_symbols_are_local() { let symbol_table = scan_program_symbol_table( @@ -14890,6 +15218,58 @@ _pathseps_with_colon = {f':{s}' for s in path_separators} ); } + #[test] + fn test_function_scope_inlined_comprehension_restore_keeps_swap_before_duplicate_store() { + let code = compile_exec( + "\ +def f(): + a = [1 for a in [0]] + return 1 +", + ); + let f = find_code(&code, "f").expect("missing f code"); + let ops: Vec<_> = f + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + assert!( + ops.windows(4).any(|window| matches!( + window, + [ + Instruction::PopIter, + Instruction::Swap { .. }, + Instruction::StoreFast { .. }, + Instruction::StoreFast { .. } + ] + )), + "expected PopIter/SWAP 2/STORE_FAST/STORE_FAST restore tail, got ops={ops:?}" + ); + } + + #[test] + fn test_single_mode_folded_multiline_constant_does_not_leave_nops() { + let code = compile_single( + "\ +(- + - + - + 1) +", + ); + + assert!( + !code + .instructions + .iter() + .any(|unit| matches!(unit.op, Instruction::Nop)), + "expected folded single-mode multiline constant to drop NOP anchors, got instructions={:?}", + code.instructions + ); + } + #[test] fn test_or_condition_in_jump_context_uses_shared_true_fallthrough() { let code = compile_exec( @@ -15771,9 +16151,9 @@ def f(msg): fn test_protected_conditional_tail_keeps_strong_load_fast() { let code = compile_exec( "\ -def f(m, klass, category, warning_base): +def f(m, class_name, category, warning_base): try: - cat = getattr(m, klass) + cat = getattr(m, class_name) except AttributeError: raise ValueError(category) if not issubclass(cat, warning_base): diff --git a/crates/codegen/src/ir.rs b/crates/codegen/src/ir.rs index 73c8028d74b..dc455c7d1a2 100644 --- a/crates/codegen/src/ir.rs +++ b/crates/codegen/src/ir.rs @@ -45,6 +45,7 @@ pub struct CodeUnitMetadata { pub cellvars: IndexSet, // u_cellvars pub freevars: IndexSet, // u_freevars pub fast_hidden: IndexMap, // u_fast_hidden + pub fast_hidden_final: IndexSet, // final CO_FAST_HIDDEN names pub argcount: u32, // u_argcount pub posonlyargcount: u32, // u_posonlyargcount pub kwonlyargcount: u32, // u_kwonlyargcount @@ -285,6 +286,7 @@ impl CodeInfo { // CPython resolves line numbers once before cold-block extraction and // again after reordering blocks. resolve_line_numbers(&mut self.blocks); + inline_single_predecessor_artificial_expr_exit_blocks(&mut self.blocks); push_cold_blocks_to_end(&mut self.blocks); reorder_conditional_chain_and_jump_back_blocks(&mut self.blocks); @@ -315,6 +317,7 @@ impl CodeInfo { inline_pop_except_return_blocks(&mut self.blocks); duplicate_named_except_cleanup_returns(&mut self.blocks, &self.metadata); self.eliminate_unreachable_blocks(); + resolve_line_numbers(&mut self.blocks); let cellfixedoffsets = build_cellfixedoffsets( &self.metadata.varnames, &self.metadata.cellvars, @@ -335,6 +338,7 @@ impl CodeInfo { self.compute_load_fast_start_depths(); // optimize_load_fast: after normalize_jumps self.optimize_load_fast_borrow(); + self.deoptimize_borrow_for_folded_nonliteral_exprs(); self.deoptimize_borrow_after_multi_handler_resume_join(); self.deoptimize_borrow_after_named_except_cleanup_join(); self.deoptimize_borrow_in_protected_conditional_tail(); @@ -375,6 +379,7 @@ impl CodeInfo { cellvars: cellvar_cache, freevars: freevar_cache, fast_hidden, + fast_hidden_final, argcount: arg_count, posonlyargcount: posonlyarg_count, kwonlyargcount: kwonlyarg_count, @@ -663,7 +668,9 @@ impl CodeInfo { } // Apply CO_FAST_HIDDEN for inlined comprehension variables for (name, &hidden) in &fast_hidden { - if hidden && let Some(idx) = varname_cache.get_index_of(name.as_str()) { + if (hidden || fast_hidden_final.contains(name)) + && let Some(idx) = varname_cache.get_index_of(name.as_str()) + { localspluskinds[idx] |= CO_FAST_HIDDEN; } } @@ -856,6 +863,11 @@ impl CodeInfo { ) => Some(ConstantData::Integer { value: BigInt::from(i32::from(*value)), }), + ( + ConstantData::Complex { value }, + Instruction::CallIntrinsic1 { .. }, + Some(oparg::IntrinsicFunction1::UnaryPositive), + ) => Some(ConstantData::Complex { value: *value }), _ => None, } } @@ -902,6 +914,18 @@ impl CodeInfo { { let (const_idx, _) = self.metadata.consts.insert_full(folded_const); set_to_nop(&mut block.instructions[operand_index]); + block.instructions[operand_index].location = block.instructions[i].location; + block.instructions[operand_index].end_location = + block.instructions[i].end_location; + let mut prev = operand_index; + while let Some(idx) = prev.checked_sub(1) { + if !matches!(block.instructions[idx].instr.real(), Some(Instruction::Nop)) { + break; + } + block.instructions[idx].location = block.instructions[i].location; + block.instructions[idx].end_location = block.instructions[i].end_location; + prev = idx; + } block.instructions[i].instr = Instruction::LoadConst { consti: Arg::marker(), } @@ -2107,10 +2131,16 @@ impl CodeInfo { for i in 0..instructions.len().saturating_sub(1) { let lhs = &instructions[i]; let rhs = &instructions[i + 1]; + let preceded_by_swap = i > 0 + && matches!( + instructions[i - 1].instr.real(), + Some(Instruction::Swap { .. }) + ); if !matches!(lhs.instr.real(), Some(Instruction::StoreFast { .. })) || !matches!(rhs.instr.real(), Some(Instruction::StoreFast { .. })) || u32::from(lhs.arg) != u32::from(rhs.arg) || instruction_lineno(lhs) != instruction_lineno(rhs) + || preceded_by_swap { continue; } @@ -3416,6 +3446,31 @@ impl CodeInfo { } } + fn deoptimize_borrow_for_folded_nonliteral_exprs(&mut self) { + for block in &mut self.blocks { + for info in &mut block.instructions { + if !info.folded_from_nonliteral_expr { + continue; + } + match info.instr.real() { + Some(Instruction::LoadFastBorrow { .. }) => { + info.instr = Instruction::LoadFast { + var_num: Arg::marker(), + } + .into(); + } + Some(Instruction::LoadFastBorrowLoadFastBorrow { .. }) => { + info.instr = Instruction::LoadFastLoadFast { + var_nums: Arg::marker(), + } + .into(); + } + _ => {} + } + } + } + } + fn deoptimize_borrow_after_push_exc_info(&mut self) { for block in &mut self.blocks { let mut in_exception_state = false; @@ -4161,6 +4216,7 @@ impl CodeInfo { self.add_checks_for_loads_of_uninitialized_variables(); self.insert_superinstructions(); resolve_line_numbers(&mut self.blocks); + inline_single_predecessor_artificial_expr_exit_blocks(&mut self.blocks); trace.push(( "after_first_resolve_line_numbers".to_owned(), self.debug_block_dump(), @@ -4219,12 +4275,38 @@ impl CodeInfo { self.debug_block_dump(), )); + resolve_line_numbers(&mut self.blocks); + trace.push(( + "after_final_resolve_line_numbers".to_owned(), + self.debug_block_dump(), + )); + + self.remove_redundant_const_pop_top_pairs(); remove_redundant_nops_and_jumps(&mut self.blocks); trace.push(( "after_remove_redundant_nops_and_jumps".to_owned(), self.debug_block_dump(), )); + jump_threading_unconditional(&mut self.blocks); + 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); + inline_pop_except_return_blocks(&mut self.blocks); + duplicate_named_except_cleanup_returns(&mut self.blocks, &self.metadata); + self.eliminate_unreachable_blocks(); + trace.push(( + "after_final_cfg_cleanup".to_owned(), + self.debug_block_dump(), + )); + + resolve_line_numbers(&mut self.blocks); + trace.push(( + "after_post_cleanup_resolve_line_numbers".to_owned(), + self.debug_block_dump(), + )); + let cellfixedoffsets = build_cellfixedoffsets( &self.metadata.varnames, &self.metadata.cellvars, @@ -4244,6 +4326,7 @@ impl CodeInfo { self.debug_block_dump(), )); self.optimize_load_fast_borrow(); + self.deoptimize_borrow_for_folded_nonliteral_exprs(); self.deoptimize_borrow_after_multi_handler_resume_join(); self.deoptimize_borrow_after_named_except_cleanup_join(); self.deoptimize_borrow_in_protected_conditional_tail(); @@ -5180,6 +5263,17 @@ fn inline_small_or_no_lineno_blocks(blocks: &mut [Block]) { }; loop { let mut changes = false; + let mut predecessors = vec![0usize; blocks.len()]; + for block in blocks.iter() { + if block.next != BlockIdx::NULL { + predecessors[block.next.idx()] += 1; + } + for info in &block.instructions { + if info.target != BlockIdx::NULL { + predecessors[info.target.idx()] += 1; + } + } + } let mut current = BlockIdx(0); while current != BlockIdx::NULL { let next = blocks[current.idx()].next; @@ -5205,7 +5299,13 @@ fn inline_small_or_no_lineno_blocks(blocks: &mut [Block]) { && blocks[target.idx()].instructions.len() <= MAX_COPY_SIZE; let no_lineno_no_fallthrough = block_has_no_lineno(&blocks[target.idx()]) && !block_has_fallthrough(&blocks[target.idx()]); - if small_exit_block || no_lineno_no_fallthrough { + let shared_artificial_expr_exit = small_exit_block + && predecessors[target.idx()] > 1 + && is_artificial_expr_stmt_exit_block(&blocks[target.idx()]) + && !instruction_has_lineno(&blocks[target.idx()].instructions[0]) + && !instruction_has_lineno(&blocks[target.idx()].instructions[1]) + && !instruction_has_lineno(&blocks[target.idx()].instructions[2]); + if !shared_artificial_expr_exit && (small_exit_block || no_lineno_no_fallthrough) { let removed_jump_had_lineno = blocks[current.idx()] .instructions .last() @@ -5232,6 +5332,74 @@ fn inline_small_or_no_lineno_blocks(blocks: &mut [Block]) { } } +fn is_artificial_expr_stmt_exit_block(block: &Block) -> bool { + matches!( + block.instructions.as_slice(), + [ + InstructionInfo { + instr: AnyInstruction::Real(Instruction::PopTop), + .. + }, + InstructionInfo { + instr: AnyInstruction::Real(Instruction::LoadConst { .. }), + .. + }, + InstructionInfo { + instr: AnyInstruction::Real(Instruction::ReturnValue), + .. + } + ] + ) +} + +fn inline_single_predecessor_artificial_expr_exit_blocks(blocks: &mut [Block]) { + let predecessors = compute_predecessors(blocks); + + for idx in 0..blocks.len() { + let Some(last) = blocks[idx].instructions.last().copied() else { + continue; + }; + if !last.instr.is_unconditional_jump() || last.target == BlockIdx::NULL { + continue; + } + + let target = next_nonempty_block(blocks, last.target); + if target == BlockIdx::NULL + || predecessors[target.idx()] != 1 + || !is_artificial_expr_stmt_exit_block(&blocks[target.idx()]) + { + continue; + } + + let is_jump_wrapper = blocks[idx] + .instructions + .split_last() + .is_some_and(|(_, prefix)| { + prefix + .iter() + .all(|ins| matches!(ins.instr.real(), Some(Instruction::Nop))) + }); + if is_jump_wrapper { + continue; + } + + if blocks[idx] + .instructions + .last() + .is_some_and(instruction_has_lineno) + { + if let Some(last_instr) = blocks[idx].instructions.last_mut() { + set_to_nop(last_instr); + } + } else { + let _ = blocks[idx].instructions.pop(); + } + blocks[idx] + .instructions + .extend(blocks[target.idx()].instructions.clone()); + } +} + fn remove_redundant_nops_in_blocks(blocks: &mut [Block]) -> usize { let mut changes = 0; let mut block_order = Vec::new(); @@ -5606,6 +5774,28 @@ fn block_has_no_lineno(block: &Block) -> bool { .all(|ins| !instruction_has_lineno(ins)) } +fn shared_jump_back_target(block: &Block) -> Option { + if !block_has_no_lineno(block) { + return None; + } + + let (last, prefix) = block.instructions.split_last()?; + if !last.instr.is_unconditional_jump() || last.target == BlockIdx::NULL { + return None; + } + + if !prefix.iter().all(|info| { + matches!( + info.instr.real(), + Some(Instruction::PopExcept) | Some(Instruction::Nop) + ) + }) { + return None; + } + + Some(last.target) +} + fn is_jump_only_block(block: &Block) -> bool { let [instr] = block.instructions.as_slice() else { return false; @@ -6243,6 +6433,42 @@ fn compute_predecessors(blocks: &[Block]) -> Vec { predecessors } +fn record_incoming_origin(origins: &mut [Vec], target: BlockIdx, source: BlockIdx) { + let incoming = &mut origins[target.idx()]; + if !incoming.contains(&source) { + incoming.push(source); + } +} + +fn compute_incoming_origins(blocks: &[Block], reachable: &[bool]) -> Vec> { + let mut origins = vec![Vec::new(); blocks.len()]; + let mut current = BlockIdx(0); + while current != BlockIdx::NULL { + if !reachable[current.idx()] { + current = blocks[current.idx()].next; + continue; + } + + let block = &blocks[current.idx()]; + if block_has_fallthrough(block) { + let next = next_nonempty_block(blocks, block.next); + if next != BlockIdx::NULL && reachable[next.idx()] { + record_incoming_origin(&mut origins, next, current); + } + } + for ins in &block.instructions { + if ins.target != BlockIdx::NULL { + let target = next_nonempty_block(blocks, ins.target); + if target != BlockIdx::NULL && reachable[target.idx()] { + record_incoming_origin(&mut origins, target, current); + } + } + } + current = block.next; + } + origins +} + fn duplicate_exits_without_lineno(blocks: &mut Vec, predecessors: &mut Vec) { let mut current = BlockIdx(0); while current != BlockIdx::NULL { @@ -6287,6 +6513,8 @@ fn duplicate_exits_without_lineno(blocks: &mut Vec, predecessors: &mut Ve current = blocks[current.idx()].next; } + let reachable = compute_reachable_blocks(blocks); + let incoming_origins = compute_incoming_origins(blocks, &reachable); current = BlockIdx(0); while current != BlockIdx::NULL { let block = &blocks[current.idx()]; @@ -6295,7 +6523,14 @@ fn duplicate_exits_without_lineno(blocks: &mut Vec, predecessors: &mut Ve { let target = next_nonempty_block(blocks, block.next); if target != BlockIdx::NULL - && predecessors[target.idx()] == 1 + && (predecessors[target.idx()] == 1 + || has_unique_fallthrough_origin( + blocks, + &reachable, + &incoming_origins, + current, + target, + )) && is_exit_without_lineno(&blocks[target.idx()]) && let Some((location, end_location)) = propagation_location(last) && let Some(first) = blocks[target.idx()].instructions.first_mut() @@ -6308,6 +6543,8 @@ fn duplicate_exits_without_lineno(blocks: &mut Vec, predecessors: &mut Ve } fn propagate_line_numbers(blocks: &mut [Block], predecessors: &[u32]) { + let reachable = compute_reachable_blocks(blocks); + let incoming_origins = compute_incoming_origins(blocks, &reachable); let mut current = BlockIdx(0); while current != BlockIdx::NULL { let last = blocks[current.idx()].instructions.last().copied(); @@ -6331,7 +6568,14 @@ fn propagate_line_numbers(blocks: &mut [Block], predecessors: &[u32]) { if has_fallthrough { let target = next_nonempty_block(blocks, next_block); if target != BlockIdx::NULL - && predecessors[target.idx()] == 1 + && (predecessors[target.idx()] == 1 + || has_unique_fallthrough_origin( + blocks, + &reachable, + &incoming_origins, + current, + target, + )) && let Some((location, end_location)) = propagation_location(&last) && let Some(first) = blocks[target.idx()].instructions.first_mut() { @@ -6380,6 +6624,42 @@ fn find_layout_predecessor(blocks: &[Block], target: BlockIdx) -> BlockIdx { BlockIdx::NULL } +fn has_unique_fallthrough_origin( + blocks: &[Block], + reachable: &[bool], + incoming_origins: &[Vec], + source: BlockIdx, + target: BlockIdx, +) -> bool { + if source == BlockIdx::NULL + || target == BlockIdx::NULL + || !reachable[source.idx()] + || !block_has_fallthrough(&blocks[source.idx()]) + || next_nonempty_block(blocks, blocks[source.idx()].next) != target + { + return false; + } + + let mut allowed = vec![false; blocks.len()]; + allowed[source.idx()] = true; + + let mut current = blocks[source.idx()].next; + while current != BlockIdx::NULL && current != target { + if !blocks[current.idx()].instructions.is_empty() { + return false; + } + allowed[current.idx()] = true; + current = blocks[current.idx()].next; + } + if current != target { + return false; + } + + incoming_origins[target.idx()] + .iter() + .all(|origin| allowed[origin.idx()]) +} + fn comes_before(blocks: &[Block], first: BlockIdx, second: BlockIdx) -> bool { let mut current = BlockIdx(0); while current != BlockIdx::NULL { @@ -6400,12 +6680,11 @@ fn duplicate_shared_jump_back_targets(blocks: &mut Vec) { for target in 0..blocks.len() { let target = BlockIdx(target as u32); - if !is_jump_only_block(&blocks[target.idx()]) || !block_has_no_lineno(&blocks[target.idx()]) - { + let Some(jump_target) = shared_jump_back_target(&blocks[target.idx()]) else { continue; - } + }; - let jump_target = next_nonempty_block(blocks, blocks[target.idx()].instructions[0].target); + let jump_target = next_nonempty_block(blocks, jump_target); if jump_target == BlockIdx::NULL || !comes_before(blocks, jump_target, target) { continue; } @@ -6425,18 +6704,15 @@ fn duplicate_shared_jump_back_targets(blocks: &mut Vec) { continue; } - let Some(instr_idx) = trailing_conditional_jump_index(&blocks[block_idx.idx()]) else { - continue; - }; - if next_nonempty_block( - blocks, - blocks[block_idx.idx()].instructions[instr_idx].target, - ) != target - { - continue; + for (instr_idx, info) in blocks[block_idx.idx()].instructions.iter().enumerate() { + if !is_jump_instruction(info) || info.target == BlockIdx::NULL { + continue; + } + if next_nonempty_block(blocks, info.target) != target { + continue; + } + clones.push((target, block_idx, instr_idx)); } - - clones.push((target, block_idx, instr_idx)); } } diff --git a/crates/codegen/src/symboltable.rs b/crates/codegen/src/symboltable.rs index 4ab2ba65548..ceaff167eee 100644 --- a/crates/codegen/src/symboltable.rs +++ b/crates/codegen/src/symboltable.rs @@ -317,26 +317,6 @@ fn drop_class_free(symbol_table: &mut SymbolTable, newfree: &mut IndexSet bool { - use ast::visitor::Visitor; - struct AwaitFinder(bool); - impl ast::visitor::Visitor<'_> for AwaitFinder { - fn visit_expr(&mut self, expr: &ast::Expr) { - if !self.0 { - if matches!(expr, ast::Expr::Await(_)) { - self.0 = true; - } else { - ast::visitor::walk_expr(self, expr); - } - } - } - } - let mut finder = AwaitFinder(false); - finder.visit_expr(expr); - finder.0 -} - /// PEP 709: Merge symbols from an inlined comprehension into the parent scope. /// Matches symtable.c inline_comprehension(). fn inline_comprehension( @@ -345,7 +325,8 @@ fn inline_comprehension( comp_free: &mut IndexSet, inlined_cells: &mut IndexSet, parent_type: CompilerScope, -) { +) -> IndexSet { + let mut removed_class_implicit = IndexSet::default(); for (name, sub_symbol) in &comp.symbols { // Skip the .0 parameter if sub_symbol.flags.contains(SymbolFlags::PARAMETER) { @@ -362,9 +343,12 @@ fn inline_comprehension( // Handle __class__ in ClassBlock let scope = if sub_symbol.scope == SymbolScope::Free && parent_type == CompilerScope::Class - && name == "__class__" - { + && matches!( + name.as_str(), + "__class__" | "__classdict__" | "__conditional_annotations__" + ) { comp_free.swap_remove(name); + removed_class_implicit.insert(name.clone()); SymbolScope::GlobalImplicit } else { sub_symbol.scope @@ -393,6 +377,7 @@ fn inline_comprehension( parent_symbols.insert(name.clone(), symbol); } } + removed_class_implicit } type SymbolMap = IndexMap; @@ -557,13 +542,18 @@ impl SymbolTableAnalyzer { let mut newfree = IndexSet::default(); for (idx, (mut child_free, is_inlined)) in child_frees.into_iter().enumerate() { if is_inlined { - inline_comprehension( + let removed_class_implicit = inline_comprehension( &mut symbol_table.symbols, - &sub_tables[idx], + &symbol_table.sub_tables[idx], &mut child_free, &mut inlined_cells, symbol_table.typ, ); + for name in removed_class_implicit { + symbol_table.sub_tables[idx] + .symbols + .shift_remove(name.as_str()); + } } newfree.extend(child_free); } @@ -579,6 +569,12 @@ impl SymbolTableAnalyzer { let sub_tables = &*symbol_table.sub_tables; + for symbol in symbol_table.symbols.values_mut() { + if inlined_cells.contains(&symbol.name) { + symbol.flags.insert(SymbolFlags::COMP_CELL); + } + } + // Analyze symbols in current scope for symbol in symbol_table.symbols.values_mut() { self.analyze_symbol(symbol, symbol_table.typ, sub_tables, class_entry)?; @@ -589,13 +585,24 @@ impl SymbolTableAnalyzer { } } - // PEP 709: Promote LOCAL to CELL and set COMP_CELL for inlined cell vars + // PEP 709 / CPython symtable.c: + // - only promote LOCAL -> CELL in function-like scopes, where + // analyze_cells() runs. Module and class scopes keep their normal + // scope and rely on DEF_COMP_CELL for comprehension-only cells. + let promote_inlined_cells_to_cell = matches!( + symbol_table.typ, + CompilerScope::Function + | CompilerScope::AsyncFunction + | CompilerScope::Lambda + | CompilerScope::Comprehension + | CompilerScope::Annotation + ); for symbol in symbol_table.symbols.values_mut() { - if inlined_cells.contains(&symbol.name) { - if symbol.scope == SymbolScope::Local { - symbol.scope = SymbolScope::Cell; - } - symbol.flags.insert(SymbolFlags::COMP_CELL); + if inlined_cells.contains(&symbol.name) + && promote_inlined_cells_to_cell + && symbol.scope == SymbolScope::Local + { + symbol.scope = SymbolScope::Cell; } } @@ -679,8 +686,19 @@ impl SymbolTableAnalyzer { SymbolScope::Unknown => { // Try hard to figure out what the scope of this symbol is. let scope = if symbol.is_bound() { - self.found_in_inner_scope(sub_tables, &symbol.name, st_typ) - .unwrap_or(SymbolScope::Local) + if symbol.flags.contains(SymbolFlags::COMP_CELL) + && matches!(st_typ, CompilerScope::Module | CompilerScope::Class) + { + // CPython keeps comprehension-only cells in + // module/class scopes as normal local/name + // bindings and uses DEF_COMP_CELL to allocate the + // synthetic cell slot. The spliced comp child + // should not force the outer name itself to CELL. + SymbolScope::Local + } else { + self.found_in_inner_scope(sub_tables, &symbol.name, st_typ) + .unwrap_or(SymbolScope::Local) + } } else if let Some(scope) = self.found_in_outer_scope(&symbol.name, st_typ) { // If found in enclosing scope (function/TypeParams), use that scope @@ -2207,10 +2225,10 @@ impl SymbolTableBuilder { self.tables.last_mut().unwrap().is_generator = is_generator; // PEP 709: Mark non-generator comprehensions for inlining. - // Excluded: generator expressions, async comprehensions, - // and annotation scopes nested in classes (can_see_class_scope). - let element_has_await = expr_contains_await(elt1) || elt2.is_some_and(expr_contains_await); - if !is_generator && !has_async_gen && !element_has_await { + // CPython's symtable marks all non-generator comprehensions for + // inlining, except annotation scopes nested in classes that can see + // class scope. + if !is_generator { let parent = self.tables.iter().rev().nth(1); let parent_can_see_class = parent.is_some_and(|t| t.can_see_class_scope); if !parent_can_see_class { diff --git a/crates/vm/src/frame.rs b/crates/vm/src/frame.rs index 180c4fad0ed..e54e25672c2 100644 --- a/crates/vm/src/frame.rs +++ b/crates/vm/src/frame.rs @@ -898,6 +898,32 @@ impl Frame { let nfrees = code.freevars.len(); let free_start = nlocalsplus - nfrees; let is_optimized = code.flags.contains(bytecode::CodeFlags::OPTIMIZED); + let has_active_hidden_locals = !is_optimized + && code.localspluskinds.iter().enumerate().any(|(i, &kind)| { + if kind & CO_FAST_HIDDEN == 0 { + return false; + } + match fastlocals[i].as_ref() { + None => false, + Some(obj) => { + if kind & (CO_FAST_CELL | CO_FAST_FREE) != 0 { + obj.downcast_ref::() + .is_none_or(|cell| cell.get().is_some()) + } else { + true + } + } + } + }); + let overlay_locals = if has_active_hidden_locals { + // Match CPython's PyEval_GetLocals() behavior for frames with + // PEP 709 hidden locals: locals() inside an inlined comprehension + // returns a snapshot of the active fast locals only, not the + // backing module/class/eval mapping. + Some(vm.ctx.new_dict()) + } else { + None + }; // Track which non-merged cellvar index we're at let mut nonmerged_cell_idx = 0; @@ -975,13 +1001,25 @@ impl Frame { fastlocals[i].clone() }; - match locals_map.ass_subscript(name, value, vm) { + let result = if let Some(dict) = &overlay_locals { + match value { + Some(value) => dict.set_item(name, value, vm), + None => dict.del_item(name, vm), + } + } else { + locals_map.ass_subscript(name, value, vm) + }; + match result { Ok(()) => {} Err(e) if e.fast_isinstance(vm.ctx.exceptions.key_error) => {} Err(e) => return Err(e), } } - Ok(locals.clone_mapping(vm)) + if let Some(dict) = overlay_locals { + Ok(ArgMapping::from_dict_exact(dict)) + } else { + Ok(locals.clone_mapping(vm)) + } } } diff --git a/scripts/compare_bytecode.py b/scripts/compare_bytecode.py index 48da77688d8..7e441b02f00 100644 --- a/scripts/compare_bytecode.py +++ b/scripts/compare_bytecode.py @@ -180,8 +180,7 @@ def _finish_one(job): proc.kill() proc.communicate() print( - " Timeout (600s), retrying %d file(s) serially" - % len(job["targets"]), + " Timeout (600s), retrying %d file(s) serially" % len(job["targets"]), file=sys.stderr, ) data = None @@ -205,7 +204,9 @@ def _finish_one(job): ) missing = [ - target for target in job["targets"] if target[0] not in data and target[0] in expected + target + for target in job["targets"] + if target[0] not in data and target[0] in expected ] if missing: print( diff --git a/scripts/dis_dump.py b/scripts/dis_dump.py index 23dfd1e0cda..813de22e658 100755 --- a/scripts/dis_dump.py +++ b/scripts/dis_dump.py @@ -392,11 +392,7 @@ def main(): sys.stderr.write(".") sys.stderr.flush() - output = ( - open(args.output, "w", encoding="utf-8") - if args.output - else sys.stdout - ) + output = open(args.output, "w", encoding="utf-8") if args.output else sys.stdout try: json.dump(results, output, ensure_ascii=False, separators=(",", ":")) finally: From a04edb12deecb7302bf85bebff141157e055eafa Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Mon, 20 Apr 2026 22:52:37 +0900 Subject: [PATCH 3/3] apply reviews --- crates/codegen/src/symboltable.rs | 4 +- crates/vm/src/builtins/frame.rs | 3 +- crates/vm/src/frame.rs | 113 ++++++++++++++++++++---------- scripts/compare_bytecode.py | 74 ++++++++++++++----- 4 files changed, 136 insertions(+), 58 deletions(-) diff --git a/crates/codegen/src/symboltable.rs b/crates/codegen/src/symboltable.rs index ceaff167eee..f19453bc59a 100644 --- a/crates/codegen/src/symboltable.rs +++ b/crates/codegen/src/symboltable.rs @@ -2226,8 +2226,8 @@ impl SymbolTableBuilder { // PEP 709: Mark non-generator comprehensions for inlining. // CPython's symtable marks all non-generator comprehensions for - // inlining, except annotation scopes nested in classes that can see - // class scope. + // inlining, except scopes nested under a parent that can see class + // scope (for example annotation scopes inside classes). if !is_generator { let parent = self.tables.iter().rev().nth(1); let parent_can_see_class = parent.is_some_and(|t| t.can_see_class_scope); diff --git a/crates/vm/src/builtins/frame.rs b/crates/vm/src/builtins/frame.rs index 4601eee4467..42615e81a7a 100644 --- a/crates/vm/src/builtins/frame.rs +++ b/crates/vm/src/builtins/frame.rs @@ -454,7 +454,7 @@ impl Frame { #[pygetset] fn f_locals(&self, vm: &VirtualMachine) -> PyResult { - let result = self.locals(vm).map(Into::into); + let result = self.f_locals_mapping(vm).map(Into::into); self.locals_dirty .store(true, core::sync::atomic::Ordering::Release); result @@ -703,6 +703,7 @@ impl Py { // Clear temporary refs self.temporary_refs.lock().clear(); + self.f_locals_hidden_overlay.lock().take(); Ok(()) } diff --git a/crates/vm/src/frame.rs b/crates/vm/src/frame.rs index e54e25672c2..93783b1e5cc 100644 --- a/crates/vm/src/frame.rs +++ b/crates/vm/src/frame.rs @@ -30,7 +30,7 @@ use crate::{ function::{ArgMapping, Either, FuncArgs, PyMethodFlags}, object::PyAtomicBorrow, object::{Traverse, TraverseFn}, - protocol::{PyIter, PyIterReturn}, + protocol::{PyIter, PyIterReturn, PyMapping}, scope::Scope, sliceable::SliceableSequenceOp, stdlib::{_typing, builtins, sys::monitoring}, @@ -609,6 +609,9 @@ pub struct InterpreterFrame { pub(crate) owner: atomic::AtomicI8, /// Set when f_locals is accessed. Cleared after locals_to_fast() sync. pub(crate) locals_dirty: atomic::AtomicBool, + /// Persistent overlay for `frame.f_locals` when hidden locals need a + /// snapshot separate from the backing locals mapping. + pub(crate) f_locals_hidden_overlay: PyMutex>, /// Number of stack entries to pop after set_f_lineno returns to the /// execution loop. set_f_lineno cannot pop directly because the /// execution loop holds the state mutex. @@ -660,6 +663,7 @@ unsafe impl Traverse for Frame { iframe.builtins.traverse(tracer_fn); iframe.trace.traverse(tracer_fn); iframe.temporary_refs.traverse(tracer_fn); + iframe.f_locals_hidden_overlay.traverse(tracer_fn); } } @@ -739,6 +743,7 @@ impl Frame { previous: AtomicPtr::new(core::ptr::null_mut()), owner: atomic::AtomicI8::new(FrameOwner::FrameObject as i8), locals_dirty: atomic::AtomicBool::new(false), + f_locals_hidden_overlay: PyMutex::new(None), pending_stack_pops: Default::default(), pending_unwind_from_stack: Default::default(), }; @@ -797,6 +802,7 @@ impl Frame { for slot in fastlocals.iter_mut() { *slot = None; } + self.f_locals_hidden_overlay.lock().take(); } /// Get cell contents by localsplus index. @@ -865,9 +871,16 @@ impl Frame { return Ok(()); } let code = &**self.code; + let overlay_locals = self + .has_active_hidden_locals() + .then(|| self.f_locals_hidden_overlay.lock().clone()) + .flatten() + .map(ArgMapping::from_dict_exact); + let locals_map = overlay_locals + .as_ref() + .map_or_else(|| self.locals.mapping(vm), ArgMapping::mapping); // SAFETY: Called before generator resume; no concurrent access. let fastlocals = unsafe { (*self.iframe.get()).localsplus.fastlocals_mut() }; - let locals_map = self.locals.mapping(vm); for (i, &varname) in code.varnames.iter().enumerate() { if i >= fastlocals.len() { break; @@ -882,23 +895,12 @@ impl Frame { Ok(()) } - pub fn locals(&self, vm: &VirtualMachine) -> PyResult { - use rustpython_compiler_core::bytecode::{ - CO_FAST_CELL, CO_FAST_FREE, CO_FAST_HIDDEN, CO_FAST_LOCAL, - }; - // SAFETY: Either the frame is not executing (caller checked owner), - // or we're in a trace callback on the same thread that's executing. - let locals = &self.locals; + fn has_active_hidden_locals(&self) -> bool { + use rustpython_compiler_core::bytecode::{CO_FAST_CELL, CO_FAST_FREE, CO_FAST_HIDDEN}; let code = &**self.code; - let locals_map = locals.mapping(vm); let fastlocals = unsafe { (*self.iframe.get()).localsplus.fastlocals() }; - - // Iterate through all localsplus slots using localspluskinds - let nlocalsplus = code.localspluskinds.len(); - let nfrees = code.freevars.len(); - let free_start = nlocalsplus - nfrees; let is_optimized = code.flags.contains(bytecode::CodeFlags::OPTIMIZED); - let has_active_hidden_locals = !is_optimized + !is_optimized && code.localspluskinds.iter().enumerate().any(|(i, &kind)| { if kind & CO_FAST_HIDDEN == 0 { return false; @@ -914,16 +916,27 @@ impl Frame { } } } - }); - let overlay_locals = if has_active_hidden_locals { - // Match CPython's PyEval_GetLocals() behavior for frames with - // PEP 709 hidden locals: locals() inside an inlined comprehension - // returns a snapshot of the active fast locals only, not the - // backing module/class/eval mapping. - Some(vm.ctx.new_dict()) - } else { - None + }) + } + + fn sync_visible_locals_to_mapping( + &self, + locals_map: PyMapping<'_>, + vm: &VirtualMachine, + ) -> PyResult<()> { + use rustpython_compiler_core::bytecode::{ + CO_FAST_CELL, CO_FAST_FREE, CO_FAST_HIDDEN, CO_FAST_LOCAL, }; + // SAFETY: Either the frame is not executing (caller checked owner), + // or we're in a trace callback on the same thread that's executing. + let code = &**self.code; + let fastlocals = unsafe { (*self.iframe.get()).localsplus.fastlocals() }; + + // Iterate through all localsplus slots using localspluskinds + let nlocalsplus = code.localspluskinds.len(); + let nfrees = code.freevars.len(); + let free_start = nlocalsplus - nfrees; + let is_optimized = code.flags.contains(bytecode::CodeFlags::OPTIMIZED); // Track which non-merged cellvar index we're at let mut nonmerged_cell_idx = 0; @@ -1001,24 +1014,52 @@ impl Frame { fastlocals[i].clone() }; - let result = if let Some(dict) = &overlay_locals { - match value { - Some(value) => dict.set_item(name, value, vm), - None => dict.del_item(name, vm), - } - } else { - locals_map.ass_subscript(name, value, vm) - }; + let result = locals_map.ass_subscript(name, value, vm); match result { Ok(()) => {} Err(e) if e.fast_isinstance(vm.ctx.exceptions.key_error) => {} Err(e) => return Err(e), } } - if let Some(dict) = overlay_locals { - Ok(ArgMapping::from_dict_exact(dict)) + Ok(()) + } + + pub fn f_locals_mapping(&self, vm: &VirtualMachine) -> PyResult { + if !self.has_active_hidden_locals() { + self.f_locals_hidden_overlay.lock().take(); + return self.locals(vm); + } + + let needs_refresh = !self.locals_dirty.load(atomic::Ordering::Acquire); + let overlay_dict = { + let mut overlay = self.f_locals_hidden_overlay.lock(); + match overlay.as_ref() { + Some(dict) => dict.clone(), + None => { + let dict = vm.ctx.new_dict(); + *overlay = Some(dict.clone()); + dict + } + } + }; + if needs_refresh { + PyDict::clear(&overlay_dict); + let overlay = ArgMapping::from_dict_exact(overlay_dict.clone()); + self.sync_visible_locals_to_mapping(overlay.mapping(), vm)?; + } + Ok(ArgMapping::from_dict_exact(overlay_dict)) + } + + pub fn locals(&self, vm: &VirtualMachine) -> PyResult { + if self.has_active_hidden_locals() { + // Match CPython's locals() behavior for frames with PEP 709 hidden + // locals: return a fresh snapshot instead of the backing mapping. + let overlay = ArgMapping::from_dict_exact(vm.ctx.new_dict()); + self.sync_visible_locals_to_mapping(overlay.mapping(), vm)?; + Ok(overlay) } else { - Ok(locals.clone_mapping(vm)) + self.sync_visible_locals_to_mapping(self.locals.mapping(vm), vm)?; + Ok(self.locals.clone_mapping(vm)) } } } diff --git a/scripts/compare_bytecode.py b/scripts/compare_bytecode.py index 7e441b02f00..56ad77ffc1d 100644 --- a/scripts/compare_bytecode.py +++ b/scripts/compare_bytecode.py @@ -25,6 +25,7 @@ PROJECT_ROOT = os.path.dirname(SCRIPT_DIR) DIS_DUMP = os.path.join(SCRIPT_DIR, "dis_dump.py") DEFAULT_REPORT = os.path.join(PROJECT_ROOT, "compare_bytecode.report") +DUMP_TIMEOUT = 600 def find_rustpython(): @@ -69,13 +70,15 @@ def _start_one(interpreter, targets, base_dir): if interpreter != sys.executable: env["RUSTPYTHONPATH"] = base_dir - files_file = tempfile.NamedTemporaryFile( - mode="w", encoding="utf-8", delete=False, dir=PROJECT_ROOT - ) - output_file = tempfile.NamedTemporaryFile( - mode="w", encoding="utf-8", delete=False, dir=PROJECT_ROOT - ) + files_file = None + output_file = None try: + files_file = tempfile.NamedTemporaryFile( + mode="w", encoding="utf-8", delete=False, dir=PROJECT_ROOT + ) + output_file = tempfile.NamedTemporaryFile( + mode="w", encoding="utf-8", delete=False, dir=PROJECT_ROOT + ) for _, path in targets: files_file.write(path) files_file.write("\n") @@ -109,8 +112,14 @@ def _start_one(interpreter, targets, base_dir): "base_dir": base_dir, } except Exception: - os.unlink(files_file.name) - os.unlink(output_file.name) + for handle in (files_file, output_file): + if handle is None: + continue + try: + handle.close() + finally: + if os.path.exists(handle.name): + os.unlink(handle.name) raise @@ -130,19 +139,18 @@ def _load_dump_output(output_file): return None -def _run_sync_dump(interpreter, targets, base_dir, timeout=600): +def _run_sync_dump(interpreter, targets, base_dir, timeout=DUMP_TIMEOUT): job = _start_one(interpreter, targets, base_dir) proc = job["proc"] + stdout = b"" + timed_out = False try: - stdout = proc.communicate(timeout=600)[0] + stdout = proc.communicate(timeout=timeout)[0] except subprocess.TimeoutExpired: proc.kill() proc.communicate() - print(" Timeout (600s)", file=sys.stderr) - stdout = b"" + print(f" Timeout ({timeout}s)", file=sys.stderr) timed_out = True - finally: - timed_out = locals().get("timed_out", False) try: data = _load_dump_output(job["output_file"]) @@ -163,24 +171,52 @@ def _run_sync_dump(interpreter, targets, base_dir, timeout=600): def _rerun_missing_targets(interpreter, targets, base_dir): recovered = {} + failed = [] + empty = [] for target in targets: + relpath = target[0] data = _run_sync_dump(interpreter, [target], base_dir) - if data: + if data is None: + failed.append(relpath) + recovered[relpath] = { + "status": "error", + "error": "dump helper failed while rerunning target", + } + elif data: recovered.update(data) + else: + empty.append(relpath) + recovered[relpath] = { + "status": "error", + "error": "dump helper produced no data while rerunning target", + } + if failed: + print( + " Warning: rerun failed for %d file(s): %s" + % (len(failed), ", ".join(failed[:5])), + file=sys.stderr, + ) + if empty: + print( + " Warning: rerun produced no data for %d file(s): %s" + % (len(empty), ", ".join(empty[:5])), + file=sys.stderr, + ) return recovered -def _finish_one(job): +def _finish_one(job, timeout=DUMP_TIMEOUT): """Wait for a single dis_dump.py process and return parsed JSON.""" proc = job["proc"] expected = {relpath for relpath, _ in job["targets"]} + stdout = b"" try: - stdout = proc.communicate(timeout=600)[0] + stdout = proc.communicate(timeout=timeout)[0] except subprocess.TimeoutExpired: proc.kill() proc.communicate() print( - " Timeout (600s), retrying %d file(s) serially" % len(job["targets"]), + f" Timeout ({timeout}s), retrying {len(job['targets'])} file(s) serially", file=sys.stderr, ) data = None @@ -194,7 +230,7 @@ def _finish_one(job): if proc.returncode != 0: print(" Warning: exited with code %d" % proc.returncode, file=sys.stderr) - stray = stdout.decode(errors="replace").strip() if "stdout" in locals() else "" + stray = stdout.decode(errors="replace").strip() if stray: print(" Warning: unexpected stdout from dump helper", file=sys.stderr)