diff --git a/Lib/test/test_compile.py b/Lib/test/test_compile.py index 9676aded5d1..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: @@ -1103,7 +1102,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_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 e20f712a31a..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 = ["{}, {}, {}", "{}, {}, *{}", "{}, *{}, {}", "*{}, {}, {}"] @@ -863,7 +861,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 +903,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/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 1a35d2c23d1..2f7510f0d44 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::{ @@ -124,8 +125,6 @@ enum SuperCallType<'a> { #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum BuiltinGeneratorCallKind { Tuple, - List, - Set, All, Any, } @@ -167,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)] @@ -439,6 +444,34 @@ 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 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, @@ -446,6 +479,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), @@ -485,9 +519,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<()> { @@ -584,12 +671,15 @@ impl Compiler { _ => n > 4, }; - // Fold all-constant collections (>= 3 elements) regardless of size - if !seen_star + 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()) - && let Some(folded) = self.try_fold_constant_collection(elts)? + && can_fold_const_collection + && let Some(folded) = self.try_fold_constant_collection(elts, collection_type)? { match collection_type { CollectionType::Tuple => { @@ -652,7 +742,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 }); @@ -1044,17 +1134,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 @@ -1203,6 +1294,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 @@ -1238,25 +1330,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 @@ -1302,11 +1376,42 @@ impl Compiler { location, end_location, except_handler, + folded_from_nonliteral_expr: false, lineno_override, cache_entries: 0, }); } + 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,16 +1933,14 @@ 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); + emit!(self, PseudoInstruction::AnnotationsPlaceholder); let (doc, statements) = split_doc(&body.body, &self.opts); if let Some(value) = doc { @@ -1854,10 +1957,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__")?; @@ -1868,6 +1969,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: @@ -1886,6 +1991,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) { @@ -1893,10 +1999,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__")?; @@ -1940,6 +2044,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(()) } @@ -2379,7 +2487,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)?; @@ -2405,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, .. @@ -2508,7 +2617,6 @@ impl Compiler { if self.opts.optimize == 0 { let after_block = self.new_block(); self.compile_jump_if(test, true, after_block)?; - emit!( self, Instruction::LoadCommonConstant { @@ -2516,9 +2624,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, @@ -2526,7 +2633,6 @@ impl Compiler { argc: bytecode::RaiseKind::Raise, } ); - self.switch_to_block(after_block); } else { // Optimized-out asserts still need to consume any nested @@ -2593,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() { @@ -3437,16 +3547,26 @@ 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(); - let orelse_block = if orelse.is_empty() { - end_block - } else { - 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!( self, PseudoInstruction::SetupFinally { @@ -3459,11 +3579,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(); @@ -3505,7 +3624,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, @@ -3519,25 +3637,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); @@ -3556,6 +3655,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)?; @@ -3591,18 +3700,10 @@ 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); + if let Some(range) = normal_exit_range { + self.set_source_range(range); + } Ok(()) } @@ -4224,18 +4325,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 = @@ -4244,12 +4360,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 { @@ -4287,24 +4404,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, @@ -4355,10 +4463,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); } @@ -4400,20 +4511,42 @@ 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 }); + + 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; + } - // 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(); + let not_set_block = has_conditional.then(|| self.new_block()); + let name = simple_name.expect("missing simple annotation name"); - // LOAD_CONST index - self.emit_load_const(ConstantData::Integer { value: idx.into() }); - // Load __conditional_annotations__ from appropriate scope - // Class scope: LoadDeref (freevars), Module scope: LoadGlobal + if has_conditional { + 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 }); @@ -4421,60 +4554,35 @@ 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); + simple_idx += 1; - // 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 @@ -5127,9 +5235,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)?; } } @@ -5146,6 +5251,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 @@ -5468,84 +5577,36 @@ impl Compiler { test: &ast::Expr, body: &[ast::Stmt], elif_else_clauses: &[ast::ElifElseClause], + _stmt_range: TextRange, ) -> CompileResult<()> { - let constant = Self::expr_constant(test); + let end_block = self.new_block(); + let next_block = if elif_else_clauses.is_empty() { + end_block + } else { + self.new_block() + }; - // 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 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)?; - // 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; + let Some((clause, rest)) = elif_else_clauses.split_first() else { + self.switch_to_block(end_block); return Ok(()); - } - - // Non-constant test: normal compilation - match elif_else_clauses { - // Only if - [] => { - let after_block = self.new_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(); - - self.compile_jump_if(test, false, next_block)?; - self.compile_statements(body)?; - emit!(self, PseudoInstruction::Jump { delta: after_block }); + }; - 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 }); - } + emit!(self, PseudoInstruction::Jump { delta: end_block }); + self.switch_to_block(next_block); - 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); - } + 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(()) } @@ -5557,37 +5618,17 @@ 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)?; + 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)?; let was_in_loop = self.ctx.loop_data.replace((while_block, after_block)); self.compile_statements(body)?; @@ -5855,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 { @@ -5907,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)?; }; @@ -5943,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__"))); @@ -6150,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()); @@ -6660,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)?; } @@ -6668,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)?; } @@ -6680,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)?; @@ -6691,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. @@ -6702,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(()) } @@ -6790,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, @@ -6798,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(()) } @@ -6826,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) } @@ -6851,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( @@ -6860,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) { @@ -6875,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)?; } @@ -6912,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)?; } @@ -7034,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); } @@ -7085,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 }); @@ -7156,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, @@ -7222,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(()) } @@ -7398,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 { @@ -7419,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, + )?; } } } @@ -7436,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. @@ -7445,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); } } @@ -7460,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, @@ -7507,10 +7669,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, @@ -7534,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<()> { @@ -7846,6 +8014,50 @@ 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")); + 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 => { + simplified_prefix += 1; + } + } + } + + 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 { + 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, .. @@ -8381,8 +8593,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 +8601,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 +8630,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 +8679,6 @@ impl Compiler { } ); } - BuiltinGeneratorCallKind::List | BuiltinGeneratorCallKind::Set => {} BuiltinGeneratorCallKind::All => { self.emit_load_const(ConstantData::Boolean { value: true }); } @@ -8574,18 +8758,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 @@ -8707,7 +8885,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 { @@ -8763,7 +8941,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; @@ -8857,20 +9035,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), } @@ -8889,6 +9063,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, @@ -8910,9 +9142,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 @@ -8930,12 +9159,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, @@ -8966,8 +9201,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(); @@ -9009,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 { @@ -9048,9 +9287,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 { - 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); } } @@ -9107,11 +9351,15 @@ impl Compiler { self.make_closure(code, bytecode::MakeFunctionFlags::new())?; // Evaluate iterated item: - self.compile_expression(&generators[0].iter)?; + self.compile_for_iterable_expression(&outermost.iter, outermost.is_async)?; + 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); @@ -9132,25 +9380,24 @@ 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.compile_for_iterable_expression( + &generators[0].iter, + has_async && generators[0].is_async, + )?; + 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. @@ -9176,22 +9423,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); + } } } @@ -9287,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 { @@ -9324,7 +9598,8 @@ 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 { self.compile_jump_if(if_condition, false, if_cleanup_block)?; } @@ -9446,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) { @@ -9516,14 +9798,85 @@ 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( &mut self, elts: &[ast::Expr], + collection_type: CollectionType, ) -> CompileResult> { let mut constants = Vec::with_capacity(elts.len()); for elt in elts { @@ -9532,9 +9885,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> { @@ -9567,25 +9926,85 @@ impl Compiler { } ConstantData::Tuple { elements } } - _ => return Ok(None), - })) - } - - fn emit_load_const(&mut self, constant: ConstantData) { - let idx = self.arg_constant(constant); - self.emit_arg(idx, |consti| Instruction::LoadConst { consti }) - } - - /// Fold constant slice: if all parts are compile-time constants, emit LOAD_CONST(slice). - fn try_fold_constant_slice( - &mut self, - lower: Option<&ast::Expr>, - upper: Option<&ast::Expr>, - step: Option<&ast::Expr>, - ) -> CompileResult { - let to_const = |expr: Option<&ast::Expr>, this: &mut Self| -> CompileResult<_> { - match expr { - None => Ok(Some(ConstantData::None)), + 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 { + 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), + })) + } + + fn emit_load_const(&mut self, constant: ConstantData) { + let idx = self.arg_constant(constant); + self.emit_arg(idx, |consti| Instruction::LoadConst { consti }) + } + + /// Fold constant slice: if all parts are compile-time constants, emit LOAD_CONST(slice). + fn try_fold_constant_slice( + &mut self, + lower: Option<&ast::Expr>, + upper: Option<&ast::Expr>, + step: Option<&ast::Expr>, + ) -> CompileResult { + let to_const = |expr: Option<&ast::Expr>, this: &mut Self| -> CompileResult<_> { + match expr { + None => Ok(Some(ConstantData::None)), Some(expr) => this.try_fold_constant_expr(expr), } }; @@ -9670,44 +10089,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) { @@ -9800,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 { @@ -9949,7 +10335,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 } @@ -10840,6 +11232,91 @@ 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 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); @@ -10891,6 +11368,195 @@ 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_trace_for_unpack_list_literal() { + let trace = compile_exec_late_cfg_trace( + "\ +result = [] +for x, in [(1,), (2,), (3,)]: + result.append(x) +", + ); + for (stage, dump) in trace { + eprintln!("=== {stage} ===\n{dump}"); + } + } + + #[test] + fn test_trace_break_in_finally_function() { + let trace = compile_single_function_late_cfg_trace( + "\ +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_trace_constant_false_elif_chain() { + let trace = compile_exec_late_cfg_trace( + "\ +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_trace_multi_pass_suite() { + let trace = compile_exec_late_cfg_trace( + "\ +if 1: + # + # + # + pass + pass + # + pass + # +", + ); + for (stage, dump) in trace { + eprintln!("=== {stage} ===\n{dump}"); + } + } + + #[test] + fn test_trace_single_compare_if() { + let trace = compile_exec_late_cfg_trace( + "\ +if 1 == 1: + pass +", + ); + 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( @@ -10930,44 +11596,165 @@ x = not True } #[test] - fn test_nested_double_async_with() { - assert_dis_snapshot!(compile_exec( - "\ -async def test(): - for stop_exc in (StopIteration('spam'), StopAsyncIteration('ham')): - with self.subTest(type=type(stop_exc)): - try: - async with egg(): - raise stop_exc - except Exception as ex: - self.assertIs(ex, stop_exc) - else: - self.fail(f'{stop_exc} was suppressed') -" - )); - } - - #[test] - fn test_scope_exit_instructions_keep_line_numbers() { + fn test_plain_constant_bool_op_folds_to_selected_operand() { let code = compile_exec( "\ -async def test(): - for stop_exc in (StopIteration('spam'), StopAsyncIteration('ham')): - with self.subTest(type=type(stop_exc)): - try: - async with egg(): - raise stop_exc - except Exception as ex: - self.assertIs(ex, stop_exc) - else: - self.fail(f'{stop_exc} was suppressed') +x = 1 or 2 or 3 ", ); - assert_scope_exit_locations(&code); + 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_attribute_ex_call_uses_plain_load_attr() { + 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( + "\ +async def test(): + for stop_exc in (StopIteration('spam'), StopAsyncIteration('ham')): + with self.subTest(type=type(stop_exc)): + try: + async with egg(): + raise stop_exc + except Exception as ex: + self.assertIs(ex, stop_exc) + else: + self.fail(f'{stop_exc} was suppressed') +" + )); + } + + #[test] + fn test_scope_exit_instructions_keep_line_numbers() { + let code = compile_exec( + "\ +async def test(): + for stop_exc in (StopIteration('spam'), StopAsyncIteration('ham')): + with self.subTest(type=type(stop_exc)): + try: + async with egg(): + raise stop_exc + except Exception as ex: + self.assertIs(ex, stop_exc) + else: + self.fail(f'{stop_exc} was suppressed') +", + ); + assert_scope_exit_locations(&code); + } + + #[test] + fn test_attribute_ex_call_uses_plain_load_attr() { let code = compile_exec( "\ def f(cls, args, kwargs): @@ -11060,7 +11847,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 +11878,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)" ); } @@ -11458,137 +12247,836 @@ def f(parts): } #[test] - fn test_assert_without_message_raises_class_directly() { + fn test_for_exit_before_elif_does_not_leave_line_anchor_nop() { let code = compile_exec( "\ -def f(x): - assert x +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 f = find_code(&code, "f").expect("missing function code"); - let call_count = f + let ops: Vec<_> = code .instructions .iter() - .filter(|unit| matches!(unit.op, Instruction::Call { .. })) - .count(); - let push_null_count = f + .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() - .filter(|unit| matches!(unit.op, Instruction::PushNull)) - .count(); + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); - assert_eq!(call_count, 0); - assert_eq!(push_null_count, 0); + 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_chained_compare_jump_uses_single_cleanup_copy() { + fn test_assert_without_message_raises_class_directly() { let code = compile_exec( "\ -def f(code): - if not 1 <= code <= 2147483647: - raise ValueError('x') +def f(x): + assert x ", ); let f = find_code(&code, "f").expect("missing function code"); - let copy_count = f + let call_count = f .instructions .iter() - .filter(|unit| matches!(unit.op, Instruction::Copy { .. })) + .filter(|unit| matches!(unit.op, Instruction::Call { .. })) .count(); - let pop_top_count = f + let push_null_count = f .instructions .iter() - .filter(|unit| matches!(unit.op, Instruction::PopTop)) + .filter(|unit| matches!(unit.op, Instruction::PushNull)) .count(); - assert_eq!(copy_count, 1); - assert_eq!(pop_top_count, 1); + assert_eq!(call_count, 0); + assert_eq!(push_null_count, 0); } #[test] - fn test_yield_from_cleanup_jumps_to_shared_end_send() { + fn test_assert_with_message_uses_common_constant_direct_call() { let code = compile_exec( "\ -def outer(): - def inner(): - yield from outer_gen - return inner +def f(x, y): + assert x, y ", ); - let inner = find_code(&code, "inner").expect("missing inner code"); - let ops: Vec<_> = inner + let f = find_code(&code, "f").expect("missing f code"); + let load_assertion = f .instructions .iter() - .map(|unit| unit.op) - .filter(|op| !matches!(op, Instruction::Cache)) - .collect(); + .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"); - let cleanup_idx = ops - .iter() - .position(|op| matches!(op, Instruction::CleanupThrow)) - .expect("missing CLEANUP_THROW"); assert!( - matches!( - ops.get(cleanup_idx + 1), - Some(Instruction::JumpBackwardNoInterrupt { .. }) - | Some(Instruction::JumpForward { .. }) + !matches!( + f.instructions.get(load_assertion + 1).map(|unit| unit.op), + Some(Instruction::PushNull) ), - "expected CLEANUP_THROW to jump to shared END_SEND block, got ops={ops:?}" + "assert message path should not use PUSH_NULL, got ops={:?}", + f.instructions + .iter() + .map(|unit| unit.op) + .collect::>() ); assert!( - !matches!(ops.get(cleanup_idx + 1), Some(Instruction::EndSend)), - "CLEANUP_THROW should not inline END_SEND directly, got ops={ops:?}" + 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_try_except_falls_through_to_post_handler_code() { + 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 f(): - try: - line = 2 - raise KeyError - except: - line = 5 - line = 6 +def f2bad(): + (no_such_global): int + print(no_such_global) ", ); - 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 f = find_code(&code, "f2bad").expect("missing f2bad code"); assert!( - !matches!( - ops.get(first_pop_except + 1), - Some(Instruction::JumpForward { .. }) - ), - "expected except body to fall through to post-handler code, got ops={ops:?}" + 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!( - matches!( - ops.get(first_pop_except + 1), - Some(Instruction::LoadSmallInt { .. }) | Some(Instruction::LoadConst { .. }) - ), - "expected line-after-except code immediately after POP_EXCEPT, got ops={ops:?}" + !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_slice_folding_handles_string_and_bigint_bounds() { + 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( "\ -def f(obj): - return obj['a':123456789012345678901234567890] +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_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( + "\ +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_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( + "\ +def f(code): + if not 1 <= code <= 2147483647: + raise ValueError('x') +", + ); + let f = find_code(&code, "f").expect("missing function code"); + let copy_count = f + .instructions + .iter() + .filter(|unit| matches!(unit.op, Instruction::Copy { .. })) + .count(); + let pop_top_count = f + .instructions + .iter() + .filter(|unit| matches!(unit.op, Instruction::PopTop)) + .count(); + + assert_eq!(copy_count, 1); + assert_eq!(pop_top_count, 1); + } + + #[test] + fn test_yield_from_cleanup_jumps_to_shared_end_send() { + let code = compile_exec( + "\ +def outer(): + def inner(): + yield from outer_gen + return inner +", + ); + let inner = find_code(&code, "inner").expect("missing inner code"); + let ops: Vec<_> = inner + .instructions + .iter() + .map(|unit| unit.op) + .filter(|op| !matches!(op, Instruction::Cache)) + .collect(); + + let cleanup_idx = ops + .iter() + .position(|op| matches!(op, Instruction::CleanupThrow)) + .expect("missing CLEANUP_THROW"); + assert!( + matches!( + ops.get(cleanup_idx + 1), + Some(Instruction::JumpBackwardNoInterrupt { .. }) + | Some(Instruction::JumpForward { .. }) + ), + "expected CLEANUP_THROW to jump to shared END_SEND block, got ops={ops:?}" + ); + assert!( + !matches!(ops.get(cleanup_idx + 1), Some(Instruction::EndSend)), + "CLEANUP_THROW should not inline END_SEND directly, got ops={ops:?}" + ); + } + + #[test] + fn test_try_except_falls_through_to_post_handler_code() { + let code = compile_exec( + "\ +def f(): + try: + line = 2 + raise KeyError + except: + line = 5 + line = 6 +", + ); + 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"); + assert!( + !matches!( + ops.get(first_pop_except + 1), + Some(Instruction::JumpForward { .. }) + ), + "expected except body to fall through to post-handler code, got ops={ops:?}" + ); + assert!( + matches!( + ops.get(first_pop_except + 1), + Some(Instruction::LoadSmallInt { .. }) | Some(Instruction::LoadConst { .. }) + ), + "expected line-after-except code immediately after POP_EXCEPT, got ops={ops:?}" + ); + } + + #[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( + "\ +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( + "\ +def f(obj): + return obj['a':123456789012345678901234567890] ", ); let f = find_code(&code, "f").expect("missing function code"); @@ -11812,7 +13300,7 @@ def f(): ); assert!(f.constants.iter().any(|constant| matches!( constant, - ConstantData::Tuple { elements } + ConstantData::Frozenset { elements } if matches!( elements.as_slice(), [ @@ -11824,6 +13312,306 @@ 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"); + + 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_uses_direct_raise_shape() { + 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 use direct raise shape, 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_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( @@ -11864,4 +13652,40 @@ 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 61e549199d5..01e6971f65b 100644 --- a/crates/codegen/src/ir.rs +++ b/crates/codegen/src/ir.rs @@ -3,15 +3,17 @@ 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::{ 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, + 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; } @@ -148,6 +152,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 +165,7 @@ impl Default for Block { preserve_lasti: false, start_depth: None, cold: false, + disable_load_fast_borrow: false, } } } @@ -170,6 +177,7 @@ pub struct CodeInfo { pub blocks: Vec, pub current_block: BlockIdx, + pub annotations_blocks: Option>, pub metadata: CodeUnitMetadata, @@ -199,20 +207,17 @@ impl CodeInfo { mut self, opts: &crate::compile::CompileOpts, ) -> crate::InternalResult { + self.splice_annotations_blocks(); // Constant folding passes self.fold_binop_constants(); - self.remove_nops(); - self.fold_unary_negative(); + 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(); @@ -223,7 +228,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) @@ -235,7 +240,11 @@ impl CodeInfo { jump_threading(&mut self.blocks); self.eliminate_unreachable_blocks(); self.remove_nops(); - // TODO: insert_superinstructions disabled pending StoreFastLoadFast VM fix + 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) @@ -247,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 @@ -260,12 +271,39 @@ impl CodeInfo { reorder_jump_over_exception_cleanup_blocks(&mut self.blocks); self.eliminate_unreachable_blocks(); remove_redundant_nops_and_jumps(&mut self.blocks); - self.add_checks_for_loads_of_uninitialized_variables(); + // 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, + &self.metadata.freevars, + ); + // 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); + 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. + 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); + self.compute_load_fast_start_depths(); // 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(); - - let max_stackdepth = self.max_stackdepth()?; + self.reorder_entry_prefix_cell_setup(); + self.remove_unused_consts(); let Self { flags, @@ -274,6 +312,7 @@ impl CodeInfo { mut blocks, current_block: _, + annotations_blocks: _, metadata, static_attributes: _, in_inlined_comp: _, @@ -308,6 +347,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(); @@ -330,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; @@ -620,6 +657,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) { @@ -668,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; } @@ -720,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 @@ -728,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; @@ -759,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; @@ -781,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, @@ -806,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 { @@ -817,6 +1039,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; @@ -886,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, @@ -914,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) => @@ -962,6 +1219,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 { @@ -972,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; @@ -1026,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; } @@ -1053,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; @@ -1101,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; } @@ -1137,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 { .. }) @@ -1186,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), @@ -1207,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; } @@ -1241,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; @@ -1282,27 +1606,35 @@ 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; 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; } @@ -1366,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!( @@ -1410,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) { @@ -1472,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); + } } } } @@ -1544,11 +1942,35 @@ 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() { 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()) @@ -1557,65 +1979,136 @@ impl CodeInfo { continue; }; - if matches!( - next_instr.into(), - Opcode::PopJumpIfFalse | Opcode::PopJumpIfTrue - ) && matches!(curr_instr.into(), Opcode::CompareOp) - { - block.instructions[i].arg = OpArg::new( - u32::from(block.instructions[i].arg) | oparg::COMPARE_OP_BOOL_MASK, - ); - 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 - } + 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 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 { .. } + ) && matches!(next_instr, Instruction::PopTop) + { + 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) { // 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 } @@ -1686,6 +2179,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) { @@ -1787,52 +2313,96 @@ 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 }); } } - /// 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)] - fn combine_store_fast_load_fast(&mut self) { + /// 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 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 { .. }), Some(Instruction::LoadFast { .. })) = - (curr.instr.real(), next.instr.real()) - else { + let line = curr.location.line; + + let mut j = i + 1; + while j < block.instructions.len() + && matches!(block.instructions[j].instr.real(), Some(Instruction::Nop)) + && block.instructions[j].location.line == line + { + j += 1; + } + if j >= block.instructions.len() { i += 1; continue; - }; - // Skip if instructions are on different lines (matching make_super_instruction) - let line1 = curr.location.line; - let line2 = next.location.line; - if line1 != line2 { + } + + let next = &block.instructions[j]; + if next.location.line != line { 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 { - i += 1; + + 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 = (idx1 << 4) | idx2; + block.instructions[i].instr = Instruction::LoadFastLoadFast { + var_nums: Arg::marker(), + } + .into(); + block.instructions[i].arg = OpArg::new(packed); + block.instructions.drain(i + 1..=j); + } + ( + 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 { + 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); + block.instructions.drain(i + 1..=j); + } + (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].arg = OpArg::new(packed); + block.instructions.drain(i + 1..=j); + } + _ => i += 1, } } } @@ -1841,86 +2411,654 @@ 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 { - if block.instructions.is_empty() { - continue; + 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], + source: BlockIdx, + target: BlockIdx, + start_depth: 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={:?}", + Some(start_depth), + blocks[source.idx()] + .instructions + .last() + .and_then(|info| info.instr.real()), + blocks[target.idx()] + .instructions + .iter() + .map(|info| info.instr) + .collect::>(), + ); + return; } + 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; - // 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()]; + while let Some(block_idx) = worklist.pop() { + let block = &self.blocks[block_idx]; - // 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 { + let instr = info.instr; + let arg_u32 = u32::from(info.arg); + + 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 { .. }) => { + let target = info.target; + if target != BlockIdx::NULL { + push_block( + &mut worklist, + &mut visited, + &self.blocks, + block_idx, + 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 { .. }) => { + let target = info.target; + if target != BlockIdx::NULL { + push_block( + &mut worklist, + &mut visited, + &self.blocks, + block_idx, + 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; + let target = info.target; + if target != BlockIdx::NULL { + let target_depth = refs + .len() + .saturating_sub(num_popped) + .saturating_add(num_pushed); + push_block( + &mut worklist, + &mut visited, + &self.blocks, + block_idx, + 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); + } + } + } + } + } + + let next = block.next; + if next != BlockIdx::NULL + && block.instructions.last().is_none_or(|term| { + !term.instr.is_unconditional_jump() && !term.instr.is_scope_exit() + }) + { + push_block( + &mut worklist, + &mut visited, + &self.blocks, + block_idx, + next, + refs.len(), + ); + } + + for r in refs { + if r.instr != DUMMY_INSTR { + instr_flags[r.instr as usize] |= REF_UNCONSUMED; + } + } + + 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; + } + 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 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 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; + }; + 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_borrow_after_push_exc_info(&mut self) { + for block in &mut self.blocks { + let mut in_exception_state = false; + for info in &mut block.instructions { + match info.instr.real() { + Some(Instruction::PushExcInfo) => { + in_exception_state = true; + } + Some(Instruction::PopExcept) | Some(Instruction::Reraise { .. }) => { + in_exception_state = false; + } + Some(Instruction::LoadFastBorrow { .. }) if in_exception_state => { + info.instr = Instruction::LoadFast { + var_num: Arg::marker(), + } + .into(); + } + Some(Instruction::LoadFastBorrowLoadFastBorrow { .. }) + if in_exception_state => + { + info.instr = Instruction::LoadFastLoadFast { + var_nums: Arg::marker(), + } + .into(); + } + _ => {} + } + } + } + } + + fn deoptimize_borrow_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 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; + 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; } - } - if underflow { - 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; } - // 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 saw_build_tuple && saw_match_keys { + to_deopt.push((block_idx, i)); } } + } - if underflow { - continue; + 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 + .instructions + .iter() + .rev() + .find_map(|info| info.instr.real()) + } - // Mark instructions whose values remain on stack at block end - for &src in &stack { - if src != NOT_LOCAL { - unconsumed[src] = true; + 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)); } } + } - // Convert LOAD_FAST to LOAD_FAST_BORROW where value is fully consumed - for (i, info) in block.instructions.iter_mut().enumerate() { - if unconsumed[i] { - continue; + let starts_after_cleanup: Vec = predecessors + .iter() + .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 { .. }) + ) + }) + }) + .collect(); + + 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 Some(instr) = info.instr.real() else { - continue; - }; - match instr.into() { - Opcode::LoadFast => { - info.instr = Opcode::LoadFastBorrow.into(); + 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]) + || in_restore_prefix); + + 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; } + + in_restore_prefix &= + matches!(info.instr.real(), Some(Instruction::StoreFast { .. })); + new_instructions.push(info); } + block.instructions = new_instructions; } } @@ -1930,6 +3068,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; @@ -1967,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)); @@ -1989,6 +3140,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 +3169,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 +3190,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); @@ -2140,7 +3311,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() { @@ -2148,8 +3322,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 { @@ -2182,6 +3357,261 @@ 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={} start_depth={}", + 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, + 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, + " [disp={} raw={} override={:?}] {:?} arg={} target={}", + lineno, + info.location.line.get(), + info.lineno_override, + 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(); + trace.push(("initial".to_owned(), self.debug_block_dump())); + + self.splice_annotations_blocks(); + self.fold_binop_constants(); + self.fold_unary_constants(); + self.fold_binop_constants(); + self.fold_tuple_constants(); + self.fold_list_constants(); + self.fold_set_constants(); + self.fold_const_iterable_for_iter(); + self.convert_to_load_small_int(); + self.remove_unused_consts(); + 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(), + )); + + 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(), + )); + + 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(), + )); + + 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(), + )); + + 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) + } +} + +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 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); + 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; @@ -2543,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, }); @@ -2624,11 +4055,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) @@ -2676,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, }) } @@ -2723,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 @@ -2830,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, }; @@ -2845,11 +4289,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 { @@ -2859,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, }); @@ -2869,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, }); @@ -2949,7 +4397,30 @@ fn inline_small_or_no_lineno_blocks(blocks: &mut [Block]) { .iter() .all(|ins| !instruction_has_lineno(ins)) }; - + let current_is_named_except_cleanup_normal_exit = |block: &Block| { + let len = block.instructions.len(); + if len < 5 { + return false; + } + let tail = &block.instructions[len - 5..]; + matches!(tail[0].instr.real(), Some(Instruction::PopExcept)) + && matches!(tail[1].instr.real(), Some(Instruction::LoadConst { .. })) + && matches!( + tail[2].instr.real(), + Some(Instruction::StoreName { .. } | Instruction::StoreFast { .. }) + ) + && matches!( + tail[3].instr.real(), + Some(Instruction::DeleteName { .. } | Instruction::DeleteFast { .. }) + ) + && tail[4].instr.is_unconditional_jump() + }; + let target_pushes_handler = |block: &Block| { + block + .instructions + .iter() + .any(|ins| ins.instr.is_block_push()) + }; loop { let mut changes = false; let mut current = BlockIdx(0); @@ -2967,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; @@ -2975,12 +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 || 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; } @@ -3018,28 +4495,40 @@ 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 { - 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); + 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; + } } } 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 { + 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; + } } } } @@ -3093,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) => { @@ -3112,6 +4628,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 @@ -3447,6 +4995,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()]) { @@ -3454,6 +5003,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; @@ -3461,6 +5011,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; @@ -3487,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, @@ -3707,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; } @@ -3785,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 @@ -3795,14 +5374,15 @@ 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); 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 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 0cdb7f0a3df..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,23 +1,8 @@ --- source: crates/codegen/src/compile.rs -assertion_line: 10890 +assertion_line: 11413 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 - - 2 13 LOAD_CONST (None) - 14 RETURN_VALUE - 15 LOAD_CONST (None) - 16 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 8e91189912d..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,27 +1,8 @@ --- source: crates/codegen/src/compile.rs -assertion_line: 10900 +assertion_line: 11423 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 - - 2 17 LOAD_CONST (None) - 18 RETURN_VALUE - 19 LOAD_CONST (None) - 20 RETURN_VALUE + 1 LOAD_CONST (None) + 2 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..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: 10936 +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) @@ -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,115 +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 (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) + 226 POP_TOP + 227 POP_TOP + 228 JUMP_BACKWARD (205) + 229 CACHE + 230 COPY (3) + 231 POP_EXCEPT + 232 RERAISE (1) - 2 230 CALL_INTRINSIC_1 (StopIterationError) - 231 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 f8d7a2b26ea..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,48 +1572,42 @@ 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)?; } } _ => { 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)?; } @@ -1639,6 +1636,10 @@ impl SymbolTableBuilder { let saved_in_conditional_block = self.in_conditional_block; self.in_conditional_block = true; self.scan_statements(body)?; + // 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_, @@ -1654,7 +1655,22 @@ impl SymbolTableBuilder { } self.scan_statements(body)?; } - self.scan_statements(orelse)?; + 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; } @@ -2160,6 +2176,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, @@ -2195,36 +2218,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(()) } 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/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/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/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/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/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) } } 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/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 e30a8955fdd..e8b9c1bf5f8 100755 --- a/scripts/dis_dump.py +++ b/scripts/dis_dump.py @@ -11,6 +11,8 @@ """ import argparse +import ast +import builtins import dis import json import os @@ -38,6 +40,7 @@ "JUMP_IF_TRUE_OR_POP", "JUMP_IF_FALSE_OR_POP", "FOR_ITER", + "END_ASYNC_FOR", "SEND", } ) @@ -93,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 @@ -100,10 +111,33 @@ 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 +147,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 +158,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", @@ -237,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