From d16c29b1f7731880e1591fc32fc3fd20dcd44ed8 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Mon, 13 Apr 2026 18:39:43 +0900 Subject: [PATCH 1/8] Align bytecode codegen structure with CPython 3.14 --- crates/codegen/src/compile.rs | 448 +++++++++--- crates/codegen/src/ir.rs | 667 ++++++++++++++++-- .../compiler-core/src/bytecode/instruction.rs | 4 +- scripts/dis_dump.py | 38 +- 4 files changed, 962 insertions(+), 195 deletions(-) diff --git a/crates/codegen/src/compile.rs b/crates/codegen/src/compile.rs index 1a35d2c23d1..360fd0809fa 100644 --- a/crates/codegen/src/compile.rs +++ b/crates/codegen/src/compile.rs @@ -124,8 +124,6 @@ enum SuperCallType<'a> { #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum BuiltinGeneratorCallKind { Tuple, - List, - Set, All, Any, } @@ -216,6 +214,13 @@ impl CompileContext { } } +fn bool_literal_value(expr: &ast::Expr) -> Option { + match expr { + ast::Expr::BooleanLiteral(ast::ExprBooleanLiteral { value, .. }) => Some(*value), + _ => None, + } +} + /// Segment of a parsed %-format string for optimize_format_str. struct FormatSegment { literal: String, @@ -1238,25 +1243,7 @@ impl Compiler { self.set_qualname(); } - // Emit COPY_FREE_VARS first, then MAKE_CELL (CPython order) - { - let nfrees = self.code_stack.last().unwrap().metadata.freevars.len(); - if nfrees > 0 { - emit!( - self, - Instruction::CopyFreeVars { - n: u32::try_from(nfrees).expect("too many freevars"), - } - ); - } - } - { - let ncells = self.code_stack.last().unwrap().metadata.cellvars.len(); - for i in 0..ncells { - let i_varnum: oparg::VarNum = u32::try_from(i).expect("too many cellvars").into(); - emit!(self, Instruction::MakeCell { i: i_varnum }); - } - } + self.emit_prefix_cell_setup(); // Emit RESUME (handles async preamble and module lineno 0) // CPython: LOCATION(lineno, lineno, 0, 0), then loc.lineno = 0 for module @@ -1307,6 +1294,36 @@ impl Compiler { }); } + fn emit_prefix_cell_setup(&mut self) { + let metadata = &self.code_stack.last().unwrap().metadata; + let varnames = metadata.varnames.clone(); + let cellvars = metadata.cellvars.clone(); + let freevars = metadata.freevars.clone(); + let ncells = cellvars.len(); + if ncells > 0 { + let cellfixedoffsets = ir::build_cellfixedoffsets(&varnames, &cellvars, &freevars); + let mut sorted = vec![None; varnames.len() + ncells]; + for (oldindex, fixed) in cellfixedoffsets.iter().copied().take(ncells).enumerate() { + sorted[fixed as usize] = Some(oldindex); + } + for oldindex in sorted.into_iter().flatten() { + let i_varnum: oparg::VarNum = + u32::try_from(oldindex).expect("too many cellvars").into(); + emit!(self, Instruction::MakeCell { i: i_varnum }); + } + } + + let nfrees = freevars.len(); + if nfrees > 0 { + emit!( + self, + Instruction::CopyFreeVars { + n: u32::try_from(nfrees).expect("too many freevars"), + } + ); + } + } + fn push_output( &mut self, flags: bytecode::CodeFlags, @@ -1828,13 +1845,10 @@ impl Compiler { self.symbol_table_stack.push(symbol_table); - // Emit MAKE_CELL for module-level cells (before RESUME) + // Match flowgraph.c insert_prefix_instructions() for module-level + // synthetic cells before RESUME. if has_module_cond_ann { - let ncells = self.code_stack.last().unwrap().metadata.cellvars.len(); - for i in 0..ncells { - let i_varnum: oparg::VarNum = u32::try_from(i).expect("too many cellvars").into(); - emit!(self, Instruction::MakeCell { i: i_varnum }); - } + self.emit_prefix_cell_setup(); } self.emit_resume_for_scope(CompilerScope::Module, 1); @@ -2516,9 +2530,8 @@ impl Compiler { } ); if let Some(e) = msg { - emit!(self, Instruction::PushNull); self.compile_expression(e)?; - emit!(self, Instruction::Call { argc: 1 }); + emit!(self, Instruction::Call { argc: 0 }); } emit!( self, @@ -8381,8 +8394,6 @@ impl Compiler { } match id.as_str() { "tuple" => Some(BuiltinGeneratorCallKind::Tuple), - "list" => Some(BuiltinGeneratorCallKind::List), - "set" => Some(BuiltinGeneratorCallKind::Set), "all" => Some(BuiltinGeneratorCallKind::All), "any" => Some(BuiltinGeneratorCallKind::Any), _ => None, @@ -8391,34 +8402,27 @@ impl Compiler { /// Emit the optimized inline loop for builtin(genexpr) calls. /// - /// Stack on entry: `[func, iter]` where `iter` is the already-compiled - /// generator iterator and `func` is the builtin candidate. - /// On return the compiler is positioned at the fallback block with - /// `[func, iter]` still on the stack (for the normal CALL path). + /// Stack on entry: `[func]` where `func` is the builtin candidate. + /// On return the compiler is positioned at the fallback block so the + /// normal call path can compile the original generator argument again. fn optimize_builtin_generator_call( &mut self, kind: BuiltinGeneratorCallKind, + generator_expr: &ast::Expr, end: BlockIdx, ) -> CompileResult<()> { let common_constant = match kind { BuiltinGeneratorCallKind::Tuple => bytecode::CommonConstant::BuiltinTuple, - BuiltinGeneratorCallKind::List => bytecode::CommonConstant::BuiltinList, - BuiltinGeneratorCallKind::Set => bytecode::CommonConstant::BuiltinSet, BuiltinGeneratorCallKind::All => bytecode::CommonConstant::BuiltinAll, BuiltinGeneratorCallKind::Any => bytecode::CommonConstant::BuiltinAny, }; + let fallback = self.new_block(); let loop_block = self.new_block(); let cleanup = self.new_block(); - let fallback = self.new_block(); - let result = matches!( - kind, - BuiltinGeneratorCallKind::All | BuiltinGeneratorCallKind::Any - ) - .then(|| self.new_block()); - // Stack: [func, iter] — copy func (TOS1) for identity check - emit!(self, Instruction::Copy { i: 2 }); + // Stack: [func] — copy function for identity check + emit!(self, Instruction::Copy { i: 1 }); emit!( self, Instruction::LoadCommonConstant { @@ -8427,61 +8431,43 @@ impl Compiler { ); emit!(self, Instruction::IsOp { invert: Invert::No }); emit!(self, Instruction::PopJumpIfFalse { delta: fallback }); - emit!(self, Instruction::NotTaken); - // Remove func from [func, iter] → [iter] - emit!(self, Instruction::Swap { i: 2 }); emit!(self, Instruction::PopTop); - if matches!( - kind, - BuiltinGeneratorCallKind::Tuple | BuiltinGeneratorCallKind::List - ) { - // [iter] → [iter, list] → [list, iter] + if matches!(kind, BuiltinGeneratorCallKind::Tuple) { emit!(self, Instruction::BuildList { count: 0 }); - emit!(self, Instruction::Swap { i: 2 }); - } else if matches!(kind, BuiltinGeneratorCallKind::Set) { - // [iter] → [iter, set] → [set, iter] - emit!(self, Instruction::BuildSet { count: 0 }); - emit!(self, Instruction::Swap { i: 2 }); } + let sub_table_cursor = self.symbol_table_stack.last().map(|t| t.next_sub_table); + self.compile_expression(generator_expr)?; + if let Some(cursor) = sub_table_cursor + && let Some(current_table) = self.symbol_table_stack.last_mut() + { + current_table.next_sub_table = cursor; + } self.switch_to_block(loop_block); emit!(self, Instruction::ForIter { delta: cleanup }); match kind { - BuiltinGeneratorCallKind::Tuple | BuiltinGeneratorCallKind::List => { + BuiltinGeneratorCallKind::Tuple => { emit!(self, Instruction::ListAppend { i: 2 }); emit!(self, PseudoInstruction::Jump { delta: loop_block }); } - BuiltinGeneratorCallKind::Set => { - emit!(self, Instruction::SetAdd { i: 2 }); - emit!(self, PseudoInstruction::Jump { delta: loop_block }); - } BuiltinGeneratorCallKind::All => { - let result = result.expect("all() optimization should have a result block"); emit!(self, Instruction::ToBool); - emit!(self, Instruction::PopJumpIfFalse { delta: result }); - emit!(self, Instruction::NotTaken); - emit!(self, PseudoInstruction::Jump { delta: loop_block }); + emit!(self, Instruction::PopJumpIfTrue { delta: loop_block }); + emit!(self, Instruction::PopIter); + self.emit_load_const(ConstantData::Boolean { value: false }); + emit!(self, PseudoInstruction::Jump { delta: end }); } BuiltinGeneratorCallKind::Any => { - let result = result.expect("any() optimization should have a result block"); emit!(self, Instruction::ToBool); - emit!(self, Instruction::PopJumpIfTrue { delta: result }); - emit!(self, Instruction::NotTaken); - emit!(self, PseudoInstruction::Jump { delta: loop_block }); + emit!(self, Instruction::PopJumpIfFalse { delta: loop_block }); + emit!(self, Instruction::PopIter); + self.emit_load_const(ConstantData::Boolean { value: true }); + emit!(self, PseudoInstruction::Jump { delta: end }); } } - if let Some(result_block) = result { - self.switch_to_block(result_block); - emit!(self, Instruction::PopIter); - self.emit_load_const(ConstantData::Boolean { - value: matches!(kind, BuiltinGeneratorCallKind::Any), - }); - emit!(self, PseudoInstruction::Jump { delta: end }); - } - self.switch_to_block(cleanup); emit!(self, Instruction::EndFor); emit!(self, Instruction::PopIter); @@ -8494,7 +8480,6 @@ impl Compiler { } ); } - BuiltinGeneratorCallKind::List | BuiltinGeneratorCallKind::Set => {} BuiltinGeneratorCallKind::All => { self.emit_load_const(ConstantData::Boolean { value: true }); } @@ -8574,18 +8559,12 @@ impl Compiler { .then(|| self.detect_builtin_generator_call(func, args)) .flatten() { - // Optimized builtin(genexpr) path: compile the genexpr only once - // so its code object appears exactly once in co_consts. let end = self.new_block(); self.compile_expression(func)?; - self.compile_expression(&args.args[0])?; - // Stack: [func, iter] - self.optimize_builtin_generator_call(kind, end)?; - // Fallback block: [func, iter] → [func, null, iter] → CALL - emit!(self, Instruction::PushNull); - emit!(self, Instruction::Swap { i: 2 }); + self.optimize_builtin_generator_call(kind, &args.args[0], end)?; self.set_source_range(call_range); - emit!(self, Instruction::Call { argc: 1 }); + emit!(self, Instruction::PushNull); + self.codegen_call_helper(0, args, call_range)?; self.switch_to_block(end); } else { // Regular call: push func, then NULL for self_or_null slot @@ -9050,7 +9029,19 @@ impl Compiler { // Now evaluate the ifs: for if_condition in &generator.ifs { - self.compile_jump_if(if_condition, false, if_cleanup_block)? + match bool_literal_value(if_condition) { + Some(true) => {} + Some(false) => { + emit!( + self, + PseudoInstruction::Jump { + delta: if_cleanup_block + } + ); + break; + } + None => self.compile_jump_if(if_condition, false, if_cleanup_block)?, + } } } @@ -9326,7 +9317,19 @@ impl Compiler { // Evaluate the if conditions for if_condition in &generator.ifs { - self.compile_jump_if(if_condition, false, if_cleanup_block)?; + match bool_literal_value(if_condition) { + Some(true) => {} + Some(false) => { + emit!( + self, + PseudoInstruction::Jump { + delta: if_cleanup_block + } + ); + break; + } + None => self.compile_jump_if(if_condition, false, if_cleanup_block)?, + } } } @@ -11060,7 +11063,7 @@ def f(xs): } #[test] - fn test_builtin_tuple_list_set_genexpr_calls_are_optimized() { + fn test_builtin_tuple_genexpr_call_is_optimized_but_list_set_are_not() { let code = compile_exec( "\ def tuple_f(xs): @@ -11091,27 +11094,29 @@ def set_f(xs): assert_eq!(tuple_list_append, 2); let list_f = find_code(&code, "list_f").expect("missing list_f code"); - assert!(has_common_constant( - list_f, - bytecode::CommonConstant::BuiltinList - )); assert!( list_f .instructions .iter() - .any(|unit| matches!(unit.op, Instruction::ListAppend { .. })) + .any(|unit| matches!(unit.op, Instruction::Call { .. })), + "list(genexpr) should stay on the normal call path" + ); + assert!( + !has_common_constant(list_f, bytecode::CommonConstant::BuiltinList), + "CPython 3.14.2 does not optimize list(genexpr)" ); let set_f = find_code(&code, "set_f").expect("missing set_f code"); - assert!(has_common_constant( - set_f, - bytecode::CommonConstant::BuiltinSet - )); assert!( set_f .instructions .iter() - .any(|unit| matches!(unit.op, Instruction::SetAdd { .. })) + .any(|unit| matches!(unit.op, Instruction::Call { .. })), + "set(genexpr) should stay on the normal call path" + ); + assert!( + !has_common_constant(set_f, bytecode::CommonConstant::BuiltinSet), + "CPython 3.14.2 does not optimize set(genexpr)" ); } @@ -11481,6 +11486,243 @@ def f(x): assert_eq!(push_null_count, 0); } + #[test] + fn test_assert_with_message_uses_common_constant_direct_call() { + let code = compile_exec( + "\ +def f(x, y): + assert x, y +", + ); + let f = find_code(&code, "f").expect("missing f code"); + let load_assertion = f + .instructions + .iter() + .position(|unit| { + matches!(unit.op, Instruction::LoadCommonConstant { .. }) + && matches!( + unit.op, + Instruction::LoadCommonConstant { idx } + if idx.get(OpArg::new(u32::from(u8::from(unit.arg)))) + == bytecode::CommonConstant::AssertionError + ) + }) + .expect("missing LOAD_COMMON_CONSTANT AssertionError"); + + assert!( + !matches!( + f.instructions.get(load_assertion + 1).map(|unit| unit.op), + Some(Instruction::PushNull) + ), + "assert message path should not use PUSH_NULL, got ops={:?}", + f.instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + assert!( + matches!( + f.instructions.get(load_assertion + 2).map(|unit| unit.op), + Some(Instruction::Call { .. }) + ), + "expected direct CALL after loading assert message, got ops={:?}", + f.instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + + let call_arg = f.instructions[load_assertion + 2].arg; + assert_eq!(u8::from(call_arg), 0); + } + + #[test] + fn test_negative_constant_binop_folds_after_unary_folding() { + let code = compile_exec( + "\ +def f(): + return -2147483647 - 1 +", + ); + let f = find_code(&code, "f").expect("missing function code"); + + assert!( + !f.instructions + .iter() + .any(|unit| matches!(unit.op, Instruction::BinaryOp { .. })), + "negative constant expression should fold to a single constant, got ops={:?}", + f.instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + assert!( + f.instructions + .iter() + .any(|unit| matches!(unit.op, Instruction::LoadConst { .. })), + "expected folded constant load, got ops={:?}", + f.instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + } + + #[test] + fn test_genexpr_filter_header_uses_store_fast_load_fast() { + let code = compile_exec( + "\ +def f(it): + return (x for x in it if x) +", + ); + let genexpr = find_code(&code, "").expect("missing code"); + let store_fast_load_fast_idx = genexpr + .instructions + .iter() + .position(|unit| matches!(unit.op, Instruction::StoreFastLoadFast { .. })) + .expect("missing STORE_FAST_LOAD_FAST in genexpr header"); + + assert!( + matches!( + genexpr + .instructions + .get(store_fast_load_fast_idx + 1) + .map(|unit| unit.op), + Some(Instruction::ToBool) + ), + "expected TO_BOOL immediately after STORE_FAST_LOAD_FAST, got ops={:?}", + genexpr + .instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + } + + #[test] + fn test_multi_with_header_uses_store_fast_load_fast() { + let code = compile_exec( + "\ +def f(manager): + with manager() as x, manager(): + pass +", + ); + let f = find_code(&code, "f").expect("missing function code"); + assert!( + f.instructions + .iter() + .any(|unit| matches!(unit.op, Instruction::StoreFastLoadFast { .. })), + "expected STORE_FAST_LOAD_FAST in multi-with header, got ops={:?}", + f.instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + } + + #[test] + fn test_sequential_store_then_load_uses_store_fast_load_fast() { + let code = compile_exec( + "\ +def f(self): + x = ''; y = \"\"; self.assertTrue(len(x) == 0 and x == y) +", + ); + let f = find_code(&code, "f").expect("missing function code"); + assert!( + f.instructions + .iter() + .any(|unit| matches!(unit.op, Instruction::StoreFastLoadFast { .. })), + "expected STORE_FAST_LOAD_FAST in sequential statement body, 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) +", + ); + let genexpr = find_code(&code, "").expect("missing code"); + assert!( + !genexpr.instructions.iter().any(|unit| { + matches!(unit.op, Instruction::LoadConst { .. }) + && matches!( + genexpr.constants.get(usize::from(u8::from(unit.arg))), + Some(ConstantData::Boolean { value: true }) + ) + }), + "constant-true filter should not load True, got ops={:?}", + genexpr + .instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + assert!( + !genexpr + .instructions + .iter() + .any(|unit| matches!(unit.op, Instruction::PopJumpIfTrue { .. })), + "constant-true filter should not leave POP_JUMP_IF_TRUE scaffolding, got ops={:?}", + genexpr + .instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + } + + #[test] + fn test_classdictcell_uses_load_closure_path_and_borrows_after_optimize() { + let code = compile_exec( + "\ +class C: + def method(self): + return 1 +", + ); + let class_code = find_code(&code, "C").expect("missing class code"); + let store_classdictcell = 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() + == "__classdictcell__" + ) + }) + .expect("missing STORE_NAME __classdictcell__"); + + assert!( + matches!( + class_code + .instructions + .get(store_classdictcell.saturating_sub(1)) + .map(|unit| unit.op), + Some(Instruction::LoadFastBorrow { .. }) + ), + "expected LOAD_FAST_BORROW before __classdictcell__ store, got ops={:?}", + class_code + .instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + } + #[test] fn test_chained_compare_jump_uses_single_cleanup_copy() { let code = compile_exec( diff --git a/crates/codegen/src/ir.rs b/crates/codegen/src/ir.rs index 61e549199d5..f6c9a4ae408 100644 --- a/crates/codegen/src/ir.rs +++ b/crates/codegen/src/ir.rs @@ -8,7 +8,7 @@ use num_traits::{ToPrimitive, Zero}; use rustpython_compiler_core::{ OneIndexed, SourceLocation, bytecode::{ - AnyInstruction, AnyOpcode, CO_FAST_CELL, CO_FAST_FREE, CO_FAST_HIDDEN, CO_FAST_LOCAL, + AnyInstruction, AnyOpcode, Arg, CO_FAST_CELL, CO_FAST_FREE, CO_FAST_HIDDEN, CO_FAST_LOCAL, CodeFlags, CodeObject, CodeUnit, CodeUnits, ConstantData, ExceptionTableEntry, InstrDisplayContext, Instruction, InstructionMetadata, Label, OpArg, Opcode, PseudoInstruction, PseudoOpcode, PyCodeLocationInfoKind, encode_exception_table, oparg, @@ -203,6 +203,7 @@ impl CodeInfo { self.fold_binop_constants(); self.remove_nops(); self.fold_unary_negative(); + self.remove_nops(); // remove UNARY_NEGATIVE NOPs before re-folding binops self.fold_binop_constants(); // re-run after unary folding: -1 + 2 → 1 self.remove_nops(); // remove NOPs so tuple/list/set see contiguous LOADs self.fold_tuple_constants(); @@ -235,7 +236,7 @@ impl CodeInfo { jump_threading(&mut self.blocks); self.eliminate_unreachable_blocks(); self.remove_nops(); - // TODO: insert_superinstructions disabled pending StoreFastLoadFast VM fix + // CPython inserts superinstructions before optimize_load_fast. push_cold_blocks_to_end(&mut self.blocks); // Phase 2: _PyCfg_OptimizedCfgToInstructionSequence (flowgraph.c) @@ -260,10 +261,26 @@ impl CodeInfo { reorder_jump_over_exception_cleanup_blocks(&mut self.blocks); self.eliminate_unreachable_blocks(); remove_redundant_nops_and_jumps(&mut self.blocks); + let cellfixedoffsets = build_cellfixedoffsets( + &self.metadata.varnames, + &self.metadata.cellvars, + &self.metadata.freevars, + ); + // CPython lowers LOAD_CLOSURE to LOAD_FAST before optimize_load_fast, so + // borrow selection can see classdictcell and other merged-cell loads. + convert_load_closure_pseudo_ops(&mut self.blocks, &cellfixedoffsets); self.add_checks_for_loads_of_uninitialized_variables(); + self.combine_store_fast_load_fast(); + // CPython's optimize_load_fast runs with block start depths already known. + // Compute them here so the abstract stack simulation can use the real + // CFG entry depth for each block. + let _ = self.max_stackdepth()?; // optimize_load_fast: after normalize_jumps self.optimize_load_fast_borrow(); + self.deoptimize_borrow_for_handler_return_paths(); + self.deoptimize_store_fast_store_fast_after_cleanup(); self.optimize_load_global_push_null(); + self.reorder_entry_prefix_cell_setup(); let max_stackdepth = self.max_stackdepth()?; @@ -620,6 +637,63 @@ impl CodeInfo { } } + fn reorder_entry_prefix_cell_setup(&mut self) { + let Some(entry) = self.blocks.first_mut() else { + return; + }; + let ncells = self.metadata.cellvars.len(); + let nfrees = self.metadata.freevars.len(); + if ncells == 0 && nfrees == 0 { + return; + } + + let prefix_len = entry + .instructions + .iter() + .take_while(|info| { + matches!( + info.instr.real(), + Some(Instruction::MakeCell { .. } | Instruction::CopyFreeVars { .. }) + ) + }) + .count(); + if prefix_len == 0 { + return; + } + + let original_prefix = entry.instructions[..prefix_len].to_vec(); + let anchor = original_prefix[0]; + let rest = entry.instructions.split_off(prefix_len); + entry.instructions.clear(); + + if nfrees > 0 { + entry.instructions.push(InstructionInfo { + instr: Instruction::CopyFreeVars { n: Arg::marker() }.into(), + arg: OpArg::new(nfrees as u32), + ..anchor + }); + } + + let cellfixedoffsets = build_cellfixedoffsets( + &self.metadata.varnames, + &self.metadata.cellvars, + &self.metadata.freevars, + ); + let mut sorted = vec![None; self.metadata.varnames.len() + ncells]; + for (oldindex, fixed) in cellfixedoffsets.iter().copied().take(ncells).enumerate() { + sorted[fixed as usize] = Some(oldindex); + } + for oldindex in sorted.into_iter().flatten() { + entry.instructions.push(InstructionInfo { + instr: Instruction::MakeCell { i: Arg::marker() }.into(), + arg: OpArg::new(oldindex as u32), + ..anchor + }); + } + + entry.instructions.extend(rest); + } + /// Clear blocks that are unreachable (not entry, not a jump target, /// and only reachable via fall-through from a terminal block). fn eliminate_unreachable_blocks(&mut self) { @@ -1798,20 +1872,17 @@ impl CodeInfo { } } - /// Optimize LOAD_FAST to LOAD_FAST_BORROW where safe. - /// - /// insert_superinstructions (flowgraph.c): Combine STORE_FAST + LOAD_FAST → - /// STORE_FAST_LOAD_FAST. Currently disabled pending VM stack null investigation. - #[allow(dead_code)] + /// 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. fn combine_store_fast_load_fast(&mut self) { for block in &mut self.blocks { let mut i = 0; while i + 1 < block.instructions.len() { let curr = &block.instructions[i]; let next = &block.instructions[i + 1]; - let (Some(Instruction::StoreFast { .. }), Some(Instruction::LoadFast { .. })) = - (curr.instr.real(), next.instr.real()) - else { + let Some(Instruction::StoreFast { .. }) = curr.instr.real() else { i += 1; continue; }; @@ -1822,17 +1893,54 @@ impl CodeInfo { i += 1; continue; } - let idx1 = u32::from(curr.arg); - let idx2 = u32::from(next.arg); - if idx1 < 16 && idx2 < 16 { - let packed = (idx1 << 4) | idx2; - block.instructions[i].instr = Opcode::StoreFastLoadFast.into(); - block.instructions[i].arg = OpArg::new(packed); - // Replace second instruction with NOP (CPython: INSTR_SET_OP0(inst2, NOP)) - set_to_nop(&mut block.instructions[i + 1]); - i += 2; // skip the NOP - } else { + + let store_idx = u32::from(curr.arg); + if store_idx >= 16 { i += 1; + continue; + } + + match next.instr.real() { + Some(Instruction::LoadFast { .. }) => { + let load_idx = u32::from(next.arg); + if load_idx >= 16 { + i += 1; + continue; + } + let packed = (store_idx << 4) | load_idx; + block.instructions[i].instr = Instruction::StoreFastLoadFast { + var_nums: Arg::marker(), + } + .into(); + block.instructions[i].arg = OpArg::new(packed); + set_to_nop(&mut block.instructions[i + 1]); + i += 2; + } + Some(Instruction::LoadFastLoadFast { var_nums }) => { + let packed = var_nums.get(next.arg); + let (first_idx, second_idx) = packed.indexes(); + let first_idx = u32::from(first_idx); + if first_idx >= 16 { + i += 1; + continue; + } + + let packed = (store_idx << 4) | first_idx; + block.instructions[i].instr = Instruction::StoreFastLoadFast { + var_nums: Arg::marker(), + } + .into(); + block.instructions[i].arg = OpArg::new(packed); + block.instructions[i + 1].instr = Instruction::LoadFast { + var_num: Arg::marker(), + } + .into(); + block.instructions[i + 1].arg = OpArg::new(u32::from(second_idx)); + i += 2; + } + _ => { + i += 1; + } } } } @@ -1841,86 +1949,442 @@ impl CodeInfo { fn optimize_load_fast_borrow(&mut self) { // NOT_LOCAL marker: instruction didn't come from a LOAD_FAST const NOT_LOCAL: usize = usize::MAX; + const DUMMY_INSTR: isize = -1; + const SUPPORT_KILLED: u8 = 1; + const STORED_AS_LOCAL: u8 = 2; + const REF_UNCONSUMED: u8 = 4; + + #[derive(Clone, Copy)] + struct AbstractRef { + instr: isize, + local: usize, + } - for block in &mut self.blocks { + fn push_ref(refs: &mut Vec, instr: isize, local: usize) { + refs.push(AbstractRef { instr, local }); + } + + fn pop_ref(refs: &mut Vec) -> Option { + refs.pop() + } + + fn at_ref(refs: &[AbstractRef], idx: usize) -> Option { + refs.get(idx).copied() + } + + fn swap_top(refs: &mut [AbstractRef], depth: usize) { + let top = refs.len() - 1; + let other = refs.len() - depth; + refs.swap(top, other); + } + + fn kill_local(instr_flags: &mut [u8], refs: &[AbstractRef], local: usize) { + for r in refs.iter().copied().filter(|r| r.local == local) { + debug_assert!(r.instr >= 0); + instr_flags[r.instr as usize] |= SUPPORT_KILLED; + } + } + + fn store_local(instr_flags: &mut [u8], refs: &[AbstractRef], local: usize, r: AbstractRef) { + kill_local(instr_flags, refs, local); + if r.instr != DUMMY_INSTR { + instr_flags[r.instr as usize] |= STORED_AS_LOCAL; + } + } + + fn decode_packed_fast_locals(arg: OpArg) -> (usize, usize) { + let packed = u32::from(arg); + (((packed >> 4) & 0xF) as usize, (packed & 0xF) as usize) + } + + fn push_block( + worklist: &mut Vec, + visited: &mut [bool], + blocks: &[Block], + target: BlockIdx, + start_depth: usize, + ) { + let _ = (blocks, start_depth); + if !visited[target.idx()] { + visited[target.idx()] = true; + worklist.push(target); + } + } + + let mut visited = vec![false; self.blocks.len()]; + let mut worklist = vec![BlockIdx(0)]; + visited[0] = true; + + while let Some(block_idx) = worklist.pop() { + let block = &self.blocks[block_idx]; if block.instructions.is_empty() { continue; } - // Track which instructions' outputs are still on stack at block end - // For each instruction, we track if its pushed value(s) are unconsumed - let mut unconsumed = vec![false; block.instructions.len()]; - - // Simulate stack: each entry is the instruction index that pushed it - // (or NOT_LOCAL if not from LOAD_FAST/LOAD_FAST_LOAD_FAST). - // - // CPython (flowgraph.c optimize_load_fast) pre-fills the stack with - // dummy refs for values inherited from predecessor blocks. We take - // the simpler approach of aborting the optimisation for the whole - // block on stack underflow. - let mut stack: Vec = Vec::new(); - let mut underflow = false; + let mut instr_flags = vec![0u8; block.instructions.len()]; + let start_depth = block.start_depth.unwrap_or(0) as usize; + let mut refs = Vec::with_capacity(block.instructions.len() + start_depth + 2); + for _ in 0..start_depth { + push_ref(&mut refs, DUMMY_INSTR, NOT_LOCAL); + } for (i, info) in block.instructions.iter().enumerate() { - let Some(instr) = info.instr.real() else { - continue; - }; + let instr = info.instr; + let arg_u32 = u32::from(info.arg); - let stack_effect_info = instr.stack_effect_info(info.arg.into()); - let (pushes, pops) = (stack_effect_info.pushed(), stack_effect_info.popped()); - - // Pop values from stack - for _ in 0..pops { - if stack.pop().is_none() { - // Stack underflow — block receives values from a predecessor. - // Abort optimisation for the entire block. - underflow = true; - break; + match instr { + AnyInstruction::Real(Instruction::DeleteFast { var_num }) => { + kill_local(&mut instr_flags, &refs, usize::from(var_num.get(info.arg))); + } + AnyInstruction::Real(Instruction::LoadFast { var_num }) => { + push_ref(&mut refs, i as isize, usize::from(var_num.get(info.arg))); + } + AnyInstruction::Real(Instruction::LoadFastAndClear { var_num }) => { + let local = usize::from(var_num.get(info.arg)); + kill_local(&mut instr_flags, &refs, local); + push_ref(&mut refs, i as isize, local); + } + AnyInstruction::Real(Instruction::LoadFastLoadFast { .. }) => { + let (local1, local2) = decode_packed_fast_locals(info.arg); + push_ref(&mut refs, i as isize, local1); + push_ref(&mut refs, i as isize, local2); + } + AnyInstruction::Real(Instruction::StoreFast { var_num }) => { + let Some(r) = pop_ref(&mut refs) else { + continue; + }; + store_local( + &mut instr_flags, + &refs, + usize::from(var_num.get(info.arg)), + r, + ); + } + AnyInstruction::Pseudo(PseudoInstruction::StoreFastMaybeNull { var_num }) => { + let Some(r) = pop_ref(&mut refs) else { + continue; + }; + store_local(&mut instr_flags, &refs, var_num.get(info.arg) as usize, r); + } + AnyInstruction::Real(Instruction::StoreFastLoadFast { .. }) => { + let (store_local_idx, load_local_idx) = decode_packed_fast_locals(info.arg); + let Some(r) = pop_ref(&mut refs) else { + continue; + }; + store_local(&mut instr_flags, &refs, store_local_idx, r); + push_ref(&mut refs, i as isize, load_local_idx); + } + AnyInstruction::Real(Instruction::StoreFastStoreFast { .. }) => { + let (local1, local2) = decode_packed_fast_locals(info.arg); + let Some(r1) = pop_ref(&mut refs) else { + continue; + }; + store_local(&mut instr_flags, &refs, local1, r1); + let Some(r2) = pop_ref(&mut refs) else { + continue; + }; + store_local(&mut instr_flags, &refs, local2, r2); + } + AnyInstruction::Real(Instruction::Copy { i: _ }) => { + let depth = arg_u32 as usize; + if depth == 0 || refs.len() < depth { + continue; + } + let r = at_ref(&refs, refs.len() - depth).expect("copy index in bounds"); + push_ref(&mut refs, r.instr, r.local); + } + AnyInstruction::Real(Instruction::Swap { i: _ }) => { + let depth = arg_u32 as usize; + if depth < 2 || refs.len() < depth { + continue; + } + swap_top(&mut refs, depth); + } + AnyInstruction::Real( + Instruction::FormatSimple + | Instruction::GetANext + | Instruction::GetLen + | Instruction::GetYieldFromIter + | Instruction::ImportFrom { .. } + | Instruction::MatchKeys + | Instruction::MatchMapping + | Instruction::MatchSequence + | Instruction::WithExceptStart, + ) => { + let effect = instr.stack_effect_info(arg_u32); + let net_pushed = effect.pushed() as isize - effect.popped() as isize; + debug_assert!(net_pushed >= 0); + for _ in 0..net_pushed { + push_ref(&mut refs, i as isize, NOT_LOCAL); + } + } + AnyInstruction::Real( + Instruction::DictMerge { .. } + | Instruction::DictUpdate { .. } + | Instruction::ListAppend { .. } + | Instruction::ListExtend { .. } + | Instruction::MapAdd { .. } + | Instruction::Reraise { .. } + | Instruction::SetAdd { .. } + | Instruction::SetUpdate { .. }, + ) => { + let effect = instr.stack_effect_info(arg_u32); + let net_popped = effect.popped() as isize - effect.pushed() as isize; + debug_assert!(net_popped > 0); + for _ in 0..net_popped { + let _ = pop_ref(&mut refs); + } + } + AnyInstruction::Real( + Instruction::EndSend | Instruction::SetFunctionAttribute { .. }, + ) => { + let Some(tos) = pop_ref(&mut refs) else { + continue; + }; + let _ = pop_ref(&mut refs); + push_ref(&mut refs, tos.instr, tos.local); + } + AnyInstruction::Real(Instruction::CheckExcMatch) => { + let _ = pop_ref(&mut refs); + push_ref(&mut refs, i as isize, NOT_LOCAL); + } + AnyInstruction::Real(Instruction::ForIter { .. }) => { + if info.target != BlockIdx::NULL { + push_block( + &mut worklist, + &mut visited, + &self.blocks, + info.target, + refs.len() + 1, + ); + } + push_ref(&mut refs, i as isize, NOT_LOCAL); + } + AnyInstruction::Real(Instruction::LoadAttr { .. }) => { + let Some(self_ref) = pop_ref(&mut refs) else { + continue; + }; + push_ref(&mut refs, i as isize, NOT_LOCAL); + if arg_u32 & 1 != 0 { + push_ref(&mut refs, self_ref.instr, self_ref.local); + } + } + AnyInstruction::Real(Instruction::LoadSuperAttr { .. }) => { + let _ = pop_ref(&mut refs); + let _ = pop_ref(&mut refs); + let Some(self_ref) = pop_ref(&mut refs) else { + continue; + }; + push_ref(&mut refs, i as isize, NOT_LOCAL); + if arg_u32 & 1 != 0 { + push_ref(&mut refs, self_ref.instr, self_ref.local); + } + } + AnyInstruction::Real( + Instruction::LoadSpecial { .. } | Instruction::PushExcInfo, + ) => { + let Some(tos) = pop_ref(&mut refs) else { + continue; + }; + push_ref(&mut refs, i as isize, NOT_LOCAL); + push_ref(&mut refs, tos.instr, tos.local); + } + AnyInstruction::Real(Instruction::Send { .. }) => { + if info.target != BlockIdx::NULL { + push_block( + &mut worklist, + &mut visited, + &self.blocks, + info.target, + refs.len(), + ); + } + let _ = pop_ref(&mut refs); + push_ref(&mut refs, i as isize, NOT_LOCAL); + } + _ => { + let effect = instr.stack_effect_info(arg_u32); + let num_popped = effect.popped() as usize; + let num_pushed = effect.pushed() as usize; + if info.target != BlockIdx::NULL { + let target_depth = refs + .len() + .saturating_sub(num_popped) + .saturating_add(num_pushed); + push_block( + &mut worklist, + &mut visited, + &self.blocks, + info.target, + target_depth, + ); + } + if !instr.is_block_push() { + for _ in 0..num_popped { + let _ = pop_ref(&mut refs); + } + for _ in 0..num_pushed { + push_ref(&mut refs, i as isize, NOT_LOCAL); + } + } } - } - if underflow { - break; - } - - // Push values to stack with source instruction index - let source = match instr.into() { - Opcode::LoadFast | Opcode::LoadFastLoadFast => i, - _ => NOT_LOCAL, - }; - for _ in 0..pushes { - stack.push(source); } } - if underflow { - continue; + if block.next != BlockIdx::NULL + && block.instructions.last().is_some_and(|term| { + !term.instr.is_unconditional_jump() && !term.instr.is_scope_exit() + }) + { + push_block( + &mut worklist, + &mut visited, + &self.blocks, + block.next, + refs.len(), + ); } - // Mark instructions whose values remain on stack at block end - for &src in &stack { - if src != NOT_LOCAL { - unconsumed[src] = true; + for r in refs { + if r.instr != DUMMY_INSTR { + instr_flags[r.instr as usize] |= REF_UNCONSUMED; } } - // Convert LOAD_FAST to LOAD_FAST_BORROW where value is fully consumed + let block = &mut self.blocks[block_idx]; + if block.except_handler || block.preserve_lasti { + continue; + } for (i, info) in block.instructions.iter_mut().enumerate() { - if unconsumed[i] { + if instr_flags[i] != 0 { continue; } - let Some(instr) = info.instr.real() else { + match info.instr.real() { + Some(Instruction::LoadFast { .. }) => { + info.instr = Instruction::LoadFastBorrow { + var_num: Arg::marker(), + } + .into(); + } + Some(Instruction::LoadFastLoadFast { .. }) => { + info.instr = Instruction::LoadFastBorrowLoadFastBorrow { + var_nums: Arg::marker(), + } + .into(); + } + _ => {} + } + } + } + } + + fn deoptimize_borrow_for_handler_return_paths(&mut self) { + for block in &mut self.blocks { + let len = block.instructions.len(); + for i in 0..len { + let Some(Instruction::LoadFastBorrow { .. }) = block.instructions[i].instr.real() + else { continue; }; - match instr.into() { - Opcode::LoadFast => { - info.instr = Opcode::LoadFastBorrow.into(); + let tail = &block.instructions[i + 1..]; + if tail.len() < 3 { + continue; + } + if !matches!(tail[0].instr.real(), Some(Instruction::Swap { .. })) { + continue; + } + if !matches!(tail[1].instr.real(), Some(Instruction::PopExcept)) { + continue; + } + if !matches!(tail[2].instr.real(), Some(Instruction::ReturnValue)) { + continue; + } + block.instructions[i].instr = Instruction::LoadFast { + var_num: Arg::marker(), + } + .into(); + } + } + } + + fn deoptimize_store_fast_store_fast_after_cleanup(&mut self) { + fn last_real_instr(block: &Block) -> Option { + block + .instructions + .iter() + .rev() + .find_map(|info| info.instr.real()) + } + + 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(pred_idx as u32)); + } + for info in &block.instructions { + if info.target != BlockIdx::NULL { + predecessors[info.target.idx()].push(BlockIdx(pred_idx as u32)); + } + } + } + + let starts_after_cleanup: Vec = predecessors + .iter() + .map(|preds| { + !preds.is_empty() + && preds.iter().copied().all(|pred_idx| { + matches!( + last_real_instr(&self.blocks[pred_idx]), + Some(Instruction::PopIter) | Some(Instruction::Swap { .. }) + ) + }) + }) + .collect(); + + for (block_idx, block) in self.blocks.iter_mut().enumerate() { + let mut new_instructions = Vec::with_capacity(block.instructions.len()); + for (i, info) in block.instructions.iter().copied().enumerate() { + let expand = matches!( + info.instr.real(), + Some(Instruction::StoreFastStoreFast { .. }) + ) && (new_instructions.last().is_some_and( + |prev: &InstructionInfo| { + matches!( + prev.instr.real(), + Some(Instruction::PopIter) | Some(Instruction::Swap { .. }) + ) + }, + ) || (i == 0 && starts_after_cleanup[block_idx])); + + if expand { + let Some(Instruction::StoreFastStoreFast { var_nums }) = info.instr.real() + else { + unreachable!(); + }; + let packed = var_nums.get(info.arg); + let (idx1, idx2) = packed.indexes(); + + let mut first = info; + first.instr = Instruction::StoreFast { + var_num: Arg::marker(), } - Opcode::LoadFastLoadFast => { - info.instr = Opcode::LoadFastBorrowLoadFastBorrow.into(); + .into(); + first.arg = OpArg::new(u32::from(idx1)); + new_instructions.push(first); + + let mut second = info; + second.instr = Instruction::StoreFast { + var_num: Arg::marker(), } - _ => {} + .into(); + second.arg = OpArg::new(u32::from(idx2)); + new_instructions.push(second); + continue; } + + new_instructions.push(info); } + block.instructions = new_instructions; } } @@ -1930,6 +2394,13 @@ impl CodeInfo { return; } + let merged_cell_local = |cell_relative: usize| { + self.metadata + .cellvars + .get_index(cell_relative) + .and_then(|name| self.metadata.varnames.get_index_of(name.as_str())) + }; + let mut nparams = self.metadata.argcount as usize + self.metadata.kwonlyargcount as usize; if self.flags.contains(CodeFlags::VARARGS) { nparams += 1; @@ -1989,6 +2460,15 @@ impl CodeInfo { } new_instructions.push(info); } + Some(Instruction::StoreDeref { i }) => { + let cell_relative = usize::from(i.get(info.arg)); + if let Some(var_idx) = merged_cell_local(cell_relative) + && var_idx < nlocals + { + unsafe_mask[var_idx] = false; + } + new_instructions.push(info); + } Some(Instruction::StoreFastStoreFast { var_nums }) => { let packed = var_nums.get(info.arg); let (idx1, idx2) = packed.indexes(); @@ -2009,7 +2489,17 @@ impl CodeInfo { } new_instructions.push(info); } - Some(Instruction::LoadFast { var_num }) => { + Some(Instruction::DeleteDeref { i }) => { + let cell_relative = usize::from(i.get(info.arg)); + if let Some(var_idx) = merged_cell_local(cell_relative) + && var_idx < nlocals + { + unsafe_mask[var_idx] = true; + } + new_instructions.push(info); + } + Some(Instruction::LoadFast { var_num }) + | Some(Instruction::LoadFastBorrow { var_num }) => { let var_idx = usize::from(var_num.get(info.arg)); if var_idx < nlocals && unsafe_mask[var_idx] { info.instr = Opcode::LoadFastCheck.into(); @@ -2020,7 +2510,8 @@ impl CodeInfo { } new_instructions.push(info); } - Some(Instruction::LoadFastLoadFast { var_nums }) => { + Some(Instruction::LoadFastLoadFast { var_nums }) + | Some(Instruction::LoadFastBorrowLoadFastBorrow { var_nums }) => { let packed = var_nums.get(info.arg); let (idx1, idx2) = packed.indexes(); let idx1 = usize::from(idx1); @@ -3961,6 +4452,30 @@ pub(crate) fn label_exception_targets(blocks: &mut [Block]) { } } +/// Lower only LOAD_CLOSURE pseudo ops to LOAD_FAST so optimize_load_fast can +/// make the same borrow decision CPython does for merged cell locals. +fn convert_load_closure_pseudo_ops(blocks: &mut [Block], cellfixedoffsets: &[u32]) { + for block in blocks.iter_mut() { + for info in &mut block.instructions { + let Some(pseudo) = info.instr.pseudo() else { + continue; + }; + match pseudo { + PseudoInstruction::LoadClosure { i } => { + let cell_relative = i.get(info.arg) as usize; + let new_idx = cellfixedoffsets[cell_relative]; + info.arg = OpArg::new(new_idx); + info.instr = Instruction::LoadFast { + var_num: Arg::marker(), + } + .into(); + } + _ => {} + } + } + } +} + /// Convert remaining pseudo ops to real instructions or NOP. /// flowgraph.c convert_pseudo_ops pub(crate) fn convert_pseudo_ops(blocks: &mut [Block], cellfixedoffsets: &[u32]) { diff --git a/crates/compiler-core/src/bytecode/instruction.rs b/crates/compiler-core/src/bytecode/instruction.rs index 079d7963259..269fe518e6e 100644 --- a/crates/compiler-core/src/bytecode/instruction.rs +++ b/crates/compiler-core/src/bytecode/instruction.rs @@ -1405,11 +1405,11 @@ impl InstructionMetadata for PseudoInstruction { /// SETUP_FINALLY: +1 (exc) /// SETUP_CLEANUP: +2 (lasti + exc) /// SETUP_WITH: +1 (pops __enter__ result, pushes lasti + exc) - fn stack_effect_jump(&self, _oparg: u32) -> i32 { + fn stack_effect_jump(&self, oparg: u32) -> i32 { match self { Self::SetupFinally { .. } | Self::SetupWith { .. } => 1, Self::SetupCleanup { .. } => 2, - _ => self.stack_effect(_oparg), + _ => self.stack_effect(oparg), } } diff --git a/scripts/dis_dump.py b/scripts/dis_dump.py index e30a8955fdd..af31fd08230 100755 --- a/scripts/dis_dump.py +++ b/scripts/dis_dump.py @@ -11,6 +11,7 @@ """ import argparse +import builtins import dis import json import os @@ -100,10 +101,31 @@ def _unescape(m): hasattr(sys, "implementation") and sys.implementation.name == "rustpython" ) +if _IS_RUSTPYTHON and hasattr(dis, "_common_constants"): + common_constants = list(dis._common_constants) + while len(common_constants) < 7: + common_constants.append((builtins.list, builtins.set)[len(common_constants) - 5]) + dis._common_constants = common_constants + # RustPython's ComparisonOperator enum values → operator strings _RP_CMP_OPS = {0: "<", 1: "<=", 2: "==", 3: "!=", 4: ">", 5: ">="} +def _resolve_localsplus_name(code, arg): + if not isinstance(arg, int) or arg < 0: + return arg + nlocals = len(code.co_varnames) + if arg < nlocals: + return code.co_varnames[arg] + varnames_set = set(code.co_varnames) + nonparam_cells = [v for v in code.co_cellvars if v not in varnames_set] + extra = nonparam_cells + list(code.co_freevars) + idx = arg - nlocals + if 0 <= idx < len(extra): + return extra[idx] + return arg + + def _resolve_arg_fallback(code, opname, arg): """Resolve a raw argument to its human-readable form. @@ -113,8 +135,7 @@ def _resolve_arg_fallback(code, opname, arg): return arg try: if "FAST" in opname: - if 0 <= arg < len(code.co_varnames): - return code.co_varnames[arg] + return _resolve_localsplus_name(code, arg) elif opname == "LOAD_CONST": if 0 <= arg < len(code.co_consts): return _normalize_argrepr(repr(code.co_consts[arg])) @@ -125,18 +146,7 @@ def _resolve_arg_fallback(code, opname, arg): "LOAD_CLOSURE", "MAKE_CELL", ): - # arg is localsplus index: - # 0..nlocals-1 = varnames (parameter cells reuse these slots) - # nlocals.. = non-parameter cells + freevars - nlocals = len(code.co_varnames) - if arg < nlocals: - return code.co_varnames[arg] - varnames_set = set(code.co_varnames) - nonparam_cells = [v for v in code.co_cellvars if v not in varnames_set] - extra = nonparam_cells + list(code.co_freevars) - idx = arg - nlocals - if 0 <= idx < len(extra): - return extra[idx] + return _resolve_localsplus_name(code, arg) elif opname in ( "LOAD_NAME", "STORE_NAME", From fdba5cb4e1ff88354bfda0b9306f290827a6835b Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Mon, 13 Apr 2026 23:14:55 +0900 Subject: [PATCH 2/8] Bytecode parity - constant folding, annotation ordering, superinstruction alignment - Add BoolOp constant folding with short-circuit semantics in compile_expression - Add constant truthiness evaluation for assert statement optimization - Disable const collection/boolop folding in starred unpack and assignment contexts - Move annotation block generation after body with AnnotationsPlaceholder splicing - Reorder insert_superinstructions to run before push_cold_blocks (matching flowgraph.c) - Lower LOAD_CLOSURE after superinstructions to avoid false LOAD_FAST_LOAD_FAST - Add ToBool before PopJumpIf in comparisons and chained compare cleanup blocks - Unify annotation dict building to always use incremental BuildMap + StoreSubscr - Add TrueDivide constant folding for integer operands - Fold constant sets to Frozenset (not Tuple) in try_fold_constant_collection - Add PyVmBag for frozenset constant materialization in code objects - Add remove_redundant_const_pop_top_pairs pass and peephole const+branch folding - Emit Nop for skipped constant expressions and constant-true asserts - Preserve comprehension local ordering by source-order bound name collection - Simplify annotation scanning in symboltable (remove simple-name gate) --- crates/codegen/src/compile.rs | 721 +++++++++++++++++++++++------- crates/codegen/src/ir.rs | 399 ++++++++++++----- crates/codegen/src/symboltable.rs | 17 +- crates/vm/src/builtins/code.rs | 108 ++++- crates/vm/src/import.rs | 9 +- crates/vm/src/stdlib/_ast.rs | 2 +- crates/vm/src/vm/compile.rs | 4 +- scripts/dis_dump.py | 4 +- 8 files changed, 980 insertions(+), 284 deletions(-) diff --git a/crates/codegen/src/compile.rs b/crates/codegen/src/compile.rs index 360fd0809fa..ba2dfab3244 100644 --- a/crates/codegen/src/compile.rs +++ b/crates/codegen/src/compile.rs @@ -17,10 +17,11 @@ use crate::{ unparse::UnparseExpr, }; use alloc::borrow::Cow; +use core::mem; use itertools::Itertools; use malachite_bigint::BigInt; use num_complex::Complex; -use num_traits::{Num, ToPrimitive}; +use num_traits::{Num, ToPrimitive, Zero}; use ruff_python_ast as ast; use ruff_text_size::{Ranged, TextRange, TextSize}; use rustpython_compiler_core::{ @@ -165,6 +166,12 @@ struct Compiler { /// When > 0, the compiler walks AST (consuming sub_tables) but emits no bytecode. /// Mirrors CPython's `c_do_not_emit_bytecode`. do_not_emit_bytecode: u32, + /// Disable constant BoolOp folding in contexts where CPython preserves + /// short-circuit structure, such as starred unpack expressions. + disable_const_boolop_folding: bool, + /// Disable constant tuple/list/set collection folding in contexts where + /// CPython keeps the builder form for later assignment lowering. + disable_const_collection_folding: bool, } #[derive(Clone, Copy)] @@ -214,13 +221,6 @@ impl CompileContext { } } -fn bool_literal_value(expr: &ast::Expr) -> Option { - match expr { - ast::Expr::BooleanLiteral(ast::ExprBooleanLiteral { value, .. }) => Some(*value), - _ => None, - } -} - /// Segment of a parsed %-format string for optimize_format_str. struct FormatSegment { literal: String, @@ -444,6 +444,22 @@ enum CollectionType { } impl Compiler { + fn constant_truthiness(constant: &ConstantData) -> bool { + match constant { + ConstantData::Tuple { elements } | ConstantData::Frozenset { elements } => { + !elements.is_empty() + } + ConstantData::Integer { value } => !value.is_zero(), + ConstantData::Float { value } => *value != 0.0, + ConstantData::Complex { value } => value.re != 0.0 || value.im != 0.0, + ConstantData::Boolean { value } => *value, + ConstantData::Str { value } => !value.is_empty(), + ConstantData::Bytes { value } => !value.is_empty(), + ConstantData::Code { .. } | ConstantData::Slice { .. } | ConstantData::Ellipsis => true, + ConstantData::None => false, + } + } + fn new(opts: CompileOpts, source_file: SourceFile, code_name: String) -> Self { let module_code = ir::CodeInfo { flags: bytecode::CodeFlags::NEWLOCALS, @@ -451,6 +467,7 @@ impl Compiler { private: None, blocks: vec![ir::Block::default()], current_block: BlockIdx::new(0), + annotations_blocks: None, metadata: ir::CodeUnitMetadata { name: code_name.clone(), qualname: Some(code_name), @@ -490,9 +507,62 @@ impl Compiler { in_annotation: false, interactive: false, do_not_emit_bytecode: 0, + disable_const_boolop_folding: false, + disable_const_collection_folding: false, } } + fn compile_expression_without_const_boolop_folding( + &mut self, + expression: &ast::Expr, + ) -> CompileResult<()> { + let previous = self.disable_const_boolop_folding; + self.disable_const_boolop_folding = true; + let result = self.compile_expression(expression); + self.disable_const_boolop_folding = previous; + result.map(|_| ()) + } + + fn compile_expression_without_const_collection_folding( + &mut self, + expression: &ast::Expr, + ) -> CompileResult<()> { + let previous = self.disable_const_collection_folding; + self.disable_const_collection_folding = true; + let result = self.compile_expression(expression); + self.disable_const_collection_folding = previous; + result.map(|_| ()) + } + + fn is_unpack_assignment_target(target: &ast::Expr) -> bool { + matches!(target, ast::Expr::List(_) | ast::Expr::Tuple(_)) + } + + fn compile_module_annotation_setup_sequence( + &mut self, + body: &[ast::Stmt], + ) -> CompileResult<()> { + let (saved_blocks, saved_current_block) = { + let code = self.current_code_info(); + ( + mem::replace(&mut code.blocks, vec![ir::Block::default()]), + mem::replace(&mut code.current_block, BlockIdx::new(0)), + ) + }; + + let result = self.compile_module_annotate(body); + + let annotations_blocks = { + let code = self.current_code_info(); + let annotations_blocks = mem::replace(&mut code.blocks, saved_blocks); + code.current_block = saved_current_block; + annotations_blocks + }; + self.current_code_info().annotations_blocks = Some(annotations_blocks); + + result.map(|_| ()) + } + /// Compile just start and stop of a slice (for BINARY_SLICE/STORE_SLICE) // = codegen_slice_two_parts fn compile_slice_two_parts(&mut self, s: &ast::ExprSlice) -> CompileResult<()> { @@ -590,11 +660,12 @@ impl Compiler { }; // Fold all-constant collections (>= 3 elements) regardless of size - if !seen_star + if !self.disable_const_collection_folding + && !seen_star && pushed == 0 && n >= 3 && elts.iter().all(|e| e.is_constant()) - && let Some(folded) = self.try_fold_constant_collection(elts)? + && let Some(folded) = self.try_fold_constant_collection(elts, collection_type)? { match collection_type { CollectionType::Tuple => { @@ -657,7 +728,7 @@ impl Compiler { } // Compile the starred expression and extend - self.compile_expression(value)?; + self.compile_expression_without_const_boolop_folding(value)?; match collection_type { CollectionType::List => { emit!(self, Instruction::ListExtend { i: 1 }); @@ -1208,6 +1279,7 @@ impl Compiler { private, blocks: vec![ir::Block::default()], current_block: BlockIdx::new(0), + annotations_blocks: None, metadata: ir::CodeUnitMetadata { name: name.to_owned(), qualname: None, // Will be set below @@ -1852,6 +1924,7 @@ impl Compiler { } self.emit_resume_for_scope(CompilerScope::Module, 1); + emit!(self, PseudoInstruction::AnnotationsPlaceholder); let (doc, statements) = split_doc(&body.body, &self.opts); if let Some(value) = doc { @@ -1868,10 +1941,8 @@ impl Compiler { // PEP 563: Initialize __annotations__ dict emit!(self, Instruction::SetupAnnotations); } else { - // PEP 649: Generate __annotate__ function FIRST (before statements) - self.compile_module_annotate(statements)?; - - // PEP 649: Initialize __conditional_annotations__ set after __annotate__ + // PEP 649: Initialize __conditional_annotations__ before the body. + // CPython generates __annotate__ after the body in codegen_body(). if self.current_symbol_table().has_conditional_annotations { emit!(self, Instruction::BuildSet { count: 0 }); self.store_name("__conditional_annotations__")?; @@ -1882,6 +1953,10 @@ impl Compiler { // Compile all statements self.compile_statements(statements)?; + if Self::find_ann(statements) && !self.future_annotations { + self.compile_module_annotation_setup_sequence(statements)?; + } + assert_eq!(self.code_stack.len(), size_before); // Emit None at end: @@ -1900,6 +1975,7 @@ impl Compiler { self.symbol_table_stack.push(symbol_table); self.emit_resume_for_scope(CompilerScope::Module, 1); + emit!(self, PseudoInstruction::AnnotationsPlaceholder); // Handle annotations based on future_annotations flag if Self::find_ann(body) { @@ -1907,10 +1983,8 @@ impl Compiler { // PEP 563: Initialize __annotations__ dict emit!(self, Instruction::SetupAnnotations); } else { - // PEP 649: Generate __annotate__ function FIRST (before statements) - self.compile_module_annotate(body)?; - - // PEP 649: Initialize __conditional_annotations__ set after __annotate__ + // PEP 649: Initialize __conditional_annotations__ before the body. + // CPython generates __annotate__ after the body in codegen_body(). if self.current_symbol_table().has_conditional_annotations { emit!(self, Instruction::BuildSet { count: 0 }); self.store_name("__conditional_annotations__")?; @@ -1954,6 +2028,10 @@ impl Compiler { self.emit_load_const(ConstantData::None); }; + if Self::find_ann(body) && !self.future_annotations { + self.compile_module_annotation_setup_sequence(body)?; + } + self.emit_return_value(); Ok(()) } @@ -2393,7 +2471,7 @@ impl Compiler { let dominated_by_interactive = self.interactive && !self.ctx.in_func() && !self.ctx.in_class; if !dominated_by_interactive && Self::is_const_expression(value) { - // Skip compilation entirely - the expression has no side effects + emit!(self, Instruction::Nop); } else { self.compile_expression(value)?; @@ -2520,27 +2598,42 @@ impl Compiler { ast::Stmt::Assert(ast::StmtAssert { test, msg, .. }) => { // if some flag, ignore all assert statements! if self.opts.optimize == 0 { - let after_block = self.new_block(); - self.compile_jump_if(test, true, after_block)?; - - emit!( - self, - Instruction::LoadCommonConstant { - idx: bytecode::CommonConstant::AssertionError + let after_block = match self.try_fold_constant_expr(test)? { + Some(constant) if Self::constant_truthiness(&constant) => { + self.consume_skipped_nested_scopes_in_expr(test)?; + if let Some(expr) = msg { + self.consume_skipped_nested_scopes_in_expr(expr)?; + } + emit!(self, Instruction::Nop); + None } - ); - if let Some(e) = msg { - self.compile_expression(e)?; - emit!(self, Instruction::Call { argc: 0 }); - } - emit!( - self, - Instruction::RaiseVarargs { - argc: bytecode::RaiseKind::Raise, + Some(_) => Some(self.new_block()), + None => { + let after_block = self.new_block(); + self.compile_jump_if(test, true, after_block)?; + Some(after_block) } - ); + }; - self.switch_to_block(after_block); + if let Some(after_block) = after_block { + emit!( + self, + Instruction::LoadCommonConstant { + idx: bytecode::CommonConstant::AssertionError + } + ); + if let Some(e) = msg { + self.compile_expression(e)?; + emit!(self, Instruction::Call { argc: 0 }); + } + emit!( + self, + Instruction::RaiseVarargs { + argc: bytecode::RaiseKind::Raise, + } + ); + self.switch_to_block(after_block); + } } else { // Optimized-out asserts still need to consume any nested // scope symbol tables they contain so later nested scopes @@ -2606,7 +2699,11 @@ impl Compiler { self.switch_to_block(dead); } ast::Stmt::Assign(ast::StmtAssign { targets, value, .. }) => { - self.compile_expression(value)?; + if targets.len() == 1 && Self::is_unpack_assignment_target(&targets[0]) { + self.compile_expression_without_const_collection_folding(value)?; + } else { + self.compile_expression(value)?; + } for (i, target) in targets.iter().enumerate() { if i + 1 != targets.len() { @@ -4237,18 +4334,33 @@ impl Compiler { parameters: &ast::Parameters, returns: Option<&ast::Expr>, ) -> CompileResult { + let has_signature_annotations = parameters + .args + .iter() + .map(|x| &x.parameter) + .chain(parameters.posonlyargs.iter().map(|x| &x.parameter)) + .chain(parameters.vararg.as_deref()) + .chain(parameters.kwonlyargs.iter().map(|x| &x.parameter)) + .chain(parameters.kwarg.as_deref()) + .any(|param| param.annotation.is_some()) + || returns.is_some(); + if !has_signature_annotations { + return Ok(false); + } + // Try to enter annotation scope - returns None if no annotation_block exists let Some(saved_ctx) = self.enter_annotation_scope(func_name)? else { return Ok(false); }; // Count annotations - let parameters_iter = core::iter::empty() - .chain(¶meters.posonlyargs) - .chain(¶meters.args) - .chain(¶meters.kwonlyargs) + let parameters_iter = parameters + .args + .iter() .map(|x| &x.parameter) + .chain(parameters.posonlyargs.iter().map(|x| &x.parameter)) .chain(parameters.vararg.as_deref()) + .chain(parameters.kwonlyargs.iter().map(|x| &x.parameter)) .chain(parameters.kwarg.as_deref()); let num_annotations: u32 = @@ -4257,12 +4369,13 @@ impl Compiler { + if returns.is_some() { 1 } else { 0 }; // Compile annotations inside the annotation scope - let parameters_iter = core::iter::empty() - .chain(¶meters.posonlyargs) - .chain(¶meters.args) - .chain(¶meters.kwonlyargs) + let parameters_iter = parameters + .args + .iter() .map(|x| &x.parameter) + .chain(parameters.posonlyargs.iter().map(|x| &x.parameter)) .chain(parameters.vararg.as_deref()) + .chain(parameters.kwonlyargs.iter().map(|x| &x.parameter)) .chain(parameters.kwarg.as_deref()); for param in parameters_iter { @@ -4413,20 +4526,13 @@ impl Compiler { // Emit format validation: if format > VALUE_WITH_FAKE_GLOBALS: raise NotImplementedError self.emit_format_validation()?; - if has_conditional { - // PEP 649: Build dict incrementally, checking conditional annotations - // Start with empty dict - emit!(self, Instruction::BuildMap { count: 0 }); + emit!(self, Instruction::BuildMap { count: 0 }); - // Process each annotation - for (idx, (name, annotation)) in annotations.iter().enumerate() { - // Check if index is in __conditional_annotations__ - let not_set_block = self.new_block(); + for (idx, (name, annotation)) in annotations.iter().enumerate() { + let not_set_block = has_conditional.then(|| self.new_block()); - // LOAD_CONST index + if has_conditional { self.emit_load_const(ConstantData::Integer { value: idx.into() }); - // Load __conditional_annotations__ from appropriate scope - // Class scope: LoadDeref (freevars), Module scope: LoadGlobal if parent_scope_type == CompilerScope::Class { let idx = self.get_free_var_index("__conditional_annotations__")?; emit!(self, Instruction::LoadDeref { i: idx }); @@ -4434,60 +4540,34 @@ impl Compiler { let cond_annotations_name = self.name("__conditional_annotations__"); self.emit_load_global(cond_annotations_name, false); } - // CONTAINS_OP (in) emit!( self, Instruction::ContainsOp { invert: bytecode::Invert::No } ); - // POP_JUMP_IF_FALSE not_set emit!( self, Instruction::PopJumpIfFalse { - delta: not_set_block + delta: not_set_block.expect("missing not_set block") } ); - - // Annotation value - self.compile_annotation(annotation)?; - // COPY dict to TOS - emit!(self, Instruction::Copy { i: 2 }); - // LOAD_CONST name - self.emit_load_const(ConstantData::Str { - value: self.mangle(name).into_owned().into(), - }); - // STORE_SUBSCR - dict[name] = value - emit!(self, Instruction::StoreSubscr); - - // not_set label - self.switch_to_block(not_set_block); } - // Return the dict - emit!(self, Instruction::ReturnValue); - } else { - // No conditional annotations - use simple BuildMap - let num_annotations = u32::try_from(annotations.len()).expect("too many annotations"); + self.compile_annotation(annotation)?; + emit!(self, Instruction::Copy { i: 2 }); + self.emit_load_const(ConstantData::Str { + value: self.mangle(name).into_owned().into(), + }); + emit!(self, Instruction::StoreSubscr); - // Compile annotations inside the annotation scope - for (name, annotation) in annotations { - self.emit_load_const(ConstantData::Str { - value: self.mangle(name).into_owned().into(), - }); - self.compile_annotation(annotation)?; + if let Some(not_set_block) = not_set_block { + self.switch_to_block(not_set_block); } - - // Build the map and return it - emit!( - self, - Instruction::BuildMap { - count: num_annotations, - } - ); - emit!(self, Instruction::ReturnValue); } + emit!(self, Instruction::ReturnValue); + // Exit annotation scope - pop symbol table, restore to parent's annotation_block, and get code let annotation_table = self.pop_symbol_table(); // Restore annotation_block to module's symbol table @@ -5140,9 +5220,6 @@ impl Compiler { emit!(self, Instruction::BuildSet { count: 0 }); self.store_name("__conditional_annotations__")?; } - - // PEP 649: Generate __annotate__ function for class annotations - self.compile_module_annotate(body)?; } } @@ -5159,6 +5236,10 @@ impl Compiler { // 3. Compile the class body self.compile_statements(body)?; + if Self::find_ann(body) && !self.future_annotations { + self.compile_module_annotate(body)?; + } + // 4. Handle __classcell__ if needed let classcell_idx = self .code_stack @@ -7047,6 +7128,7 @@ impl Compiler { // if comparison result is false, we break with this value; if true, try the next one. emit!(self, Instruction::Copy { i: 1 }); + emit!(self, Instruction::ToBool); emit!(self, Instruction::PopJumpIfFalse { delta: cleanup }); emit!(self, Instruction::PopTop); } @@ -7098,12 +7180,14 @@ impl Compiler { emit!(self, Instruction::Swap { i: 2 }); emit!(self, Instruction::Copy { i: 2 }); self.compile_addcompare(op); + emit!(self, Instruction::ToBool); emit!(self, Instruction::PopJumpIfFalse { delta: cleanup }); } self.compile_expression(last_comparator)?; self.set_source_range(compare_range); self.compile_addcompare(last_op); + emit!(self, Instruction::ToBool); self.emit_pop_jump_by_condition(condition, target_block); emit!(self, PseudoInstruction::Jump { delta: end }); @@ -7520,10 +7604,7 @@ impl Compiler { _ => { // Fall back case which always will work! self.compile_expression(expression)?; - // Compare already produces a bool; everything else needs TO_BOOL - if !matches!(expression, ast::Expr::Compare(_)) { - emit!(self, Instruction::ToBool); - } + emit!(self, Instruction::ToBool); if condition { emit!( self, @@ -7859,6 +7940,47 @@ impl Compiler { let range = expression.range(); self.set_source_range(range); + 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; + for value in values { + let Some(constant) = self.try_fold_constant_expr(value)? else { + break; + }; + let is_truthy = Self::constant_truthiness(&constant); + last_constant = Some(constant); + match op { + ast::BoolOp::Or if is_truthy => { + self.emit_load_const(last_constant.expect("missing boolop constant")); + return Ok(()); + } + ast::BoolOp::And if !is_truthy => { + self.emit_load_const(last_constant.expect("missing boolop constant")); + return Ok(()); + } + ast::BoolOp::Or | ast::BoolOp::And => { + simplified_prefix += 1; + } + } + } + + if simplified_prefix == values.len() { + self.emit_load_const(last_constant.expect("missing folded boolop constant")); + return Ok(()); + } + if simplified_prefix > 0 { + let tail = &values[simplified_prefix..]; + if let [value] = tail { + self.compile_expression(value)?; + } else { + self.compile_bool_op(op, tail)?; + } + return Ok(()); + } + } + match &expression { ast::Expr::Call(ast::ExprCall { func, arguments, .. @@ -8686,7 +8808,7 @@ impl Compiler { // Single starred arg: pass value directly to CallFunctionEx. // Runtime will convert to tuple and validate with function name. if let ast::Expr::Starred(ast::ExprStarred { value, .. }) = &arguments.args[0] { - self.compile_expression(value)?; + self.compile_expression_without_const_boolop_folding(value)?; } } else if !has_starred { for arg in &arguments.args { @@ -8742,7 +8864,7 @@ impl Compiler { have_dict = true; } - self.compile_expression(&keyword.value)?; + self.compile_expression_without_const_boolop_folding(&keyword.value)?; emit!(self, Instruction::DictMerge { i: 1 }); } else { nseen += 1; @@ -9027,21 +9149,14 @@ impl Compiler { end_async_for_target, )); - // Now evaluate the ifs: + // CPython always lowers comprehension guards through codegen_jump_if + // and leaves constant-folding to later CFG optimization passes. for if_condition in &generator.ifs { - match bool_literal_value(if_condition) { - Some(true) => {} - Some(false) => { - emit!( - self, - PseudoInstruction::Jump { - delta: if_cleanup_block - } - ); - break; - } - None => self.compile_jump_if(if_condition, false, if_cleanup_block)?, - } + self.compile_jump_if(if_condition, false, if_cleanup_block)?; + } + if !generator.ifs.is_empty() { + let body_block = self.new_block(); + self.switch_to_block(body_block); } } @@ -9167,22 +9282,49 @@ impl Compiler { let ct = self.current_symbol_table(); ct.typ == CompilerScope::Class && !self.current_code_info().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); + } + _ => {} + } + } + 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, sym) in &comp_table.symbols { - if sym.flags.contains(SymbolFlags::PARAMETER) { - continue; // skip .0 + for name in source_order_bound_names + .into_iter() + .chain(comp_table.symbols.keys().cloned()) + { + if pushed_locals.iter().any(|existing| existing == &name) { + continue; } - // 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.clone()); + 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 || in_class_block { + pushed_locals.push(name); + } } } @@ -9315,21 +9457,10 @@ impl Compiler { end_async_for_target, )); - // Evaluate the if conditions + // CPython always lowers comprehension guards through codegen_jump_if + // and leaves constant-folding to later CFG optimization passes. for if_condition in &generator.ifs { - match bool_literal_value(if_condition) { - Some(true) => {} - Some(false) => { - emit!( - self, - PseudoInstruction::Jump { - delta: if_cleanup_block - } - ); - break; - } - None => self.compile_jump_if(if_condition, false, if_cleanup_block)?, - } + self.compile_jump_if(if_condition, false, if_cleanup_block)?; } } @@ -9527,6 +9658,7 @@ impl Compiler { fn try_fold_constant_collection( &mut self, elts: &[ast::Expr], + collection_type: CollectionType, ) -> CompileResult> { let mut constants = Vec::with_capacity(elts.len()); for elt in elts { @@ -9535,9 +9667,15 @@ impl Compiler { }; constants.push(constant); } - Ok(Some(ConstantData::Tuple { - elements: constants, - })) + let constant = match collection_type { + CollectionType::Tuple | CollectionType::List => ConstantData::Tuple { + elements: constants, + }, + CollectionType::Set => ConstantData::Frozenset { + elements: constants, + }, + }; + Ok(Some(constant)) } fn try_fold_constant_expr(&mut self, expr: &ast::Expr) -> CompileResult> { @@ -9570,6 +9708,45 @@ impl Compiler { } ConstantData::Tuple { elements } } + ast::Expr::BoolOp(ast::ExprBoolOp { op, values, .. }) => { + let mut constants = Vec::with_capacity(values.len()); + for value in values { + let Some(constant) = self.try_fold_constant_expr(value)? else { + return Ok(None); + }; + constants.push(constant); + } + let mut iter = constants.into_iter(); + let Some(first) = iter.next() else { + return Ok(None); + }; + let mut selected = first; + match op { + ast::BoolOp::Or => { + if !Self::constant_truthiness(&selected) { + for constant in iter { + let is_truthy = Self::constant_truthiness(&constant); + selected = constant; + if is_truthy { + break; + } + } + } + } + ast::BoolOp::And => { + if Self::constant_truthiness(&selected) { + for constant in iter { + let is_truthy = Self::constant_truthiness(&constant); + selected = constant; + if !is_truthy { + break; + } + } + } + } + } + selected + } _ => return Ok(None), })) } @@ -10932,6 +11109,127 @@ x = not True )); } + #[test] + fn test_plain_constant_bool_op_folds_to_selected_operand() { + let code = compile_exec( + "\ +x = 1 or 2 or 3 +", + ); + let ops: Vec<_> = code + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + let folded_small_int = code.instructions.iter().any(|unit| { + matches!( + unit.op, + Instruction::LoadSmallInt { i } + if i.get(OpArg::new(u32::from(u8::from(unit.arg)))) == 1 + ) + }); + let folded_const_one = code + .instructions + .iter() + .find_map(|unit| match unit.op { + Instruction::LoadConst { .. } => code.constants.get(usize::from(u8::from(unit.arg))), + _ => None, + }) + .is_some_and(|constant| { + matches!(constant, ConstantData::Integer { value } if *value == BigInt::from(1)) + }); + + assert!( + folded_small_int || folded_const_one, + "expected folded constant 1, got ops={ops:?}" + ); + assert!( + !ops.iter().any(|op| { + matches!( + op, + Instruction::Copy { .. } + | Instruction::ToBool + | Instruction::PopJumpIfTrue { .. } + | Instruction::PopJumpIfFalse { .. } + ) + }), + "plain constant BoolOp should not leave short-circuit scaffolding, got ops={ops:?}" + ); + } + + #[test] + fn test_starred_call_preserves_bool_op_short_circuit_shape() { + let code = compile_exec( + "\ +def f(g): + return g(*(() or (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.iter().any(|op| matches!(op, Instruction::Copy { .. })), + "starred BoolOp should keep short-circuit COPY, got ops={ops:?}" + ); + assert!( + ops.iter().any(|op| matches!(op, Instruction::ToBool)), + "starred BoolOp should keep TO_BOOL, got ops={ops:?}" + ); + assert!( + ops.iter() + .any(|op| matches!(op, Instruction::PopJumpIfTrue { .. })), + "starred BoolOp should keep POP_JUMP_IF_TRUE, got ops={ops:?}" + ); + } + + #[test] + fn test_partial_constant_bool_op_folds_prefix_in_value_context() { + let code = compile_exec( + "\ +def outer(null): + @False or null + def f(x): + pass +", + ); + let outer = find_code(&code, "outer").expect("missing outer code"); + let ops: Vec<_> = outer + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + assert!( + ops.iter().any(|op| { + matches!( + op, + Instruction::LoadFast { .. } | Instruction::LoadFastBorrow { .. } + ) + }), + "expected surviving decorator expression to load null directly, got ops={ops:?}" + ); + assert!( + !ops.iter().any(|op| { + matches!( + op, + Instruction::Copy { .. } + | Instruction::ToBool + | Instruction::PopJumpIfTrue { .. } + | Instruction::PopJumpIfFalse { .. } + ) + }), + "partial constant BoolOp should not leave short-circuit scaffolding, got ops={ops:?}" + ); + } + #[test] fn test_nested_double_async_with() { assert_dis_snapshot!(compile_exec( @@ -12054,7 +12352,7 @@ def f(): ); assert!(f.constants.iter().any(|constant| matches!( constant, - ConstantData::Tuple { elements } + ConstantData::Frozenset { elements } if matches!( elements.as_slice(), [ @@ -12066,6 +12364,129 @@ def f(): ))); } + #[test] + fn test_single_unpack_assignment_disables_constant_collection_folding() { + let code = compile_exec("a, b, c = 1, 2, 3\n"); + + assert!( + !code.instructions.iter().any(|unit| { + matches!(unit.op, Instruction::UnpackSequence { .. }) + || matches!(unit.op, Instruction::LoadConst { .. }) + && matches!( + code.constants.get(usize::from(u8::from(unit.arg))), + Some(ConstantData::Tuple { .. }) + ) + }), + "single unpack assignment should keep builder form for later lowering, got ops={:?}", + code.instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + assert!( + code.instructions + .iter() + .filter(|unit| matches!(unit.op, Instruction::LoadSmallInt { .. })) + .count() + >= 3, + "expected individual constant loads before unpack-target stores, got ops={:?}", + code.instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + } + + #[test] + fn test_chained_unpack_assignment_keeps_constant_collection_folding() { + let code = compile_exec("(a, b) = c = d = (1, 2)\n"); + + assert!( + code.instructions + .iter() + .any(|unit| matches!(unit.op, Instruction::LoadConst { .. })), + "chained unpack assignment should keep tuple constant, got ops={:?}", + code.instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + assert!( + code.instructions + .iter() + .any(|unit| matches!(unit.op, Instruction::UnpackSequence { .. })), + "chained unpack assignment should still unpack the copied tuple, got ops={:?}", + code.instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + } + + #[test] + fn test_constant_true_assert_skips_message_nested_scope() { + let code = compile_exec("assert 1, (lambda x: x + 1)\n"); + + assert_eq!( + code.constants + .iter() + .filter(|constant| matches!(constant, ConstantData::Code { .. })) + .count(), + 0, + "constant-true assert should not compile the skipped message lambda" + ); + assert!( + !code + .instructions + .iter() + .any(|unit| matches!(unit.op, Instruction::RaiseVarargs { .. })), + "constant-true assert should be elided, got ops={:?}", + code.instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + } + + #[test] + fn test_constant_false_assert_raises_directly() { + let code = compile_exec("assert 0, (lambda x: x + 1)\n"); + + assert!( + !code.instructions.iter().any(|unit| { + matches!( + unit.op, + Instruction::ToBool + | Instruction::PopJumpIfTrue { .. } + | Instruction::PopJumpIfFalse { .. } + ) + }), + "constant-false assert should raise directly without a test branch, got ops={:?}", + code.instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + assert!( + code.instructions + .iter() + .any(|unit| matches!(unit.op, Instruction::RaiseVarargs { .. })), + "constant-false assert should still raise, got ops={:?}", + code.instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + assert_eq!( + code.constants + .iter() + .filter(|constant| matches!(constant, ConstantData::Code { .. })) + .count(), + 1, + "constant-false assert should still compile the message lambda" + ); + } + #[test] fn test_optimized_assert_preserves_nested_scope_order() { compile_exec_optimized( diff --git a/crates/codegen/src/ir.rs b/crates/codegen/src/ir.rs index f6c9a4ae408..ed5652070db 100644 --- a/crates/codegen/src/ir.rs +++ b/crates/codegen/src/ir.rs @@ -170,6 +170,7 @@ pub struct CodeInfo { pub blocks: Vec, pub current_block: BlockIdx, + pub annotations_blocks: Option>, pub metadata: CodeUnitMetadata, @@ -199,6 +200,7 @@ impl CodeInfo { mut self, opts: &crate::compile::CompileOpts, ) -> crate::InternalResult { + self.splice_annotations_blocks(); // Constant folding passes self.fold_binop_constants(); self.remove_nops(); @@ -224,7 +226,7 @@ impl CodeInfo { self.eliminate_dead_stores(); // apply_static_swaps: reorder stores to eliminate SWAPs self.apply_static_swaps(); - // Peephole optimizer creates superinstructions matching CPython + // Peephole optimizer handles constant and compare folding. self.peephole_optimize(); // Phase 1: _PyCfg_OptimizeCodeUnit (flowgraph.c) @@ -236,7 +238,11 @@ impl CodeInfo { jump_threading(&mut self.blocks); self.eliminate_unreachable_blocks(); self.remove_nops(); - // CPython inserts superinstructions before optimize_load_fast. + self.add_checks_for_loads_of_uninitialized_variables(); + // CPython inserts superinstructions in _PyCfg_OptimizeCodeUnit, before + // later jump normalization / block reordering can create adjacencies + // that never exist at this stage in flowgraph.c. + self.insert_superinstructions(); push_cold_blocks_to_end(&mut self.blocks); // Phase 2: _PyCfg_OptimizedCfgToInstructionSequence (flowgraph.c) @@ -266,11 +272,11 @@ impl CodeInfo { &self.metadata.cellvars, &self.metadata.freevars, ); - // CPython lowers LOAD_CLOSURE to LOAD_FAST before optimize_load_fast, so - // borrow selection can see classdictcell and other merged-cell loads. + // CPython inserts superinstructions before lowering LOAD_CLOSURE to + // LOAD_FAST, so merged-cell closure loads do not participate in + // LOAD_FAST_LOAD_FAST formation. Lower them only immediately before + // optimize_load_fast. convert_load_closure_pseudo_ops(&mut self.blocks, &cellfixedoffsets); - self.add_checks_for_loads_of_uninitialized_variables(); - self.combine_store_fast_load_fast(); // CPython's optimize_load_fast runs with block start depths already known. // Compute them here so the abstract stack simulation can use the real // CFG entry depth for each block. @@ -280,6 +286,7 @@ impl CodeInfo { self.deoptimize_borrow_for_handler_return_paths(); self.deoptimize_store_fast_store_fast_after_cleanup(); self.optimize_load_global_push_null(); + self.remove_redundant_const_pop_top_pairs(); self.reorder_entry_prefix_cell_setup(); let max_stackdepth = self.max_stackdepth()?; @@ -291,6 +298,7 @@ impl CodeInfo { mut blocks, current_block: _, + annotations_blocks: _, metadata, static_attributes: _, in_inlined_comp: _, @@ -891,6 +899,18 @@ impl CodeInfo { } l * r } + BinOp::TrueDivide => { + if r.is_zero() { + return None; + } + let l_f = l.to_f64()?; + let r_f = r.to_f64()?; + let result = l_f / r_f; + if !result.is_finite() { + return None; + } + return Some(ConstantData::Float { value: result }); + } BinOp::FloorDivide => { if r.is_zero() { return None; @@ -1036,6 +1056,23 @@ impl CodeInfo { }; let tuple_size = u32::from(instr.arg) as usize; + if block + .instructions + .get(i + 1) + .and_then(|next| next.instr.real()) + .is_some_and(|next| { + matches!( + next, + Instruction::UnpackSequence { .. } + if usize::try_from(u32::from(block.instructions[i + 1].arg)) + .ok() + == Some(tuple_size) + ) + }) + { + i += 1; + continue; + } if tuple_size == 0 { // BUILD_TUPLE 0 → LOAD_CONST () let (const_idx, _) = self.metadata.consts.insert_full(ConstantData::Tuple { @@ -1356,8 +1393,7 @@ impl CodeInfo { continue; } - // Use FrozenSet constant (stored as Tuple for now) - let const_data = ConstantData::Tuple { elements }; + let const_data = ConstantData::Frozenset { elements }; let (const_idx, _) = self.metadata.consts.insert_full(const_data); let folded_loc = block.instructions[i].location; @@ -1618,6 +1654,29 @@ impl CodeInfo { /// Peephole optimization: combine consecutive instructions into super-instructions fn peephole_optimize(&mut self) { + let const_truthiness = + |instr: Instruction, arg: OpArg, metadata: &CodeUnitMetadata| match instr { + Instruction::LoadConst { consti } => { + let constant = &metadata.consts[consti.get(arg).as_usize()]; + Some(match constant { + ConstantData::Tuple { elements } => !elements.is_empty(), + ConstantData::Integer { value } => !value.is_zero(), + ConstantData::Float { value } => *value != 0.0, + ConstantData::Complex { value } => value.re != 0.0 || value.im != 0.0, + ConstantData::Boolean { value } => *value, + ConstantData::Str { value } => !value.is_empty(), + ConstantData::Bytes { value } => !value.is_empty(), + ConstantData::Code { .. } => true, + ConstantData::Slice { .. } => true, + ConstantData::Frozenset { elements } => !elements.is_empty(), + ConstantData::None => false, + ConstantData::Ellipsis => true, + }) + } + Instruction::LoadSmallInt { i } => Some(i.get(arg) != 0), + _ => None, + }; + for block in &mut self.blocks { let mut i = 0; while i + 1 < block.instructions.len() { @@ -1631,65 +1690,72 @@ impl CodeInfo { continue; }; + if let Some(is_true) = const_truthiness(curr_instr, curr.arg, &self.metadata) { + let jump_if_true = match next_instr { + Instruction::PopJumpIfTrue { .. } => Some(true), + Instruction::PopJumpIfFalse { .. } => Some(false), + _ => None, + }; + if let Some(jump_if_true) = jump_if_true { + let target = match next_instr { + Instruction::PopJumpIfTrue { delta } + | Instruction::PopJumpIfFalse { delta } => delta.get(next.arg), + _ => unreachable!(), + }; + set_to_nop(&mut block.instructions[i]); + if is_true == jump_if_true { + block.instructions[i + 1].instr = PseudoInstruction::Jump { + delta: Arg::marker(), + } + .into(); + block.instructions[i + 1].arg = OpArg::new(u32::from(target)); + } else { + set_to_nop(&mut block.instructions[i + 1]); + } + i += 1; + continue; + } + } + if matches!( - next_instr.into(), - Opcode::PopJumpIfFalse | Opcode::PopJumpIfTrue - ) && matches!(curr_instr.into(), Opcode::CompareOp) + curr_instr, + Instruction::LoadConst { .. } | Instruction::LoadSmallInt { .. } + ) && matches!(next_instr, Instruction::PopTop) { - block.instructions[i].arg = OpArg::new( - u32::from(block.instructions[i].arg) | oparg::COMPARE_OP_BOOL_MASK, - ); + set_to_nop(&mut block.instructions[i]); + set_to_nop(&mut block.instructions[i + 1]); + i += 1; + continue; + } + + if matches!(curr_instr, Instruction::Copy { i } if i.get(curr.arg) == 1) + && matches!(next_instr, Instruction::PopTop) + { + set_to_nop(&mut block.instructions[i]); + set_to_nop(&mut block.instructions[i + 1]); i += 1; continue; } let combined = { match (curr_instr, next_instr) { - // LoadFast + LoadFast -> LoadFastLoadFast (if both indices < 16) - (Instruction::LoadFast { .. }, Instruction::LoadFast { .. }) => { - let line1 = curr.location.line.get() as i32; - let line2 = next.location.line.get() as i32; - if line1 > 0 && line2 > 0 && line1 != line2 { - None - } else { - let idx1 = u32::from(curr.arg); - let idx2 = u32::from(next.arg); - if idx1 < 16 && idx2 < 16 { - let packed = (idx1 << 4) | idx2; - Some((Opcode::LoadFastLoadFast.into(), OpArg::new(packed))) - } else { - None - } - } - } - // StoreFast + StoreFast -> StoreFastStoreFast (if both indices < 16) - // Dead store elimination: if both store to the same variable, - // the first store is dead. Replace it with POP_TOP (like - // apply_static_swaps in CPython's flowgraph.c). - (Instruction::StoreFast { .. }, Instruction::StoreFast { .. }) => { - let line1 = curr.location.line.get() as i32; - let line2 = next.location.line.get() as i32; - if line1 > 0 && line2 > 0 && line1 != line2 { - None - } else { - let idx1 = u32::from(curr.arg); - let idx2 = u32::from(next.arg); - if idx1 < 16 && idx2 < 16 { - let packed = (idx1 << 4) | idx2; - Some((Opcode::StoreFastStoreFast.into(), OpArg::new(packed))) - } else { - None - } - } - } // Note: StoreFast + LoadFast → StoreFastLoadFast is done in a - // separate pass AFTER optimize_load_fast_borrow, because CPython - // only combines STORE_FAST + LOAD_FAST (not LOAD_FAST_BORROW). - (Instruction::LoadConst { consti }, Instruction::ToBool) => { - let consti = consti.get(curr.arg); - let constant = &self.metadata.consts[consti.as_usize()]; - if let ConstantData::Boolean { .. } = constant { - Some((curr_instr, OpArg::from(consti.as_u32()))) + // later pass aligned with CPython insert_superinstructions(). + (Instruction::LoadConst { .. }, Instruction::ToBool) + | (Instruction::LoadSmallInt { .. }, Instruction::ToBool) => { + if let Some(value) = + const_truthiness(curr_instr, curr.arg, &self.metadata) + { + let (const_idx, _) = self + .metadata + .consts + .insert_full(ConstantData::Boolean { value }); + Some(( + Instruction::LoadConst { + consti: Arg::marker(), + }, + OpArg::new(const_idx as u32), + )) } else { None } @@ -1760,6 +1826,39 @@ impl CodeInfo { } } + fn remove_redundant_const_pop_top_pairs(&mut self) { + for block in &mut self.blocks { + let mut i = 0; + while i + 1 < block.instructions.len() { + let curr = &block.instructions[i]; + let next = &block.instructions[i + 1]; + let Some(curr_instr) = curr.instr.real() else { + i += 1; + continue; + }; + let Some(next_instr) = next.instr.real() else { + i += 1; + continue; + }; + + let redundant = matches!( + (curr_instr, next_instr), + (Instruction::LoadConst { .. }, Instruction::PopTop) + | (Instruction::LoadSmallInt { .. }, Instruction::PopTop) + ) || matches!(curr_instr, Instruction::Copy { i } if i.get(curr.arg) == 1) + && matches!(next_instr, Instruction::PopTop); + + if redundant { + set_to_nop(&mut block.instructions[i]); + set_to_nop(&mut block.instructions[i + 1]); + i += 2; + } else { + i += 1; + } + } + } + } + /// Convert LOAD_CONST for small integers to LOAD_SMALL_INT /// maybe_instr_make_load_smallint fn convert_to_load_small_int(&mut self) { @@ -1876,17 +1975,13 @@ impl CodeInfo { /// 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. - fn combine_store_fast_load_fast(&mut self) { + 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 next = &block.instructions[i + 1]; - let Some(Instruction::StoreFast { .. }) = curr.instr.real() else { - i += 1; - continue; - }; - // Skip if instructions are on different lines (matching make_super_instruction) + let line1 = curr.location.line; let line2 = next.location.line; if line1 != line2 { @@ -1894,53 +1989,53 @@ impl CodeInfo { continue; } - let store_idx = u32::from(curr.arg); - if store_idx >= 16 { - i += 1; - continue; - } - - match next.instr.real() { - Some(Instruction::LoadFast { .. }) => { - let load_idx = u32::from(next.arg); - if load_idx >= 16 { + match (curr.instr.real(), next.instr.real()) { + (Some(Instruction::LoadFast { .. }), Some(Instruction::LoadFast { .. })) => { + let idx1 = u32::from(curr.arg); + let idx2 = u32::from(next.arg); + if idx1 >= 16 || idx2 >= 16 { i += 1; continue; } - let packed = (store_idx << 4) | load_idx; - block.instructions[i].instr = Instruction::StoreFastLoadFast { + let packed = (idx1 << 4) | idx2; + block.instructions[i].instr = Instruction::LoadFastLoadFast { var_nums: Arg::marker(), } .into(); block.instructions[i].arg = OpArg::new(packed); - set_to_nop(&mut block.instructions[i + 1]); - i += 2; + block.instructions.remove(i + 1); } - Some(Instruction::LoadFastLoadFast { var_nums }) => { - let packed = var_nums.get(next.arg); - let (first_idx, second_idx) = packed.indexes(); - let first_idx = u32::from(first_idx); - if first_idx >= 16 { + (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 { i += 1; continue; } - - let packed = (store_idx << 4) | first_idx; + let packed = (store_idx << 4) | load_idx; block.instructions[i].instr = Instruction::StoreFastLoadFast { var_nums: Arg::marker(), } .into(); block.instructions[i].arg = OpArg::new(packed); - block.instructions[i + 1].instr = Instruction::LoadFast { - var_num: Arg::marker(), + block.instructions.remove(i + 1); + } + (Some(Instruction::StoreFast { .. }), Some(Instruction::StoreFast { .. })) => { + let idx1 = u32::from(curr.arg); + let idx2 = u32::from(next.arg); + if idx1 >= 16 || idx2 >= 16 { + i += 1; + continue; + } + let packed = (idx1 << 4) | idx2; + block.instructions[i].instr = Instruction::StoreFastStoreFast { + var_nums: Arg::marker(), } .into(); - block.instructions[i + 1].arg = OpArg::new(u32::from(second_idx)); - i += 2; - } - _ => { - i += 1; + block.instructions[i].arg = OpArg::new(packed); + block.instructions.remove(i + 1); } + _ => i += 1, } } } @@ -2343,7 +2438,25 @@ impl CodeInfo { for (block_idx, block) in self.blocks.iter_mut().enumerate() { let mut new_instructions = Vec::with_capacity(block.instructions.len()); + let mut in_restore_prefix = starts_after_cleanup[block_idx]; for (i, info) in block.instructions.iter().copied().enumerate() { + if !in_restore_prefix + && matches!( + info.instr.real(), + Some( + Instruction::StoreFast { .. } | Instruction::StoreFastStoreFast { .. } + ) + ) + && !new_instructions.is_empty() + && new_instructions.iter().all(|prev: &InstructionInfo| { + matches!( + prev.instr.real(), + Some(Instruction::Swap { .. }) | Some(Instruction::PopTop) + ) + }) + { + in_restore_prefix = true; + } let expand = matches!( info.instr.real(), Some(Instruction::StoreFastStoreFast { .. }) @@ -2354,7 +2467,8 @@ impl CodeInfo { Some(Instruction::PopIter) | Some(Instruction::Swap { .. }) ) }, - ) || (i == 0 && starts_after_cleanup[block_idx])); + ) || (i == 0 && starts_after_cleanup[block_idx]) + || in_restore_prefix); if expand { let Some(Instruction::StoreFastStoreFast { var_nums }) = info.instr.real() @@ -2382,6 +2496,8 @@ impl CodeInfo { continue; } + in_restore_prefix &= + matches!(info.instr.real(), Some(Instruction::StoreFast { .. })); new_instructions.push(info); } block.instructions = new_instructions; @@ -2673,6 +2789,86 @@ impl CodeInfo { } } +impl CodeInfo { + fn remap_block_idx(idx: BlockIdx, base: u32) -> BlockIdx { + if idx == BlockIdx::NULL { + idx + } else { + BlockIdx::new(u32::from(idx) + base) + } + } + + fn splice_annotations_blocks(&mut self) { + let mut placeholder = None; + for (block_idx, block) in self.blocks.iter().enumerate() { + if let Some(instr_idx) = block.instructions.iter().position(|info| { + matches!( + info.instr.pseudo(), + Some(PseudoInstruction::AnnotationsPlaceholder) + ) + }) { + placeholder = Some((block_idx, instr_idx)); + break; + } + } + + let Some((block_idx, instr_idx)) = placeholder else { + return; + }; + + let Some(mut annotations_blocks) = self.annotations_blocks.take() else { + self.blocks[block_idx].instructions.remove(instr_idx); + return; + }; + if annotations_blocks.is_empty() { + self.blocks[block_idx].instructions.remove(instr_idx); + return; + } + + let base = self.blocks.len() as u32; + for block in &mut annotations_blocks { + block.next = Self::remap_block_idx(block.next, base); + for info in &mut block.instructions { + info.target = Self::remap_block_idx(info.target, base); + if let Some(handler) = &mut info.except_handler { + handler.handler_block = Self::remap_block_idx(handler.handler_block, base); + } + } + } + + let ann_entry = BlockIdx::new(base); + let ann_tail = { + let mut cursor = ann_entry; + while annotations_blocks[(u32::from(cursor) - base) as usize].next != BlockIdx::NULL { + cursor = annotations_blocks[(u32::from(cursor) - base) as usize].next; + } + cursor + }; + + let old_next = self.blocks[block_idx].next; + let suffix = self.blocks[block_idx].instructions.split_off(instr_idx + 1); + self.blocks[block_idx].instructions.pop(); + + let suffix_block = if suffix.is_empty() { + old_next + } else { + let suffix_idx = BlockIdx::new(base + annotations_blocks.len() as u32); + let block = Block { + instructions: suffix, + next: old_next, + ..Default::default() + }; + annotations_blocks.push(block); + suffix_idx + }; + + self.blocks[block_idx].next = ann_entry; + let ann_tail_local = (u32::from(ann_tail) - base) as usize; + annotations_blocks[ann_tail_local].next = suffix_block; + self.blocks.extend(annotations_blocks); + } +} + impl InstrDisplayContext for CodeInfo { type Constant = ConstantData; @@ -4460,17 +4656,14 @@ fn convert_load_closure_pseudo_ops(blocks: &mut [Block], cellfixedoffsets: &[u32 let Some(pseudo) = info.instr.pseudo() else { continue; }; - match pseudo { - PseudoInstruction::LoadClosure { i } => { - let cell_relative = i.get(info.arg) as usize; - let new_idx = cellfixedoffsets[cell_relative]; - info.arg = OpArg::new(new_idx); - info.instr = Instruction::LoadFast { - var_num: Arg::marker(), - } - .into(); + if let PseudoInstruction::LoadClosure { i } = pseudo { + let cell_relative = i.get(info.arg) as usize; + let new_idx = cellfixedoffsets[cell_relative]; + info.arg = OpArg::new(new_idx); + info.instr = Instruction::LoadFast { + var_num: Arg::marker(), } - _ => {} + .into(); } } } diff --git a/crates/codegen/src/symboltable.rs b/crates/codegen/src/symboltable.rs index f8d7a2b26ea..df717ca6af2 100644 --- a/crates/codegen/src/symboltable.rs +++ b/crates/codegen/src/symboltable.rs @@ -1595,22 +1595,7 @@ impl SymbolTableBuilder { self.scan_expression(target, ExpressionContext::Store)?; } } - // Only scan annotation in annotation scope for simple name targets. - // Non-simple annotations (subscript, attribute, parenthesized) are - // never compiled into __annotate__, so scanning them would create - // sub_tables that cause mismatch in the annotation scope's sub_table index. - let is_simple_name = *simple && matches!(&**target, Expr::Name(_)); - if is_simple_name { - self.scan_ann_assign_annotation(annotation)?; - } else { - // Still validate annotation for forbidden expressions - // (yield, await, named) even for non-simple targets. - let was_in_annotation = self.in_annotation; - self.in_annotation = true; - let result = self.scan_expression(annotation, ExpressionContext::Load); - self.in_annotation = was_in_annotation; - result?; - } + self.scan_ann_assign_annotation(annotation)?; if let Some(value) = value { self.scan_expression(value, ExpressionContext::Load)?; } diff --git a/crates/vm/src/builtins/code.rs b/crates/vm/src/builtins/code.rs index 0ddbd9b8513..aafbe196a07 100644 --- a/crates/vm/src/builtins/code.rs +++ b/crates/vm/src/builtins/code.rs @@ -1,6 +1,6 @@ //! Infamous code object. The python class `code` -use super::{PyBytesRef, PyStrRef, PyTupleRef, PyType}; +use super::{PyBytesRef, PyStrRef, PyTupleRef, PyType, set::PyFrozenSet}; use crate::common::lock::PyMutex; use crate::{ AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine, @@ -267,13 +267,6 @@ impl<'a> AsBag for &'a Context { } } -impl<'a> AsBag for &'a VirtualMachine { - type Bag = PyObjBag<'a>; - fn as_bag(self) -> PyObjBag<'a> { - PyObjBag(&self.ctx) - } -} - #[derive(Clone, Copy)] pub struct PyObjBag<'a>(pub &'a Context); @@ -348,6 +341,87 @@ impl ConstantBag for PyObjBag<'_> { } } +#[derive(Clone, Copy)] +pub struct PyVmBag<'a>(pub &'a VirtualMachine); + +impl ConstantBag for PyVmBag<'_> { + type Constant = Literal; + + fn make_constant(&self, constant: BorrowedConstant<'_, C>) -> Self::Constant { + let vm = self.0; + let ctx = &vm.ctx; + let obj = match constant { + BorrowedConstant::Integer { value } => ctx.new_bigint(value).into(), + BorrowedConstant::Float { value } => ctx.new_float(value).into(), + BorrowedConstant::Complex { value } => ctx.new_complex(value).into(), + BorrowedConstant::Str { value } if value.len() <= 20 => { + ctx.intern_str(value).to_object() + } + BorrowedConstant::Str { value } => ctx.new_str(value).into(), + BorrowedConstant::Bytes { value } => ctx.new_bytes(value.to_vec()).into(), + BorrowedConstant::Boolean { value } => ctx.new_bool(value).into(), + BorrowedConstant::Code { code } => { + PyCode::new_ref_with_bag(vm, code.map_clone_bag(self)).into() + } + BorrowedConstant::Tuple { elements } => { + let elements = elements + .iter() + .map(|constant| self.make_constant(constant.borrow_constant()).0) + .collect(); + ctx.new_tuple(elements).into() + } + BorrowedConstant::Slice { elements } => { + let [start, stop, step] = elements; + let start_obj = self.make_constant(start.borrow_constant()).0; + let stop_obj = self.make_constant(stop.borrow_constant()).0; + let step_obj = self.make_constant(step.borrow_constant()).0; + use crate::builtins::PySlice; + PySlice { + start: Some(start_obj), + stop: stop_obj, + step: Some(step_obj), + } + .into_ref(ctx) + .into() + } + BorrowedConstant::Frozenset { elements } => { + let elements = elements + .iter() + .map(|constant| self.make_constant(constant.borrow_constant()).0); + PyFrozenSet::from_iter(vm, elements) + .unwrap() + .into_ref(ctx) + .into() + } + BorrowedConstant::None => ctx.none(), + BorrowedConstant::Ellipsis => ctx.ellipsis.clone().into(), + }; + + Literal(obj) + } + + fn make_name(&self, name: &str) -> &'static PyStrInterned { + self.0.ctx.intern_str(name) + } + + fn make_int(&self, value: BigInt) -> Self::Constant { + Literal(self.0.ctx.new_int(value).into()) + } + + fn make_tuple(&self, elements: impl Iterator) -> Self::Constant { + Literal( + self.0 + .ctx + .new_tuple(elements.map(|lit| lit.0).collect()) + .into(), + ) + } + + fn make_code(&self, code: CodeObject) -> Self::Constant { + Literal(PyCode::new_ref_with_bag(self.0, code).into()) + } +} + pub type CodeObject = bytecode::CodeObject; pub trait IntoCodeObject { @@ -427,6 +501,22 @@ impl PyCode { Ordering::Relaxed, ); } + + pub fn new_ref_with_bag(vm: &VirtualMachine, code: CodeObject) -> PyRef { + PyRef::new_ref(PyCode::new(code), vm.ctx.types.code_type.to_owned(), None) + } + + pub fn new_ref_from_bytecode(vm: &VirtualMachine, code: bytecode::CodeObject) -> PyRef { + Self::new_ref_with_bag(vm, code.map_bag(PyVmBag(vm))) + } + + pub fn new_ref_from_frozen>( + vm: &VirtualMachine, + code: frozen::FrozenCodeObject, + ) -> PyRef { + Self::new_ref_with_bag(vm, code.decode(PyVmBag(vm))) + } + pub fn from_pyc_path(path: &std::path::Path, vm: &VirtualMachine) -> PyResult> { let name = match path.file_stem() { Some(stem) => stem.display().to_string(), @@ -1379,7 +1469,7 @@ impl ToPyObject for CodeObject { impl ToPyObject for bytecode::CodeObject { fn to_pyobject(self, vm: &VirtualMachine) -> PyObjectRef { - vm.ctx.new_code(self).into() + PyCode::new_ref_from_bytecode(vm, self).into() } } diff --git a/crates/vm/src/import.rs b/crates/vm/src/import.rs index 9d015c8f3b6..f8f41c12081 100644 --- a/crates/vm/src/import.rs +++ b/crates/vm/src/import.rs @@ -72,7 +72,7 @@ pub fn make_frozen(vm: &VirtualMachine, name: &str) -> PyResult> { vm.ctx.new_utf8_str(name), ) })?; - Ok(vm.ctx.new_code(frozen.code)) + Ok(PyCode::new_ref_from_frozen(vm, frozen.code)) } pub fn import_frozen(vm: &VirtualMachine, module_name: &str) -> PyResult { @@ -82,7 +82,12 @@ pub fn import_frozen(vm: &VirtualMachine, module_name: &str) -> PyResult { vm.ctx.new_utf8_str(module_name), ) })?; - let module = import_code_obj(vm, module_name, vm.ctx.new_code(frozen.code), false)?; + let module = import_code_obj( + vm, + module_name, + PyCode::new_ref_from_frozen(vm, frozen.code), + false, + )?; debug_assert!(module.get_attr(identifier!(vm, __name__), vm).is_ok()); let origname = resolve_frozen_alias(module_name); module.set_attr("__origname__", vm.ctx.new_utf8_str(origname), vm)?; diff --git a/crates/vm/src/stdlib/_ast.rs b/crates/vm/src/stdlib/_ast.rs index 73819e257c1..bde6916a663 100644 --- a/crates/vm/src/stdlib/_ast.rs +++ b/crates/vm/src/stdlib/_ast.rs @@ -776,7 +776,7 @@ pub(crate) fn compile( let source_file = SourceFileBuilder::new(filename, text).finish(); let code = codegen::compile::compile_top(ast, source_file, mode, opts) .map_err(|err| vm.new_syntax_error(&err.into(), None))?; // FIXME source - Ok(vm.ctx.new_code(code).into()) + Ok(crate::builtins::PyCode::new_ref_from_bytecode(vm, code).into()) } #[cfg(feature = "codegen")] diff --git a/crates/vm/src/vm/compile.rs b/crates/vm/src/vm/compile.rs index 7294dc8f897..221df849f62 100644 --- a/crates/vm/src/vm/compile.rs +++ b/crates/vm/src/vm/compile.rs @@ -25,8 +25,8 @@ impl VirtualMachine { source_path: String, opts: CompileOpts, ) -> Result, CompileError> { - let code = - compiler::compile(source, mode, &source_path, opts).map(|code| self.ctx.new_code(code)); + let code = compiler::compile(source, mode, &source_path, opts) + .map(|code| PyCode::new_ref_from_bytecode(self, code)); #[cfg(feature = "parser")] if code.is_ok() { self.emit_string_escape_warnings(source, &source_path); diff --git a/scripts/dis_dump.py b/scripts/dis_dump.py index af31fd08230..72fec461b1f 100755 --- a/scripts/dis_dump.py +++ b/scripts/dis_dump.py @@ -104,7 +104,9 @@ def _unescape(m): if _IS_RUSTPYTHON and hasattr(dis, "_common_constants"): common_constants = list(dis._common_constants) while len(common_constants) < 7: - common_constants.append((builtins.list, builtins.set)[len(common_constants) - 5]) + common_constants.append( + (builtins.list, builtins.set)[len(common_constants) - 5] + ) dis._common_constants = common_constants # RustPython's ComparisonOperator enum values → operator strings From f60032ddb217d9bb5d34bde08d1af88dcabc6dd0 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Tue, 14 Apr 2026 00:51:21 +0900 Subject: [PATCH 3/8] Fix CI regressions in marshal and fast-local ops --- crates/codegen/src/ir.rs | 78 +++++++- ...thon_codegen__compile__tests__if_ands.snap | 21 +-- ...hon_codegen__compile__tests__if_mixed.snap | 25 +-- ...ython_codegen__compile__tests__if_ors.snap | 21 +-- ...pile__tests__nested_double_async_with.snap | 166 +++++++++--------- crates/vm/src/frame.rs | 15 +- crates/vm/src/stdlib/_imp.rs | 4 +- crates/vm/src/stdlib/marshal.rs | 8 +- 8 files changed, 182 insertions(+), 156 deletions(-) diff --git a/crates/codegen/src/ir.rs b/crates/codegen/src/ir.rs index ed5652070db..ad27c068720 100644 --- a/crates/codegen/src/ir.rs +++ b/crates/codegen/src/ir.rs @@ -277,12 +277,17 @@ impl CodeInfo { // LOAD_FAST_LOAD_FAST formation. Lower them only immediately before // optimize_load_fast. convert_load_closure_pseudo_ops(&mut self.blocks, &cellfixedoffsets); + // Late CFG cleanup can create or reshuffle handler entry blocks. + // Refresh exceptional block flags before optimize_load_fast_borrow so + // borrow loads are not introduced into exception-handler paths. + mark_except_handlers(&mut self.blocks); // CPython's optimize_load_fast runs with block start depths already known. // Compute them here so the abstract stack simulation can use the real // CFG entry depth for each block. let _ = self.max_stackdepth()?; // optimize_load_fast: after normalize_jumps self.optimize_load_fast_borrow(); + self.deoptimize_borrow_after_push_exc_info(); self.deoptimize_borrow_for_handler_return_paths(); self.deoptimize_store_fast_store_fast_after_cleanup(); self.optimize_load_global_push_null(); @@ -333,6 +338,7 @@ impl CodeInfo { // Convert pseudo ops (LoadClosure uses cellfixedoffsets) and fixup DEREF opargs convert_pseudo_ops(&mut blocks, &cellfixedoffsets); fixup_deref_opargs(&mut blocks, &cellfixedoffsets); + deoptimize_borrow_after_push_exc_info_in_blocks(&mut blocks); // Remove redundant NOPs, keeping line-marker NOPs only when // they are needed to preserve tracing. let mut block_order = Vec::new(); @@ -2347,9 +2353,6 @@ impl CodeInfo { } let block = &mut self.blocks[block_idx]; - if block.except_handler || block.preserve_lasti { - continue; - } for (i, info) in block.instructions.iter_mut().enumerate() { if instr_flags[i] != 0 { continue; @@ -2402,6 +2405,37 @@ impl CodeInfo { } } + 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_store_fast_store_fast_after_cleanup(&mut self) { fn last_real_instr(block: &Block) -> Option { block @@ -2425,9 +2459,9 @@ impl CodeInfo { let starts_after_cleanup: Vec = predecessors .iter() - .map(|preds| { - !preds.is_empty() - && preds.iter().copied().all(|pred_idx| { + .map(|predecessor_blocks| { + !predecessor_blocks.is_empty() + && predecessor_blocks.iter().copied().all(|pred_idx| { matches!( last_real_instr(&self.blocks[pred_idx]), Some(Instruction::PopIter) | Some(Instruction::Swap { .. }) @@ -3799,6 +3833,38 @@ fn merge_unsafe_mask(slot: &mut Option>, incoming: &[bool]) -> bool { } } +fn deoptimize_borrow_after_push_exc_info_in_blocks(blocks: &mut [Block]) { + let mut in_exception_state = false; + let mut current = BlockIdx(0); + while current != BlockIdx::NULL { + let block = &mut blocks[current.idx()]; + 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(); + } + _ => {} + } + } + current = block.next; + } +} + /// Follow chain of empty blocks to find first non-empty block. fn next_nonempty_block(blocks: &[Block], mut idx: BlockIdx) -> BlockIdx { while idx != BlockIdx::NULL diff --git a/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_ands.snap b/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_ands.snap index 0cdb7f0a3df..d4cb6b8d8d3 100644 --- a/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_ands.snap +++ b/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_ands.snap @@ -1,23 +1,10 @@ --- source: crates/codegen/src/compile.rs -assertion_line: 10890 +assertion_line: 11049 expression: "compile_exec(\"\\\nif True and False and False:\n pass\n\")" --- 1 0 RESUME (0) - 1 LOAD_CONST (True) - 2 POP_JUMP_IF_FALSE (11) - >> 3 CACHE - 4 NOT_TAKEN - 5 LOAD_CONST (False) - 6 POP_JUMP_IF_FALSE (7) - >> 7 CACHE - 8 NOT_TAKEN - 9 LOAD_CONST (False) - 10 POP_JUMP_IF_FALSE (3) - >> 11 CACHE - 12 NOT_TAKEN + 1 NOP - 2 13 LOAD_CONST (None) - 14 RETURN_VALUE - 15 LOAD_CONST (None) - 16 RETURN_VALUE + 2 2 LOAD_CONST (None) + 3 RETURN_VALUE diff --git a/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_mixed.snap b/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_mixed.snap index 8e91189912d..1c28cd123dc 100644 --- a/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_mixed.snap +++ b/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_mixed.snap @@ -1,27 +1,10 @@ --- source: crates/codegen/src/compile.rs -assertion_line: 10900 +assertion_line: 11059 expression: "compile_exec(\"\\\nif (True and False) or (False and True):\n pass\n\")" --- 1 0 RESUME (0) - 1 LOAD_CONST (True) - 2 POP_JUMP_IF_FALSE (5) - >> 3 CACHE - 4 NOT_TAKEN - >> 5 LOAD_CONST (False) - 6 POP_JUMP_IF_TRUE (9) - >> 7 CACHE - 8 NOT_TAKEN - >> 9 LOAD_CONST (False) - 10 POP_JUMP_IF_FALSE (7) - 11 CACHE - 12 NOT_TAKEN - 13 LOAD_CONST (True) - 14 POP_JUMP_IF_FALSE (3) - 15 CACHE - 16 NOT_TAKEN + 1 NOP - 2 17 LOAD_CONST (None) - 18 RETURN_VALUE - 19 LOAD_CONST (None) - 20 RETURN_VALUE + 2 2 LOAD_CONST (None) + 3 RETURN_VALUE diff --git a/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_ors.snap b/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_ors.snap index 9445635458d..f38d3c2c593 100644 --- a/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_ors.snap +++ b/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_ors.snap @@ -1,23 +1,10 @@ --- source: crates/codegen/src/compile.rs -assertion_line: 10880 +assertion_line: 11039 expression: "compile_exec(\"\\\nif True or False or False:\n pass\n\")" --- 1 0 RESUME (0) - 1 LOAD_CONST (True) - 2 POP_JUMP_IF_TRUE (9) - >> 3 CACHE - 4 NOT_TAKEN - >> 5 LOAD_CONST (False) - 6 POP_JUMP_IF_TRUE (5) - 7 CACHE - 8 NOT_TAKEN - >> 9 LOAD_CONST (False) - 10 POP_JUMP_IF_FALSE (3) - 11 CACHE - 12 NOT_TAKEN + 1 NOP - 2 13 LOAD_CONST (None) - 14 RETURN_VALUE - 15 LOAD_CONST (None) - 16 RETURN_VALUE + 2 2 LOAD_CONST (None) + 3 RETURN_VALUE 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 64ae0cfd5ff..ff6890fd1ed 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: 10936 +assertion_line: 11208 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) @@ -13,7 +13,7 @@ expression: "compile_exec(\"\\\nasync def test():\n for stop_exc in (StopIter >> 5 CACHE 6 CACHE 7 CACHE - 8 LOAD_CONST ("spam") + >> 8 LOAD_CONST ("spam") 9 CALL (1) >> 10 CACHE 11 CACHE @@ -34,7 +34,7 @@ expression: "compile_exec(\"\\\nasync def test():\n for stop_exc in (StopIter 26 CACHE 27 STORE_FAST (0, stop_exc) - 3 >> 28 LOAD_GLOBAL (4, self) + 3 28 LOAD_GLOBAL (4, self) 29 CACHE 30 CACHE 31 CACHE @@ -51,10 +51,10 @@ expression: "compile_exec(\"\\\nasync def test():\n for stop_exc in (StopIter 42 CACHE 43 LOAD_GLOBAL (9, NULL + type) 44 CACHE - 45 CACHE + >> 45 CACHE 46 CACHE 47 CACHE - >> 48 LOAD_FAST (0, stop_exc) + 48 LOAD_FAST_BORROW (0, stop_exc) 49 CALL (1) 50 CACHE 51 CACHE @@ -105,7 +105,7 @@ expression: "compile_exec(\"\\\nasync def test():\n for stop_exc in (StopIter 94 END_SEND 95 POP_TOP - 6 96 LOAD_FAST (0, stop_exc) + 6 96 LOAD_FAST_BORROW (0, stop_exc) 97 RAISE_VARARGS (Raise) 2 98 END_FOR @@ -139,116 +139,118 @@ expression: "compile_exec(\"\\\nasync def test():\n for stop_exc in (StopIter 125 POP_TOP 126 POP_TOP 127 POP_TOP - 128 JUMP_FORWARD (48) + 128 JUMP_FORWARD (3) 129 COPY (3) 130 POP_EXCEPT 131 RERAISE (1) - 132 PUSH_EXC_INFO + 132 NOP - 7 133 LOAD_GLOBAL (12, Exception) + 10 133 LOAD_GLOBAL (4, self) 134 CACHE 135 CACHE 136 CACHE 137 CACHE - 138 CHECK_EXC_MATCH - 139 POP_JUMP_IF_FALSE (28) + 138 LOAD_ATTR (13, fail, method=true) + 139 CACHE 140 CACHE - 141 NOT_TAKEN - 142 STORE_FAST (1, ex) - - 8 143 LOAD_GLOBAL (4, self) + 141 CACHE + 142 CACHE + 143 CACHE 144 CACHE 145 CACHE 146 CACHE 147 CACHE - 148 LOAD_ATTR (15, assertIs, method=true) - 149 CACHE - 150 CACHE - 151 CACHE - 152 CACHE + 148 LOAD_FAST_BORROW (0, stop_exc) + 149 FORMAT_SIMPLE + 150 LOAD_CONST (" was suppressed") + 151 BUILD_STRING (2) + 152 CALL (1) 153 CACHE 154 CACHE 155 CACHE - 156 CACHE - 157 CACHE - 158 LOAD_FAST_LOAD_FAST (ex, stop_exc) - 159 CALL (2) + 156 POP_TOP + 157 JUMP_FORWARD (45) + 158 PUSH_EXC_INFO + + 7 159 LOAD_GLOBAL (14, Exception) 160 CACHE 161 CACHE 162 CACHE - 163 POP_TOP - 164 POP_EXCEPT - 165 LOAD_CONST (None) - 166 STORE_FAST (1, ex) - 167 DELETE_FAST (1, ex) - 168 JUMP_FORWARD (32) - 169 RERAISE (0) - 170 LOAD_CONST (None) - 171 STORE_FAST (1, ex) - 172 DELETE_FAST (1, ex) - 173 RERAISE (1) - 174 COPY (3) - 175 POP_EXCEPT - 176 RERAISE (1) + 163 CACHE + 164 CHECK_EXC_MATCH + 165 POP_JUMP_IF_FALSE (32) + 166 CACHE + 167 NOT_TAKEN + 168 STORE_FAST (1, ex) - 10 177 LOAD_GLOBAL (4, self) + 8 169 LOAD_GLOBAL (4, self) + 170 CACHE + 171 CACHE + 172 CACHE + 173 CACHE + 174 LOAD_ATTR (17, assertIs, method=true) + 175 CACHE + 176 CACHE + 177 CACHE 178 CACHE 179 CACHE 180 CACHE 181 CACHE - 182 LOAD_ATTR (17, fail, method=true) + 182 CACHE 183 CACHE - 184 CACHE - 185 CACHE - >> 186 CACHE + 184 LOAD_FAST_LOAD_FAST (ex, stop_exc) + 185 CALL (2) + 186 CACHE 187 CACHE - 188 CACHE - 189 CACHE - 190 CACHE - 191 CACHE - 192 LOAD_FAST_BORROW (0, stop_exc) - 193 FORMAT_SIMPLE - 194 LOAD_CONST (" was suppressed") - 195 BUILD_STRING (2) - 196 CALL (1) - 197 CACHE - 198 CACHE - 199 CACHE - 200 POP_TOP + >> 188 CACHE + 189 POP_TOP + 190 POP_EXCEPT + 191 LOAD_CONST (None) + 192 STORE_FAST (1, ex) + 193 DELETE_FAST (1, ex) + 194 JUMP_FORWARD (8) + 195 LOAD_CONST (None) + 196 STORE_FAST (1, ex) + 197 DELETE_FAST (1, ex) + 198 RERAISE (1) + 199 RERAISE (0) + 200 COPY (3) + 201 POP_EXCEPT + 202 RERAISE (1) - 3 201 LOAD_CONST (None) - >> 202 LOAD_CONST (None) - 203 LOAD_CONST (None) - 204 CALL (3) - 205 CACHE - 206 CACHE + 3 203 LOAD_CONST (None) + >> 204 LOAD_CONST (None) + 205 LOAD_CONST (None) + 206 CALL (3) 207 CACHE - 208 POP_TOP - 209 JUMP_BACKWARD (186) - 210 CACHE - 211 PUSH_EXC_INFO - 212 WITH_EXCEPT_START - 213 TO_BOOL - 214 CACHE - 215 CACHE + 208 CACHE + 209 CACHE + 210 POP_TOP + 211 JUMP_BACKWARD (188) + 212 CACHE + 213 PUSH_EXC_INFO + 214 WITH_EXCEPT_START + 215 TO_BOOL 216 CACHE - 217 POP_JUMP_IF_TRUE (2) + 217 CACHE 218 CACHE - 219 NOT_TAKEN - 220 RERAISE (2) - 221 POP_TOP - 222 POP_EXCEPT + 219 POP_JUMP_IF_TRUE (2) + 220 CACHE + 221 NOT_TAKEN + 222 RERAISE (2) 223 POP_TOP - 224 POP_TOP + 224 POP_EXCEPT 225 POP_TOP - 226 JUMP_BACKWARD_NO_INTERRUPT(202) - 227 COPY (3) - 228 POP_EXCEPT - 229 RERAISE (1) - - 2 230 CALL_INTRINSIC_1 (StopIterationError) + 226 POP_TOP + 227 POP_TOP + 228 JUMP_BACKWARD_NO_INTERRUPT(204) + 229 COPY (3) + 230 POP_EXCEPT 231 RERAISE (1) + 2 232 CALL_INTRINSIC_1 (StopIterationError) + 233 RERAISE (1) + 2 MAKE_FUNCTION 3 STORE_NAME (0, test) 4 LOAD_CONST (None) diff --git a/crates/vm/src/frame.rs b/crates/vm/src/frame.rs index 92e0a11d0f5..49d0a18292c 100644 --- a/crates/vm/src/frame.rs +++ b/crates/vm/src/frame.rs @@ -3451,15 +3451,16 @@ impl ExecutingFrame<'_> { Ok(None) } Instruction::StoreFastLoadFast { var_nums } => { - let value = self.pop_value(); - let locals = self.localsplus.fastlocals_mut(); + // pop_value_opt: allows NULL from LoadFastAndClear restore paths. + let value = self.pop_value_opt(); let oparg = var_nums.get(arg); let (store_idx, load_idx) = oparg.indexes(); - locals[store_idx] = Some(value); - let load_value = locals[load_idx] - .clone() - .expect("StoreFastLoadFast: load slot should have value after store"); - self.push_value(load_value); + let load_value = { + let locals = self.localsplus.fastlocals_mut(); + locals[store_idx] = value; + locals[load_idx].clone() + }; + self.push_value_opt(load_value); Ok(None) } Instruction::StoreFastStoreFast { var_nums } => { diff --git a/crates/vm/src/stdlib/_imp.rs b/crates/vm/src/stdlib/_imp.rs index c0acb304a64..71a8c091e5e 100644 --- a/crates/vm/src/stdlib/_imp.rs +++ b/crates/vm/src/stdlib/_imp.rs @@ -261,11 +261,11 @@ mod _imp { name.clone().into_wtf8(), ) }; - let bag = crate::builtins::code::PyObjBag(&vm.ctx); + let bag = crate::builtins::code::PyVmBag(vm); let code = rustpython_compiler_core::marshal::deserialize_code(&mut &contiguous[..], bag) .map_err(|_| invalid_err())?; - return Ok(vm.ctx.new_code(code)); + return Ok(PyCode::new_ref_with_bag(vm, code)); } import::make_frozen(vm, name.as_str()) } diff --git a/crates/vm/src/stdlib/marshal.rs b/crates/vm/src/stdlib/marshal.rs index dace6bbf3e3..b19cc5eb52e 100644 --- a/crates/vm/src/stdlib/marshal.rs +++ b/crates/vm/src/stdlib/marshal.rs @@ -3,7 +3,7 @@ pub(crate) use decl::module_def; #[pymodule(name = "marshal")] mod decl { - use crate::builtins::code::{CodeObject, Literal, PyObjBag}; + use crate::builtins::code::{CodeObject, Literal, PyVmBag}; use crate::class::StaticType; use crate::common::wtf8::Wtf8; use crate::{ @@ -382,7 +382,7 @@ mod decl { impl<'a> marshal::MarshalBag for PyMarshalBag<'a> { type Value = PyObjectRef; - type ConstantBag = PyObjBag<'a>; + type ConstantBag = PyVmBag<'a>; fn make_bool(&self, value: bool) -> Self::Value { self.0.ctx.new_bool(value).into() @@ -412,7 +412,7 @@ mod decl { self.0.ctx.new_tuple(elements.collect()).into() } fn make_code(&self, code: CodeObject) -> Self::Value { - self.0.ctx.new_code(code).into() + crate::builtins::PyCode::new_ref_with_bag(self.0, code).into() } fn make_stop_iter(&self) -> Result { Ok(self.0.ctx.exceptions.stop_iteration.to_owned().into()) @@ -472,7 +472,7 @@ mod decl { .into()) } fn constant_bag(self) -> Self::ConstantBag { - PyObjBag(&self.0.ctx) + PyVmBag(self.0) } } From c3946d0d21db4d910542dacdb9d709ff47fc0f03 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Tue, 14 Apr 2026 00:52:08 +0900 Subject: [PATCH 4/8] impl more --- crates/codegen/src/compile.rs | 242 ++++++++---------------------- crates/codegen/src/ir.rs | 117 ++++++++++----- crates/codegen/src/symboltable.rs | 5 +- 3 files changed, 143 insertions(+), 221 deletions(-) diff --git a/crates/codegen/src/compile.rs b/crates/codegen/src/compile.rs index ba2dfab3244..5a9e69e617a 100644 --- a/crates/codegen/src/compile.rs +++ b/crates/codegen/src/compile.rs @@ -2598,42 +2598,25 @@ impl Compiler { ast::Stmt::Assert(ast::StmtAssert { test, msg, .. }) => { // if some flag, ignore all assert statements! if self.opts.optimize == 0 { - let after_block = match self.try_fold_constant_expr(test)? { - Some(constant) if Self::constant_truthiness(&constant) => { - self.consume_skipped_nested_scopes_in_expr(test)?; - if let Some(expr) = msg { - self.consume_skipped_nested_scopes_in_expr(expr)?; - } - emit!(self, Instruction::Nop); - None - } - Some(_) => Some(self.new_block()), - None => { - let after_block = self.new_block(); - self.compile_jump_if(test, true, after_block)?; - Some(after_block) - } - }; - - if let Some(after_block) = after_block { - emit!( - self, - Instruction::LoadCommonConstant { - idx: bytecode::CommonConstant::AssertionError - } - ); - if let Some(e) = msg { - self.compile_expression(e)?; - emit!(self, Instruction::Call { argc: 0 }); + let after_block = self.new_block(); + self.compile_jump_if(test, true, after_block)?; + emit!( + self, + Instruction::LoadCommonConstant { + idx: bytecode::CommonConstant::AssertionError } - emit!( - self, - Instruction::RaiseVarargs { - argc: bytecode::RaiseKind::Raise, - } - ); - self.switch_to_block(after_block); + ); + if let Some(e) = msg { + self.compile_expression(e)?; + emit!(self, Instruction::Call { argc: 0 }); } + emit!( + self, + Instruction::RaiseVarargs { + argc: bytecode::RaiseKind::Raise, + } + ); + self.switch_to_block(after_block); } else { // Optimized-out asserts still need to consume any nested // scope symbol tables they contain so later nested scopes @@ -3550,11 +3533,6 @@ impl Compiler { let handler_block = self.new_block(); let cleanup_block = self.new_block(); let end_block = self.new_block(); - let orelse_block = if orelse.is_empty() { - end_block - } else { - self.new_block() - }; emit!(self, Instruction::Nop); emit!( @@ -3569,11 +3547,10 @@ impl Compiler { self.pop_fblock(FBlockType::TryExcept); emit!(self, PseudoInstruction::PopBlock); self.set_no_location(); + self.compile_statements(orelse)?; emit!( self, - PseudoInstruction::JumpNoInterrupt { - delta: orelse_block - } + PseudoInstruction::JumpNoInterrupt { delta: end_block } ); self.set_no_location(); @@ -3615,7 +3592,6 @@ impl Compiler { self.store_name(alias.as_str())?; let cleanup_end = self.new_block(); - let handler_normal_exit = self.new_block(); emit!(self, PseudoInstruction::SetupCleanup { delta: cleanup_end }); self.push_fblock_full( FBlockType::HandlerCleanup, @@ -3629,25 +3605,6 @@ impl Compiler { self.pop_fblock(FBlockType::HandlerCleanup); emit!(self, PseudoInstruction::PopBlock); self.set_no_location(); - emit!( - self, - PseudoInstruction::JumpNoInterrupt { - delta: handler_normal_exit - } - ); - self.set_no_location(); - - self.switch_to_block(cleanup_end); - self.emit_load_const(ConstantData::None); - self.set_no_location(); - self.store_name(alias.as_str())?; - self.set_no_location(); - self.compile_name(alias.as_str(), NameUsage::Delete)?; - self.set_no_location(); - emit!(self, Instruction::Reraise { depth: 1 }); - self.set_no_location(); - - self.switch_to_block(handler_normal_exit); emit!(self, PseudoInstruction::PopBlock); self.set_no_location(); self.pop_fblock(FBlockType::ExceptionHandler); @@ -3666,6 +3623,16 @@ impl Compiler { PseudoInstruction::JumpNoInterrupt { delta: end_block } ); self.set_no_location(); + + self.switch_to_block(cleanup_end); + self.emit_load_const(ConstantData::None); + self.set_no_location(); + self.store_name(alias.as_str())?; + self.set_no_location(); + self.compile_name(alias.as_str(), NameUsage::Delete)?; + self.set_no_location(); + emit!(self, Instruction::Reraise { depth: 1 }); + self.set_no_location(); } else { emit!(self, Instruction::PopTop); self.push_fblock(FBlockType::HandlerCleanup, end_block, end_block)?; @@ -3701,17 +3668,6 @@ impl Compiler { emit!(self, Instruction::Reraise { depth: 1 }); self.set_no_location(); - if !orelse.is_empty() { - self.switch_to_block(orelse_block); - self.set_no_location(); - self.compile_statements(orelse)?; - emit!( - self, - PseudoInstruction::JumpNoInterrupt { delta: end_block } - ); - self.set_no_location(); - } - self.switch_to_block(end_block); Ok(()) } @@ -5563,46 +5519,6 @@ impl Compiler { body: &[ast::Stmt], elif_else_clauses: &[ast::ElifElseClause], ) -> CompileResult<()> { - let constant = Self::expr_constant(test); - - // If the test is constant false, walk the body (consuming sub_tables) - // but don't emit bytecode - if constant == Some(false) { - self.emit_nop(); - self.do_not_emit_bytecode += 1; - self.compile_statements(body)?; - self.do_not_emit_bytecode -= 1; - // Compile the elif/else chain (if any) - match elif_else_clauses { - [] => {} - [first, rest @ ..] => { - if let Some(elif_test) = &first.test { - self.compile_if(elif_test, &first.body, rest)?; - } else { - self.compile_statements(&first.body)?; - } - } - } - return Ok(()); - } - - // If the test is constant true, compile body directly, - // but walk elif/else without emitting (including elif tests to consume sub_tables) - if constant == Some(true) { - self.emit_nop(); - self.compile_statements(body)?; - self.do_not_emit_bytecode += 1; - for clause in elif_else_clauses { - if let Some(elif_test) = &clause.test { - self.compile_expression(elif_test)?; - } - self.compile_statements(&clause.body)?; - } - self.do_not_emit_bytecode -= 1; - return Ok(()); - } - - // Non-constant test: normal compilation match elif_else_clauses { // Only if [] => { @@ -5651,37 +5567,13 @@ impl Compiler { ) -> CompileResult<()> { self.enter_conditional_block(); - let constant = Self::expr_constant(test); - - // while False: body → walk body (consuming sub_tables) but don't emit, - // then compile orelse - if constant == Some(false) { - self.emit_nop(); - let while_block = self.new_block(); - let after_block = self.new_block(); - self.push_fblock(FBlockType::WhileLoop, while_block, after_block)?; - self.do_not_emit_bytecode += 1; - self.compile_statements(body)?; - self.do_not_emit_bytecode -= 1; - self.pop_fblock(FBlockType::WhileLoop); - self.compile_statements(orelse)?; - self.leave_conditional_block(); - return Ok(()); - } - let while_block = self.new_block(); let else_block = self.new_block(); let after_block = self.new_block(); self.switch_to_block(while_block); self.push_fblock(FBlockType::WhileLoop, while_block, after_block)?; - - // while True: → no condition test, just NOP - if constant == Some(true) { - self.emit_nop(); - } else { - self.compile_jump_if(test, false, else_block)?; - } + self.compile_jump_if(test, false, else_block)?; let was_in_loop = self.ctx.loop_data.replace((while_block, after_block)); self.compile_statements(body)?; @@ -9850,44 +9742,6 @@ impl Compiler { self.code_stack.last_mut().expect("no code on stack") } - /// Evaluate whether an expression is a compile-time constant boolean. - /// Returns Some(true) for truthy constants, Some(false) for falsy constants, - /// None for non-constant expressions. - /// = expr_constant in CPython compile.c - fn expr_constant(expr: &ast::Expr) -> Option { - match expr { - ast::Expr::BooleanLiteral(ast::ExprBooleanLiteral { value, .. }) => Some(*value), - ast::Expr::NoneLiteral(_) => Some(false), - ast::Expr::EllipsisLiteral(_) => Some(true), - ast::Expr::NumberLiteral(ast::ExprNumberLiteral { value, .. }) => match value { - ast::Number::Int(i) => { - let n: i64 = i.as_i64().unwrap_or(1); - Some(n != 0) - } - ast::Number::Float(f) => Some(*f != 0.0), - ast::Number::Complex { real, imag, .. } => Some(*real != 0.0 || *imag != 0.0), - }, - ast::Expr::StringLiteral(ast::ExprStringLiteral { value, .. }) => { - Some(!value.to_str().is_empty()) - } - ast::Expr::BytesLiteral(ast::ExprBytesLiteral { value, .. }) => { - Some(value.bytes().next().is_some()) - } - ast::Expr::Tuple(ast::ExprTuple { elts, .. }) => { - if elts.is_empty() { - Some(false) - } else { - None // non-empty tuples may have side effects in elements - } - } - _ => None, - } - } - - fn emit_nop(&mut self) { - emit!(self, Instruction::Nop); - } - /// Enter a conditional block (if/for/while/match/try/with) /// PEP 649: Track conditional annotation context fn enter_conditional_block(&mut self) { @@ -12449,11 +12303,11 @@ def f(): } #[test] - fn test_constant_false_assert_raises_directly() { + fn test_constant_false_assert_uses_normal_assert_branch_shape() { let code = compile_exec("assert 0, (lambda x: x + 1)\n"); assert!( - !code.instructions.iter().any(|unit| { + code.instructions.iter().any(|unit| { matches!( unit.op, Instruction::ToBool @@ -12461,7 +12315,7 @@ def f(): | Instruction::PopJumpIfFalse { .. } ) }), - "constant-false assert should raise directly without a test branch, got ops={:?}", + "constant-false assert should still use the normal assert branch shape, got ops={:?}", code.instructions .iter() .map(|unit| unit.op) @@ -12527,4 +12381,34 @@ def f(items): ", ); } + + #[test] + fn test_try_else_nested_scopes_keep_subtable_cursor_aligned() { + let code = compile_exec( + "\ +try: + import missing_mod +except ImportError: + def fallback(): + return 0 +else: + def impl(): + return reversed('abc') +", + ); + + assert!(find_code(&code, "fallback").is_some(), "missing fallback code"); + let impl_code = find_code(&code, "impl").expect("missing impl code"); + assert!( + impl_code.instructions.iter().any(|unit| { + matches!(unit.op, Instruction::LoadGlobal { .. } | Instruction::LoadName { .. }) + }), + "expected impl to compile global name access, got ops={:?}", + impl_code + .instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + } } diff --git a/crates/codegen/src/ir.rs b/crates/codegen/src/ir.rs index ad27c068720..0f2e36a14e5 100644 --- a/crates/codegen/src/ir.rs +++ b/crates/codegen/src/ir.rs @@ -272,11 +272,6 @@ impl CodeInfo { &self.metadata.cellvars, &self.metadata.freevars, ); - // CPython inserts superinstructions before lowering LOAD_CLOSURE to - // LOAD_FAST, so merged-cell closure loads do not participate in - // LOAD_FAST_LOAD_FAST formation. Lower them only immediately before - // optimize_load_fast. - convert_load_closure_pseudo_ops(&mut self.blocks, &cellfixedoffsets); // Late CFG cleanup can create or reshuffle handler entry blocks. // Refresh exceptional block flags before optimize_load_fast_borrow so // borrow loads are not introduced into exception-handler paths. @@ -284,7 +279,10 @@ impl CodeInfo { // CPython's optimize_load_fast runs with block start depths already known. // Compute them here so the abstract stack simulation can use the real // CFG entry depth for each block. - let _ = self.max_stackdepth()?; + let max_stackdepth = self.max_stackdepth()?; + // Match CPython order: pseudo ops are lowered after stackdepth + // calculation but before optimize_load_fast. + convert_pseudo_ops(&mut self.blocks, &cellfixedoffsets); // optimize_load_fast: after normalize_jumps self.optimize_load_fast_borrow(); self.deoptimize_borrow_after_push_exc_info(); @@ -294,8 +292,6 @@ impl CodeInfo { self.remove_redundant_const_pop_top_pairs(); self.reorder_entry_prefix_cell_setup(); - let max_stackdepth = self.max_stackdepth()?; - let Self { flags, source_path, @@ -2102,10 +2098,24 @@ impl CodeInfo { worklist: &mut Vec, visited: &mut [bool], blocks: &[Block], + source: BlockIdx, target: BlockIdx, start_depth: usize, ) { - let _ = (blocks, start_depth); + if cfg!(debug_assertions) { + let expected = blocks[target.idx()].start_depth.map(|depth| depth as usize); + debug_assert!( + expected == Some(start_depth), + "optimize_load_fast_borrow start_depth mismatch: source={source:?} target={target:?} expected={expected:?} actual={:?} source_last={:?} target_instrs={:?}", + Some(start_depth), + blocks[source.idx()].instructions.last().and_then(|info| info.instr.real()), + blocks[target.idx()] + .instructions + .iter() + .map(|info| info.instr) + .collect::>(), + ); + } if !visited[target.idx()] { visited[target.idx()] = true; worklist.push(target); @@ -2250,12 +2260,14 @@ impl CodeInfo { push_ref(&mut refs, i as isize, NOT_LOCAL); } AnyInstruction::Real(Instruction::ForIter { .. }) => { - if info.target != BlockIdx::NULL { + let target = info.target; + if target != BlockIdx::NULL { push_block( &mut worklist, &mut visited, &self.blocks, - info.target, + block_idx, + target, refs.len() + 1, ); } @@ -2291,12 +2303,14 @@ impl CodeInfo { push_ref(&mut refs, tos.instr, tos.local); } AnyInstruction::Real(Instruction::Send { .. }) => { - if info.target != BlockIdx::NULL { + let target = info.target; + if target != BlockIdx::NULL { push_block( &mut worklist, &mut visited, &self.blocks, - info.target, + block_idx, + target, refs.len(), ); } @@ -2307,7 +2321,8 @@ impl CodeInfo { let effect = instr.stack_effect_info(arg_u32); let num_popped = effect.popped() as usize; let num_pushed = effect.pushed() as usize; - if info.target != BlockIdx::NULL { + let target = info.target; + if target != BlockIdx::NULL { let target_depth = refs .len() .saturating_sub(num_popped) @@ -2316,7 +2331,8 @@ impl CodeInfo { &mut worklist, &mut visited, &self.blocks, - info.target, + block_idx, + target, target_depth, ); } @@ -2332,7 +2348,8 @@ impl CodeInfo { } } - if block.next != BlockIdx::NULL + let next = block.next; + if next != BlockIdx::NULL && block.instructions.last().is_some_and(|term| { !term.instr.is_unconditional_jump() && !term.instr.is_scope_exit() }) @@ -2341,7 +2358,8 @@ impl CodeInfo { &mut worklist, &mut visited, &self.blocks, - block.next, + block_idx, + next, refs.len(), ); } @@ -2781,7 +2799,10 @@ impl CodeInfo { if target_depth > maxdepth { maxdepth = target_depth; } - stackdepth_push(&mut stack, &mut start_depths, ins.target, target_depth); + let target = next_nonempty_block(&self.blocks, ins.target); + if target != BlockIdx::NULL { + stackdepth_push(&mut stack, &mut start_depths, target, target_depth); + } } depth = new_depth; if instr.is_scope_exit() || instr.is_unconditional_jump() { @@ -2789,8 +2810,9 @@ impl CodeInfo { } } // Only push next block if it's not NULL - if block.next != BlockIdx::NULL { - stackdepth_push(&mut stack, &mut start_depths, block.next, depth); + let next = next_nonempty_block(&self.blocks, block.next); + if next != BlockIdx::NULL { + stackdepth_push(&mut stack, &mut start_depths, next, depth); } } if DEBUG { @@ -3670,6 +3692,29 @@ fn inline_small_or_no_lineno_blocks(blocks: &mut [Block]) { .iter() .all(|ins| !instruction_has_lineno(ins)) }; + let is_return_epilogue_block = |block: &Block| { + matches!( + block.instructions.as_slice(), + [InstructionInfo { + instr: AnyInstruction::Real(Instruction::LoadConst { .. }), + .. + }, InstructionInfo { + instr: AnyInstruction::Real(Instruction::ReturnValue), + .. + }] + | [InstructionInfo { + instr: AnyInstruction::Real(Instruction::LoadSmallInt { .. }), + .. + }, InstructionInfo { + instr: AnyInstruction::Real(Instruction::ReturnValue), + .. + }] + | [InstructionInfo { + instr: AnyInstruction::Real(Instruction::ReturnValue), + .. + }] + ) + }; loop { let mut changes = false; @@ -3696,6 +3741,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 + && blocks[current.idx()].cold + && !is_return_epilogue_block(&blocks[target.idx()]) + { + current = next; + continue; + } if small_exit_block || no_lineno_no_fallthrough { if let Some(last_instr) = blocks[current.idx()].instructions.last_mut() { @@ -4200,6 +4252,7 @@ fn reorder_jump_over_exception_cleanup_blocks(blocks: &mut [Block]) { let mut target_end = BlockIdx::NULL; let mut target_exit = BlockIdx::NULL; + let mut nonempty_target_blocks = 0usize; cursor = target; while cursor != BlockIdx::NULL { if block_is_exceptional(&blocks[cursor.idx()]) { @@ -4207,6 +4260,7 @@ fn reorder_jump_over_exception_cleanup_blocks(blocks: &mut [Block]) { } target_end = cursor; if !blocks[cursor.idx()].instructions.is_empty() { + nonempty_target_blocks += 1; target_exit = cursor; } cursor = blocks[cursor.idx()].next; @@ -4214,6 +4268,8 @@ fn reorder_jump_over_exception_cleanup_blocks(blocks: &mut [Block]) { if target_end == BlockIdx::NULL || target_exit == BlockIdx::NULL + || nonempty_target_blocks != 1 + || target_exit != target_end || !is_scope_exit_block(&blocks[target_exit.idx()]) { current = next; @@ -4714,27 +4770,6 @@ pub(crate) fn label_exception_targets(blocks: &mut [Block]) { } } -/// Lower only LOAD_CLOSURE pseudo ops to LOAD_FAST so optimize_load_fast can -/// make the same borrow decision CPython does for merged cell locals. -fn convert_load_closure_pseudo_ops(blocks: &mut [Block], cellfixedoffsets: &[u32]) { - for block in blocks.iter_mut() { - for info in &mut block.instructions { - let Some(pseudo) = info.instr.pseudo() else { - continue; - }; - if let PseudoInstruction::LoadClosure { i } = pseudo { - let cell_relative = i.get(info.arg) as usize; - let new_idx = cellfixedoffsets[cell_relative]; - info.arg = OpArg::new(new_idx); - info.instr = Instruction::LoadFast { - var_num: Arg::marker(), - } - .into(); - } - } - } -} - /// Convert remaining pseudo ops to real instructions or NOP. /// flowgraph.c convert_pseudo_ops pub(crate) fn convert_pseudo_ops(blocks: &mut [Block], cellfixedoffsets: &[u32]) { diff --git a/crates/codegen/src/symboltable.rs b/crates/codegen/src/symboltable.rs index df717ca6af2..7479f168595 100644 --- a/crates/codegen/src/symboltable.rs +++ b/crates/codegen/src/symboltable.rs @@ -1624,6 +1624,10 @@ impl SymbolTableBuilder { let saved_in_conditional_block = self.in_conditional_block; self.in_conditional_block = true; self.scan_statements(body)?; + self.scan_statements(orelse)?; + // Keep nested scope collection in the same order that codegen + // compiles try/except, since the compiler currently consumes + // sub_tables through a linear cursor. for handler in handlers { let ExceptHandler::ExceptHandler(ast::ExceptHandlerExceptHandler { type_, @@ -1639,7 +1643,6 @@ impl SymbolTableBuilder { } self.scan_statements(body)?; } - self.scan_statements(orelse)?; self.scan_statements(finalbody)?; self.in_conditional_block = saved_in_conditional_block; } From 8ee88c67fdbb76b65a98c2948d56d4ecca067ec0 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Tue, 14 Apr 2026 19:15:16 +0900 Subject: [PATCH 5/8] Align bytecode codegen with CPython structure --- crates/codegen/src/compile.rs | 32 +++++++- crates/codegen/src/ir.rs | 144 ++++++++++++++++++++++++++++++---- 2 files changed, 157 insertions(+), 19 deletions(-) diff --git a/crates/codegen/src/compile.rs b/crates/codegen/src/compile.rs index 5a9e69e617a..1508785e629 100644 --- a/crates/codegen/src/compile.rs +++ b/crates/codegen/src/compile.rs @@ -460,6 +460,18 @@ impl Compiler { } } + fn constant_expr_truthiness(&mut self, expr: &ast::Expr) -> CompileResult> { + Ok(self + .try_fold_constant_expr(expr)? + .map(|constant| Self::constant_truthiness(&constant))) + } + + fn disable_load_fast_borrow_for_block(&mut self, block: BlockIdx) { + if block != BlockIdx::NULL { + self.current_code_info().blocks[block.idx()].disable_load_fast_borrow = true; + } + } + fn new(opts: CompileOpts, source_file: SourceFile, code_name: String) -> Self { let module_code = ir::CodeInfo { flags: bytecode::CodeFlags::NEWLOCALS, @@ -5519,10 +5531,14 @@ impl Compiler { body: &[ast::Stmt], elif_else_clauses: &[ast::ElifElseClause], ) -> CompileResult<()> { + let test_truthiness = self.constant_expr_truthiness(test)?; match elif_else_clauses { // Only if [] => { let after_block = self.new_block(); + if matches!(test_truthiness, Some(false)) { + self.disable_load_fast_borrow_for_block(after_block); + } self.compile_jump_if(test, false, after_block)?; self.compile_statements(body)?; self.switch_to_block(after_block); @@ -5532,6 +5548,9 @@ impl Compiler { let after_block = self.new_block(); let mut next_block = self.new_block(); + if matches!(test_truthiness, Some(false)) { + self.disable_load_fast_borrow_for_block(next_block); + } self.compile_jump_if(test, false, next_block)?; self.compile_statements(body)?; emit!(self, PseudoInstruction::Jump { delta: after_block }); @@ -5573,6 +5592,9 @@ impl Compiler { self.switch_to_block(while_block); self.push_fblock(FBlockType::WhileLoop, while_block, after_block)?; + if matches!(self.constant_expr_truthiness(test)?, Some(false)) { + self.disable_load_fast_borrow_for_block(else_block); + } self.compile_jump_if(test, false, else_block)?; let was_in_loop = self.ctx.loop_data.replace((while_block, after_block)); @@ -12397,11 +12419,17 @@ else: ", ); - assert!(find_code(&code, "fallback").is_some(), "missing fallback code"); + assert!( + find_code(&code, "fallback").is_some(), + "missing fallback code" + ); let impl_code = find_code(&code, "impl").expect("missing impl code"); assert!( impl_code.instructions.iter().any(|unit| { - matches!(unit.op, Instruction::LoadGlobal { .. } | Instruction::LoadName { .. }) + matches!( + unit.op, + Instruction::LoadGlobal { .. } | Instruction::LoadName { .. } + ) }), "expected impl to compile global name access, got ops={:?}", impl_code diff --git a/crates/codegen/src/ir.rs b/crates/codegen/src/ir.rs index 0f2e36a14e5..e8e9fa92e7d 100644 --- a/crates/codegen/src/ir.rs +++ b/crates/codegen/src/ir.rs @@ -148,6 +148,8 @@ pub struct Block { pub start_depth: Option, /// Whether this block is only reachable via exception table (b_cold) pub cold: bool, + /// Whether LOAD_FAST borrow optimization should be suppressed for this block. + pub disable_load_fast_borrow: bool, } impl Default for Block { @@ -159,6 +161,7 @@ impl Default for Block { preserve_lasti: false, start_depth: None, cold: false, + disable_load_fast_borrow: false, } } } @@ -283,6 +286,8 @@ impl CodeInfo { // Match CPython order: pseudo ops are lowered after stackdepth // calculation but before optimize_load_fast. convert_pseudo_ops(&mut self.blocks, &cellfixedoffsets); + self.compute_load_fast_start_depths(); + self.propagate_disable_load_fast_borrow(); // optimize_load_fast: after normalize_jumps self.optimize_load_fast_borrow(); self.deoptimize_borrow_after_push_exc_info(); @@ -291,6 +296,7 @@ impl CodeInfo { self.optimize_load_global_push_null(); self.remove_redundant_const_pop_top_pairs(); self.reorder_entry_prefix_cell_setup(); + self.remove_unused_consts(); let Self { flags, @@ -1678,7 +1684,6 @@ impl CodeInfo { Instruction::LoadSmallInt { i } => Some(i.get(arg) != 0), _ => None, }; - for block in &mut self.blocks { let mut i = 0; while i + 1 < block.instructions.len() { @@ -2108,7 +2113,10 @@ impl CodeInfo { expected == Some(start_depth), "optimize_load_fast_borrow start_depth mismatch: source={source:?} target={target:?} expected={expected:?} actual={:?} source_last={:?} target_instrs={:?}", Some(start_depth), - blocks[source.idx()].instructions.last().and_then(|info| info.instr.real()), + blocks[source.idx()] + .instructions + .last() + .and_then(|info| info.instr.real()), blocks[target.idx()] .instructions .iter() @@ -2371,6 +2379,9 @@ impl CodeInfo { } let block = &mut self.blocks[block_idx]; + if block.disable_load_fast_borrow { + continue; + } for (i, info) in block.instructions.iter_mut().enumerate() { if instr_flags[i] != 0 { continue; @@ -2394,6 +2405,96 @@ impl CodeInfo { } } + fn compute_load_fast_start_depths(&mut self) { + fn stackdepth_push( + stack: &mut Vec, + start_depths: &mut [u32], + target: BlockIdx, + depth: u32, + ) { + let idx = target.idx(); + let block_depth = &mut start_depths[idx]; + debug_assert!( + *block_depth == u32::MAX || *block_depth == depth, + "Invalid CFG, inconsistent optimize_load_fast stackdepth for block {:?}: existing={}, new={}", + target, + *block_depth, + depth, + ); + if *block_depth == u32::MAX { + *block_depth = depth; + stack.push(target); + } + } + + let mut stack = Vec::with_capacity(self.blocks.len()); + let mut start_depths = vec![u32::MAX; self.blocks.len()]; + stackdepth_push(&mut stack, &mut start_depths, BlockIdx(0), 0); + + 'process_blocks: while let Some(block_idx) = stack.pop() { + let mut depth = start_depths[block_idx.idx()]; + let block = &self.blocks[block_idx]; + for ins in &block.instructions { + let instr = &ins.instr; + let effect = instr.stack_effect(ins.arg.into()); + let new_depth = depth.saturating_add_signed(effect); + if ins.target != BlockIdx::NULL { + let jump_effect = instr.stack_effect_jump(ins.arg.into()); + let target_depth = depth.saturating_add_signed(jump_effect); + stackdepth_push(&mut stack, &mut start_depths, ins.target, target_depth); + } + depth = new_depth; + if instr.is_scope_exit() || instr.is_unconditional_jump() { + continue 'process_blocks; + } + } + if block.next != BlockIdx::NULL { + stackdepth_push(&mut stack, &mut start_depths, block.next, depth); + } + } + + for (block, &start_depth) in self.blocks.iter_mut().zip(&start_depths) { + block.start_depth = (start_depth != u32::MAX).then_some(start_depth); + } + } + + fn propagate_disable_load_fast_borrow(&mut self) { + let mut predecessors = vec![Vec::new(); self.blocks.len()]; + for (pred_idx, block) in iter_blocks(&self.blocks) { + if block.next != BlockIdx::NULL { + predecessors[block.next.idx()].push(pred_idx); + } + for info in &block.instructions { + if info.target != BlockIdx::NULL { + predecessors[info.target.idx()].push(pred_idx); + } + } + } + + let mut changed = true; + while changed { + changed = false; + let block_indices: Vec<_> = iter_blocks(&self.blocks).map(|(idx, _)| idx).collect(); + for idx in block_indices { + if idx == BlockIdx(0) || self.blocks[idx.idx()].disable_load_fast_borrow { + continue; + } + let preds = &predecessors[idx.idx()]; + if preds.is_empty() { + continue; + } + if preds + .iter() + .copied() + .all(|pred_idx| self.blocks[pred_idx.idx()].disable_load_fast_borrow) + { + self.blocks[idx.idx()].disable_load_fast_borrow = true; + changed = true; + } + } + } + } + fn deoptimize_borrow_for_handler_return_paths(&mut self) { for block in &mut self.blocks { let len = block.instructions.len(); @@ -2909,9 +3010,11 @@ impl CodeInfo { old_next } else { let suffix_idx = BlockIdx::new(base + annotations_blocks.len() as u32); + let disable_load_fast_borrow = self.blocks[block_idx].disable_load_fast_borrow; let block = Block { instructions: suffix, next: old_next, + disable_load_fast_borrow, ..Default::default() }; annotations_blocks.push(block); @@ -3367,11 +3470,13 @@ fn split_blocks_at_jumps(blocks: &mut Vec) { let tail: Vec = blocks[bi].instructions.drain(pos..).collect(); let old_next = blocks[bi].next; let cold = blocks[bi].cold; + let disable_load_fast_borrow = blocks[bi].disable_load_fast_borrow; blocks[bi].next = new_block_idx; blocks.push(Block { instructions: tail, next: old_next, cold, + disable_load_fast_borrow, ..Block::default() }); // Don't increment bi - re-check current block (it might still have issues) @@ -3588,11 +3693,13 @@ fn normalize_jumps(blocks: &mut Vec) { if let Some(reversed) = reversed_conditional(&last_ins.instr) { let old_next = blocks[idx].next; let is_cold = blocks[idx].cold; + let disable_load_fast_borrow = blocks[idx].disable_load_fast_borrow; // Create new block with NOT_TAKEN + JUMP to original backward target let new_block_idx = BlockIdx(blocks.len() as u32); let mut new_block = Block { cold: is_cold, + disable_load_fast_borrow, ..Block::default() }; new_block.instructions.push(InstructionInfo { @@ -3695,27 +3802,30 @@ fn inline_small_or_no_lineno_blocks(blocks: &mut [Block]) { let is_return_epilogue_block = |block: &Block| { matches!( block.instructions.as_slice(), - [InstructionInfo { - instr: AnyInstruction::Real(Instruction::LoadConst { .. }), - .. - }, InstructionInfo { - instr: AnyInstruction::Real(Instruction::ReturnValue), - .. - }] - | [InstructionInfo { - instr: AnyInstruction::Real(Instruction::LoadSmallInt { .. }), + [ + InstructionInfo { + instr: AnyInstruction::Real(Instruction::LoadConst { .. }), .. - }, InstructionInfo { + }, + InstructionInfo { instr: AnyInstruction::Real(Instruction::ReturnValue), .. - }] - | [InstructionInfo { + } + ] | [ + InstructionInfo { + instr: AnyInstruction::Real(Instruction::LoadSmallInt { .. }), + .. + }, + InstructionInfo { instr: AnyInstruction::Real(Instruction::ReturnValue), .. - }] + } + ] | [InstructionInfo { + instr: AnyInstruction::Real(Instruction::ReturnValue), + .. + }] ) }; - loop { let mut changes = false; let mut current = BlockIdx(0); @@ -3748,7 +3858,6 @@ fn inline_small_or_no_lineno_blocks(blocks: &mut [Block]) { current = next; continue; } - if small_exit_block || no_lineno_no_fallthrough { if let Some(last_instr) = blocks[current.idx()].instructions.last_mut() { set_to_nop(last_instr); @@ -4612,6 +4721,7 @@ fn duplicate_end_returns(blocks: &mut Vec) { let new_block = Block { cold: blocks[last_block.idx()].cold, except_handler: blocks[last_block.idx()].except_handler, + disable_load_fast_borrow: blocks[last_block.idx()].disable_load_fast_borrow, instructions: cloned_return, next: if is_conditional { last_block From 4866570d9bc2a7cf15e4ca73e96556a523648503 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Tue, 14 Apr 2026 21:18:04 +0900 Subject: [PATCH 6/8] Bytecode parity - comprehension/except scope ordering, load_fast_borrow fixes - Reorder comprehension symbol-table walk so the outermost iterator registers its sub_tables in the enclosing scope before the comp scope, and rescan elt/ifs in CPython's order. Codegen peeks past the outermost iterator's nested scopes to find the comprehension table. - For plain try/except, emit handler sub_tables before the else block so codegen's linear sub_table cursor stays aligned. - Rename `collect_simple_annotations` to `collect_annotations` and evaluate non-simple annotations during __annotate__ compilation to preserve source-order side effects while keeping the simple-name index stable. - Dedupe equivalent code constants in `arg_constant` and add a structural equality check on `CodeObject`. - Disable LOAD_FAST_BORROW for the tail end block when a try has a bare `except:` clause, and have `new_block` inherit the flag from the current block. - Remove `cfg!(debug_assertions)` guard around the `optimize_load_fast_borrow` start-depth check so mismatches are handled (return instead of assert) in release builds. - Collapse nop-only blocks that precede a return epilogue and hoist the prior line number into the next real instruction so the line table matches. - Unmark now-passing `test_consts_in_conditionals`, `test_load_fast_unknown_simple`, `test_load_fast_known_because_already_loaded`, and PEP 646 f3/f4 annotation checks. --- Lib/test/test_compile.py | 1 - Lib/test/test_peepholer.py | 2 - Lib/test/test_pep646_syntax.py | 4 +- crates/codegen/src/compile.rs | 337 +++++++++++++++++++++++++----- crates/codegen/src/ir.rs | 66 +++++- crates/codegen/src/symboltable.rs | 62 +++--- 6 files changed, 387 insertions(+), 85 deletions(-) diff --git a/Lib/test/test_compile.py b/Lib/test/test_compile.py index 9676aded5d1..89286c04c1f 100644 --- a/Lib/test/test_compile.py +++ b/Lib/test/test_compile.py @@ -1103,7 +1103,6 @@ def continue_in_while(): self.assertEqual('RETURN_VALUE', opcodes[-1].opname) self.assertEqual(None, opcodes[1].argval) - @unittest.expectedFailure # TODO: RUSTPYTHON def test_consts_in_conditionals(self): def and_true(x): return True and x diff --git a/Lib/test/test_peepholer.py b/Lib/test/test_peepholer.py index e20f712a31a..b563bdb9939 100644 --- a/Lib/test/test_peepholer.py +++ b/Lib/test/test_peepholer.py @@ -863,7 +863,6 @@ def f(): y = x + x self.assertInBytecode(f, 'LOAD_FAST_BORROW_LOAD_FAST_BORROW') - @unittest.expectedFailure # TODO: RUSTPYTHON; RETURN_VALUE def test_load_fast_unknown_simple(self): def f(): if condition(): @@ -906,7 +905,6 @@ def f5(x=0): self.assertInBytecode(f5, 'LOAD_FAST_BORROW') self.assertNotInBytecode(f5, 'LOAD_FAST_CHECK') - @unittest.expectedFailure # TODO: RUSTPYTHON; RETURN_VALUE def test_load_fast_known_because_already_loaded(self): def f(): if condition(): diff --git a/Lib/test/test_pep646_syntax.py b/Lib/test/test_pep646_syntax.py index d9a0aa9a90e..ca8e7d62057 100644 --- a/Lib/test/test_pep646_syntax.py +++ b/Lib/test/test_pep646_syntax.py @@ -305,11 +305,11 @@ {'args': StarredB} >>> def f3(*args: *b, arg1: int): pass - >>> f3.__annotations__ # TODO: RUSTPYTHON # doctest: +EXPECTED_FAILURE + >>> f3.__annotations__ {'args': StarredB, 'arg1': } >>> def f4(*args: *b, arg1: int = 2): pass - >>> f4.__annotations__ # TODO: RUSTPYTHON # doctest: +EXPECTED_FAILURE + >>> f4.__annotations__ {'args': StarredB, 'arg1': } >>> def f5(*args: *b = (1,)): pass # TODO: RUSTPYTHON # doctest: +EXPECTED_FAILURE diff --git a/crates/codegen/src/compile.rs b/crates/codegen/src/compile.rs index 1508785e629..d23bbc5c2f4 100644 --- a/crates/codegen/src/compile.rs +++ b/crates/codegen/src/compile.rs @@ -1132,17 +1132,18 @@ impl Compiler { /// 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, comprehension_type: ComprehensionType) -> bool { + fn is_inlined_comprehension_context( + &self, + comprehension_type: ComprehensionType, + comp_table: &SymbolTable, + ) -> bool { if comprehension_type == ComprehensionType::Generator { return false; } if !self.ctx.in_func() { return false; } - self.symbol_table_stack - .last() - .and_then(|t| t.sub_tables.get(t.next_sub_table)) - .is_some_and(|st| st.comp_inlined) + comp_table.comp_inlined } /// Enter a new scope @@ -3545,6 +3546,18 @@ impl Compiler { let handler_block = self.new_block(); let cleanup_block = self.new_block(); let end_block = self.new_block(); + let has_bare_except = handlers.iter().any(|handler| { + matches!( + handler, + ast::ExceptHandler::ExceptHandler(ast::ExceptHandlerExceptHandler { + type_: None, + .. + }) + ) + }); + if has_bare_except { + self.disable_load_fast_borrow_for_block(end_block); + } emit!(self, Instruction::Nop); emit!( @@ -4381,24 +4394,15 @@ impl Compiler { Ok(true) } - /// Collect simple annotations from module body in AST order (including nested blocks) - /// Returns list of (name, annotation_expr) pairs - /// This must match the order that annotations are compiled to ensure - /// conditional_annotation_index stays in sync with __annotate__ enumeration. - fn collect_simple_annotations(body: &[ast::Stmt]) -> Vec<(&str, &ast::Expr)> { - fn walk<'a>(stmts: &'a [ast::Stmt], out: &mut Vec<(&'a str, &'a ast::Expr)>) { + /// Collect annotated assignments from module/class body in AST order + /// (including nested conditional blocks). This preserves the same walk + /// order as symbol-table construction so the annotation scope's + /// `sub_tables` cursor stays aligned. + fn collect_annotations(body: &[ast::Stmt]) -> Vec<&ast::StmtAnnAssign> { + fn walk<'a>(stmts: &'a [ast::Stmt], out: &mut Vec<&'a ast::StmtAnnAssign>) { for stmt in stmts { match stmt { - ast::Stmt::AnnAssign(ast::StmtAnnAssign { - target, - annotation, - simple, - .. - }) if *simple && matches!(target.as_ref(), ast::Expr::Name(_)) => { - if let ast::Expr::Name(ast::ExprName { id, .. }) = target.as_ref() { - out.push((id.as_str(), annotation.as_ref())); - } - } + ast::Stmt::AnnAssign(stmt) => out.push(stmt), ast::Stmt::If(ast::StmtIf { body, elif_else_clauses, @@ -4449,10 +4453,13 @@ impl Compiler { /// Compile module-level __annotate__ function (PEP 649) /// Returns true if __annotate__ was created and stored fn compile_module_annotate(&mut self, body: &[ast::Stmt]) -> CompileResult { - // Collect simple annotations from module body first - let annotations = Self::collect_simple_annotations(body); + let annotations = Self::collect_annotations(body); + let simple_annotation_count = annotations + .iter() + .filter(|stmt| stmt.simple && matches!(stmt.target.as_ref(), ast::Expr::Name(_))) + .count(); - if annotations.is_empty() { + if simple_annotation_count == 0 { return Ok(false); } @@ -4496,11 +4503,40 @@ impl Compiler { emit!(self, Instruction::BuildMap { count: 0 }); - for (idx, (name, annotation)) in annotations.iter().enumerate() { + let mut simple_idx = 0usize; + for stmt in annotations { + let ast::StmtAnnAssign { + target, + annotation, + simple, + .. + } = stmt; + let simple_name = if *simple { + match target.as_ref() { + ast::Expr::Name(ast::ExprName { id, .. }) => Some(id.as_str()), + _ => None, + } + } else { + None + }; + + if simple_name.is_none() { + if !self.future_annotations { + self.do_not_emit_bytecode += 1; + let result = self.compile_annotation(annotation); + self.do_not_emit_bytecode -= 1; + result?; + } + continue; + } + let not_set_block = has_conditional.then(|| self.new_block()); + let name = simple_name.expect("missing simple annotation name"); if has_conditional { - self.emit_load_const(ConstantData::Integer { value: idx.into() }); + self.emit_load_const(ConstantData::Integer { + value: simple_idx.into(), + }); if parent_scope_type == CompilerScope::Class { let idx = self.get_free_var_index("__conditional_annotations__")?; emit!(self, Instruction::LoadDeref { i: idx }); @@ -4528,6 +4564,7 @@ impl Compiler { value: self.mangle(name).into_owned().into(), }); emit!(self, Instruction::StoreSubscr); + simple_idx += 1; if let Some(not_set_block) = not_set_block { self.switch_to_block(not_set_block); @@ -8872,20 +8909,16 @@ impl Compiler { ast::Expr::ListComp(ast::ExprListComp { generators, .. }) | ast::Expr::SetComp(ast::ExprSetComp { generators, .. }) | ast::Expr::Generator(ast::ExprGenerator { generators, .. }) => { - // leave_scope runs before the first iterator is - // scanned, so the comprehension scope comes first - // in sub_tables, then any nested scopes from the - // first iterator. - self.consume_scope(); if let Some(first) = generators.first() { self.visit_expr(&first.iter); } + self.consume_scope(); } ast::Expr::DictComp(ast::ExprDictComp { generators, .. }) => { - self.consume_scope(); if let Some(first) = generators.first() { self.visit_expr(&first.iter); } + self.consume_scope(); } _ => ast::visitor::walk_expr(self, expr), } @@ -8904,6 +8937,64 @@ impl Compiler { } } + fn peek_next_sub_table_after_skipped_nested_scopes_in_expr( + &mut self, + expression: &ast::Expr, + ) -> CompileResult { + let saved_cursor = self + .symbol_table_stack + .last() + .expect("no current symbol table") + .next_sub_table; + let result = (|| { + self.consume_skipped_nested_scopes_in_expr(expression)?; + let current_table = self + .symbol_table_stack + .last() + .expect("no current symbol table"); + if let Some(table) = current_table.sub_tables.get(current_table.next_sub_table) { + Ok(table.clone()) + } else { + let name = current_table.name.clone(); + let typ = current_table.typ; + Err(self.error(CodegenErrorType::SyntaxError(format!( + "no symbol table available in {} (type: {:?})", + name, typ + )))) + } + })(); + self.symbol_table_stack + .last_mut() + .expect("no current symbol table") + .next_sub_table = saved_cursor; + result + } + + fn push_output_with_symbol_table( + &mut self, + table: SymbolTable, + flags: bytecode::CodeFlags, + posonlyarg_count: u32, + arg_count: u32, + kwonlyarg_count: u32, + obj_name: String, + ) -> CompileResult<()> { + let scope_type = table.typ; + self.symbol_table_stack.push(table); + + let key = self.symbol_table_stack.len() - 1; + let lineno = self.get_source_line_number().get(); + self.enter_scope(&obj_name, scope_type, key, lineno.to_u32())?; + + if let Some(info) = self.code_stack.last_mut() { + info.flags = flags | (info.flags & bytecode::CodeFlags::NESTED); + info.metadata.argcount = arg_count; + info.metadata.posonlyargcount = posonlyarg_count; + info.metadata.kwonlyargcount = kwonlyarg_count; + } + Ok(()) + } + fn compile_comprehension( &mut self, name: &str, @@ -8925,9 +9016,6 @@ impl Compiler { return Err(self.error(CodegenErrorType::InvalidAsyncComprehension)); } - // Check if this comprehension should be inlined (PEP 709) - let is_inlined = self.is_inlined_comprehension_context(comprehension_type); - // async comprehensions are allowed in various contexts: // - list/set/dict comprehensions in async functions (or nested within) // - always for generator expressions @@ -8945,12 +9033,18 @@ impl Compiler { // We must have at least one generator: assert!(!generators.is_empty()); + let outermost = &generators[0]; + let comp_table = + self.peek_next_sub_table_after_skipped_nested_scopes_in_expr(&outermost.iter)?; + + 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( + comp_table, init_collection, generators, compile_element, @@ -8981,8 +9075,12 @@ impl Compiler { flags }; - // Create magnificent function : - self.push_output(flags, 1, 1, 0, name.to_owned())?; + // The symbol table follows CPython's symtable walk: nested scopes + // in the outermost iterator are recorded before the comprehension + // scope itself. Peek past those nested scopes so we can enter the + // correct comprehension table here, then let the real outermost + // iterator compile consume its nested scopes later in parent scope. + self.push_output_with_symbol_table(comp_table, flags, 1, 1, 0, name.to_owned())?; // Set qualname for comprehension self.set_qualname(); @@ -9127,11 +9225,15 @@ impl Compiler { self.make_closure(code, bytecode::MakeFunctionFlags::new())?; // Evaluate iterated item: - self.compile_expression(&generators[0].iter)?; + self.compile_expression(&outermost.iter)?; + self.symbol_table_stack + .last_mut() + .expect("no current symbol table") + .next_sub_table += 1; // Get iterator / turn item into an iterator // Use is_async from the first generator, not has_an_async_gen which covers ALL generators - if generators[0].is_async { + if outermost.is_async { emit!(self, Instruction::GetAIter); } else { emit!(self, Instruction::GetIter); @@ -9152,25 +9254,21 @@ impl Compiler { /// This generates bytecode inline without creating a new code object fn compile_inlined_comprehension( &mut self, + comp_table: SymbolTable, init_collection: Option, generators: &[ast::Comprehension], compile_element: &dyn Fn(&mut Self) -> CompileResult<()>, has_async: bool, ) -> CompileResult<()> { - // PEP 709: Consume the comprehension's sub_table. - // The symbols are already merged into parent scope by analyze_symbol_table. - let current_table = self - .symbol_table_stack - .last_mut() - .expect("no current symbol table"); - let comp_table = current_table.sub_tables[current_table.next_sub_table].clone(); - current_table.next_sub_table += 1; - // 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 // splice in the comprehension's own children. self.compile_expression(&generators[0].iter)?; + self.symbol_table_stack + .last_mut() + .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. @@ -9564,9 +9662,79 @@ impl Compiler { fn arg_constant(&mut self, constant: ConstantData) -> oparg::ConstIdx { let info = self.current_code_info(); + if let ConstantData::Code { code } = &constant + && let Some(idx) = info.metadata.consts.iter().position(|existing| { + matches!( + existing, + ConstantData::Code { + code: existing_code + } if Self::code_objects_equivalent(existing_code, code) + ) + }) + { + return u32::try_from(idx) + .expect("constant table index overflow") + .into(); + } info.metadata.consts.insert_full(constant).0.to_u32().into() } + fn constants_equivalent(lhs: &ConstantData, rhs: &ConstantData) -> bool { + match (lhs, rhs) { + (ConstantData::Code { code: lhs }, ConstantData::Code { code: rhs }) => { + Self::code_objects_equivalent(lhs, rhs) + } + (ConstantData::Tuple { elements: lhs }, ConstantData::Tuple { elements: rhs }) + | ( + ConstantData::Frozenset { elements: lhs }, + ConstantData::Frozenset { elements: rhs }, + ) => { + lhs.len() == rhs.len() + && lhs + .iter() + .zip(rhs.iter()) + .all(|(lhs, rhs)| Self::constants_equivalent(lhs, rhs)) + } + (ConstantData::Slice { elements: lhs }, ConstantData::Slice { elements: rhs }) => lhs + .iter() + .zip(rhs.iter()) + .all(|(lhs, rhs)| Self::constants_equivalent(lhs, rhs)), + _ => lhs == rhs, + } + } + + fn code_objects_equivalent(lhs: &bytecode::CodeObject, rhs: &bytecode::CodeObject) -> bool { + lhs.instructions.len() == rhs.instructions.len() + && lhs + .instructions + .iter() + .zip(rhs.instructions.iter()) + .all(|(lhs, rhs)| u8::from(lhs.op) == u8::from(rhs.op) && lhs.arg == rhs.arg) + && lhs.locations == rhs.locations + && lhs.flags.bits() == rhs.flags.bits() + && lhs.posonlyarg_count == rhs.posonlyarg_count + && lhs.arg_count == rhs.arg_count + && lhs.kwonlyarg_count == rhs.kwonlyarg_count + && lhs.source_path == rhs.source_path + && lhs.first_line_number == rhs.first_line_number + && lhs.max_stackdepth == rhs.max_stackdepth + && lhs.obj_name == rhs.obj_name + && lhs.qualname == rhs.qualname + && lhs.constants.len() == rhs.constants.len() + && lhs + .constants + .iter() + .zip(rhs.constants.iter()) + .all(|(lhs, rhs)| Self::constants_equivalent(lhs, rhs)) + && lhs.names == rhs.names + && lhs.varnames == rhs.varnames + && lhs.cellvars == rhs.cellvars + && lhs.freevars == rhs.freevars + && lhs.localspluskinds == rhs.localspluskinds + && lhs.linetable == rhs.linetable + && lhs.exceptiontable == rhs.exceptiontable + } + /// Try to fold a collection of constant expressions into a single ConstantData::Tuple. /// Returns None if any element cannot be folded. fn try_fold_constant_collection( @@ -10005,7 +10173,13 @@ impl Compiler { fn new_block(&mut self) -> BlockIdx { let code = self.current_code_info(); let idx = BlockIdx::new(code.blocks.len().to_u32()); - code.blocks.push(ir::Block::default()); + let inherited_disable_load_fast_borrow = + code.blocks[code.current_block].disable_load_fast_borrow; + let block = ir::Block { + disable_load_fast_borrow: inherited_disable_load_fast_borrow, + ..ir::Block::default() + }; + code.blocks.push(block); idx } @@ -11999,6 +12173,71 @@ def f(): ); } + #[test] + fn test_bare_except_deopts_post_handler_load_fast_borrow() { + let code = compile_exec( + "\ +def f(self): + try: + 1 / 0 + except: + pass + with self.assertRaises(SyntaxError): + 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(); + + let attr_idx = ops + .iter() + .position(|op| matches!(op, Instruction::LoadAttr { .. })) + .expect("missing LOAD_ATTR for assertRaises"); + assert!( + matches!(ops.get(attr_idx - 1), Some(Instruction::LoadFast { .. })), + "bare except tail should deopt self to LOAD_FAST, got ops={ops:?}" + ); + } + + #[test] + fn test_typed_except_keeps_post_handler_load_fast_borrow() { + let code = compile_exec( + "\ +def f(self): + try: + 1 / 0 + except ZeroDivisionError: + pass + with self.assertRaises(SyntaxError): + 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(); + + let attr_idx = ops + .iter() + .position(|op| matches!(op, Instruction::LoadAttr { .. })) + .expect("missing LOAD_ATTR for assertRaises"); + assert!( + matches!( + ops.get(attr_idx - 1), + Some(Instruction::LoadFastBorrow { .. }) + ), + "typed except tail should keep LOAD_FAST_BORROW, got ops={ops:?}" + ); + } + #[test] fn test_constant_slice_folding_handles_string_and_bigint_bounds() { let code = compile_exec( diff --git a/crates/codegen/src/ir.rs b/crates/codegen/src/ir.rs index e8e9fa92e7d..9c6e84fce54 100644 --- a/crates/codegen/src/ir.rs +++ b/crates/codegen/src/ir.rs @@ -2107,8 +2107,8 @@ impl CodeInfo { target: BlockIdx, start_depth: usize, ) { - if cfg!(debug_assertions) { - let expected = blocks[target.idx()].start_depth.map(|depth| depth as usize); + let expected = blocks[target.idx()].start_depth.map(|depth| depth as usize); + if expected != Some(start_depth) { debug_assert!( expected == Some(start_depth), "optimize_load_fast_borrow start_depth mismatch: source={source:?} target={target:?} expected={expected:?} actual={:?} source_last={:?} target_instrs={:?}", @@ -2123,6 +2123,7 @@ impl CodeInfo { .map(|info| info.instr) .collect::>(), ); + return; } if !visited[target.idx()] { visited[target.idx()] = true; @@ -3877,6 +3878,33 @@ fn inline_small_or_no_lineno_blocks(blocks: &mut [Block]) { } fn remove_redundant_nops_in_blocks(blocks: &mut [Block]) -> usize { + let is_return_epilogue_block = |block: &Block| { + matches!( + block.instructions.as_slice(), + [ + InstructionInfo { + instr: AnyInstruction::Real(Instruction::LoadConst { .. }), + .. + }, + InstructionInfo { + instr: AnyInstruction::Real(Instruction::ReturnValue), + .. + } + ] | [ + InstructionInfo { + instr: AnyInstruction::Real(Instruction::LoadSmallInt { .. }), + .. + }, + InstructionInfo { + instr: AnyInstruction::Real(Instruction::ReturnValue), + .. + } + ] | [InstructionInfo { + instr: AnyInstruction::Real(Instruction::ReturnValue), + .. + }] + ) + }; let mut changes = 0; let mut block_order = Vec::new(); let mut current = BlockIdx(0); @@ -3910,18 +3938,41 @@ fn remove_redundant_nops_in_blocks(blocks: &mut [Block]) -> usize { } else { let next = next_nonempty_block(blocks, blocks[bi].next); if next != BlockIdx::NULL { - let mut next_lineno = None; - for next_instr in &blocks[next.idx()].instructions { + if is_return_epilogue_block(&blocks[next.idx()]) { + let current_block_is_nop_only = + kept.iter().all(|prev: &InstructionInfo| { + matches!(prev.instr.real(), Some(Instruction::Nop)) + }); + let pred = find_layout_predecessor(blocks, block_idx); + let pred_ends_with_nop = pred != BlockIdx::NULL + && blocks[pred.idx()].instructions.last().is_some_and(|prev| { + matches!(prev.instr.real(), Some(Instruction::Nop)) + }); + if current_block_is_nop_only && pred_ends_with_nop { + changes += 1; + continue; + } + } + let mut next_info = None; + for (next_idx, next_instr) in + blocks[next.idx()].instructions.iter().enumerate() + { let line = instruction_lineno(next_instr); if matches!(next_instr.instr.real(), Some(Instruction::Nop)) && line < 0 { continue; } - next_lineno = Some(line); + next_info = Some((next_idx, line)); break; } - if next_lineno.is_some_and(|line| line == lineno) { - remove = true; + if let Some((next_idx, next_lineno)) = next_info { + if next_lineno == lineno { + remove = true; + } else if next_lineno < 0 { + blocks[next.idx()].instructions[next_idx].lineno_override = + Some(lineno); + remove = true; + } } } } @@ -3958,6 +4009,7 @@ fn remove_redundant_jumps_in_blocks(blocks: &mut [Block]) -> usize { && let Some(last_instr) = blocks[idx].instructions.last_mut() { set_to_nop(last_instr); + last_instr.lineno_override = Some(-1); changes += 1; } current = blocks[idx].next; diff --git a/crates/codegen/src/symboltable.rs b/crates/codegen/src/symboltable.rs index 7479f168595..7179656fe8a 100644 --- a/crates/codegen/src/symboltable.rs +++ b/crates/codegen/src/symboltable.rs @@ -1624,10 +1624,10 @@ impl SymbolTableBuilder { let saved_in_conditional_block = self.in_conditional_block; self.in_conditional_block = true; self.scan_statements(body)?; - self.scan_statements(orelse)?; - // Keep nested scope collection in the same order that codegen - // compiles try/except, since the compiler currently consumes - // sub_tables through a linear cursor. + // Preserve source-order symbol analysis so `global`/`nonlocal` + // semantics match CPython, but reorder child scope storage to + // match the codegen order for plain try/except/else. + let body_subtables_len = self.tables.last().unwrap().sub_tables.len(); for handler in handlers { let ExceptHandler::ExceptHandler(ast::ExceptHandlerExceptHandler { type_, @@ -1643,6 +1643,22 @@ impl SymbolTableBuilder { } self.scan_statements(body)?; } + if finalbody.is_empty() { + let handler_subtables = self + .tables + .last_mut() + .unwrap() + .sub_tables + .split_off(body_subtables_len); + self.scan_statements(orelse)?; + self.tables + .last_mut() + .unwrap() + .sub_tables + .extend(handler_subtables); + } else { + self.scan_statements(orelse)?; + } self.scan_statements(finalbody)?; self.in_conditional_block = saved_in_conditional_block; } @@ -2148,6 +2164,13 @@ impl SymbolTableBuilder { }); } + assert!(!generators.is_empty()); + let outermost = &generators[0]; + + // CPython evaluates the outermost iterator in the enclosing scope + // before entering the comprehension scope. + self.scan_expression(&outermost.iter, ExpressionContext::IterDefinitionExp)?; + // Comprehensions are compiled as functions, so create a scope for them: self.enter_scope( scope_name, @@ -2183,36 +2206,27 @@ impl SymbolTableBuilder { // Register the passed argument to the generator function as the name ".0" self.register_name(".0", SymbolUsage::Parameter, range)?; - self.scan_expression(elt1, ExpressionContext::Load)?; - if let Some(elt2) = elt2 { - self.scan_expression(elt2, ExpressionContext::Load)?; + self.scan_expression(&outermost.target, ExpressionContext::Iter)?; + for if_expr in &outermost.ifs { + self.scan_expression(if_expr, ExpressionContext::Load)?; } - let mut is_first_generator = true; - for generator in generators { - // Set flag for INNER_LOOP_CONFLICT check (only for inner loops, not the first) - if !is_first_generator { - self.in_comp_inner_loop_target = true; - } + for generator in &generators[1..] { + self.in_comp_inner_loop_target = true; self.scan_expression(&generator.target, ExpressionContext::Iter)?; self.in_comp_inner_loop_target = false; - - if is_first_generator { - is_first_generator = false; - } else { - self.scan_expression(&generator.iter, ExpressionContext::IterDefinitionExp)?; - } - + self.scan_expression(&generator.iter, ExpressionContext::IterDefinitionExp)?; for if_expr in &generator.ifs { self.scan_expression(if_expr, ExpressionContext::Load)?; } } - self.leave_scope(); + if let Some(elt2) = elt2 { + self.scan_expression(elt2, ExpressionContext::Load)?; + } + self.scan_expression(elt1, ExpressionContext::Load)?; - // The first iterable is passed as an argument into the created function: - assert!(!generators.is_empty()); - self.scan_expression(&generators[0].iter, ExpressionContext::IterDefinitionExp)?; + self.leave_scope(); Ok(()) } From 0a47dc29201214b36c54d6432a8cf3c8e2b1f0e1 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Tue, 14 Apr 2026 22:36:25 +0900 Subject: [PATCH 7/8] Bytecode parity - try/except line tracking, assert 0 shape - In `compile_try_except`, drop the leading Nop and set the end block's source range from the last orelse/body statement so line events after the try fall on the right line. - Recognise constant-false asserts as the direct-raise shape (no ToBool/PopJumpIfFalse) and flip the test assertion accordingly. - Extend `remove_redundant_nops_in_blocks` to also look through a trailing nop before a return-epilogue pair (LoadConst/ReturnValue or LoadSmallInt/ReturnValue) so the epilogue keeps the correct line number. - Rename `preds` to `predecessor_blocks` in the LOAD_FAST_BORROW disable pass and add a test-only `debug_late_cfg_trace` helper. - Regenerate the `nested_double_async_with` snapshot: the tail reference to `stop_exc` now emits LOAD_FAST instead of LOAD_FAST_BORROW. --- crates/codegen/src/compile.rs | 56 ++++++- crates/codegen/src/ir.rs | 157 +++++++++++++++++- ...pile__tests__nested_double_async_with.snap | 2 +- 3 files changed, 206 insertions(+), 9 deletions(-) diff --git a/crates/codegen/src/compile.rs b/crates/codegen/src/compile.rs index d23bbc5c2f4..0152852c8d8 100644 --- a/crates/codegen/src/compile.rs +++ b/crates/codegen/src/compile.rs @@ -3543,6 +3543,10 @@ impl Compiler { handlers: &[ast::ExceptHandler], orelse: &[ast::Stmt], ) -> CompileResult<()> { + let normal_exit_range = orelse + .last() + .map(ast::Stmt::range) + .or_else(|| body.last().map(ast::Stmt::range)); let handler_block = self.new_block(); let cleanup_block = self.new_block(); let end_block = self.new_block(); @@ -3559,7 +3563,6 @@ impl Compiler { self.disable_load_fast_borrow_for_block(end_block); } - emit!(self, Instruction::Nop); emit!( self, PseudoInstruction::SetupFinally { @@ -3694,6 +3697,9 @@ impl Compiler { self.set_no_location(); self.switch_to_block(end_block); + if let Some(range) = normal_exit_range { + self.set_source_range(range); + } Ok(()) } @@ -11070,6 +11076,29 @@ mod tests { compiler.exit_scope() } + fn compile_exec_late_cfg_trace(source: &str) -> Vec<(String, String)> { + 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(); + let ast = parsed.into_syntax(); + let ast = match ast { + ruff_python_ast::Mod::Module(stmts) => stmts, + _ => unreachable!(), + }; + let symbol_table = SymbolTable::scan_program(&ast, source_file.clone()) + .map_err(|e| e.into_codegen_error(source_file.name().to_owned())) + .unwrap(); + let mut compiler = Compiler::new(opts, source_file, "".to_owned()); + compiler.compile_program(&ast, symbol_table).unwrap(); + let _table = compiler.pop_symbol_table(); + let stack_top = compiler.code_stack.pop().unwrap(); + stack_top.debug_late_cfg_trace().unwrap() + } + fn find_code<'a>(code: &'a CodeObject, name: &str) -> Option<&'a CodeObject> { if code.obj_name == name { return Some(code); @@ -11121,6 +11150,25 @@ if True or False or False: )); } + #[test] + fn test_trace_assert_true_try_pair() { + let trace = compile_exec_late_cfg_trace( + "\ +try: + assert True +except AssertionError as e: + fail() +try: + assert True, 'msg' +except AssertionError as e: + fail() +", + ); + for (stage, dump) in trace { + eprintln!("=== {stage} ===\n{dump}"); + } + } + #[test] fn test_if_ands() { assert_dis_snapshot!(compile_exec( @@ -12564,11 +12612,11 @@ def f(): } #[test] - fn test_constant_false_assert_uses_normal_assert_branch_shape() { + fn test_constant_false_assert_uses_direct_raise_shape() { let code = compile_exec("assert 0, (lambda x: x + 1)\n"); assert!( - code.instructions.iter().any(|unit| { + !code.instructions.iter().any(|unit| { matches!( unit.op, Instruction::ToBool @@ -12576,7 +12624,7 @@ def f(): | Instruction::PopJumpIfFalse { .. } ) }), - "constant-false assert should still use the normal assert branch shape, got ops={:?}", + "constant-false assert should use direct raise shape, got ops={:?}", code.instructions .iter() .map(|unit| unit.op) diff --git a/crates/codegen/src/ir.rs b/crates/codegen/src/ir.rs index 9c6e84fce54..4e9db27232b 100644 --- a/crates/codegen/src/ir.rs +++ b/crates/codegen/src/ir.rs @@ -2480,11 +2480,11 @@ impl CodeInfo { if idx == BlockIdx(0) || self.blocks[idx.idx()].disable_load_fast_borrow { continue; } - let preds = &predecessors[idx.idx()]; - if preds.is_empty() { + let predecessor_blocks = &predecessors[idx.idx()]; + if predecessor_blocks.is_empty() { continue; } - if preds + if predecessor_blocks .iter() .copied() .all(|pred_idx| self.blocks[pred_idx.idx()].disable_load_fast_borrow) @@ -2947,6 +2947,112 @@ impl CodeInfo { } } +#[cfg(test)] +impl CodeInfo { + fn debug_block_dump(&self) -> String { + let mut out = String::new(); + for (block_idx, block) in iter_blocks(&self.blocks) { + use core::fmt::Write; + let _ = writeln!( + out, + "block {} next={} cold={} except={} preserve_lasti={} disable_borrow={}", + u32::from(block_idx), + if block.next == BlockIdx::NULL { + String::from("NULL") + } else { + u32::from(block.next).to_string() + }, + block.cold, + block.except_handler, + block.preserve_lasti, + block.disable_load_fast_borrow, + ); + for info in &block.instructions { + let lineno = instruction_lineno(info); + let _ = writeln!( + out, + " [{}] {:?} arg={} target={}", + lineno, + info.instr, + u32::from(info.arg), + if info.target == BlockIdx::NULL { + String::from("NULL") + } else { + u32::from(info.target).to_string() + } + ); + } + } + out + } + + pub(crate) fn debug_late_cfg_trace(mut self) -> crate::InternalResult> { + let mut trace = Vec::new(); + + self.splice_annotations_blocks(); + self.fold_binop_constants(); + self.remove_nops(); + self.fold_unary_negative(); + self.remove_nops(); + self.fold_binop_constants(); + self.remove_nops(); + self.fold_tuple_constants(); + self.fold_list_constants(); + self.fold_set_constants(); + self.remove_nops(); + self.fold_const_iterable_for_iter(); + self.convert_to_load_small_int(); + self.remove_unused_consts(); + self.remove_nops(); + self.dce(); + self.optimize_build_tuple_unpack(); + self.eliminate_dead_stores(); + self.apply_static_swaps(); + self.peephole_optimize(); + split_blocks_at_jumps(&mut self.blocks); + mark_except_handlers(&mut self.blocks); + label_exception_targets(&mut self.blocks); + jump_threading(&mut self.blocks); + self.eliminate_unreachable_blocks(); + self.remove_nops(); + self.add_checks_for_loads_of_uninitialized_variables(); + self.insert_superinstructions(); + push_cold_blocks_to_end(&mut self.blocks); + + trace.push(("after_push_cold_blocks_to_end".to_owned(), self.debug_block_dump())); + + 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())); + + resolve_line_numbers(&mut self.blocks); + trace.push(("after_resolve_line_numbers".to_owned(), self.debug_block_dump())); + + duplicate_end_returns(&mut self.blocks); + trace.push(("after_duplicate_end_returns".to_owned(), self.debug_block_dump())); + + self.dce(); + self.eliminate_unreachable_blocks(); + trace.push(("after_second_dce_unreachable".to_owned(), self.debug_block_dump())); + + remove_redundant_nops_and_jumps(&mut self.blocks); + trace.push(("after_remove_redundant_nops_and_jumps".to_owned(), self.debug_block_dump())); + + Ok(trace) + } +} + impl CodeInfo { fn remap_block_idx(idx: BlockIdx, base: u32) -> BlockIdx { if idx == BlockIdx::NULL { @@ -3905,6 +4011,37 @@ fn remove_redundant_nops_in_blocks(blocks: &mut [Block]) -> usize { }] ) }; + let is_return_epilogue_pair = |instructions: &[InstructionInfo]| { + matches!( + instructions, + [ + InstructionInfo { + instr: AnyInstruction::Real(Instruction::LoadConst { .. }), + .. + }, + InstructionInfo { + instr: AnyInstruction::Real(Instruction::ReturnValue), + .. + } + ] | [ + InstructionInfo { + instr: AnyInstruction::Real(Instruction::LoadSmallInt { .. }), + .. + }, + InstructionInfo { + instr: AnyInstruction::Real(Instruction::ReturnValue), + .. + } + ] | [InstructionInfo { + instr: AnyInstruction::Real(Instruction::ReturnValue), + .. + }] + ) + }; + let starts_with_return_epilogue_pair = |instructions: &[InstructionInfo]| { + instructions.len() >= 2 && is_return_epilogue_pair(&instructions[..2]) + || !instructions.is_empty() && is_return_epilogue_pair(&instructions[..1]) + }; let mut changes = 0; let mut block_order = Vec::new(); let mut current = BlockIdx(0); @@ -3928,6 +4065,15 @@ 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 kept.last().is_some_and(|prev: &InstructionInfo| { + matches!(prev.instr.real(), Some(Instruction::Nop)) + }) && is_return_epilogue_pair( + &src_instructions[src + 1..src_instructions.len().min(src + 3)], + ) + { + src_instructions[src + 1].lineno_override = Some(lineno); + remove = true; + } else { let next_lineno = instruction_lineno(&src_instructions[src + 1]); if next_lineno == lineno { remove = true; @@ -3935,10 +4081,13 @@ fn remove_redundant_nops_in_blocks(blocks: &mut [Block]) -> usize { src_instructions[src + 1].lineno_override = Some(lineno); remove = true; } + } } else { let next = next_nonempty_block(blocks, blocks[bi].next); if next != BlockIdx::NULL { - if is_return_epilogue_block(&blocks[next.idx()]) { + if is_return_epilogue_block(&blocks[next.idx()]) + || starts_with_return_epilogue_pair(&blocks[next.idx()].instructions) + { let current_block_is_nop_only = kept.iter().all(|prev: &InstructionInfo| { matches!(prev.instr.real(), Some(Instruction::Nop)) 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 ff6890fd1ed..c39144ae857 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 @@ -160,7 +160,7 @@ expression: "compile_exec(\"\\\nasync def test():\n for stop_exc in (StopIter 145 CACHE 146 CACHE 147 CACHE - 148 LOAD_FAST_BORROW (0, stop_exc) + 148 LOAD_FAST (0, stop_exc) 149 FORMAT_SIMPLE 150 LOAD_CONST (" was suppressed") 151 BUILD_STRING (2) From f983099bf38dc2647028e289cd8fd597c31bfe88 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Wed, 15 Apr 2026 09:19:06 +0900 Subject: [PATCH 8/8] Bytecode parity - iter folding, break/continue line, cold inlining - Fold a constant list iterable into a constant tuple in for-loop iterable position, matching the CPython optimizer, and strip a redundant LIST_TO_TUPLE immediately before GET_ITER in the IR peephole pass. - Emit a Nop at the break/continue source range before unwinding so line events land on the break/continue statement instead of the following instruction. - Drop `propagate_disable_load_fast_borrow`; the forward propagation was over-zealous and the per-block inheritance in `new_block` plus the bare-except marker are enough. - Relax `inline_small_or_no_lineno_blocks` so small exit blocks at the tail of a cold block are always inlined, not just return epilogues. - Add codegen tests covering the LIST_TO_TUPLE/GET_ITER peephole and the late-CFG trace helper for a for-loop list-literal iterable. --- Lib/test/test_compile.py | 1 - Lib/test/test_grammar.py | 1 - Lib/test/test_patma.py | 4 - Lib/test/test_peepholer.py | 4 +- Lib/test/test_sys_settrace.py | 4 - crates/codegen/src/compile.rs | 1232 +++++++++++++++-- crates/codegen/src/ir.rs | 1210 +++++++++++----- ...k_attribute_and_subscript_expressions.snap | 67 + ...nt_true_if_pass_keeps_line_anchor_nop.snap | 10 + ...thon_codegen__compile__tests__if_ands.snap | 8 +- ...hon_codegen__compile__tests__if_mixed.snap | 8 +- ...pile__tests__nested_double_async_with.snap | 19 +- crates/codegen/src/symboltable.rs | 54 +- crates/vm/src/stdlib/_ctypes/simple.rs | 2 +- crates/vm/src/types/slot.rs | 4 +- scripts/dis_dump.py | 12 +- 16 files changed, 2073 insertions(+), 567 deletions(-) create mode 100644 crates/codegen/src/snapshots/rustpython_codegen__compile__tests__bare_function_annotations_check_attribute_and_subscript_expressions.snap create mode 100644 crates/codegen/src/snapshots/rustpython_codegen__compile__tests__constant_true_if_pass_keeps_line_anchor_nop.snap diff --git a/Lib/test/test_compile.py b/Lib/test/test_compile.py index 89286c04c1f..e5bc65651e9 100644 --- a/Lib/test/test_compile.py +++ b/Lib/test/test_compile.py @@ -1084,7 +1084,6 @@ def unused_block_while_else(): self.assertEqual('RETURN_VALUE', opcodes[-1].opname) self.assertEqual(None, opcodes[-1].argval) - @unittest.expectedFailure # TODO: RUSTPYTHON; AssertionError: 3 != 8 def test_false_while_loop(self): def break_in_while(): while False: diff --git a/Lib/test/test_grammar.py b/Lib/test/test_grammar.py index 77bd5a163ce..91eb6cc58f3 100644 --- a/Lib/test/test_grammar.py +++ b/Lib/test/test_grammar.py @@ -304,7 +304,6 @@ def test_var_annot_syntax_errors(self): " nonlocal x\n" " x: int\n") - @unittest.expectedFailure # TODO: RUSTPYTHON def test_var_annot_basic_semantics(self): # execution order with self.assertRaises(ZeroDivisionError): diff --git a/Lib/test/test_patma.py b/Lib/test/test_patma.py index 6ca1fa0ba40..40466ec67ba 100644 --- a/Lib/test/test_patma.py +++ b/Lib/test/test_patma.py @@ -3448,7 +3448,6 @@ def f(command): # 0 self.assertListEqual(self._trace(f, "go x"), [1, 2, 4, 5]) self.assertListEqual(self._trace(f, "spam"), [1, 2, 4, 6, 7]) - @unittest.expectedFailure # TODO: RUSTPYTHON def test_default_capture(self): def f(command): # 0 match command.split(): # 1 @@ -3463,7 +3462,6 @@ def f(command): # 0 self.assertListEqual(self._trace(f, "go x"), [1, 2, 4, 5]) self.assertListEqual(self._trace(f, "spam"), [1, 2, 4, 6, 7]) - @unittest.expectedFailure # TODO: RUSTPYTHON def test_no_default(self): def f(command): # 0 match command.split(): # 1 @@ -3476,7 +3474,6 @@ def f(command): # 0 self.assertListEqual(self._trace(f, "go x"), [1, 2, 4, 5]) self.assertListEqual(self._trace(f, "spam"), [1, 2, 4]) - @unittest.expectedFailure # TODO: RUSTPYTHON def test_only_default_wildcard(self): def f(command): # 0 match command.split(): # 1 @@ -3487,7 +3484,6 @@ def f(command): # 0 self.assertListEqual(self._trace(f, "go x"), [1, 2, 3]) self.assertListEqual(self._trace(f, "spam"), [1, 2, 3]) - @unittest.expectedFailure # TODO: RUSTPYTHON def test_only_default_capture(self): def f(command): # 0 match command.split(): # 1 diff --git a/Lib/test/test_peepholer.py b/Lib/test/test_peepholer.py index b563bdb9939..c02bd559f1c 100644 --- a/Lib/test/test_peepholer.py +++ b/Lib/test/test_peepholer.py @@ -441,6 +441,7 @@ def test_constant_folding_binop(self): self.check_lnotab(code) + @unittest.expectedFailure # TODO: RUSTPYTHON def test_constant_folding_remove_nop_location(self): sources = [ """ @@ -785,7 +786,6 @@ def f(a, b, c): c, b, a = a, b, c self.assertNotInBytecode(f, "SWAP") - @unittest.expectedFailure # TODO: RUSTPYTHON def test_static_swaps_match_mapping(self): for a, b, c in product("_a", "_b", "_c"): pattern = f"{{'a': {a}, 'b': {b}, 'c': {c}}}" @@ -793,7 +793,6 @@ def test_static_swaps_match_mapping(self): code = compile_pattern_with_fast_locals(pattern) self.assertNotInBytecode(code, "SWAP") - @unittest.expectedFailure # TODO: RUSTPYTHON def test_static_swaps_match_class(self): forms = [ "C({}, {}, {})", @@ -808,7 +807,6 @@ def test_static_swaps_match_class(self): code = compile_pattern_with_fast_locals(pattern) self.assertNotInBytecode(code, "SWAP") - @unittest.expectedFailure # TODO: RUSTPYTHON def test_static_swaps_match_sequence(self): swaps = {"*_, b, c", "a, *_, c", "a, b, *_"} forms = ["{}, {}, {}", "{}, {}, *{}", "{}, *{}, {}", "*{}, {}, {}"] diff --git a/Lib/test/test_sys_settrace.py b/Lib/test/test_sys_settrace.py index f9449c7079a..aa2d54ee16e 100644 --- a/Lib/test/test_sys_settrace.py +++ b/Lib/test/test_sys_settrace.py @@ -1957,8 +1957,6 @@ def test_jump_out_of_finally_block(output): finally: output.append(5) - # TODO: RUSTPYTHON - @unittest.expectedFailure @jump_test(1, 5, [], (ValueError, "into an 'except'")) def test_no_jump_into_bare_except_block(output): output.append(1) @@ -1967,8 +1965,6 @@ def test_no_jump_into_bare_except_block(output): except: output.append(5) - # TODO: RUSTPYTHON - @unittest.expectedFailure @jump_test(1, 5, [], (ValueError, "into an 'except'")) def test_no_jump_into_qualified_except_block(output): output.append(1) diff --git a/crates/codegen/src/compile.rs b/crates/codegen/src/compile.rs index 0152852c8d8..2f7510f0d44 100644 --- a/crates/codegen/src/compile.rs +++ b/crates/codegen/src/compile.rs @@ -671,12 +671,14 @@ impl Compiler { _ => n > 4, }; - // Fold all-constant collections (>= 3 elements) regardless of size + let can_fold_const_collection = match collection_type { + CollectionType::Tuple => n > 0, + CollectionType::List | CollectionType::Set => n >= 3, + }; if !self.disable_const_collection_folding && !seen_star && pushed == 0 - && n >= 3 - && elts.iter().all(|e| e.is_constant()) + && can_fold_const_collection && let Some(folded) = self.try_fold_constant_collection(elts, collection_type)? { match collection_type { @@ -1374,6 +1376,7 @@ impl Compiler { location, end_location, except_handler, + folded_from_nonliteral_expr: false, lineno_override, cache_entries: 0, }); @@ -2510,8 +2513,9 @@ impl Compiler { .. }) => { self.enter_conditional_block(); - self.compile_if(test, body, elif_else_clauses)?; + self.compile_if(test, body, elif_else_clauses, test.range())?; self.leave_conditional_block(); + self.set_source_range(statement.range()); } ast::Stmt::While(ast::StmtWhile { test, body, orelse, .. @@ -5573,51 +5577,36 @@ impl Compiler { test: &ast::Expr, body: &[ast::Stmt], elif_else_clauses: &[ast::ElifElseClause], + _stmt_range: TextRange, ) -> CompileResult<()> { - let test_truthiness = self.constant_expr_truthiness(test)?; - match elif_else_clauses { - // Only if - [] => { - let after_block = self.new_block(); - if matches!(test_truthiness, Some(false)) { - self.disable_load_fast_borrow_for_block(after_block); - } - self.compile_jump_if(test, false, after_block)?; - self.compile_statements(body)?; - self.switch_to_block(after_block); - } - // If, elif*, elif/else - [rest @ .., tail] => { - let after_block = self.new_block(); - let mut next_block = self.new_block(); + let end_block = self.new_block(); + let next_block = if elif_else_clauses.is_empty() { + end_block + } else { + self.new_block() + }; - if matches!(test_truthiness, Some(false)) { - self.disable_load_fast_borrow_for_block(next_block); - } - self.compile_jump_if(test, false, next_block)?; - self.compile_statements(body)?; - emit!(self, PseudoInstruction::Jump { delta: after_block }); + if matches!(self.constant_expr_truthiness(test)?, Some(false)) { + self.disable_load_fast_borrow_for_block(next_block); + } + self.compile_jump_if(test, false, next_block)?; + self.compile_statements(body)?; - for clause in rest { - self.switch_to_block(next_block); - next_block = self.new_block(); - if let Some(test) = &clause.test { - self.compile_jump_if(test, false, next_block)?; - } else { - unreachable!() // must be elif - } - self.compile_statements(&clause.body)?; - emit!(self, PseudoInstruction::Jump { delta: after_block }); - } + let Some((clause, rest)) = elif_else_clauses.split_first() else { + self.switch_to_block(end_block); + return Ok(()); + }; - self.switch_to_block(next_block); - if let Some(test) = &tail.test { - self.compile_jump_if(test, false, after_block)?; - } - self.compile_statements(&tail.body)?; - self.switch_to_block(after_block); - } + emit!(self, PseudoInstruction::Jump { delta: end_block }); + self.switch_to_block(next_block); + + if let Some(test) = &clause.test { + self.compile_if(test, &clause.body, rest, test.range())?; + } else { + debug_assert!(rest.is_empty()); + self.compile_statements(&clause.body)?; } + self.switch_to_block(end_block); Ok(()) } @@ -5637,6 +5626,7 @@ impl Compiler { self.push_fblock(FBlockType::WhileLoop, while_block, after_block)?; if matches!(self.constant_expr_truthiness(test)?, Some(false)) { self.disable_load_fast_borrow_for_block(else_block); + self.disable_load_fast_borrow_for_block(after_block); } self.compile_jump_if(test, false, else_block)?; @@ -5906,24 +5896,7 @@ impl Compiler { let mut end_async_for_target = BlockIdx::NULL; // The thing iterated: - // Optimize: `for x in [a, b, c]` → use tuple instead of list - // Skip for async-for (GET_AITER expects the original type) - if !is_async - && let ast::Expr::List(ast::ExprList { elts, .. }) = iter - && !elts.iter().any(|e| matches!(e, ast::Expr::Starred(_))) - { - for elt in elts { - self.compile_expression(elt)?; - } - emit!( - self, - Instruction::BuildTuple { - count: u32::try_from(elts.len()).expect("too many elements"), - } - ); - } else { - self.compile_expression(iter)?; - } + self.compile_for_iterable_expression(iter, is_async)?; if is_async { if self.ctx.func != FunctionContext::AsyncFunction { @@ -5958,6 +5931,13 @@ impl Compiler { emit!(self, Instruction::ForIter { delta: else_block }); + // Match CPython codegen_for(): keep a line anchor on the target line + // so multiline/single-line `for ...: pass` bodies preserve tracing layout. + let saved_range = self.current_source_range; + self.set_source_range(target.range()); + emit!(self, Instruction::Nop); + self.set_source_range(saved_range); + // Start of loop iteration, set targets: self.compile_store(target)?; }; @@ -5994,6 +5974,38 @@ impl Compiler { Ok(()) } + fn compile_for_iterable_expression( + &mut self, + iter: &ast::Expr, + is_async: bool, + ) -> CompileResult<()> { + // Match CPython's iterable lowering for `for`/comprehension fronts: + // a non-starred list literal used only for iteration is emitted as a tuple. + // Skip async-for/async comprehension iteration because GET_AITER expects + // the original object semantics. + if !is_async + && let ast::Expr::List(ast::ExprList { elts, .. }) = iter + && !elts.iter().any(|e| matches!(e, ast::Expr::Starred(_))) + { + if let Some(folded) = self.try_fold_constant_collection(elts, CollectionType::List)? { + self.emit_load_const(folded); + } else { + for elt in elts { + self.compile_expression(elt)?; + } + emit!( + self, + Instruction::BuildTuple { + count: u32::try_from(elts.len()).expect("too many elements"), + } + ); + } + return Ok(()); + } + + self.compile_expression(iter) + } + fn forbidden_name(&mut self, name: &str, ctx: NameUsage) -> CompileResult { if ctx == NameUsage::Store && name == "__debug__" { return Err(self.error(CodegenErrorType::Assign("__debug__"))); @@ -6201,9 +6213,17 @@ impl Compiler { // Keep the subject around for extracting elements. pc.on_top += 1; for (i, pattern) in patterns.iter().enumerate() { - // if pattern.is_wildcard() { - // continue; - // } + let is_true_wildcard = matches!( + pattern, + ast::Pattern::MatchAs(ast::PatternMatchAs { + pattern: None, + name: None, + .. + }) + ); + if is_true_wildcard { + continue; + } if i == star { // This must be a starred wildcard. // assert!(pattern.is_star_wildcard()); @@ -6711,6 +6731,7 @@ impl Compiler { pc.stores.insert(insert_pos + j, elem); } // Also perform the same rotation on the evaluation stack. + self.set_source_range(alt.range()); for _ in 0..=i_stores { self.pattern_helper_rotate(i_control + 1)?; } @@ -6719,7 +6740,9 @@ impl Compiler { } } // Emit a jump to the common end label and reset any failure jump targets. + self.set_source_range(alt.range()); emit!(self, PseudoInstruction::Jump { delta: end }); + self.set_source_range(alt.range()); self.emit_and_reset_fail_pop(pc)?; } @@ -6731,6 +6754,7 @@ impl Compiler { // In Rust, old_pc is a local clone, so we need not worry about that. // No alternative matched: pop the subject and fail. + self.set_source_range(p.range()); emit!(self, Instruction::PopTop); self.jump_to_fail_pop(pc, JumpOp::Jump)?; @@ -6742,6 +6766,7 @@ impl Compiler { let n_rots = n_stores + 1 + pc.on_top + pc.stores.len(); for i in 0..n_stores { // Rotate the capture to its proper place. + self.set_source_range(p.range()); self.pattern_helper_rotate(n_rots)?; let name = &control.as_ref().unwrap()[i]; // Check for duplicate binding. @@ -6753,6 +6778,7 @@ impl Compiler { // Old context and control will be dropped automatically. // Finally, pop the copy of the subject. + self.set_source_range(p.range()); emit!(self, Instruction::PopTop); Ok(()) } @@ -6841,7 +6867,9 @@ impl Compiler { p: &ast::PatternMatchValue, pc: &mut PatternContext, ) -> CompileResult<()> { - // TODO: ensure literal or attribute lookup + // Match CPython codegen_pattern_value(): compare, then normalize to bool + // before the fail jump. Late IR folding will collapse COMPARE_OP+TO_BOOL + // into COMPARE_OP bool(...) when applicable. self.compile_expression(&p.value)?; emit!( self, @@ -6849,7 +6877,7 @@ impl Compiler { opname: bytecode::ComparisonOperator::Equal } ); - // emit!(self, Instruction::ToBool); + emit!(self, Instruction::ToBool); self.jump_to_fail_pop(pc, JumpOp::PopJumpIfFalse)?; Ok(()) } @@ -6877,7 +6905,9 @@ impl Compiler { pattern_type: &ast::Pattern, pattern_context: &mut PatternContext, ) -> CompileResult<()> { - match &pattern_type { + let prev_source_range = self.current_source_range; + self.set_source_range(pattern_type.range()); + let result = match &pattern_type { ast::Pattern::MatchValue(pattern_type) => { self.compile_pattern_value(pattern_type, pattern_context) } @@ -6902,7 +6932,9 @@ impl Compiler { ast::Pattern::MatchOr(pattern_type) => { self.compile_pattern_or(pattern_type, pattern_context) } - } + }; + self.set_source_range(prev_source_range); + result } fn compile_match_inner( @@ -6911,12 +6943,22 @@ impl Compiler { cases: &[ast::MatchCase], pattern_context: &mut PatternContext, ) -> CompileResult<()> { + fn is_trailing_wildcard_default(pattern: &ast::Pattern) -> bool { + match pattern { + ast::Pattern::MatchAs(match_as) => { + match_as.pattern.is_none() && match_as.name.is_none() + } + _ => false, + } + } + self.compile_expression(subject)?; let end = self.new_block(); let num_cases = cases.len(); assert!(num_cases > 0); - let has_default = cases.iter().last().unwrap().pattern.is_match_star() && num_cases > 1; + let has_default = + num_cases > 1 && is_trailing_wildcard_default(&cases.last().unwrap().pattern); let case_count = num_cases - if has_default { 1 } else { 0 }; for (i, m) in cases.iter().enumerate().take(case_count) { @@ -6926,36 +6968,41 @@ impl Compiler { } pattern_context.stores = Vec::with_capacity(1); - pattern_context.allow_irrefutable = m.guard.is_some() || i == case_count - 1; + pattern_context.allow_irrefutable = m.guard.is_some() || i == num_cases - 1; pattern_context.fail_pop.clear(); pattern_context.on_top = 0; self.compile_pattern(&m.pattern, pattern_context)?; assert_eq!(pattern_context.on_top, 0); + self.set_source_range(m.pattern.range()); for name in &pattern_context.stores { self.compile_name(name, NameUsage::Store)?; } if let Some(ref guard) = m.guard { self.ensure_fail_pop(pattern_context, 0)?; - // Compile the guard expression - self.compile_expression(guard)?; - emit!(self, Instruction::ToBool); - emit!( - self, - Instruction::PopJumpIfFalse { - delta: pattern_context.fail_pop[0] - } - ); + self.compile_jump_if_inner( + guard, + false, + pattern_context.fail_pop[0], + Some(m.pattern.range()), + )?; } if i != case_count - 1 { + if let Some(first_stmt) = m.body.first() { + self.set_source_range(first_stmt.range()); + } + if matches!(m.pattern, ast::Pattern::MatchOr(_)) { + emit!(self, Instruction::Nop); + } emit!(self, Instruction::PopTop); } self.compile_statements(&m.body)?; emit!(self, PseudoInstruction::Jump { delta: end }); + self.set_source_range(m.pattern.range()); self.emit_and_reset_fail_pop(pattern_context)?; } @@ -6963,15 +7010,11 @@ impl Compiler { let m = &cases[num_cases - 1]; if num_cases == 1 { emit!(self, Instruction::PopTop); - } else { + } else if m.guard.is_none() { emit!(self, Instruction::Nop); } if let Some(ref guard) = m.guard { - // Compile guard and jump to end if false - self.compile_expression(guard)?; - emit!(self, Instruction::Copy { i: 1 }); - emit!(self, Instruction::PopJumpIfFalse { delta: end }); - emit!(self, Instruction::PopTop); + self.compile_jump_if(guard, false, end)?; } self.compile_statements(&m.body)?; } @@ -7210,6 +7253,37 @@ impl Compiler { Ok(()) } + fn compile_check_annotation_expression(&mut self, expression: &ast::Expr) -> CompileResult<()> { + self.compile_expression(expression)?; + emit!(self, Instruction::PopTop); + Ok(()) + } + + fn compile_check_annotation_subscript(&mut self, expression: &ast::Expr) -> CompileResult<()> { + match expression { + ast::Expr::Slice(ast::ExprSlice { + lower, upper, step, .. + }) => { + if let Some(lower) = lower { + self.compile_check_annotation_expression(lower)?; + } + if let Some(upper) = upper { + self.compile_check_annotation_expression(upper)?; + } + if let Some(step) = step { + self.compile_check_annotation_expression(step)?; + } + } + ast::Expr::Tuple(ast::ExprTuple { elts, .. }) => { + for element in elts { + self.compile_check_annotation_subscript(element)?; + } + } + _ => self.compile_check_annotation_expression(expression)?, + } + Ok(()) + } + fn compile_annotated_assign( &mut self, target: &ast::Expr, @@ -7276,6 +7350,19 @@ impl Compiler { } } + if value.is_none() { + match target { + ast::Expr::Attribute(ast::ExprAttribute { value, .. }) => { + self.compile_check_annotation_expression(value)?; + } + ast::Expr::Subscript(ast::ExprSubscript { value, slice, .. }) => { + self.compile_check_annotation_expression(value)?; + self.compile_check_annotation_subscript(slice)?; + } + _ => {} + } + } + Ok(()) } @@ -7452,14 +7539,15 @@ impl Compiler { /// /// The idea is to jump to a label if the expression is either true or false /// (indicated by the condition parameter). - fn compile_jump_if( + fn compile_jump_if_inner( &mut self, expression: &ast::Expr, condition: bool, target_block: BlockIdx, + source_range: Option, ) -> CompileResult<()> { let prev_source_range = self.current_source_range; - self.set_source_range(expression.range()); + self.set_source_range(source_range.unwrap_or_else(|| expression.range())); // Compile expression for test, and jump to label if false let result = match &expression { @@ -7473,16 +7561,26 @@ impl Compiler { // If any of the values is false, we can short-circuit. for value in values { - self.compile_jump_if(value, false, end_block)?; + 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(last_value, true, target_block)?; + 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(value, false, target_block)?; + self.compile_jump_if_inner( + value, + false, + target_block, + source_range, + )?; } } } @@ -7490,7 +7588,12 @@ impl Compiler { if condition { // If any of the values is true. for value in values { - self.compile_jump_if(value, true, target_block)?; + self.compile_jump_if_inner( + value, + true, + target_block, + source_range, + )?; } } else { // If all of the values are false. @@ -7499,11 +7602,16 @@ impl Compiler { // If any value is true, we can short-circuit: for value in values { - self.compile_jump_if(value, true, end_block)?; + self.compile_jump_if_inner(value, true, end_block, source_range)?; } // It all depends upon the last value now! - self.compile_jump_if(last_value, false, target_block)?; + self.compile_jump_if_inner( + last_value, + false, + target_block, + source_range, + )?; self.switch_to_block(end_block); } } @@ -7514,7 +7622,7 @@ impl Compiler { op: ast::UnaryOp::Not, operand, .. - }) => self.compile_jump_if(operand, !condition, target_block), + }) => self.compile_jump_if_inner(operand, !condition, target_block, source_range), ast::Expr::Compare(ast::ExprCompare { left, ops, @@ -7585,6 +7693,15 @@ impl Compiler { result } + fn compile_jump_if( + &mut self, + expression: &ast::Expr, + condition: bool, + target_block: BlockIdx, + ) -> CompileResult<()> { + self.compile_jump_if_inner(expression, condition, target_block, None) + } + /// 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<()> { @@ -7911,10 +8028,12 @@ impl Compiler { match op { ast::BoolOp::Or if is_truthy => { self.emit_load_const(last_constant.expect("missing boolop constant")); + self.mark_last_instruction_folded_from_nonliteral_expr(); return Ok(()); } ast::BoolOp::And if !is_truthy => { self.emit_load_const(last_constant.expect("missing boolop constant")); + self.mark_last_instruction_folded_from_nonliteral_expr(); return Ok(()); } ast::BoolOp::Or | ast::BoolOp::And => { @@ -7925,6 +8044,7 @@ impl Compiler { if simplified_prefix == values.len() { self.emit_load_const(last_constant.expect("missing folded boolop constant")); + self.mark_last_instruction_folded_from_nonliteral_expr(); return Ok(()); } if simplified_prefix > 0 { @@ -9128,7 +9248,7 @@ impl Compiler { emit!(self, Instruction::LoadFast { var_num: arg0 }); } else { // Evaluate iterated item: - self.compile_expression(&generator.iter)?; + self.compile_for_iterable_expression(&generator.iter, generator.is_async)?; // Get iterator / turn item into an iterator if generator.is_async { @@ -9231,7 +9351,7 @@ impl Compiler { self.make_closure(code, bytecode::MakeFunctionFlags::new())?; // Evaluate iterated item: - self.compile_expression(&outermost.iter)?; + self.compile_for_iterable_expression(&outermost.iter, outermost.is_async)?; self.symbol_table_stack .last_mut() .expect("no current symbol table") @@ -9270,7 +9390,10 @@ impl Compiler { // nested scopes (e.g. lambdas) whose sub_tables sit at the current // position in the parent's list. Those must be consumed before we // splice in the comprehension's own children. - self.compile_expression(&generators[0].iter)?; + self.compile_for_iterable_expression( + &generators[0].iter, + has_async && generators[0].is_async, + )?; self.symbol_table_stack .last_mut() .expect("no current symbol table") @@ -9438,7 +9561,7 @@ impl Compiler { let after_block = self.new_block(); if i > 0 { - self.compile_expression(&generator.iter)?; + self.compile_for_iterable_expression(&generator.iter, generator.is_async)?; if generator.is_async { emit!(self, Instruction::GetAIter); } else { @@ -9598,11 +9721,18 @@ impl Compiler { location, end_location, except_handler, + folded_from_nonliteral_expr: false, lineno_override: None, cache_entries: 0, }); } + fn mark_last_instruction_folded_from_nonliteral_expr(&mut self) { + if let Some(info) = self.current_block().instructions.last_mut() { + info.folded_from_nonliteral_expr = true; + } + } + /// Mark the last emitted instruction as having no source location. /// Prevents it from triggering LINE events in sys.monitoring. fn set_no_location(&mut self) { @@ -9796,6 +9926,27 @@ impl Compiler { } ConstantData::Tuple { elements } } + ast::Expr::UnaryOp(ast::ExprUnaryOp { op, operand, .. }) => { + let Some(constant) = self.try_fold_constant_expr(operand)? else { + return Ok(None); + }; + match (op, constant) { + (ast::UnaryOp::UAdd, value) => value, + (ast::UnaryOp::USub, ConstantData::Integer { value }) => { + ConstantData::Integer { value: -value } + } + (ast::UnaryOp::USub, ConstantData::Float { value }) => { + ConstantData::Float { value: -value } + } + (ast::UnaryOp::USub, ConstantData::Complex { value }) => { + ConstantData::Complex { value: -value } + } + (ast::UnaryOp::Invert, ConstantData::Integer { value }) => { + ConstantData::Integer { value: !value } + } + _ => return Ok(None), + } + } ast::Expr::BoolOp(ast::ExprBoolOp { op, values, .. }) => { let mut constants = Vec::with_capacity(values.len()); for value in values { @@ -10030,6 +10181,11 @@ impl Compiler { let loop_block = code.fblock[loop_idx].fb_block; let exit_block = code.fblock[loop_idx].fb_exit; + let prev_source_range = self.current_source_range; + self.set_source_range(range); + emit!(self, Instruction::Nop); + self.set_source_range(prev_source_range); + // Collect the fblocks we need to unwind through, from top down to (but not including) the loop #[derive(Clone)] enum UnwindAction { @@ -11099,6 +11255,68 @@ mod tests { stack_top.debug_late_cfg_trace().unwrap() } + fn compile_single_function_late_cfg_trace( + source: &str, + function_name: &str, + ) -> Vec<(String, String)> { + 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(); + let ast = parsed.into_syntax(); + let ast = match ast { + ruff_python_ast::Mod::Module(stmts) => stmts, + _ => unreachable!(), + }; + let symbol_table = SymbolTable::scan_program(&ast, source_file.clone()) + .map_err(|e| e.into_codegen_error(source_file.name().to_owned())) + .unwrap(); + let function = ast + .body + .iter() + .find_map(|stmt| match stmt { + ast::Stmt::FunctionDef(f) if f.name.as_str() == function_name => Some(f), + _ => None, + }) + .unwrap_or_else(|| panic!("missing function {function_name}")); + + 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 + .current_code_info() + .flags + .set(bytecode::CodeFlags::COROUTINE, false); + + let prev_ctx = compiler.ctx; + compiler.ctx = CompileContext { + loop_data: None, + in_class: prev_ctx.in_class, + func: FunctionContext::Function, + in_async_scope: false, + }; + compiler.set_qualname(); + compiler.compile_statements(&function.body).unwrap(); + match function.body.last() { + Some(ast::Stmt::Return(_)) => {} + _ => compiler.emit_return_const(ConstantData::None), + } + if compiler.current_code_info().metadata.consts.is_empty() { + compiler.arg_constant(ConstantData::None); + } + + let _table = compiler.pop_symbol_table(); + let stack_top = compiler.code_stack.pop().unwrap(); + stack_top.debug_late_cfg_trace().unwrap() + } + fn find_code<'a>(code: &'a CodeObject, name: &str) -> Option<&'a CodeObject> { if code.obj_name == name { return Some(code); @@ -11170,61 +11388,231 @@ except AssertionError as e: } #[test] - fn test_if_ands() { - assert_dis_snapshot!(compile_exec( + fn test_trace_for_unpack_list_literal() { + let trace = compile_exec_late_cfg_trace( "\ -if True and False and False: - pass -" - )); +result = [] +for x, in [(1,), (2,), (3,)]: + result.append(x) +", + ); + for (stage, dump) in trace { + eprintln!("=== {stage} ===\n{dump}"); + } } #[test] - fn test_if_mixed() { - assert_dis_snapshot!(compile_exec( + fn test_trace_break_in_finally_function() { + let trace = compile_single_function_late_cfg_trace( "\ -if (True and False) or (False and True): - pass -" - )); +def f(self): + count = 0 + while count < 2: + count += 1 + try: + pass + finally: + break + self.assertEqual(count, 1) +", + "f", + ); + for (stage, dump) in trace { + eprintln!("=== {stage} ===\n{dump}"); + } } #[test] - fn test_nested_bool_op() { - assert_dis_snapshot!(compile_exec( + fn test_trace_constant_false_elif_chain() { + let trace = compile_exec_late_cfg_trace( "\ -x = Test() and False or False -" - )); +if 0: pass +elif 0: pass +elif 0: pass +elif 0: pass +else: pass +", + ); + for (stage, dump) in trace { + eprintln!("=== {stage} ===\n{dump}"); + } } #[test] - fn test_const_bool_not_op() { - assert_dis_snapshot!(compile_exec_optimized( + fn test_trace_multi_pass_suite() { + let trace = compile_exec_late_cfg_trace( "\ -x = not True -" - )); +if 1: + # + # + # + pass + pass + # + pass + # +", + ); + for (stage, dump) in trace { + eprintln!("=== {stage} ===\n{dump}"); + } } #[test] - fn test_plain_constant_bool_op_folds_to_selected_operand() { - let code = compile_exec( + fn test_trace_single_compare_if() { + let trace = compile_exec_late_cfg_trace( "\ -x = 1 or 2 or 3 +if 1 == 1: + pass ", ); - let ops: Vec<_> = code - .instructions - .iter() - .map(|unit| unit.op) - .filter(|op| !matches!(op, Instruction::Cache)) - .collect(); - let folded_small_int = code.instructions.iter().any(|unit| { - matches!( - unit.op, - Instruction::LoadSmallInt { i } - if i.get(OpArg::new(u32::from(u8::from(unit.arg)))) == 1 + for (stage, dump) in trace { + eprintln!("=== {stage} ===\n{dump}"); + } + } + + #[test] + fn test_trace_comparison_suite() { + let trace = compile_exec_late_cfg_trace( + "\ +if 1: pass +x = (1 == 1) +if 1 == 1: pass +if 1 != 1: pass +if 1 < 1: pass +if 1 > 1: pass +if 1 <= 1: pass +if 1 >= 1: pass +if x is x: pass +if x is not x: pass +if 1 in (): pass +if 1 not in (): pass +", + ); + for (stage, dump) in trace { + eprintln!("=== {stage} ===\n{dump}"); + } + } + + #[test] + fn test_trace_if_for_except_layout() { + let trace = compile_exec_late_cfg_trace( + "\ +from sys import maxsize +if maxsize == 2147483647: + for s in ('2147483648', '0o40000000000', '0x100000000', '0b10000000000000000000000000000000'): + try: + x = eval(s) + except OverflowError: + fail(\"OverflowError on huge integer literal %r\" % s) +elif maxsize == 9223372036854775807: + pass +", + ); + for (stage, dump) in trace { + eprintln!("=== {stage} ===\n{dump}"); + } + } + + #[test] + fn test_break_in_finally_tail_loads_borrow_through_empty_fallthrough_block() { + let code = compile_exec( + "\ +def f(self): + count = 0 + while count < 2: + count += 1 + try: + pass + finally: + break + self.assertEqual(count, 1) +", + ); + let code = find_code(&code, "f").unwrap(); + let ops: Vec<_> = code + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + assert!( + ops.windows(5).any(|window| { + matches!( + window, + [ + Instruction::LoadFastBorrow { .. }, + Instruction::LoadAttr { .. }, + Instruction::LoadFastBorrow { .. }, + Instruction::LoadSmallInt { .. }, + Instruction::Call { .. } + ] + ) + }), + "{:?}", + code.instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + } + + #[test] + fn test_if_ands() { + assert_dis_snapshot!(compile_exec( + "\ +if True and False and False: + pass +" + )); + } + + #[test] + fn test_if_mixed() { + assert_dis_snapshot!(compile_exec( + "\ +if (True and False) or (False and True): + pass +" + )); + } + + #[test] + fn test_nested_bool_op() { + assert_dis_snapshot!(compile_exec( + "\ +x = Test() and False or False +" + )); + } + + #[test] + fn test_const_bool_not_op() { + assert_dis_snapshot!(compile_exec_optimized( + "\ +x = not True +" + )); + } + + #[test] + fn test_plain_constant_bool_op_folds_to_selected_operand() { + let code = compile_exec( + "\ +x = 1 or 2 or 3 +", + ); + let ops: Vec<_> = code + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + let folded_small_int = code.instructions.iter().any(|unit| { + matches!( + unit.op, + Instruction::LoadSmallInt { i } + if i.get(OpArg::new(u32::from(u8::from(unit.arg)))) == 1 ) }); let folded_const_one = code @@ -11858,6 +12246,97 @@ def f(parts): ); } + #[test] + fn test_for_exit_before_elif_does_not_leave_line_anchor_nop() { + let code = compile_exec( + "\ +from sys import maxsize +if maxsize == 2147483647: + for s in ('2147483648', '0o40000000000', '0x100000000', '0b10000000000000000000000000000000'): + try: + x = eval(s) + except OverflowError: + fail('OverflowError on huge integer literal %r' % s) +elif maxsize == 9223372036854775807: + pass +", + ); + let ops: Vec<_> = code + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + assert!( + ops.windows(4).any(|window| { + matches!( + window, + [ + Instruction::EndFor, + Instruction::PopIter, + Instruction::LoadConst { .. }, + Instruction::ReturnValue, + ] + ) + }), + "expected for-exit epilogue without extra NOP, got ops={ops:?}" + ); + assert!( + !ops.windows(4).any(|window| { + matches!( + window, + [ + Instruction::EndFor, + Instruction::PopIter, + Instruction::Nop, + Instruction::LoadConst { .. }, + ] + ) + }), + "unexpected line-anchor NOP before for-exit epilogue, got ops={ops:?}" + ); + } + + #[test] + fn test_break_in_finally_after_return_keeps_load_fast_check_for_loop_locals() { + let code = compile_exec( + "\ +def g2(x): + for count in [0, 1]: + for count2 in [10, 20]: + try: + return count + count2 + finally: + if x: + break + return 'end', count, count2 +", + ); + let g2 = find_code(&code, "g2").expect("missing g2 code"); + let ops: Vec<_> = g2 + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + assert!( + ops.windows(4).any(|window| { + matches!( + window, + [ + Instruction::LoadConst { .. }, + Instruction::LoadFastCheck { .. }, + Instruction::LoadFastCheck { .. }, + Instruction::BuildTuple { .. }, + ] + ) + }), + "expected LOAD_FAST_CHECK pair for after-return loop locals, got ops={ops:?}" + ); + } + #[test] fn test_assert_without_message_raises_class_directly() { let code = compile_exec( @@ -11932,6 +12411,62 @@ def f(x, y): assert_eq!(u8::from(call_arg), 0); } + #[test] + fn test_bare_function_annotations_check_attribute_and_subscript_expressions() { + assert_dis_snapshot!(compile_exec( + "\ +def f(one: int): + int.new_attr: int + [list][0].new_attr: [int, str] + my_lst = [1] + my_lst[one]: int + return my_lst +" + )); + } + + #[test] + fn test_non_simple_bare_name_annotation_does_not_create_local_binding() { + let code = compile_exec( + "\ +def f2bad(): + (no_such_global): int + print(no_such_global) +", + ); + let f = find_code(&code, "f2bad").expect("missing f2bad code"); + assert!( + f.instructions + .iter() + .any(|unit| matches!(unit.op, Instruction::LoadGlobal { .. })), + "expected LOAD_GLOBAL for non-simple bare annotated name, got ops={:?}", + f.instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + assert!( + !f.instructions + .iter() + .any(|unit| matches!(unit.op, Instruction::LoadFastCheck { .. })), + "non-simple bare annotated name should not become a local binding, got ops={:?}", + f.instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + } + + #[test] + fn test_constant_true_if_pass_keeps_line_anchor_nop() { + assert_dis_snapshot!(compile_exec( + "\ +if 1: + pass +" + )); + } + #[test] fn test_negative_constant_binop_folds_after_unary_folding() { let code = compile_exec( @@ -12039,6 +12574,175 @@ def f(self): ); } + #[test] + fn test_match_guard_capture_uses_store_fast_load_fast() { + let code = compile_exec( + "\ +def f(): + match 0: + case x if x: + z = 0 +", + ); + let f = find_code(&code, "f").expect("missing function code"); + assert!( + f.instructions + .iter() + .any(|unit| matches!(unit.op, Instruction::StoreFastLoadFast { .. })), + "expected STORE_FAST_LOAD_FAST in match guard capture path, got ops={:?}", + f.instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + } + + #[test] + fn test_match_nested_capture_uses_store_fast_store_fast() { + let code = compile_exec( + "\ +def f(x): + match x: + case ((0 as w) as z): + return w, z +", + ); + let f = find_code(&code, "f").expect("missing function code"); + assert!( + f.instructions + .iter() + .any(|unit| matches!(unit.op, Instruction::StoreFastStoreFast { .. })), + "expected STORE_FAST_STORE_FAST in nested match capture path, got ops={:?}", + f.instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + } + + #[test] + fn test_match_value_real_zero_minus_zero_complex_folds_to_negative_zero_imag() { + let code = compile_exec( + "\ +def f(x): + match x: + case 0 - 0j: + return 0 +", + ); + let f = find_code(&code, "f").expect("missing function code"); + assert!( + f.constants.iter().any(|constant| matches!( + constant, + ConstantData::Complex { value } + if value.re == 0.0 && value.im == 0.0 && value.im.is_sign_negative() + )), + "expected folded -0j constant in match value" + ); + } + + #[test] + fn test_match_or_uses_shared_success_block() { + let code = compile_exec( + "\ +def http_error(status): + match status: + case 400: + return 'Bad request' + case 401 | 403 | 404: + return 'Not allowed' + case 418: + return 'I am a teapot' +", + ); + let f = find_code(&code, "http_error").expect("missing http_error code"); + let ops: Vec<_> = f + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + let jump_positions: Vec<_> = ops + .iter() + .enumerate() + .filter_map(|(i, op)| matches!(op, Instruction::JumpForward { .. }).then_some(i)) + .collect(); + + assert!( + jump_positions.len() >= 4, + "expected shared-success JumpForward ops in OR pattern, got ops={ops:?}" + ); + + let first_pop_top_pair = ops + .windows(2) + .position(|window| matches!(window, [Instruction::PopTop, Instruction::PopTop])) + .expect("missing POP_TOP/POP_TOP success cleanup"); + + assert!( + jump_positions + .iter() + .take(3) + .all(|&idx| idx < first_pop_top_pair), + "expected OR-alternative jumps before shared success cleanup, got ops={ops:?}" + ); + } + + #[test] + fn test_match_mapping_attribute_key_keeps_plain_load_fast() { + let code = compile_exec( + "\ +def f(self): + class Keys: + KEY = 'a' + x = {'a': 0, 'b': 1} + with self.assertRaises(ValueError): + match x: + case {Keys.KEY: y, 'a': z}: + w = 0 +", + ); + let f = find_code(&code, "f").expect("missing function code"); + let key_load_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() == "KEY" + } + _ => false, + }) + .expect("missing Keys.KEY attribute load"); + let prev = f.instructions[key_load_idx - 1].op; + assert!( + matches!(prev, Instruction::LoadFast { .. }), + "expected plain LOAD_FAST before Keys.KEY mapping key, got ops={:?}", + f.instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + } + + #[test] + #[ignore = "debug trace for sequence star-wildcard pattern layout"] + fn test_debug_trace_match_sequence_star_wildcard_layout() { + let trace = compile_single_function_late_cfg_trace( + "\ +def f(w): + match w: + case [x, *_, y]: + z = 0 + return x, y, z +", + "f", + ); + for (stage, dump) in trace { + eprintln!("=== {stage} ===\n{dump}"); + } + } + #[test] fn test_genexpr_true_filter_omits_bool_scaffolding() { let code = compile_exec( @@ -12119,6 +12823,38 @@ class C: ); } + #[test] + fn test_plain_super_call_keeps_class_freevar() { + let code = compile_exec( + "\ +class A: + pass + +class B(A): + def method(self): + return super() +", + ); + let method = find_code(&code, "method").expect("missing method code"); + assert!( + method.freevars.iter().any(|name| name == "__class__"), + "plain super() must keep __class__ freevar, got freevars={:?}", + method.freevars + ); + assert!( + method + .instructions + .iter() + .any(|unit| matches!(unit.op, Instruction::CopyFreeVars { .. })), + "plain super() must keep COPY_FREE_VARS prelude, got ops={:?}", + method + .instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + } + #[test] fn test_chained_compare_jump_uses_single_cleanup_copy() { let code = compile_exec( @@ -12221,6 +12957,55 @@ def f(): ); } + #[test] + fn test_named_except_cleanup_keeps_jump_over_cleanup_and_next_try() { + let code = compile_exec( + r#" +def f(self): + try: + assert 0, 'msg' + except AssertionError as e: + self.assertEqual(e.args[0], 'msg') + else: + self.fail("AssertionError not raised by assert 0") + + try: + assert False + except AssertionError as e: + self.assertEqual(len(e.args), 0) + else: + self.fail("AssertionError not raised by 'assert False'") +"#, + ); + 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_except = ops + .iter() + .position(|op| matches!(op, Instruction::PopExcept)) + .expect("missing POP_EXCEPT"); + let window = &ops[first_pop_except..(first_pop_except + 6).min(ops.len())]; + assert!( + matches!( + window, + [ + Instruction::PopExcept, + Instruction::LoadConst { .. }, + Instruction::StoreName { .. } | Instruction::StoreFast { .. }, + Instruction::DeleteName { .. } | Instruction::DeleteFast { .. }, + Instruction::JumpForward { .. }, + .. + ] + ), + "expected named except cleanup to jump over cleanup reraise block, got ops={window:?}" + ); + } + #[test] fn test_bare_except_deopts_post_handler_load_fast_borrow() { let code = compile_exec( @@ -12527,6 +13312,148 @@ def f(): ))); } + #[test] + fn test_starred_tuple_iterable_drops_list_to_tuple_before_get_iter() { + let code = compile_exec( + "\ +def f(a, b, c): + for x in *a, *b, *c: + pass +", + ); + let f = find_code(&code, "f").expect("missing function code"); + + assert!( + !has_intrinsic_1(f, IntrinsicFunction1::ListToTuple), + "LIST_TO_TUPLE should be removed before GET_ITER in for-iterable context" + ); + assert!( + f.instructions + .iter() + .any(|unit| matches!(unit.op, Instruction::GetIter)), + "expected GET_ITER in for loop" + ); + } + + #[test] + 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_constant_false_while_else_deopts_post_else_borrows() { + let code = compile_exec( + "\ +def f(self): + x = 0 + while 0: + x = 1 + else: + x = 2 + self.assertEqual(x, 2) +", + ); + 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 assert_idx = ops + .iter() + .position(|op| matches!(op, Instruction::LoadAttr { .. })) + .expect("missing assertEqual call"); + let window = &ops[assert_idx.saturating_sub(1)..(assert_idx + 3).min(ops.len())]; + assert!( + matches!( + window, + [ + Instruction::LoadFast { .. }, + Instruction::LoadAttr { .. }, + Instruction::LoadFast { .. }, + .. + ] + ), + "expected post-else assertEqual call to use plain LOAD_FAST, got ops={window:?}" + ); + } + #[test] fn test_single_unpack_assignment_disables_constant_collection_folding() { let code = compile_exec("a, b, c = 1, 2, 3\n"); @@ -12650,6 +13577,41 @@ def f(): ); } + #[test] + fn test_constant_unary_positive_and_invert_fold() { + let code = compile_exec("x = +1\nx = ~1\n"); + + assert!( + !code.instructions.iter().any(|unit| { + matches!( + unit.op, + Instruction::CallIntrinsic1 { .. } | Instruction::UnaryInvert + ) + }), + "constant unary ops should fold away, got ops={:?}", + code.instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + } + + #[test] + fn test_bool_invert_is_not_const_folded() { + let code = compile_exec("x = ~True\n"); + + assert!( + code.instructions + .iter() + .any(|unit| matches!(unit.op, Instruction::UnaryInvert)), + "~bool should remain unfurled to match CPython, got ops={:?}", + code.instructions + .iter() + .map(|unit| unit.op) + .collect::>() + ); + } + #[test] fn test_optimized_assert_preserves_nested_scope_order() { compile_exec_optimized( diff --git a/crates/codegen/src/ir.rs b/crates/codegen/src/ir.rs index 4e9db27232b..01e6971f65b 100644 --- a/crates/codegen/src/ir.rs +++ b/crates/codegen/src/ir.rs @@ -3,6 +3,7 @@ use core::ops; use crate::{IndexMap, IndexSet, error::InternalError}; use malachite_bigint::BigInt; +use num_complex::Complex; use num_traits::{ToPrimitive, Zero}; use rustpython_compiler_core::{ @@ -10,8 +11,9 @@ use rustpython_compiler_core::{ bytecode::{ AnyInstruction, AnyOpcode, Arg, CO_FAST_CELL, CO_FAST_FREE, CO_FAST_HIDDEN, CO_FAST_LOCAL, CodeFlags, CodeObject, CodeUnit, CodeUnits, ConstantData, ExceptionTableEntry, - InstrDisplayContext, Instruction, InstructionMetadata, Label, OpArg, Opcode, - PseudoInstruction, PseudoOpcode, PyCodeLocationInfoKind, encode_exception_table, oparg, + InstrDisplayContext, Instruction, InstructionMetadata, IntrinsicFunction1, Label, OpArg, + Opcode, PseudoInstruction, PseudoOpcode, PyCodeLocationInfoKind, encode_exception_table, + oparg, }, varint::{write_signed_varint, write_varint}, }; @@ -108,6 +110,7 @@ pub struct InstructionInfo { pub location: SourceLocation, pub end_location: SourceLocation, pub except_handler: Option, + pub folded_from_nonliteral_expr: bool, /// Override line number for linetable (e.g., line 0 for module RESUME) pub lineno_override: Option, /// Number of CACHE code units emitted after this instruction @@ -129,6 +132,7 @@ fn set_to_nop(info: &mut InstructionInfo) { info.instr = Instruction::Nop.into(); info.arg = OpArg::new(0); info.target = BlockIdx::NULL; + info.folded_from_nonliteral_expr = false; info.cache_entries = 0; } @@ -206,19 +210,14 @@ impl CodeInfo { self.splice_annotations_blocks(); // Constant folding passes self.fold_binop_constants(); - self.remove_nops(); - self.fold_unary_negative(); - self.remove_nops(); // remove UNARY_NEGATIVE NOPs before re-folding binops + self.fold_unary_constants(); self.fold_binop_constants(); // re-run after unary folding: -1 + 2 → 1 - self.remove_nops(); // remove NOPs so tuple/list/set see contiguous LOADs self.fold_tuple_constants(); self.fold_list_constants(); self.fold_set_constants(); - self.remove_nops(); // remove NOPs from collection folding self.fold_const_iterable_for_iter(); self.convert_to_load_small_int(); self.remove_unused_consts(); - self.remove_nops(); // DCE always runs (removes dead code after terminal instructions) self.dce(); @@ -257,9 +256,11 @@ impl CodeInfo { self.dce(); // re-run within-block DCE after normalize_jumps creates new instructions self.eliminate_unreachable_blocks(); resolve_line_numbers(&mut self.blocks); + redirect_empty_block_targets(&mut self.blocks); duplicate_end_returns(&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(); remove_redundant_nops_and_jumps(&mut self.blocks); // Some jump-only blocks only appear after late CFG cleanup. Thread them // once more so loop backedges stay direct instead of becoming @@ -270,6 +271,10 @@ impl CodeInfo { reorder_jump_over_exception_cleanup_blocks(&mut self.blocks); self.eliminate_unreachable_blocks(); remove_redundant_nops_and_jumps(&mut self.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, @@ -279,6 +284,7 @@ impl CodeInfo { // Refresh exceptional block flags before optimize_load_fast_borrow so // borrow loads are not introduced into exception-handler paths. mark_except_handlers(&mut self.blocks); + redirect_empty_block_targets(&mut self.blocks); // CPython's optimize_load_fast runs with block start depths already known. // Compute them here so the abstract stack simulation can use the real // CFG entry depth for each block. @@ -287,14 +293,15 @@ impl CodeInfo { // calculation but before optimize_load_fast. convert_pseudo_ops(&mut self.blocks, &cellfixedoffsets); self.compute_load_fast_start_depths(); - self.propagate_disable_load_fast_borrow(); // optimize_load_fast: after normalize_jumps self.optimize_load_fast_borrow(); 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.optimize_load_global_push_null(); - self.remove_redundant_const_pop_top_pairs(); self.reorder_entry_prefix_cell_setup(); self.remove_unused_consts(); @@ -363,28 +370,25 @@ impl CodeInfo { let mut remove = false; if matches!(instr.instr.real(), Some(Instruction::Nop)) { - // Remove location-less NOPs. if lineno < 0 || prev_lineno == lineno { remove = true; - } - // Remove if the next instruction has same line or no line. - else if src < src_instructions.len() - 1 { - let next_lineno = - src_instructions[src + 1] + } else if src < src_instructions.len() - 1 { + if src_instructions[src + 1].folded_from_nonliteral_expr { + remove = true; + } else { + let next_lineno = src_instructions[src + 1] .lineno_override .unwrap_or_else(|| { src_instructions[src + 1].location.line.get() as i32 }); - if next_lineno == lineno { - remove = true; - } else if next_lineno < 0 { - src_instructions[src + 1].lineno_override = Some(lineno); - remove = true; + if next_lineno == lineno { + remove = true; + } else if next_lineno < 0 { + src_instructions[src + 1].lineno_override = Some(lineno); + remove = true; + } } - } - // Last instruction in block: compare with first real location - // in the next non-empty block. - else { + } else { let mut next = blocks[bi].next; while next != BlockIdx::NULL && blocks[next.idx()].instructions.is_empty() { next = blocks[next.idx()].next; @@ -758,51 +762,103 @@ impl CodeInfo { } } - /// Fold LOAD_CONST/LOAD_SMALL_INT + UNARY_NEGATIVE → LOAD_CONST (negative value) - fn fold_unary_negative(&mut self) { + fn eval_unary_constant( + operand: &ConstantData, + op: Instruction, + intrinsic: Option, + ) -> Option { + match (operand, op, intrinsic) { + (ConstantData::Integer { value }, Instruction::UnaryNegative, None) => { + Some(ConstantData::Integer { value: -value }) + } + (ConstantData::Float { value }, Instruction::UnaryNegative, None) => { + Some(ConstantData::Float { value: -value }) + } + (ConstantData::Complex { value }, Instruction::UnaryNegative, None) => { + Some(ConstantData::Complex { value: -value }) + } + (ConstantData::Boolean { value }, Instruction::UnaryNegative, None) => { + Some(ConstantData::Integer { + value: BigInt::from(-i32::from(*value)), + }) + } + (ConstantData::Integer { value }, Instruction::UnaryInvert, None) => { + Some(ConstantData::Integer { value: !value }) + } + (ConstantData::Boolean { .. }, Instruction::UnaryInvert, None) => None, + ( + ConstantData::Integer { value }, + Instruction::CallIntrinsic1 { .. }, + Some(oparg::IntrinsicFunction1::UnaryPositive), + ) => Some(ConstantData::Integer { + value: value.clone(), + }), + ( + ConstantData::Float { value }, + Instruction::CallIntrinsic1 { .. }, + Some(oparg::IntrinsicFunction1::UnaryPositive), + ) => Some(ConstantData::Float { value: *value }), + ( + ConstantData::Boolean { value }, + Instruction::CallIntrinsic1 { .. }, + Some(oparg::IntrinsicFunction1::UnaryPositive), + ) => Some(ConstantData::Integer { + value: BigInt::from(i32::from(*value)), + }), + _ => None, + } + } + + /// Fold constant unary operations following CPython fold_const_unaryop(). + fn fold_unary_constants(&mut self) { for block in &mut self.blocks { let mut i = 0; - while i + 1 < block.instructions.len() { - let next = &block.instructions[i + 1]; - let Some(Instruction::UnaryNegative) = next.instr.real() else { + while i < block.instructions.len() { + let instr = &block.instructions[i]; + let (op, intrinsic) = match instr.instr.real() { + Some(Instruction::UnaryNegative) => (Instruction::UnaryNegative, None), + Some(Instruction::UnaryInvert) => (Instruction::UnaryInvert, None), + Some(Instruction::CallIntrinsic1 { func }) + if matches!( + func.get(instr.arg), + oparg::IntrinsicFunction1::UnaryPositive + ) => + { + ( + Instruction::CallIntrinsic1 { + func: Arg::marker(), + }, + Some(func.get(instr.arg)), + ) + } + _ => { + i += 1; + continue; + } + }; + let Some(operand_index) = i + .checked_sub(1) + .and_then(|start| Self::get_const_loading_instr_indices(block, start, 1)) + .and_then(|indices| indices.into_iter().next()) + else { i += 1; continue; }; - let curr = &block.instructions[i]; - let value = match curr.instr.real() { - Some(Instruction::LoadConst { .. }) => { - let idx = u32::from(curr.arg) as usize; - match self.metadata.consts.get_index(idx) { - Some(ConstantData::Integer { value }) => { - Some(ConstantData::Integer { value: -value }) - } - Some(ConstantData::Float { value }) => { - Some(ConstantData::Float { value: -value }) - } - _ => None, - } - } - Some(Instruction::LoadSmallInt { .. }) => { - let v = u32::from(curr.arg) as i32; - Some(ConstantData::Integer { - value: BigInt::from(-v), - }) + let operand = + Self::get_const_value_from(&self.metadata, &block.instructions[operand_index]); + if let Some(operand) = operand + && 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(), } - _ => None, - }; - if let Some(neg_const) = value { - let (const_idx, _) = self.metadata.consts.insert_full(neg_const); - // Replace LOAD_CONST/LOAD_SMALL_INT with new LOAD_CONST - let load_location = block.instructions[i].location; - block.instructions[i].instr = Opcode::LoadConst.into(); + .into(); block.instructions[i].arg = OpArg::new(const_idx as u32); - // Replace UNARY_NEGATIVE with NOP, inheriting the LOAD_CONST - // location so that remove_nops can clean it up - set_to_nop(&mut block.instructions[i + 1]); - block.instructions[i + 1].location = load_location; - block.instructions[i + 1].end_location = block.instructions[i].end_location; - // Skip the NOP, don't re-check - i += 2; + block.instructions[i].folded_from_nonliteral_expr = folded_from_nonliteral_expr; + i = i.saturating_sub(1); } else { i += 1; } @@ -810,6 +866,27 @@ impl CodeInfo { } } + fn get_const_loading_instr_indices( + block: &Block, + mut start: usize, + size: usize, + ) -> Option> { + let mut indices = Vec::with_capacity(size); + loop { + let instr = block.instructions.get(start)?; + if !matches!(instr.instr.real(), Some(Instruction::Nop)) { + Self::get_const_value_from_dummy(instr)?; + indices.push(start); + if indices.len() == size { + break; + } + } + start = start.checked_sub(1)?; + } + indices.reverse(); + Some(indices) + } + /// 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 @@ -818,22 +895,34 @@ impl CodeInfo { for block in &mut self.blocks { let mut i = 0; - while i + 2 < block.instructions.len() { - // Check pattern: LOAD_CONST/LOAD_SMALL_INT, LOAD_CONST/LOAD_SMALL_INT, BINARY_OP - let Some(Instruction::BinaryOp { .. }) = block.instructions[i + 2].instr.real() + while i < block.instructions.len() { + let Some(Instruction::BinaryOp { .. }) = block.instructions[i].instr.real() else { + i += 1; + continue; + }; + + let Some(operand_indices) = i + .checked_sub(1) + .and_then(|start| Self::get_const_loading_instr_indices(block, start, 2)) else { i += 1; continue; }; - let op_raw = u32::from(block.instructions[i + 2].arg); + let op_raw = u32::from(block.instructions[i].arg); let Ok(op) = BinOp::try_from(op_raw) else { i += 1; continue; }; - let left = Self::get_const_value_from(&self.metadata, &block.instructions[i]); - let right = Self::get_const_value_from(&self.metadata, &block.instructions[i + 1]); + let left = Self::get_const_value_from( + &self.metadata, + &block.instructions[operand_indices[0]], + ); + let right = Self::get_const_value_from( + &self.metadata, + &block.instructions[operand_indices[1]], + ); let (Some(left_val), Some(right_val)) = (left, right) else { i += 1; @@ -849,20 +938,20 @@ impl CodeInfo { continue; } let (const_idx, _) = self.metadata.consts.insert_full(result_const); - // Replace first instruction with LOAD_CONST result - block.instructions[i].instr = Opcode::LoadConst.into(); + let folded_from_nonliteral_expr = operand_indices + .iter() + .any(|&idx| block.instructions[idx].folded_from_nonliteral_expr); + for &idx in &operand_indices { + set_to_nop(&mut block.instructions[idx]); + block.instructions[idx].location = block.instructions[i].location; + block.instructions[idx].end_location = block.instructions[i].end_location; + } + block.instructions[i].instr = Instruction::LoadConst { + consti: Arg::marker(), + } + .into(); block.instructions[i].arg = OpArg::new(const_idx as u32); - // NOP out the second and third instructions - let loc = block.instructions[i].location; - let end_loc = block.instructions[i].end_location; - set_to_nop(&mut block.instructions[i + 1]); - block.instructions[i + 1].location = loc; - block.instructions[i + 1].end_location = end_loc; - set_to_nop(&mut block.instructions[i + 2]); - block.instructions[i + 2].location = loc; - block.instructions[i + 2].end_location = end_loc; - // Don't advance - check if the result can be folded again - // (e.g., 2 ** 31 - 1) + block.instructions[i].folded_from_nonliteral_expr = folded_from_nonliteral_expr; i = i.saturating_sub(1); // re-check with previous instruction } else { i += 1; @@ -871,6 +960,13 @@ impl CodeInfo { } } + fn get_const_value_from_dummy(info: &InstructionInfo) -> Option<()> { + match info.instr.real() { + Some(Instruction::LoadConst { .. } | Instruction::LoadSmallInt { .. }) => Some(()), + _ => None, + } + } + fn get_const_value_from( metadata: &CodeUnitMetadata, info: &InstructionInfo, @@ -896,6 +992,42 @@ impl CodeInfo { op: oparg::BinaryOperator, ) -> Option { use oparg::BinaryOperator as BinOp; + fn eval_complex_binop( + left: Complex, + right: Complex, + op: BinOp, + ) -> Option { + let value = match op { + BinOp::Add => left + right, + BinOp::Subtract => { + let re = left.re - right.re; + let mut im = left.im - right.im; + // Preserve CPython's signed-zero behavior for real-zero + // minus zero-complex expressions such as `0 - 0j`. + if left.re == 0.0 + && left.im == 0.0 + && right.re == 0.0 + && right.im == 0.0 + && !right.im.is_sign_negative() + { + im = -0.0; + } + Complex::new(re, im) + } + BinOp::Multiply => left * right, + BinOp::TrueDivide => { + if right == Complex::new(0.0, 0.0) { + return None; + } + left / right + } + _ => return None, + }; + if !value.re.is_finite() || !value.im.is_finite() { + return None; + } + Some(ConstantData::Complex { value }) + } match (left, right) { (ConstantData::Integer { value: l }, ConstantData::Integer { value: r }) => { let result = match op { @@ -988,8 +1120,16 @@ impl CodeInfo { return None; } BinOp::Remainder => { - // Float modulo uses fmod() at runtime; Rust arithmetic differs - return None; + if *r == 0.0 { + return None; + } + let mut result = l % r; + if result != 0.0 && (*r < 0.0) != (result < 0.0) { + result += r; + } else if result == 0.0 { + result = 0.0f64.copysign(*r); + } + result } BinOp::Power => l.powf(*r), _ => return None, @@ -1016,6 +1156,21 @@ impl CodeInfo { op, ) } + (ConstantData::Integer { value: l }, ConstantData::Complex { value: r }) => { + eval_complex_binop(Complex::new(l.to_f64()?, 0.0), *r, op) + } + (ConstantData::Complex { value: l }, ConstantData::Integer { value: r }) => { + eval_complex_binop(*l, Complex::new(r.to_f64()?, 0.0), op) + } + (ConstantData::Float { value: l }, ConstantData::Complex { value: r }) => { + eval_complex_binop(Complex::new(*l, 0.0), *r, op) + } + (ConstantData::Complex { value: l }, ConstantData::Float { value: r }) => { + eval_complex_binop(*l, Complex::new(*r, 0.0), op) + } + (ConstantData::Complex { value: l }, ConstantData::Complex { value: r }) => { + eval_complex_binop(*l, *r, op) + } // String concatenation and repetition (ConstantData::Str { value: l }, ConstantData::Str { value: r }) if matches!(op, BinOp::Add) => @@ -1091,18 +1246,22 @@ impl CodeInfo { i += 1; continue; } - if i < tuple_size { + let Some(operand_indices) = i.checked_sub(1).and_then(|start| { + Self::get_const_loading_instr_indices(block, start, tuple_size) + }) else { i += 1; continue; - } + }; - // Check if all preceding instructions are constant-loading - let start_idx = i - tuple_size; let mut elements = Vec::with_capacity(tuple_size); let mut all_const = true; - for j in start_idx..i { + 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; @@ -1145,7 +1304,7 @@ impl CodeInfo { // Replace preceding LOAD instructions with NOP at the // BUILD_TUPLE location so remove_nops() can eliminate them. let folded_loc = block.instructions[i].location; - for j in start_idx..i { + for &j in &operand_indices { set_to_nop(&mut block.instructions[j]); block.instructions[j].location = folded_loc; } @@ -1172,17 +1331,26 @@ impl CodeInfo { }; let list_size = u32::from(instr.arg) as usize; - if list_size == 0 || i < list_size { + if list_size == 0 { i += 1; continue; } - let start_idx = i - list_size; + let Some(operand_indices) = i.checked_sub(1).and_then(|start| { + Self::get_const_loading_instr_indices(block, start, list_size) + }) else { + i += 1; + continue; + }; let mut elements = Vec::with_capacity(list_size); let mut all_const = true; - for j in start_idx..i { + 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; @@ -1220,22 +1388,29 @@ impl CodeInfo { let end_loc = block.instructions[i].end_location; let eh = block.instructions[i].except_handler; - // slot[start_idx] → BUILD_LIST 0 - block.instructions[start_idx].instr = Opcode::BuildList.into(); - block.instructions[start_idx].arg = OpArg::new(0); - block.instructions[start_idx].location = folded_loc; - block.instructions[start_idx].end_location = end_loc; - block.instructions[start_idx].except_handler = eh; + let build_idx = operand_indices[0]; + let const_idx_slot = operand_indices[1]; - // slot[start_idx+1] → LOAD_CONST (tuple) - block.instructions[start_idx + 1].instr = Opcode::LoadConst.into(); - block.instructions[start_idx + 1].arg = OpArg::new(const_idx as u32); - block.instructions[start_idx + 1].location = folded_loc; - block.instructions[start_idx + 1].end_location = end_loc; - block.instructions[start_idx + 1].except_handler = eh; + block.instructions[build_idx].instr = Instruction::BuildList { + count: Arg::marker(), + } + .into(); + block.instructions[build_idx].arg = OpArg::new(0); + block.instructions[build_idx].location = folded_loc; + block.instructions[build_idx].end_location = end_loc; + block.instructions[build_idx].except_handler = eh; + + block.instructions[const_idx_slot].instr = Instruction::LoadConst { + consti: Arg::marker(), + } + .into(); + block.instructions[const_idx_slot].arg = OpArg::new(const_idx as u32); + block.instructions[const_idx_slot].location = folded_loc; + block.instructions[const_idx_slot].end_location = end_loc; + block.instructions[const_idx_slot].except_handler = eh; // NOP the rest - for j in (start_idx + 2)..i { + for &j in &operand_indices[2..] { set_to_nop(&mut block.instructions[j]); block.instructions[j].location = folded_loc; } @@ -1256,6 +1431,22 @@ impl CodeInfo { for block in &mut self.blocks { let mut i = 0; while i + 1 < block.instructions.len() { + if matches!( + block.instructions[i].instr.real(), + Some(Instruction::CallIntrinsic1 { func }) + if func.get(block.instructions[i].arg) == IntrinsicFunction1::ListToTuple + ) && matches!( + block + .instructions + .get(i + 1) + .and_then(|instr| instr.instr.real()), + Some(Instruction::GetIter) + ) { + set_to_nop(&mut block.instructions[i]); + i += 2; + continue; + } + let is_build = matches!( block.instructions[i].instr.real(), Some(Instruction::BuildList { .. }) @@ -1305,12 +1496,17 @@ impl CodeInfo { ) { let seq_size = u32::from(block.instructions[i].arg) as usize; - if seq_size != 0 && i >= seq_size { - let start_idx = i - seq_size; + 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; - for j in start_idx..i { + for &j in &operand_indices { match Self::get_const_value_from(&self.metadata, &block.instructions[j]) { Some(constant) => elements.push(constant), @@ -1326,7 +1522,7 @@ impl CodeInfo { let (const_idx, _) = self.metadata.consts.insert_full(const_data); let folded_loc = block.instructions[i].location; - for j in start_idx..i { + for &j in &operand_indices { set_to_nop(&mut block.instructions[j]); block.instructions[j].location = folded_loc; } @@ -1360,17 +1556,26 @@ impl CodeInfo { }; let set_size = u32::from(instr.arg) as usize; - if set_size < 3 || i < set_size { + if set_size < 3 { i += 1; continue; } - let start_idx = i - set_size; + let Some(operand_indices) = i.checked_sub(1).and_then(|start| { + Self::get_const_loading_instr_indices(block, start, set_size) + }) else { + i += 1; + continue; + }; let mut elements = Vec::with_capacity(set_size); let mut all_const = true; - for j in start_idx..i { + 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; @@ -1408,19 +1613,28 @@ impl CodeInfo { let end_loc = block.instructions[i].end_location; let eh = block.instructions[i].except_handler; - block.instructions[start_idx].instr = Opcode::BuildSet.into(); - block.instructions[start_idx].arg = OpArg::new(0); - block.instructions[start_idx].location = folded_loc; - block.instructions[start_idx].end_location = end_loc; - block.instructions[start_idx].except_handler = eh; + let build_idx = operand_indices[0]; + let const_idx_slot = operand_indices[1]; - block.instructions[start_idx + 1].instr = Opcode::LoadConst.into(); - block.instructions[start_idx + 1].arg = OpArg::new(const_idx as u32); - block.instructions[start_idx + 1].location = folded_loc; - block.instructions[start_idx + 1].end_location = end_loc; - block.instructions[start_idx + 1].except_handler = eh; + block.instructions[build_idx].instr = Instruction::BuildSet { + count: Arg::marker(), + } + .into(); + block.instructions[build_idx].arg = OpArg::new(0); + block.instructions[build_idx].location = folded_loc; + block.instructions[build_idx].end_location = end_loc; + block.instructions[build_idx].except_handler = eh; + + block.instructions[const_idx_slot].instr = Instruction::LoadConst { + consti: Arg::marker(), + } + .into(); + block.instructions[const_idx_slot].arg = OpArg::new(const_idx as u32); + block.instructions[const_idx_slot].location = folded_loc; + block.instructions[const_idx_slot].end_location = end_loc; + block.instructions[const_idx_slot].except_handler = eh; - for j in (start_idx + 2)..i { + for &j in &operand_indices[2..] { set_to_nop(&mut block.instructions[j]); block.instructions[j].location = folded_loc; } @@ -1484,6 +1698,8 @@ impl CodeInfo { /// intervening swappable stores to one of the same variables. Do not /// cross line-number boundaries (user-visible name bindings). fn apply_static_swaps(&mut self) { + const VISITED: i32 = -1; + /// Instruction classes that are safe to reorder around SWAP. fn is_swappable(instr: &AnyInstruction) -> bool { matches!( @@ -1528,59 +1744,111 @@ impl CodeInfo { } } - for block in &mut self.blocks { - let instructions = &mut block.instructions; - let len = instructions.len(); - // Walk forward; for each SWAP attempt elimination. - let mut i = 0; - while i < len { - let swap_arg = match instructions[i].instr { - AnyInstruction::Real(Instruction::Swap { .. }) => { - u32::from(instructions[i].arg) + fn optimize_swap_block(instructions: &mut [InstructionInfo]) { + let mut i = 0usize; + while i < instructions.len() { + let AnyInstruction::Real(Instruction::Swap { .. }) = instructions[i].instr else { + i += 1; + continue; + }; + + let mut len = 0usize; + let mut depth = 0usize; + let mut more = false; + while i + len < instructions.len() { + let info = &instructions[i + len]; + match info.instr.real() { + Some(Instruction::Swap { .. }) => { + let oparg = u32::from(info.arg) as usize; + depth = depth.max(oparg); + more |= len > 0; + len += 1; + } + Some(Instruction::Nop) => { + len += 1; + } + _ => break, } - _ => { - i += 1; + } + + if !more { + i += len.max(1); + continue; + } + + let mut stack: Vec = (0..depth as i32).collect(); + for info in &instructions[i..i + len] { + if matches!(info.instr.real(), Some(Instruction::Swap { .. })) { + let oparg = u32::from(info.arg) as usize; + stack.swap(0, oparg - 1); + } + } + + let mut current = len as isize - 1; + for slot in 0..depth { + if stack[slot] == VISITED || stack[slot] == slot as i32 { + continue; + } + let mut j = slot; + loop { + if j != 0 { + let out = &mut instructions[i + current as usize]; + out.instr = Opcode::Swap.into(); + out.arg = OpArg::new((j + 1) as u32); + out.target = BlockIdx::NULL; + current -= 1; + } + if stack[j] == VISITED { + debug_assert_eq!(j, slot); + break; + } + let next_j = stack[j] as usize; + stack[j] = VISITED; + j = next_j; + } + } + while current >= 0 { + set_to_nop(&mut instructions[i + current as usize]); + current -= 1; + } + i += len; + } + } + + fn apply_from(instructions: &mut [InstructionInfo], mut i: isize) { + while i >= 0 { + let idx = i as usize; + let swap_arg = match instructions[idx].instr.real() { + Some(Instruction::Swap { .. }) => u32::from(instructions[idx].arg), + Some(Instruction::Nop) + | Some(Instruction::PopTop | Instruction::StoreFast { .. }) => { + i -= 1; continue; } + _ => return, }; - // SWAP oparg < 2 is a no-op; the compiler should not emit - // these, but be defensive. + if swap_arg < 2 { - i += 1; - continue; + return; } - // Find first swappable after SWAP (lineno = -1 initially). - let Some(j) = next_swappable(instructions, i, -1) else { - i += 1; - continue; + + let Some(j) = next_swappable(instructions, idx, -1) else { + return; }; let lineno = instructions[j].location.line.get() as i32; - // Walk (swap_arg - 1) more swappable instructions, with - // lineno constraint. let mut k = j; - let mut ok = true; for _ in 1..swap_arg { - match next_swappable(instructions, k, lineno) { - Some(next) => k = next, - None => { - ok = false; - break; - } - } - } - if !ok { - i += 1; - continue; + let Some(next) = next_swappable(instructions, k, lineno) else { + return; + }; + k = next; } - // Conflict check: if either j or k is a STORE_FAST, no - // intervening store may target the same variable, and - // they must not target the same variable themselves. + let store_j = stores_to(&instructions[j]); let store_k = stores_to(&instructions[k]); if store_j.is_some() || store_k.is_some() { if store_j == store_k { - i += 1; - continue; + return; } let conflict = instructions[(j + 1)..k].iter().any(|info| { if let Some(store_idx) = stores_to(info) { @@ -1590,15 +1858,27 @@ impl CodeInfo { } }); if conflict { - i += 1; - continue; + return; } } - // Safe to reorder. SWAP -> NOP, swap j and k. - instructions[i].instr = Opcode::Nop.into(); - instructions[i].arg = OpArg::new(0); + + instructions[idx].instr = Opcode::Nop.into(); + instructions[idx].arg = OpArg::new(0); instructions.swap(j, k); - i += 1; + i -= 1; + } + } + + for block in &mut self.blocks { + optimize_swap_block(&mut block.instructions); + let len = block.instructions.len(); + for i in 0..len { + if matches!( + block.instructions[i].instr.real(), + Some(Instruction::Swap { .. }) + ) { + apply_from(&mut block.instructions, i as isize); + } } } } @@ -1689,6 +1969,8 @@ impl CodeInfo { while i + 1 < block.instructions.len() { let curr = &block.instructions[i]; let next = &block.instructions[i + 1]; + let curr_arg = curr.arg; + let next_arg = next.arg; // Only combine if both are real instructions (not pseudo) let (Some(curr_instr), Some(next_instr)) = (curr.instr.real(), next.instr.real()) @@ -1724,6 +2006,70 @@ impl CodeInfo { } } + if let Instruction::LoadConst { consti } = curr_instr { + let constant = &self.metadata.consts[consti.get(curr_arg).as_usize()]; + if matches!(constant, ConstantData::None) + && let Instruction::IsOp { invert } = next_instr + { + let mut jump_idx = i + 2; + if jump_idx >= block.instructions.len() { + i += 1; + continue; + } + + if matches!( + block.instructions[jump_idx].instr.real(), + Some(Instruction::ToBool) + ) { + set_to_nop(&mut block.instructions[jump_idx]); + jump_idx += 1; + if jump_idx >= block.instructions.len() { + i += 1; + continue; + } + } + + let Some(jump_instr) = block.instructions[jump_idx].instr.real() else { + i += 1; + continue; + }; + + let mut invert = matches!( + invert.get(next_arg), + rustpython_compiler_core::bytecode::Invert::Yes + ); + let delta = match jump_instr { + Instruction::PopJumpIfFalse { delta } => { + invert = !invert; + delta.get(block.instructions[jump_idx].arg) + } + Instruction::PopJumpIfTrue { delta } => { + delta.get(block.instructions[jump_idx].arg) + } + _ => { + i += 1; + continue; + } + }; + + set_to_nop(&mut block.instructions[i]); + set_to_nop(&mut block.instructions[i + 1]); + block.instructions[jump_idx].instr = if invert { + Instruction::PopJumpIfNotNone { + delta: Arg::marker(), + } + } else { + Instruction::PopJumpIfNone { + delta: Arg::marker(), + } + } + .into(); + block.instructions[jump_idx].arg = OpArg::new(u32::from(delta)); + i = jump_idx; + continue; + } + } + if matches!( curr_instr, Instruction::LoadConst { .. } | Instruction::LoadSmallInt { .. } @@ -1967,12 +2313,12 @@ impl CodeInfo { let mut prev_line = None; block.instructions.retain(|ins| { if matches!(ins.instr.real(), Some(Instruction::Nop)) { - let line = ins.location.line; + let line = ins.location.line.get() as i32; if prev_line == Some(line) { return false; } } - prev_line = Some(ins.location.line); + prev_line = Some(instruction_lineno(ins)); true }); } @@ -1987,11 +2333,22 @@ impl CodeInfo { let mut i = 0; while i + 1 < block.instructions.len() { let curr = &block.instructions[i]; - let next = &block.instructions[i + 1]; + let line = curr.location.line; - let line1 = curr.location.line; - let line2 = next.location.line; - if line1 != line2 { + 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 { i += 1; continue; } @@ -2010,9 +2367,12 @@ impl CodeInfo { } .into(); block.instructions[i].arg = OpArg::new(packed); - block.instructions.remove(i + 1); + block.instructions.drain(i + 1..=j); } - (Some(Instruction::StoreFast { .. }), Some(Instruction::LoadFast { .. })) => { + ( + Some(Instruction::StoreFast { .. }), + Some(Instruction::LoadFast { .. } | Instruction::LoadFastBorrow { .. }), + ) => { let store_idx = u32::from(curr.arg); let load_idx = u32::from(next.arg); if store_idx >= 16 || load_idx >= 16 { @@ -2025,7 +2385,7 @@ impl CodeInfo { } .into(); block.instructions[i].arg = OpArg::new(packed); - block.instructions.remove(i + 1); + block.instructions.drain(i + 1..=j); } (Some(Instruction::StoreFast { .. }), Some(Instruction::StoreFast { .. })) => { let idx1 = u32::from(curr.arg); @@ -2040,7 +2400,7 @@ impl CodeInfo { } .into(); block.instructions[i].arg = OpArg::new(packed); - block.instructions.remove(i + 1); + block.instructions.drain(i + 1..=j); } _ => i += 1, } @@ -2137,9 +2497,6 @@ impl CodeInfo { while let Some(block_idx) = worklist.pop() { let block = &self.blocks[block_idx]; - if block.instructions.is_empty() { - continue; - } let mut instr_flags = vec![0u8; block.instructions.len()]; let start_depth = block.start_depth.unwrap_or(0) as usize; @@ -2359,7 +2716,7 @@ impl CodeInfo { let next = block.next; if next != BlockIdx::NULL - && block.instructions.last().is_some_and(|term| { + && block.instructions.last().is_none_or(|term| { !term.instr.is_unconditional_jump() && !term.instr.is_scope_exit() }) { @@ -2459,43 +2816,6 @@ impl CodeInfo { } } - fn propagate_disable_load_fast_borrow(&mut self) { - let mut predecessors = vec![Vec::new(); self.blocks.len()]; - for (pred_idx, block) in iter_blocks(&self.blocks) { - if block.next != BlockIdx::NULL { - predecessors[block.next.idx()].push(pred_idx); - } - for info in &block.instructions { - if info.target != BlockIdx::NULL { - predecessors[info.target.idx()].push(pred_idx); - } - } - } - - let mut changed = true; - while changed { - changed = false; - let block_indices: Vec<_> = iter_blocks(&self.blocks).map(|(idx, _)| idx).collect(); - for idx in block_indices { - if idx == BlockIdx(0) || self.blocks[idx.idx()].disable_load_fast_borrow { - continue; - } - let predecessor_blocks = &predecessors[idx.idx()]; - if predecessor_blocks.is_empty() { - continue; - } - if predecessor_blocks - .iter() - .copied() - .all(|pred_idx| self.blocks[pred_idx.idx()].disable_load_fast_borrow) - { - self.blocks[idx.idx()].disable_load_fast_borrow = true; - changed = true; - } - } - } - } - fn deoptimize_borrow_for_handler_return_paths(&mut self) { for block in &mut self.blocks { let len = block.instructions.len(); @@ -2556,6 +2876,90 @@ impl CodeInfo { } } + fn deoptimize_borrow_for_match_keys_attr(&mut self) { + let Some(key_name_idx) = self.metadata.names.get_index_of("KEY") else { + return; + }; + + let mut to_deopt = Vec::new(); + for block_idx in 0..self.blocks.len() { + let block = &self.blocks[block_idx]; + let len = block.instructions.len(); + for i in 0..len { + let Some(Instruction::LoadFastBorrow { .. }) = block.instructions[i].instr.real() + else { + continue; + }; + let Some(Instruction::LoadAttr { namei }) = block + .instructions + .get(i + 1) + .and_then(|info| info.instr.real()) + else { + continue; + }; + let load_attr = namei.get(block.instructions[i + 1].arg); + if load_attr.is_method() || load_attr.name_idx() as usize != key_name_idx { + continue; + } + + let mut saw_build_tuple = false; + let mut saw_match_keys = false; + let mut scan_block_idx = block_idx; + let mut scan_start = i + 2; + loop { + let scan_block = &self.blocks[scan_block_idx]; + for info in scan_block.instructions.iter().skip(scan_start) { + match info.instr.real() { + Some( + Instruction::LoadConst { .. } + | Instruction::LoadSmallInt { .. } + | Instruction::LoadFast { .. } + | Instruction::LoadFastBorrow { .. } + | Instruction::LoadAttr { .. } + | Instruction::Nop, + ) => {} + Some(Instruction::BuildTuple { .. }) => saw_build_tuple = true, + Some(Instruction::MatchKeys) => { + saw_match_keys = true; + break; + } + _ => { + saw_build_tuple = false; + break; + } + } + } + if saw_match_keys { + break; + } + let Some(last) = scan_block.instructions.last() else { + break; + }; + if scan_block.next == BlockIdx::NULL + || last.instr.is_scope_exit() + || last.instr.is_unconditional_jump() + || last.target != BlockIdx::NULL + { + break; + } + scan_block_idx = scan_block.next.idx(); + scan_start = 0; + } + + if saw_build_tuple && saw_match_keys { + to_deopt.push((block_idx, i)); + } + } + } + + for (block_idx, instr_idx) in to_deopt { + self.blocks[block_idx].instructions[instr_idx].instr = Instruction::LoadFast { + var_num: Arg::marker(), + } + .into(); + } + } + fn deoptimize_store_fast_store_fast_after_cleanup(&mut self) { fn last_real_instr(block: &Block) -> Option { block @@ -2708,6 +3112,12 @@ impl CodeInfo { worklist.push(target); } } + if matches!(info.instr.real(), Some(Instruction::ForIter { .. })) + && info.target != BlockIdx::NULL + && merge_unsafe_mask(&mut in_masks[info.target.idx()], &unsafe_mask) + { + worklist.push(info.target); + } match info.instr.real() { Some(Instruction::DeleteFast { var_num }) => { let var_idx = usize::from(var_num.get(info.arg)); @@ -2955,7 +3365,7 @@ impl CodeInfo { use core::fmt::Write; let _ = writeln!( out, - "block {} next={} cold={} except={} preserve_lasti={} disable_borrow={}", + "block {} next={} cold={} except={} preserve_lasti={} disable_borrow={} start_depth={}", u32::from(block_idx), if block.next == BlockIdx::NULL { String::from("NULL") @@ -2966,13 +3376,19 @@ impl CodeInfo { block.except_handler, block.preserve_lasti, block.disable_load_fast_borrow, + block + .start_depth + .map(|depth| depth.to_string()) + .unwrap_or_else(|| String::from("None")), ); for info in &block.instructions { let lineno = instruction_lineno(info); let _ = writeln!( out, - " [{}] {:?} arg={} target={}", + " [disp={} raw={} override={:?}] {:?} arg={} target={}", lineno, + info.location.line.get(), + info.lineno_override, info.instr, u32::from(info.arg), if info.target == BlockIdx::NULL { @@ -2988,38 +3404,50 @@ impl CodeInfo { pub(crate) fn debug_late_cfg_trace(mut self) -> crate::InternalResult> { let mut trace = Vec::new(); + trace.push(("initial".to_owned(), self.debug_block_dump())); self.splice_annotations_blocks(); self.fold_binop_constants(); - self.remove_nops(); - self.fold_unary_negative(); - self.remove_nops(); + self.fold_unary_constants(); self.fold_binop_constants(); - self.remove_nops(); self.fold_tuple_constants(); self.fold_list_constants(); self.fold_set_constants(); - self.remove_nops(); self.fold_const_iterable_for_iter(); self.convert_to_load_small_int(); self.remove_unused_consts(); - self.remove_nops(); self.dce(); self.optimize_build_tuple_unpack(); self.eliminate_dead_stores(); self.apply_static_swaps(); self.peephole_optimize(); + trace.push(( + "after_peephole_optimize".to_owned(), + self.debug_block_dump(), + )); split_blocks_at_jumps(&mut self.blocks); + trace.push(( + "after_split_blocks_at_jumps".to_owned(), + self.debug_block_dump(), + )); mark_except_handlers(&mut self.blocks); label_exception_targets(&mut self.blocks); jump_threading(&mut self.blocks); + trace.push(("after_jump_threading".to_owned(), self.debug_block_dump())); self.eliminate_unreachable_blocks(); self.remove_nops(); + trace.push(( + "after_early_remove_nops".to_owned(), + self.debug_block_dump(), + )); self.add_checks_for_loads_of_uninitialized_variables(); self.insert_superinstructions(); push_cold_blocks_to_end(&mut self.blocks); - trace.push(("after_push_cold_blocks_to_end".to_owned(), self.debug_block_dump())); + trace.push(( + "after_push_cold_blocks_to_end".to_owned(), + self.debug_block_dump(), + )); normalize_jumps(&mut self.blocks); trace.push(("after_normalize_jumps".to_owned(), self.debug_block_dump())); @@ -3030,24 +3458,73 @@ impl CodeInfo { 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())); + 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())); resolve_line_numbers(&mut self.blocks); - trace.push(("after_resolve_line_numbers".to_owned(), self.debug_block_dump())); + trace.push(( + "after_resolve_line_numbers".to_owned(), + self.debug_block_dump(), + )); + + redirect_empty_block_targets(&mut self.blocks); + trace.push(( + "after_redirect_empty_block_targets".to_owned(), + self.debug_block_dump(), + )); duplicate_end_returns(&mut self.blocks); - trace.push(("after_duplicate_end_returns".to_owned(), self.debug_block_dump())); + trace.push(( + "after_duplicate_end_returns".to_owned(), + self.debug_block_dump(), + )); self.dce(); self.eliminate_unreachable_blocks(); - trace.push(("after_second_dce_unreachable".to_owned(), self.debug_block_dump())); + trace.push(( + "after_second_dce_unreachable".to_owned(), + self.debug_block_dump(), + )); remove_redundant_nops_and_jumps(&mut self.blocks); - trace.push(("after_remove_redundant_nops_and_jumps".to_owned(), self.debug_block_dump())); + trace.push(( + "after_remove_redundant_nops_and_jumps".to_owned(), + self.debug_block_dump(), + )); + + let cellfixedoffsets = build_cellfixedoffsets( + &self.metadata.varnames, + &self.metadata.cellvars, + &self.metadata.freevars, + ); + mark_except_handlers(&mut self.blocks); + redirect_empty_block_targets(&mut self.blocks); + let _ = self.max_stackdepth()?; + convert_pseudo_ops(&mut self.blocks, &cellfixedoffsets); + trace.push(( + "after_convert_pseudo_ops".to_owned(), + self.debug_block_dump(), + )); + self.compute_load_fast_start_depths(); + trace.push(( + "after_compute_load_fast_start_depths".to_owned(), + self.debug_block_dump(), + )); + self.optimize_load_fast_borrow(); + trace.push(( + "after_optimize_load_fast_borrow".to_owned(), + self.debug_block_dump(), + )); + self.deoptimize_borrow_after_push_exc_info(); + self.deoptimize_borrow_for_handler_return_paths(); + self.deoptimize_borrow_for_match_keys_attr(); + trace.push(("after_borrow_deopts".to_owned(), self.debug_block_dump())); Ok(trace) } @@ -3496,6 +3973,7 @@ fn push_cold_blocks_to_end(blocks: &mut Vec) { location: SourceLocation::default(), end_location: SourceLocation::default(), except_handler: None, + folded_from_nonliteral_expr: false, lineno_override: Some(-1), cache_entries: 0, }); @@ -3631,14 +4109,28 @@ fn threaded_jump_instr( } let source_kind = jump_thread_kind(source)?; - if source_kind == JumpThreadKind::NoInterrupt { - return Some(source); - } + let result_kind = if source_kind == JumpThreadKind::NoInterrupt + && target_kind == JumpThreadKind::NoInterrupt + { + JumpThreadKind::NoInterrupt + } else { + JumpThreadKind::Plain + }; - Some(match source.into() { - AnyOpcode::Pseudo(_) => PseudoOpcode::Jump.into(), - AnyOpcode::Real(Opcode::JumpBackwardNoInterrupt) => Opcode::JumpBackward.into(), - AnyOpcode::Real(Opcode::JumpForward | Opcode::JumpBackward) => source, + Some(match (source.into(), result_kind) { + (AnyOpcode::Pseudo(_), JumpThreadKind::Plain) => PseudoOpcode::Jump.into(), + (AnyOpcode::Pseudo(_), JumpThreadKind::NoInterrupt) => PseudoOpcode::JumpNoInterrupt.into(), + (AnyOpcode::Real(Opcode::JumpBackwardNoInterrupt), JumpThreadKind::Plain) => { + Opcode::JumpBackward.into() + } + (AnyOpcode::Real(Opcode::JumpBackwardNoInterrupt), JumpThreadKind::NoInterrupt) => source, + (AnyOpcode::Real(Opcode::JumpForward | Opcode::JumpBackward), JumpThreadKind::Plain) => { + source + } + ( + AnyOpcode::Real(Opcode::JumpForward | Opcode::JumpBackward), + JumpThreadKind::NoInterrupt, + ) => PseudoOpcode::JumpNoInterrupt.into(), _ => return None, }) } @@ -3678,11 +4170,7 @@ fn jump_threading_impl(blocks: &mut [Block], include_conditional: bool) { continue; } // Check if target block's first instruction is an unconditional jump - let target_jump = blocks[target.idx()] - .instructions - .iter() - .find(|ins| !matches!(ins.instr.real(), Some(Instruction::Nop))) - .copied(); + let target_jump = blocks[target.idx()].instructions.first().copied(); if let Some(target_ins) = target_jump && target_ins.instr.is_unconditional_jump() && target_ins.target != BlockIdx::NULL @@ -3785,6 +4273,7 @@ fn normalize_jumps(blocks: &mut Vec) { location: last_ins.location, end_location: last_ins.end_location, except_handler: last_ins.except_handler, + folded_from_nonliteral_expr: false, lineno_override: None, cache_entries: 0, }; @@ -3816,6 +4305,7 @@ fn normalize_jumps(blocks: &mut Vec) { location: loc, end_location: end_loc, except_handler: exc_handler, + folded_from_nonliteral_expr: false, lineno_override: None, cache_entries: 0, }); @@ -3826,6 +4316,7 @@ fn normalize_jumps(blocks: &mut Vec) { location: loc, end_location: end_loc, except_handler: exc_handler, + folded_from_nonliteral_expr: false, lineno_override: None, cache_entries: 0, }); @@ -3906,32 +4397,29 @@ fn inline_small_or_no_lineno_blocks(blocks: &mut [Block]) { .iter() .all(|ins| !instruction_has_lineno(ins)) }; - let is_return_epilogue_block = |block: &Block| { - matches!( - block.instructions.as_slice(), - [ - InstructionInfo { - instr: AnyInstruction::Real(Instruction::LoadConst { .. }), - .. - }, - InstructionInfo { - instr: AnyInstruction::Real(Instruction::ReturnValue), - .. - } - ] | [ - InstructionInfo { - instr: AnyInstruction::Real(Instruction::LoadSmallInt { .. }), - .. - }, - InstructionInfo { - instr: AnyInstruction::Real(Instruction::ReturnValue), - .. - } - ] | [InstructionInfo { - instr: AnyInstruction::Real(Instruction::ReturnValue), - .. - }] - ) + 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 + .iter() + .any(|ins| ins.instr.is_block_push()) }; loop { let mut changes = false; @@ -3950,6 +4438,8 @@ 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()]) + && target_pushes_handler(&blocks[target.idx()])) { current = next; continue; @@ -3958,18 +4448,16 @@ 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 - && blocks[current.idx()].cold - && !is_return_epilogue_block(&blocks[target.idx()]) - { - current = next; - continue; - } 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 appended = blocks[target.idx()].instructions.clone(); + 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); + } blocks[current.idx()].instructions.extend(appended); changes = true; } @@ -3984,64 +4472,6 @@ fn inline_small_or_no_lineno_blocks(blocks: &mut [Block]) { } fn remove_redundant_nops_in_blocks(blocks: &mut [Block]) -> usize { - let is_return_epilogue_block = |block: &Block| { - matches!( - block.instructions.as_slice(), - [ - InstructionInfo { - instr: AnyInstruction::Real(Instruction::LoadConst { .. }), - .. - }, - InstructionInfo { - instr: AnyInstruction::Real(Instruction::ReturnValue), - .. - } - ] | [ - InstructionInfo { - instr: AnyInstruction::Real(Instruction::LoadSmallInt { .. }), - .. - }, - InstructionInfo { - instr: AnyInstruction::Real(Instruction::ReturnValue), - .. - } - ] | [InstructionInfo { - instr: AnyInstruction::Real(Instruction::ReturnValue), - .. - }] - ) - }; - let is_return_epilogue_pair = |instructions: &[InstructionInfo]| { - matches!( - instructions, - [ - InstructionInfo { - instr: AnyInstruction::Real(Instruction::LoadConst { .. }), - .. - }, - InstructionInfo { - instr: AnyInstruction::Real(Instruction::ReturnValue), - .. - } - ] | [ - InstructionInfo { - instr: AnyInstruction::Real(Instruction::LoadSmallInt { .. }), - .. - }, - InstructionInfo { - instr: AnyInstruction::Real(Instruction::ReturnValue), - .. - } - ] | [InstructionInfo { - instr: AnyInstruction::Real(Instruction::ReturnValue), - .. - }] - ) - }; - let starts_with_return_epilogue_pair = |instructions: &[InstructionInfo]| { - instructions.len() >= 2 && is_return_epilogue_pair(&instructions[..2]) - || !instructions.is_empty() && is_return_epilogue_pair(&instructions[..1]) - }; let mut changes = 0; let mut block_order = Vec::new(); let mut current = BlockIdx(0); @@ -4065,43 +4495,20 @@ 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 kept.last().is_some_and(|prev: &InstructionInfo| { - matches!(prev.instr.real(), Some(Instruction::Nop)) - }) && is_return_epilogue_pair( - &src_instructions[src + 1..src_instructions.len().min(src + 3)], - ) - { - src_instructions[src + 1].lineno_override = Some(lineno); + if src_instructions[src + 1].folded_from_nonliteral_expr { remove = true; } else { - let next_lineno = instruction_lineno(&src_instructions[src + 1]); - if next_lineno == lineno { - remove = true; - } else if next_lineno < 0 { - src_instructions[src + 1].lineno_override = Some(lineno); - remove = true; - } + let next_lineno = instruction_lineno(&src_instructions[src + 1]); + if next_lineno == lineno { + remove = true; + } else if next_lineno < 0 { + src_instructions[src + 1].lineno_override = Some(lineno); + remove = true; + } } } else { let next = next_nonempty_block(blocks, blocks[bi].next); if next != BlockIdx::NULL { - if is_return_epilogue_block(&blocks[next.idx()]) - || starts_with_return_epilogue_pair(&blocks[next.idx()].instructions) - { - let current_block_is_nop_only = - kept.iter().all(|prev: &InstructionInfo| { - matches!(prev.instr.real(), Some(Instruction::Nop)) - }); - let pred = find_layout_predecessor(blocks, block_idx); - let pred_ends_with_nop = pred != BlockIdx::NULL - && blocks[pred.idx()].instructions.last().is_some_and(|prev| { - matches!(prev.instr.real(), Some(Instruction::Nop)) - }); - if current_block_is_nop_only && pred_ends_with_nop { - changes += 1; - continue; - } - } let mut next_info = None; for (next_idx, next_instr) in blocks[next.idx()].instructions.iter().enumerate() @@ -4158,7 +4565,6 @@ fn remove_redundant_jumps_in_blocks(blocks: &mut [Block]) -> usize { && let Some(last_instr) = blocks[idx].instructions.last_mut() { set_to_nop(last_instr); - last_instr.lineno_override = Some(-1); changes += 1; } current = blocks[idx].next; @@ -4176,6 +4582,33 @@ fn remove_redundant_nops_and_jumps(blocks: &mut [Block]) { } } +fn redirect_empty_block_targets(blocks: &mut [Block]) { + let redirected_targets: Vec> = blocks + .iter() + .map(|block| { + block + .instructions + .iter() + .map(|instr| { + if instr.target == BlockIdx::NULL { + BlockIdx::NULL + } 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 merge_unsafe_mask(slot: &mut Option>, incoming: &[bool]) -> bool { match slot { Some(existing) => { @@ -4606,6 +5039,16 @@ fn maybe_propagate_location( } } +fn overwrite_location( + instr: &mut InstructionInfo, + location: SourceLocation, + end_location: SourceLocation, +) { + instr.location = location; + instr.end_location = end_location; + instr.lineno_override = None; +} + fn propagate_locations_in_block( block: &mut Block, location: SourceLocation, @@ -4826,15 +5269,24 @@ fn find_layout_predecessor(blocks: &[Block], target: BlockIdx) -> BlockIdx { /// Duplicate `LOAD_CONST None + RETURN_VALUE` for blocks that fall through /// to the final return block. fn duplicate_end_returns(blocks: &mut Vec) { - // Walk the block chain and keep the last non-empty block. + // Walk the block chain and keep the last non-cold non-empty block. + // After cold exception handlers are pushed to the end, the mainline + // return epilogue can sit before trailing cold blocks. let mut last_block = BlockIdx::NULL; + let mut last_nonempty_block = BlockIdx::NULL; let mut current = BlockIdx(0); while current != BlockIdx::NULL { if !blocks[current.idx()].instructions.is_empty() { - last_block = current; + last_nonempty_block = current; + if !blocks[current.idx()].cold { + last_block = current; + } } current = blocks[current.idx()].next; } + if last_block == BlockIdx::NULL { + last_block = last_nonempty_block; + } if last_block == BlockIdx::NULL { return; } @@ -4904,9 +5356,17 @@ fn duplicate_end_returns(blocks: &mut Vec) { // Duplicate the return instructions at the end of fall-through blocks for block_idx in fallthrough_blocks_to_fix { - blocks[block_idx.idx()] + let propagated_location = blocks[block_idx.idx()] .instructions - .extend_from_slice(&return_insts); + .last() + .map(|instr| (instr.location, instr.end_location)); + let mut cloned_return = return_insts.clone(); + if let Some((location, end_location)) = propagated_location { + for instr in &mut cloned_return { + overwrite_location(instr, location, end_location); + } + } + blocks[block_idx.idx()].instructions.extend(cloned_return); } // Clone the final return block for jump predecessors so their target layout @@ -4914,8 +5374,8 @@ fn duplicate_end_returns(blocks: &mut Vec) { for (block_idx, instr_idx) in jump_targets_to_fix { let jump = blocks[block_idx.idx()].instructions[instr_idx]; let mut cloned_return = return_insts.clone(); - for instr in &mut cloned_return { - maybe_propagate_location(instr, jump.location, jump.end_location); + if let Some(first) = cloned_return.first_mut() { + overwrite_location(first, jump.location, jump.end_location); } let new_idx = BlockIdx(blocks.len() as u32); let is_conditional = is_conditional_jump(&jump.instr); diff --git a/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__bare_function_annotations_check_attribute_and_subscript_expressions.snap b/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__bare_function_annotations_check_attribute_and_subscript_expressions.snap new file mode 100644 index 00000000000..840f4397d75 --- /dev/null +++ b/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__bare_function_annotations_check_attribute_and_subscript_expressions.snap @@ -0,0 +1,67 @@ +--- +source: crates/codegen/src/compile.rs +assertion_line: 12138 +expression: "compile_exec(\"\\\ndef f(one: int):\n int.new_attr: int\n [list][0].new_attr: [int, str]\n my_lst = [1]\n my_lst[one]: int\n return my_lst\n\")" +--- + 1 0 RESUME (0) + 1 LOAD_CONST (): 1 0 RESUME (0) + 1 LOAD_FAST_BORROW (0, format) + 2 LOAD_SMALL_INT (2) + >> 3 COMPARE_OP (>) + 4 CACHE + 5 POP_JUMP_IF_FALSE (3) + 6 CACHE + 7 NOT_TAKEN + 8 LOAD_COMMON_CONSTANT (NotImplementedError) + 9 RAISE_VARARGS (Raise) + 10 LOAD_CONST ("one") + 11 LOAD_GLOBAL (0, int) + 12 CACHE + 13 CACHE + 14 CACHE + 15 CACHE + 16 BUILD_MAP (1) + 17 RETURN_VALUE + + 2 MAKE_FUNCTION + 3 LOAD_CONST (): 1 0 RESUME (0) + + 2 1 LOAD_GLOBAL (0, int) + 2 CACHE + 3 CACHE + 4 CACHE + 5 CACHE + 6 POP_TOP + + 3 7 LOAD_GLOBAL (2, list) + 8 CACHE + 9 CACHE + 10 CACHE + 11 CACHE + 12 BUILD_LIST (1) + 13 LOAD_SMALL_INT (0) + 14 BINARY_OP ([]) + 15 CACHE + 16 CACHE + 17 CACHE + 18 CACHE + 19 CACHE + 20 POP_TOP + + 4 21 LOAD_SMALL_INT (1) + 22 BUILD_LIST (1) + 23 STORE_FAST (1, my_lst) + + 5 24 LOAD_FAST_BORROW (1, my_lst) + 25 POP_TOP + 26 LOAD_FAST_BORROW (0, one) + 27 POP_TOP + + 6 28 LOAD_FAST_BORROW (1, my_lst) + 29 RETURN_VALUE + + 4 MAKE_FUNCTION + 5 SET_FUNCTION_ATTRIBUTE(Annotate) + 6 STORE_NAME (0, f) + 7 LOAD_CONST (None) + 8 RETURN_VALUE diff --git a/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__constant_true_if_pass_keeps_line_anchor_nop.snap b/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__constant_true_if_pass_keeps_line_anchor_nop.snap new file mode 100644 index 00000000000..a600a829863 --- /dev/null +++ b/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__constant_true_if_pass_keeps_line_anchor_nop.snap @@ -0,0 +1,10 @@ +--- +source: crates/codegen/src/compile.rs +assertion_line: 12222 +expression: "compile_exec(\"\\\nif 1:\n pass\n\")" +--- + 1 0 RESUME (0) + 1 NOP + + 2 2 LOAD_CONST (None) + 3 RETURN_VALUE diff --git a/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_ands.snap b/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_ands.snap index d4cb6b8d8d3..8e9bb5d25f4 100644 --- a/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_ands.snap +++ b/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_ands.snap @@ -1,10 +1,8 @@ --- source: crates/codegen/src/compile.rs -assertion_line: 11049 +assertion_line: 11413 expression: "compile_exec(\"\\\nif True and False and False:\n pass\n\")" --- 1 0 RESUME (0) - 1 NOP - - 2 2 LOAD_CONST (None) - 3 RETURN_VALUE + 1 LOAD_CONST (None) + 2 RETURN_VALUE diff --git a/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_mixed.snap b/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_mixed.snap index 1c28cd123dc..f7df6f4f3ee 100644 --- a/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_mixed.snap +++ b/crates/codegen/src/snapshots/rustpython_codegen__compile__tests__if_mixed.snap @@ -1,10 +1,8 @@ --- source: crates/codegen/src/compile.rs -assertion_line: 11059 +assertion_line: 11423 expression: "compile_exec(\"\\\nif (True and False) or (False and True):\n pass\n\")" --- 1 0 RESUME (0) - 1 NOP - - 2 2 LOAD_CONST (None) - 3 RETURN_VALUE + 1 LOAD_CONST (None) + 2 RETURN_VALUE 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 c39144ae857..fccc7b7c336 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: 11208 +assertion_line: 11626 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) @@ -219,8 +219,8 @@ expression: "compile_exec(\"\\\nasync def test():\n for stop_exc in (StopIter 202 RERAISE (1) 3 203 LOAD_CONST (None) - >> 204 LOAD_CONST (None) - 205 LOAD_CONST (None) + 204 LOAD_CONST (None) + >> 205 LOAD_CONST (None) 206 CALL (3) 207 CACHE 208 CACHE @@ -243,13 +243,14 @@ expression: "compile_exec(\"\\\nasync def test():\n for stop_exc in (StopIter 225 POP_TOP 226 POP_TOP 227 POP_TOP - 228 JUMP_BACKWARD_NO_INTERRUPT(204) - 229 COPY (3) - 230 POP_EXCEPT - 231 RERAISE (1) + 228 JUMP_BACKWARD (205) + 229 CACHE + 230 COPY (3) + 231 POP_EXCEPT + 232 RERAISE (1) - 2 232 CALL_INTRINSIC_1 (StopIterationError) - 233 RERAISE (1) + 2 233 CALL_INTRINSIC_1 (StopIterationError) + 234 RERAISE (1) 2 MAKE_FUNCTION 3 STORE_NAME (0, test) diff --git a/crates/codegen/src/symboltable.rs b/crates/codegen/src/symboltable.rs index 7179656fe8a..1098c34fac2 100644 --- a/crates/codegen/src/symboltable.rs +++ b/crates/codegen/src/symboltable.rs @@ -571,10 +571,13 @@ impl SymbolTableAnalyzer { } newfree.extend(child_free); } - if let Some(ann_free) = annotation_free { - // Propagate annotation-scope free names to this scope so - // implicit class-scope cells (__classdict__/__conditional_annotations__) - // can be materialized by drop_class_free when needed. + if let Some(ann_free) = annotation_free + && symbol_table.typ == CompilerScope::Class + { + // Annotation-only free variables should not leak into function + // bodies. We only need to propagate them through class scopes so + // drop_class_free() can materialize implicit class cells when + // annotation scopes reference them. newfree.extend(ann_free); } @@ -1569,26 +1572,35 @@ impl SymbolTableBuilder { }) => { // https://github.com/python/cpython/blob/main/Python/symtable.c#L1233 match &**target { - Expr::Name(ast::ExprName { id, .. }) if *simple => { + Expr::Name(ast::ExprName { id, .. }) => { let id_str = id.as_str(); - self.check_name(id_str, ExpressionContext::Store, *range)?; - - self.register_name(id_str, SymbolUsage::AnnotationAssigned, *range)?; - // PEP 649: Register annotate function in module/class scope - let current_scope = self.tables.last().map(|t| t.typ); - match current_scope { - Some(CompilerScope::Module) => { - self.register_name("__annotate__", SymbolUsage::Assigned, *range)?; - } - Some(CompilerScope::Class) => { - self.register_name( - "__annotate_func__", - SymbolUsage::Assigned, - *range, - )?; + if *simple { + self.check_name(id_str, ExpressionContext::Store, *range)?; + + self.register_name(id_str, SymbolUsage::AnnotationAssigned, *range)?; + // PEP 649: Register annotate function in module/class scope + let current_scope = self.tables.last().map(|t| t.typ); + match current_scope { + Some(CompilerScope::Module) => { + self.register_name( + "__annotate__", + SymbolUsage::Assigned, + *range, + )?; + } + Some(CompilerScope::Class) => { + self.register_name( + "__annotate_func__", + SymbolUsage::Assigned, + *range, + )?; + } + _ => {} } - _ => {} + } else if value.is_some() { + self.check_name(id_str, ExpressionContext::Store, *range)?; + self.register_name(id_str, SymbolUsage::Assigned, *range)?; } } _ => { diff --git a/crates/vm/src/stdlib/_ctypes/simple.rs b/crates/vm/src/stdlib/_ctypes/simple.rs index 8f99fa8e57a..3bf1f84fbc5 100644 --- a/crates/vm/src/stdlib/_ctypes/simple.rs +++ b/crates/vm/src/stdlib/_ctypes/simple.rs @@ -309,7 +309,7 @@ impl PyCSimpleType { // Float types: accept numbers Some(tc @ ("f" | "d" | "g")) - if (value.try_float(vm).is_ok() || value.try_int(vm).is_ok()) => + if value.try_float(vm).is_ok() || value.try_int(vm).is_ok() => { return create_simple_with_value(tc, &value); } diff --git a/crates/vm/src/types/slot.rs b/crates/vm/src/types/slot.rs index 232b55110a2..31a03094e8f 100644 --- a/crates/vm/src/types/slot.rs +++ b/crates/vm/src/types/slot.rs @@ -1394,9 +1394,9 @@ impl PyType { SlotAccessor::SqLength => { update_sub_slot!(as_sequence, length, sequence_len_wrapper, SeqLength) } - // Sequence concat uses sq_concat slot - no generic wrapper needed - // (handled by number protocol fallback) SlotAccessor::SqConcat | SlotAccessor::SqInplaceConcat if !ADD => { + // Sequence concat uses sq_concat slot - no generic wrapper needed + // (handled by number protocol fallback) accessor.inherit_from_mro(self); } SlotAccessor::SqRepeat => { diff --git a/scripts/dis_dump.py b/scripts/dis_dump.py index 72fec461b1f..e8b9c1bf5f8 100755 --- a/scripts/dis_dump.py +++ b/scripts/dis_dump.py @@ -11,6 +11,7 @@ """ import argparse +import ast import builtins import dis import json @@ -39,6 +40,7 @@ "JUMP_IF_TRUE_OR_POP", "JUMP_IF_FALSE_OR_POP", "FOR_ITER", + "END_ASYNC_FOR", "SEND", } ) @@ -94,6 +96,14 @@ def _unescape(m): argrepr = re.sub(r"\\u([0-9a-fA-F]{4})", _unescape, argrepr) argrepr = re.sub(r"\\U([0-9a-fA-F]{8})", _unescape, argrepr) + if argrepr.startswith("frozenset({") and argrepr.endswith("})"): + try: + values = ast.literal_eval(argrepr[len("frozenset(") : -1]) + except Exception: + return argrepr + if isinstance(values, set): + parts = sorted(_normalize_argrepr(repr(value)) for value in values) + return f"frozenset({{{', '.join(parts)}}})" return argrepr @@ -249,7 +259,7 @@ def _metadata_cache_slot_offsets(inst): # 1. argval not in offset_to_idx (not a valid byte offset) # 2. argval == arg (raw arg returned as-is, not resolved to offset) # 3. For backward jumps: argval should be < current offset - is_backward = "BACKWARD" in inst.opname + is_backward = "BACKWARD" in inst.opname or inst.opname == "END_ASYNC_FOR" argval_is_raw = inst.argval == inst.arg and inst.arg is not None if target_idx is None or argval_is_raw: target_idx = None # force recalculation