From 9be8a94248e1de40495745fd4156d6b34ae7770c Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Thu, 12 Mar 2026 23:48:38 +0900 Subject: [PATCH 1/4] Enable PEP 709 inlined comprehensions for function-like scopes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Activate the existing compile_inlined_comprehension() implementation by fixing 6 bugs that prevented it from working: - LoadFastAndClear: push NULL (not None) when slot is empty so StoreFast can restore empty state after comprehension - StoreFast: accept NULL from stack for the restore path - sub_tables.remove(0) replaced with next_sub_table cursor to match the pattern used elsewhere in the compiler - in_inlined_comp flag moved from non-inlined to inlined path - is_inlined_comprehension_context() now checks comp_inlined flag and restricts inlining to function-like scopes - comp_inlined set only when parent scope uses fastlocals Symbol table analysis handles conflict detection: - Nested scopes in comprehension → skip inlining - Bound name conflicts with parent symbol → skip inlining - Cross-comprehension reference conflicts → skip inlining - Splice comprehension sub_tables into parent for nested scope tracking --- Lib/_opcode_metadata.py | 308 ++++++------------ crates/codegen/src/compile.rs | 241 +++++++++----- crates/codegen/src/symboltable.rs | 176 +++++++--- .../compiler-core/src/bytecode/instruction.rs | 11 +- crates/vm/src/frame.rs | 53 ++- 5 files changed, 433 insertions(+), 356 deletions(-) diff --git a/Lib/_opcode_metadata.py b/Lib/_opcode_metadata.py index 4da6e507736..f1656a35cb1 100644 --- a/Lib/_opcode_metadata.py +++ b/Lib/_opcode_metadata.py @@ -1,223 +1,16 @@ # This file is generated by scripts/generate_opcode_metadata.py -# for RustPython bytecode format (CPython 3.14 compatible opcode numbers). +# for RustPython bytecode format (CPython 3.13 compatible opcode numbers). # Do not edit! -_specializations = { - "RESUME": [ - "RESUME_CHECK", - ], - "LOAD_CONST": [ - "LOAD_CONST_MORTAL", - "LOAD_CONST_IMMORTAL", - ], - "TO_BOOL": [ - "TO_BOOL_ALWAYS_TRUE", - "TO_BOOL_BOOL", - "TO_BOOL_INT", - "TO_BOOL_LIST", - "TO_BOOL_NONE", - "TO_BOOL_STR", - ], - "BINARY_OP": [ - "BINARY_OP_MULTIPLY_INT", - "BINARY_OP_ADD_INT", - "BINARY_OP_SUBTRACT_INT", - "BINARY_OP_MULTIPLY_FLOAT", - "BINARY_OP_ADD_FLOAT", - "BINARY_OP_SUBTRACT_FLOAT", - "BINARY_OP_ADD_UNICODE", - "BINARY_OP_SUBSCR_LIST_INT", - "BINARY_OP_SUBSCR_LIST_SLICE", - "BINARY_OP_SUBSCR_TUPLE_INT", - "BINARY_OP_SUBSCR_STR_INT", - "BINARY_OP_SUBSCR_DICT", - "BINARY_OP_SUBSCR_GETITEM", - "BINARY_OP_EXTEND", - "BINARY_OP_INPLACE_ADD_UNICODE", - ], - "STORE_SUBSCR": [ - "STORE_SUBSCR_DICT", - "STORE_SUBSCR_LIST_INT", - ], - "SEND": [ - "SEND_GEN", - ], - "UNPACK_SEQUENCE": [ - "UNPACK_SEQUENCE_TWO_TUPLE", - "UNPACK_SEQUENCE_TUPLE", - "UNPACK_SEQUENCE_LIST", - ], - "STORE_ATTR": [ - "STORE_ATTR_INSTANCE_VALUE", - "STORE_ATTR_SLOT", - "STORE_ATTR_WITH_HINT", - ], - "LOAD_GLOBAL": [ - "LOAD_GLOBAL_MODULE", - "LOAD_GLOBAL_BUILTIN", - ], - "LOAD_SUPER_ATTR": [ - "LOAD_SUPER_ATTR_ATTR", - "LOAD_SUPER_ATTR_METHOD", - ], - "LOAD_ATTR": [ - "LOAD_ATTR_INSTANCE_VALUE", - "LOAD_ATTR_MODULE", - "LOAD_ATTR_WITH_HINT", - "LOAD_ATTR_SLOT", - "LOAD_ATTR_CLASS", - "LOAD_ATTR_CLASS_WITH_METACLASS_CHECK", - "LOAD_ATTR_PROPERTY", - "LOAD_ATTR_GETATTRIBUTE_OVERRIDDEN", - "LOAD_ATTR_METHOD_WITH_VALUES", - "LOAD_ATTR_METHOD_NO_DICT", - "LOAD_ATTR_METHOD_LAZY_DICT", - "LOAD_ATTR_NONDESCRIPTOR_WITH_VALUES", - "LOAD_ATTR_NONDESCRIPTOR_NO_DICT", - ], - "COMPARE_OP": [ - "COMPARE_OP_FLOAT", - "COMPARE_OP_INT", - "COMPARE_OP_STR", - ], - "CONTAINS_OP": [ - "CONTAINS_OP_SET", - "CONTAINS_OP_DICT", - ], - "JUMP_BACKWARD": [ - "JUMP_BACKWARD_NO_JIT", - "JUMP_BACKWARD_JIT", - ], - "FOR_ITER": [ - "FOR_ITER_LIST", - "FOR_ITER_TUPLE", - "FOR_ITER_RANGE", - "FOR_ITER_GEN", - ], - "CALL": [ - "CALL_BOUND_METHOD_EXACT_ARGS", - "CALL_PY_EXACT_ARGS", - "CALL_TYPE_1", - "CALL_STR_1", - "CALL_TUPLE_1", - "CALL_BUILTIN_CLASS", - "CALL_BUILTIN_O", - "CALL_BUILTIN_FAST", - "CALL_BUILTIN_FAST_WITH_KEYWORDS", - "CALL_LEN", - "CALL_ISINSTANCE", - "CALL_LIST_APPEND", - "CALL_METHOD_DESCRIPTOR_O", - "CALL_METHOD_DESCRIPTOR_FAST_WITH_KEYWORDS", - "CALL_METHOD_DESCRIPTOR_NOARGS", - "CALL_METHOD_DESCRIPTOR_FAST", - "CALL_ALLOC_AND_ENTER_INIT", - "CALL_PY_GENERAL", - "CALL_BOUND_METHOD_GENERAL", - "CALL_NON_PY_GENERAL", - ], - "CALL_KW": [ - "CALL_KW_BOUND_METHOD", - "CALL_KW_PY", - "CALL_KW_NON_PY", - ], -} +_specializations = {} -_specialized_opmap = { - 'BINARY_OP_ADD_FLOAT': 129, - 'BINARY_OP_ADD_INT': 130, - 'BINARY_OP_ADD_UNICODE': 131, - 'BINARY_OP_EXTEND': 132, - 'BINARY_OP_INPLACE_ADD_UNICODE': 3, - 'BINARY_OP_MULTIPLY_FLOAT': 133, - 'BINARY_OP_MULTIPLY_INT': 134, - 'BINARY_OP_SUBSCR_DICT': 135, - 'BINARY_OP_SUBSCR_GETITEM': 136, - 'BINARY_OP_SUBSCR_LIST_INT': 137, - 'BINARY_OP_SUBSCR_LIST_SLICE': 138, - 'BINARY_OP_SUBSCR_STR_INT': 139, - 'BINARY_OP_SUBSCR_TUPLE_INT': 140, - 'BINARY_OP_SUBTRACT_FLOAT': 141, - 'BINARY_OP_SUBTRACT_INT': 142, - 'CALL_ALLOC_AND_ENTER_INIT': 143, - 'CALL_BOUND_METHOD_EXACT_ARGS': 144, - 'CALL_BOUND_METHOD_GENERAL': 145, - 'CALL_BUILTIN_CLASS': 146, - 'CALL_BUILTIN_FAST': 147, - 'CALL_BUILTIN_FAST_WITH_KEYWORDS': 148, - 'CALL_BUILTIN_O': 149, - 'CALL_ISINSTANCE': 150, - 'CALL_KW_BOUND_METHOD': 151, - 'CALL_KW_NON_PY': 152, - 'CALL_KW_PY': 153, - 'CALL_LEN': 154, - 'CALL_LIST_APPEND': 155, - 'CALL_METHOD_DESCRIPTOR_FAST': 156, - 'CALL_METHOD_DESCRIPTOR_FAST_WITH_KEYWORDS': 157, - 'CALL_METHOD_DESCRIPTOR_NOARGS': 158, - 'CALL_METHOD_DESCRIPTOR_O': 159, - 'CALL_NON_PY_GENERAL': 160, - 'CALL_PY_EXACT_ARGS': 161, - 'CALL_PY_GENERAL': 162, - 'CALL_STR_1': 163, - 'CALL_TUPLE_1': 164, - 'CALL_TYPE_1': 165, - 'COMPARE_OP_FLOAT': 166, - 'COMPARE_OP_INT': 167, - 'COMPARE_OP_STR': 168, - 'CONTAINS_OP_DICT': 169, - 'CONTAINS_OP_SET': 170, - 'FOR_ITER_GEN': 171, - 'FOR_ITER_LIST': 172, - 'FOR_ITER_RANGE': 173, - 'FOR_ITER_TUPLE': 174, - 'JUMP_BACKWARD_JIT': 175, - 'JUMP_BACKWARD_NO_JIT': 176, - 'LOAD_ATTR_CLASS': 177, - 'LOAD_ATTR_CLASS_WITH_METACLASS_CHECK': 178, - 'LOAD_ATTR_GETATTRIBUTE_OVERRIDDEN': 179, - 'LOAD_ATTR_INSTANCE_VALUE': 180, - 'LOAD_ATTR_METHOD_LAZY_DICT': 181, - 'LOAD_ATTR_METHOD_NO_DICT': 182, - 'LOAD_ATTR_METHOD_WITH_VALUES': 183, - 'LOAD_ATTR_MODULE': 184, - 'LOAD_ATTR_NONDESCRIPTOR_NO_DICT': 185, - 'LOAD_ATTR_NONDESCRIPTOR_WITH_VALUES': 186, - 'LOAD_ATTR_PROPERTY': 187, - 'LOAD_ATTR_SLOT': 188, - 'LOAD_ATTR_WITH_HINT': 189, - 'LOAD_CONST_IMMORTAL': 190, - 'LOAD_CONST_MORTAL': 191, - 'LOAD_GLOBAL_BUILTIN': 192, - 'LOAD_GLOBAL_MODULE': 193, - 'LOAD_SUPER_ATTR_ATTR': 194, - 'LOAD_SUPER_ATTR_METHOD': 195, - 'RESUME_CHECK': 196, - 'SEND_GEN': 197, - 'STORE_ATTR_INSTANCE_VALUE': 198, - 'STORE_ATTR_SLOT': 199, - 'STORE_ATTR_WITH_HINT': 200, - 'STORE_SUBSCR_DICT': 201, - 'STORE_SUBSCR_LIST_INT': 202, - 'TO_BOOL_ALWAYS_TRUE': 203, - 'TO_BOOL_BOOL': 204, - 'TO_BOOL_INT': 205, - 'TO_BOOL_LIST': 206, - 'TO_BOOL_NONE': 207, - 'TO_BOOL_STR': 208, - 'UNPACK_SEQUENCE_LIST': 209, - 'UNPACK_SEQUENCE_TUPLE': 210, - 'UNPACK_SEQUENCE_TWO_TUPLE': 211, -} +_specialized_opmap = {} opmap = { 'CACHE': 0, - 'RESERVED': 17, - 'RESUME': 128, - 'INSTRUMENTED_LINE': 254, - 'ENTER_EXECUTOR': 255, 'BINARY_SLICE': 1, 'BUILD_TEMPLATE': 2, + 'BINARY_OP_INPLACE_ADD_UNICODE': 3, 'CALL_FUNCTION_EX': 4, 'CHECK_EG_MATCH': 5, 'CHECK_EXC_MATCH': 6, @@ -231,6 +24,7 @@ 'GET_AITER': 14, 'GET_ANEXT': 15, 'GET_ITER': 16, + 'RESERVED': 17, 'GET_LEN': 18, 'GET_YIELD_FROM_ITER': 19, 'INTERPRETER_EXIT': 20, @@ -334,6 +128,91 @@ 'UNPACK_EX': 118, 'UNPACK_SEQUENCE': 119, 'YIELD_VALUE': 120, + 'RESTORE_CELL': 121, + 'RESUME': 128, + 'BINARY_OP_ADD_FLOAT': 129, + 'BINARY_OP_ADD_INT': 130, + 'BINARY_OP_ADD_UNICODE': 131, + 'BINARY_OP_EXTEND': 132, + 'BINARY_OP_MULTIPLY_FLOAT': 133, + 'BINARY_OP_MULTIPLY_INT': 134, + 'BINARY_OP_SUBSCR_DICT': 135, + 'BINARY_OP_SUBSCR_GETITEM': 136, + 'BINARY_OP_SUBSCR_LIST_INT': 137, + 'BINARY_OP_SUBSCR_LIST_SLICE': 138, + 'BINARY_OP_SUBSCR_STR_INT': 139, + 'BINARY_OP_SUBSCR_TUPLE_INT': 140, + 'BINARY_OP_SUBTRACT_FLOAT': 141, + 'BINARY_OP_SUBTRACT_INT': 142, + 'CALL_ALLOC_AND_ENTER_INIT': 143, + 'CALL_BOUND_METHOD_EXACT_ARGS': 144, + 'CALL_BOUND_METHOD_GENERAL': 145, + 'CALL_BUILTIN_CLASS': 146, + 'CALL_BUILTIN_FAST': 147, + 'CALL_BUILTIN_FAST_WITH_KEYWORDS': 148, + 'CALL_BUILTIN_O': 149, + 'CALL_ISINSTANCE': 150, + 'CALL_KW_BOUND_METHOD': 151, + 'CALL_KW_NON_PY': 152, + 'CALL_KW_PY': 153, + 'CALL_LEN': 154, + 'CALL_LIST_APPEND': 155, + 'CALL_METHOD_DESCRIPTOR_FAST': 156, + 'CALL_METHOD_DESCRIPTOR_FAST_WITH_KEYWORDS': 157, + 'CALL_METHOD_DESCRIPTOR_NOARGS': 158, + 'CALL_METHOD_DESCRIPTOR_O': 159, + 'CALL_NON_PY_GENERAL': 160, + 'CALL_PY_EXACT_ARGS': 161, + 'CALL_PY_GENERAL': 162, + 'CALL_STR_1': 163, + 'CALL_TUPLE_1': 164, + 'CALL_TYPE_1': 165, + 'COMPARE_OP_FLOAT': 166, + 'COMPARE_OP_INT': 167, + 'COMPARE_OP_STR': 168, + 'CONTAINS_OP_DICT': 169, + 'CONTAINS_OP_SET': 170, + 'FOR_ITER_GEN': 171, + 'FOR_ITER_LIST': 172, + 'FOR_ITER_RANGE': 173, + 'FOR_ITER_TUPLE': 174, + 'JUMP_BACKWARD_JIT': 175, + 'JUMP_BACKWARD_NO_JIT': 176, + 'LOAD_ATTR_CLASS': 177, + 'LOAD_ATTR_CLASS_WITH_METACLASS_CHECK': 178, + 'LOAD_ATTR_GETATTRIBUTE_OVERRIDDEN': 179, + 'LOAD_ATTR_INSTANCE_VALUE': 180, + 'LOAD_ATTR_METHOD_LAZY_DICT': 181, + 'LOAD_ATTR_METHOD_NO_DICT': 182, + 'LOAD_ATTR_METHOD_WITH_VALUES': 183, + 'LOAD_ATTR_MODULE': 184, + 'LOAD_ATTR_NONDESCRIPTOR_NO_DICT': 185, + 'LOAD_ATTR_NONDESCRIPTOR_WITH_VALUES': 186, + 'LOAD_ATTR_PROPERTY': 187, + 'LOAD_ATTR_SLOT': 188, + 'LOAD_ATTR_WITH_HINT': 189, + 'LOAD_CONST_IMMORTAL': 190, + 'LOAD_CONST_MORTAL': 191, + 'LOAD_GLOBAL_BUILTIN': 192, + 'LOAD_GLOBAL_MODULE': 193, + 'LOAD_SUPER_ATTR_ATTR': 194, + 'LOAD_SUPER_ATTR_METHOD': 195, + 'RESUME_CHECK': 196, + 'SEND_GEN': 197, + 'STORE_ATTR_INSTANCE_VALUE': 198, + 'STORE_ATTR_SLOT': 199, + 'STORE_ATTR_WITH_HINT': 200, + 'STORE_SUBSCR_DICT': 201, + 'STORE_SUBSCR_LIST_INT': 202, + 'TO_BOOL_ALWAYS_TRUE': 203, + 'TO_BOOL_BOOL': 204, + 'TO_BOOL_INT': 205, + 'TO_BOOL_LIST': 206, + 'TO_BOOL_NONE': 207, + 'TO_BOOL_STR': 208, + 'UNPACK_SEQUENCE_LIST': 209, + 'UNPACK_SEQUENCE_TUPLE': 210, + 'UNPACK_SEQUENCE_TWO_TUPLE': 211, 'INSTRUMENTED_END_FOR': 234, 'INSTRUMENTED_POP_ITER': 235, 'INSTRUMENTED_END_SEND': 236, @@ -354,6 +233,8 @@ 'INSTRUMENTED_CALL_KW': 251, 'INSTRUMENTED_CALL_FUNCTION_EX': 252, 'INSTRUMENTED_JUMP_BACKWARD': 253, + 'INSTRUMENTED_LINE': 254, + 'ENTER_EXECUTOR': 255, 'ANNOTATIONS_PLACEHOLDER': 256, 'JUMP': 257, 'JUMP_IF_FALSE': 258, @@ -367,5 +248,6 @@ 'STORE_FAST_MAYBE_NULL': 266, } -HAVE_ARGUMENT = 43 -MIN_INSTRUMENTED_OPCODE = 234 +# CPython 3.13 compatible: opcodes < 44 have no argument +HAVE_ARGUMENT = 44 +MIN_INSTRUMENTED_OPCODE = 236 diff --git a/crates/codegen/src/compile.rs b/crates/codegen/src/compile.rs index b0cbbf501f8..4e95d1b576a 100644 --- a/crates/codegen/src/compile.rs +++ b/crates/codegen/src/compile.rs @@ -13,7 +13,7 @@ use crate::{ IndexMap, IndexSet, ToPythonName, error::{CodegenError, CodegenErrorType, InternalError, PatternUnreachableReason}, ir::{self, BlockIdx}, - symboltable::{self, CompilerScope, SymbolFlags, SymbolScope, SymbolTable}, + symboltable::{self, CompilerScope, Symbol, SymbolFlags, SymbolScope, SymbolTable}, unparse::UnparseExpr, }; use alloc::borrow::Cow; @@ -968,12 +968,21 @@ impl Compiler { Ok(()) } - /// Check if this is an inlined comprehension context (PEP 709) - /// Currently disabled - always returns false to avoid stack issues - fn is_inlined_comprehension_context(&self, _comprehension_type: ComprehensionType) -> bool { - // TODO: Implement PEP 709 inlined comprehensions properly - // For now, disabled to avoid stack underflow issues - false + /// Check if this is an inlined comprehension context (PEP 709). + /// Only inline in function-like scopes (fastlocals) — module/class + /// level uses STORE_NAME which is incompatible with LOAD_FAST_AND_CLEAR. + /// Generator expressions are never inlined. + fn is_inlined_comprehension_context(&self, comprehension_type: ComprehensionType) -> 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) } /// Enter a new scope @@ -1003,12 +1012,14 @@ impl Compiler { // Use varnames from symbol table (already collected in definition order) let varname_cache: IndexSet = ste.varnames.iter().cloned().collect(); - // Build cellvars using dictbytype (CELL scope, sorted) + // Build cellvars using dictbytype (CELL scope or COMP_CELL flag, sorted) let mut cellvar_cache = IndexSet::default(); let mut cell_names: Vec<_> = ste .symbols .iter() - .filter(|(_, s)| s.scope == SymbolScope::Cell) + .filter(|(_, s)| { + s.scope == SymbolScope::Cell || s.flags.contains(SymbolFlags::COMP_CELL) + }) .map(|(name, _)| name.clone()) .collect(); cell_names.sort(); @@ -7777,12 +7788,16 @@ impl Compiler { if is_inlined { // PEP 709: Inlined comprehension - compile inline without new scope - return self.compile_inlined_comprehension( + 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( init_collection, generators, compile_element, has_an_async_gen, ); + self.current_code_info().in_inlined_comp = was_in_inlined_comp; + return result; } // Non-inlined path: create a new code object (generator expressions, etc.) @@ -7809,9 +7824,6 @@ impl Compiler { // Create magnificent function : self.push_output(flags, 1, 1, 0, name.to_owned())?; - // Mark that we're in an inlined comprehension - self.current_code_info().in_inlined_comp = true; - // Set qualname for comprehension self.set_qualname(); @@ -7937,34 +7949,6 @@ impl Compiler { Ok(()) } - /// Collect variable names from an assignment target expression - fn collect_target_names(&self, target: &ast::Expr, names: &mut Vec) { - match target { - ast::Expr::Name(name) => { - let name_str = name.id.to_string(); - if !names.contains(&name_str) { - names.push(name_str); - } - } - ast::Expr::Tuple(tuple) => { - for elt in &tuple.elts { - self.collect_target_names(elt, names); - } - } - ast::Expr::List(list) => { - for elt in &list.elts { - self.collect_target_names(elt, names); - } - } - ast::Expr::Starred(starred) => { - self.collect_target_names(&starred.value, names); - } - _ => { - // Other targets (attribute, subscript) don't bind local names - } - } - } - /// Compile an inlined comprehension (PEP 709) /// This generates bytecode inline without creating a new code object fn compile_inlined_comprehension( @@ -7972,52 +7956,129 @@ impl Compiler { init_collection: Option, generators: &[ast::Comprehension], compile_element: &dyn Fn(&mut Self) -> CompileResult<()>, - _has_an_async_gen: bool, + has_async: bool, ) -> CompileResult<()> { - // PEP 709: Consume the comprehension's sub_table (but we won't use it as a separate scope) - // We need to consume it to keep sub_tables in sync with AST traversal order. + // PEP 709: Consume the comprehension's sub_table. // The symbols are already merged into parent scope by analyze_symbol_table. - let _comp_table = self + // Splice the comprehension's children into the parent so nested scopes + // (e.g. inner comprehensions, lambdas) can be found by the compiler. + let current_table = self .symbol_table_stack .last_mut() - .expect("no current symbol table") - .sub_tables - .remove(0); - - // Collect local variables that need to be saved/restored - // These are variables bound in the comprehension (iteration vars from targets) - let mut pushed_locals: Vec = Vec::new(); - for generator in generators { - self.collect_target_names(&generator.target, &mut pushed_locals); + .expect("no current symbol table"); + let comp_table = current_table.sub_tables[current_table.next_sub_table].clone(); + current_table.next_sub_table += 1; + if !comp_table.sub_tables.is_empty() { + let insert_pos = current_table.next_sub_table; + for (i, st) in comp_table.sub_tables.iter().enumerate() { + current_table.sub_tables.insert(insert_pos + i, st.clone()); + } } - // Step 1: Compile the outermost iterator + // Step 1: Compile the outermost iterator BEFORE tweaking scopes self.compile_expression(&generators[0].iter)?; - // Use is_async from the first generator, not has_an_async_gen which covers ALL generators - if generators[0].is_async { + if has_async && generators[0].is_async { emit!(self, Instruction::GetAIter); } else { emit!(self, Instruction::GetIter); } - // Step 2: Save local variables that will be shadowed by the comprehension + // Collect local variables that need to be saved/restored. + // All DEF_LOCAL && !DEF_NONLOCAL names from the comp table, plus class block names. + let in_class_block = { + let ct = self.current_symbol_table(); + ct.typ == CompilerScope::Class && !self.current_code_info().in_inlined_comp + }; + let mut pushed_locals: Vec = Vec::new(); + for (name, sym) in &comp_table.symbols { + 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.clone()); + } + } + + // TweakInlinedComprehensionScopes: temporarily override parent symbols + // with comp scopes where they differ. + let mut temp_symbols: IndexMap = IndexMap::default(); + for (name, comp_sym) in &comp_table.symbols { + if comp_sym.flags.contains(SymbolFlags::PARAMETER) { + continue; // skip .0 + } + let comp_scope = comp_sym.scope; + + let current_table = self.symbol_table_stack.last().expect("no symbol table"); + if let Some(outer_sym) = current_table.symbols.get(name) { + let outer_scope = outer_sym.scope; + if (comp_scope != outer_scope + && comp_scope != SymbolScope::Free + && !(comp_scope == SymbolScope::Cell && outer_scope == SymbolScope::Free)) + || in_class_block + { + temp_symbols.insert(name.clone(), outer_sym.clone()); + let current_table = + self.symbol_table_stack.last_mut().expect("no symbol table"); + current_table.symbols.insert(name.clone(), comp_sym.clone()); + } + } + } + + // Step 2: Save local variables that will be shadowed by the comprehension. + // For each variable, we push the fast local value via LoadFastAndClear. + // For CELL variables, we additionally push the cell content via MakeCell + // (which saves the old cell value and clears the cell for the comprehension). + // Track cell indices to restore them later. + let mut cell_indices: Vec> = Vec::new(); + let mut total_stack_items: usize = 0; for name in &pushed_locals { let var_num = self.varname(name)?; emit!(self, Instruction::LoadFastAndClear { var_num }); + total_stack_items += 1; + // If the comp symbol is CELL, emit MAKE_CELL to save cell value + if let Some(comp_sym) = comp_table.symbols.get(name) { + if comp_sym.scope == SymbolScope::Cell { + let i = if self + .current_symbol_table() + .symbols + .get(name) + .is_some_and(|s| s.scope == SymbolScope::Free) + { + self.get_free_var_index(name)? + } else { + self.get_cell_var_index(name)? + }; + emit!(self, Instruction::MakeCell { i }); + cell_indices.push(Some(i)); + total_stack_items += 1; + } else { + cell_indices.push(None); + } + } else { + cell_indices.push(None); + } } - // Step 3: SWAP iterator to TOS (above saved locals) - if !pushed_locals.is_empty() { + // Step 3: SWAP iterator to TOS (above saved locals + cell values) + if total_stack_items > 0 { emit!( self, Instruction::Swap { - i: u32::try_from(pushed_locals.len() + 1).unwrap() + i: u32::try_from(total_stack_items + 1).unwrap() } ); } // Step 4: Create the collection (list/set/dict) - // For generator expressions, init_collection is None if let Some(init_collection) = init_collection { self._emit(init_collection, OpArg::new(0), BlockIdx::NULL); // SWAP to get iterator on top @@ -8039,14 +8100,13 @@ impl Compiler { } // Step 5: Compile the comprehension loop(s) - let mut loop_labels = vec![]; + let mut loop_labels: Vec<(BlockIdx, BlockIdx, BlockIdx, bool, BlockIdx)> = vec![]; for (i, generator) in generators.iter().enumerate() { let loop_block = self.new_block(); let if_cleanup_block = self.new_block(); let after_block = self.new_block(); if i > 0 { - // For nested loops, compile the iterator expression self.compile_expression(&generator.iter)?; if generator.is_async { emit!(self, Instruction::GetAIter); @@ -8056,17 +8116,26 @@ impl Compiler { } self.switch_to_block(loop_block); - let mut end_async_for_target = BlockIdx::NULL; + let mut end_async_for_target = BlockIdx::NULL; if generator.is_async { + emit!(self, PseudoInstruction::SetupFinally { delta: after_block }); emit!(self, Instruction::GetANext); + self.push_fblock( + FBlockType::AsyncComprehensionGenerator, + loop_block, + after_block, + )?; self.emit_load_const(ConstantData::None); end_async_for_target = self.compile_yield_from_sequence(true)?; + emit!(self, PseudoInstruction::PopBlock); + self.pop_fblock(FBlockType::AsyncComprehensionGenerator); self.compile_store(&generator.target)?; } else { emit!(self, Instruction::ForIter { delta: after_block }); self.compile_store(&generator.target)?; } + loop_labels.push(( loop_block, if_cleanup_block, @@ -8085,8 +8154,8 @@ impl Compiler { compile_element(self)?; // Step 7: Close all loops - for (loop_block, if_cleanup_block, after_block, is_async, end_async_for_target) in - loop_labels.iter().rev().copied() + for &(loop_block, if_cleanup_block, after_block, is_async, end_async_for_target) in + loop_labels.iter().rev() { emit!(self, PseudoInstruction::Jump { delta: loop_block }); @@ -8096,17 +8165,14 @@ impl Compiler { self.switch_to_block(after_block); if is_async { self.emit_end_async_for(end_async_for_target); - // Pop the iterator - emit!(self, Instruction::PopTop); } else { - // END_FOR + POP_ITER pattern (CPython 3.14) emit!(self, Instruction::EndFor); emit!(self, Instruction::PopIter); } } - // Step 8: Clean up - restore saved locals - if !pushed_locals.is_empty() { + // Step 8: Clean up - restore saved locals (and cell values) + if total_stack_items > 0 { emit!(self, PseudoInstruction::PopBlock); self.pop_fblock(FBlockType::TryExcept); @@ -8115,19 +8181,21 @@ impl Compiler { // Exception cleanup path self.switch_to_block(cleanup_block); - // Stack: [saved_locals..., collection, exception] - // Swap to get collection out from under exception + // Stack: [saved_values..., collection, exception] emit!(self, Instruction::Swap { i: 2 }); emit!(self, Instruction::PopTop); // Pop incomplete collection - // Restore locals + // Restore locals and cell values emit!( self, Instruction::Swap { - i: u32::try_from(pushed_locals.len() + 1).unwrap() + i: u32::try_from(total_stack_items + 1).unwrap() } ); - for name in pushed_locals.iter().rev() { + for (name, cell_idx) in pushed_locals.iter().rev().zip(cell_indices.iter().rev()) { + if let Some(i) = cell_idx { + emit!(self, Instruction::RestoreCell { i: *i }); + } let var_num = self.varname(name)?; emit!(self, Instruction::StoreFast { var_num }); } @@ -8138,22 +8206,31 @@ impl Compiler { self.switch_to_block(end_block); } - // SWAP result to TOS (above saved locals) - if !pushed_locals.is_empty() { + // SWAP result to TOS (above saved values) + if total_stack_items > 0 { emit!( self, Instruction::Swap { - i: u32::try_from(pushed_locals.len() + 1).unwrap() + i: u32::try_from(total_stack_items + 1).unwrap() } ); } - // Restore saved locals - for name in pushed_locals.iter().rev() { + // Restore saved locals and cell values + for (name, cell_idx) in pushed_locals.iter().rev().zip(cell_indices.iter().rev()) { + if let Some(i) = cell_idx { + emit!(self, Instruction::RestoreCell { i: *i }); + } let var_num = self.varname(name)?; emit!(self, Instruction::StoreFast { var_num }); } + // RevertInlinedComprehensionScopes: restore original symbols + let current_table = self.symbol_table_stack.last_mut().expect("no symbol table"); + for (name, original_sym) in temp_symbols { + current_table.symbols.insert(name, original_sym); + } + Ok(()) } diff --git a/crates/codegen/src/symboltable.rs b/crates/codegen/src/symboltable.rs index fdbdac2b2a7..b86e0f87b21 100644 --- a/crates/codegen/src/symboltable.rs +++ b/crates/codegen/src/symboltable.rs @@ -299,6 +299,68 @@ fn drop_class_free(symbol_table: &mut SymbolTable, newfree: &mut IndexSet, + inlined_cells: &mut IndexSet, + parent_type: CompilerScope, +) { + for (name, sub_symbol) in &comp.symbols { + // Skip the .0 parameter + if sub_symbol.flags.contains(SymbolFlags::PARAMETER) { + continue; + } + + // Track inlined cells + if sub_symbol.scope == SymbolScope::Cell + || sub_symbol.flags.contains(SymbolFlags::COMP_CELL) + { + inlined_cells.insert(name.clone()); + } + + // Handle __class__ in ClassBlock + let scope = if sub_symbol.scope == SymbolScope::Free + && parent_type == CompilerScope::Class + && name == "__class__" + { + comp_free.swap_remove(name); + SymbolScope::GlobalImplicit + } else { + sub_symbol.scope + }; + + if let Some(existing) = parent_symbols.get(name) { + // Name exists in parent + if existing.is_bound() && parent_type != CompilerScope::Class { + // Check if the name is free in any child of the comprehension + let is_free_in_child = comp.sub_tables.iter().any(|child| { + child + .symbols + .get(name) + .is_some_and(|s| s.scope == SymbolScope::Free) + }); + if !is_free_in_child { + comp_free.swap_remove(name); + } + } + } else { + // Name doesn't exist in parent, copy from comprehension. + // Reset scope to Unknown so analyze_symbol will resolve it + // in the parent's context. + let mut symbol = sub_symbol.clone(); + symbol.scope = if sub_symbol.is_bound() { + SymbolScope::Unknown + } else { + scope + }; + parent_symbols.insert(name.clone(), symbol); + } + } +} + type SymbolMap = IndexMap; mod stack { @@ -392,14 +454,9 @@ impl SymbolTableAnalyzer { let symbols = core::mem::take(&mut symbol_table.symbols); let sub_tables = &mut *symbol_table.sub_tables; - // Collect free variables from all child scopes - let mut newfree = IndexSet::default(); - let annotation_block = &mut symbol_table.annotation_block; // PEP 649: Determine class_entry to pass to children - // If current scope is a class with annotation block that can_see_class_scope, - // we need to pass class symbols to the annotation scope let is_class = symbol_table.typ == CompilerScope::Class; // Clone class symbols if needed for child scopes with can_see_class_scope @@ -418,12 +475,16 @@ impl SymbolTableAnalyzer { None }; + // Collect (child_free, is_inlined) pairs from child scopes. + // We need to process inlined comprehensions after the closure + // when we have access to symbol_table.symbols. + let mut child_frees: Vec<(IndexSet, bool)> = Vec::new(); + let mut annotation_free: Option> = None; + let mut info = (symbols, symbol_table.typ); self.tables.with_append(&mut info, |list| { let inner_scope = unsafe { &mut *(list as *mut _ as *mut Self) }; - // Analyze sub scopes and collect their free variables for sub_table in sub_tables.iter_mut() { - // Pass class_entry to sub-scopes that can see the class scope let child_class_entry = if sub_table.can_see_class_scope { if is_class { class_symbols_clone.as_ref() @@ -434,12 +495,10 @@ impl SymbolTableAnalyzer { None }; let child_free = inner_scope.analyze_symbol_table(sub_table, child_class_entry)?; - // Propagate child's free variables to this scope - newfree.extend(child_free); + child_frees.push((child_free, sub_table.comp_inlined)); } // PEP 649: Analyze annotation block if present if let Some(annotation_table) = annotation_block { - // Pass class symbols to annotation scope if can_see_class_scope let ann_class_entry = if annotation_table.can_see_class_scope { if is_class { class_symbols_clone.as_ref() @@ -451,59 +510,56 @@ impl SymbolTableAnalyzer { }; let child_free = inner_scope.analyze_symbol_table(annotation_table, ann_class_entry)?; - // Propagate annotation's free variables to this scope - newfree.extend(child_free); + annotation_free = Some(child_free); } Ok(()) })?; symbol_table.symbols = info.0; - // PEP 709: Merge symbols from inlined comprehensions into parent scope - // Only merge symbols that are actually bound in the comprehension, - // not references to outer scope variables (Free symbols). - const BOUND_FLAGS: SymbolFlags = SymbolFlags::ASSIGNED - .union(SymbolFlags::PARAMETER) - .union(SymbolFlags::ITER) - .union(SymbolFlags::ASSIGNED_IN_COMPREHENSION); - - for sub_table in sub_tables.iter() { - if sub_table.comp_inlined { - for (name, sub_symbol) in &sub_table.symbols { - // Skip the .0 parameter - it's internal to the comprehension - if name == ".0" { - continue; - } - // Only merge symbols that are bound in the comprehension - // Skip Free references to outer scope variables - if !sub_symbol.flags.intersects(BOUND_FLAGS) { - continue; - } - // If the symbol doesn't exist in parent, add it - if !symbol_table.symbols.contains_key(name) { - let mut symbol = sub_symbol.clone(); - // Mark as local in parent scope - symbol.scope = SymbolScope::Local; - symbol_table.symbols.insert(name.clone(), symbol); - } - } + // PEP 709: Process inlined comprehensions. + // Merge symbols from inlined comps into parent scope without bail-out. + let mut inlined_cells: IndexSet = IndexSet::default(); + let mut newfree = IndexSet::default(); + for (idx, (mut child_free, is_inlined)) in child_frees.into_iter().enumerate() { + if is_inlined { + inline_comprehension( + &mut symbol_table.symbols, + &sub_tables[idx], + &mut child_free, + &mut inlined_cells, + symbol_table.typ, + ); } + newfree.extend(child_free); + } + if let Some(ann_free) = annotation_free { + newfree.extend(ann_free); } + let sub_tables = &*symbol_table.sub_tables; + // Analyze symbols in current scope for symbol in symbol_table.symbols.values_mut() { self.analyze_symbol(symbol, symbol_table.typ, sub_tables, class_entry)?; // Collect free variables from this scope - // These will be propagated to the parent scope if symbol.scope == SymbolScope::Free || symbol.flags.contains(SymbolFlags::FREE_CLASS) { newfree.insert(symbol.name.clone()); } } + // PEP 709: Promote LOCAL to CELL and set COMP_CELL for inlined cell vars + for symbol in symbol_table.symbols.values_mut() { + if inlined_cells.contains(&symbol.name) { + if symbol.scope == SymbolScope::Local { + symbol.scope = SymbolScope::Cell; + } + symbol.flags.insert(SymbolFlags::COMP_CELL); + } + } + // Handle class-specific implicit cells - // This removes __class__ and __classdict__ from newfree if present - // and sets the corresponding flags on the symbol table if symbol_table.typ == CompilerScope::Class { drop_class_free(symbol_table, &mut newfree); } @@ -694,6 +750,11 @@ impl SymbolTableAnalyzer { st_typ: CompilerScope, ) -> Option { sub_tables.iter().find_map(|st| { + // PEP 709: For inlined comprehensions, check their children + // instead of the comp itself (its symbols are merged into parent). + if st.comp_inlined { + return self.found_in_inner_scope(&st.sub_tables, name, st_typ); + } let sym = st.symbols.get(name)?; if sym.scope == SymbolScope::Free || sym.flags.contains(SymbolFlags::FREE_CLASS) { if st_typ == CompilerScope::Class && name != "__class__" { @@ -2037,13 +2098,28 @@ impl SymbolTableBuilder { self.line_index_start(range), ); - // PEP 709: inlined comprehensions are not yet implemented in the - // compiler (is_inlined_comprehension_context always returns false), - // so do NOT mark comp_inlined here. Setting it would cause the - // symbol-table analyzer to merge comprehension-local symbols into - // the parent scope, while the compiler still emits a separate code - // object — leading to the merged symbols being missing from the - // comprehension's own symbol table lookup. + // PEP 709: Mark non-generator comprehensions for inlining, + // but only inside function-like scopes (fastlocals). + // Module/class scope uses STORE_NAME which is incompatible + // with LOAD_FAST_AND_CLEAR / STORE_FAST save/restore. + // Note: tables.last() is the comprehension scope we just pushed, + // so we check the second-to-last for the parent scope. + if !is_generator { + let parent = self.tables.iter().rev().nth(1); + let parent_is_func = parent.is_some_and(|t| { + matches!( + t.typ, + CompilerScope::Function + | CompilerScope::AsyncFunction + | CompilerScope::Lambda + | CompilerScope::Comprehension + ) + }); + let parent_can_see_class = parent.is_some_and(|t| t.can_see_class_scope); + if parent_is_func && !parent_can_see_class { + self.tables.last_mut().unwrap().comp_inlined = true; + } + } // Register the passed argument to the generator function as the name ".0" self.register_name(".0", SymbolUsage::Parameter, range)?; diff --git a/crates/compiler-core/src/bytecode/instruction.rs b/crates/compiler-core/src/bytecode/instruction.rs index cd85e2458d4..e36c3b9cb56 100644 --- a/crates/compiler-core/src/bytecode/instruction.rs +++ b/crates/compiler-core/src/bytecode/instruction.rs @@ -302,6 +302,13 @@ pub enum Instruction { YieldValue { arg: Arg, } = 120, + /// PEP 709: Restore a saved cell object after an inlined comprehension. + /// Pops the old cell object from the stack and writes it back to the + /// cell slot at `fastlocals[nlocals + i]`, replacing the temporary cell + /// created by `MakeCell`. + RestoreCell { + i: Arg, + } = 121, // CPython 3.14 RESUME (128) Resume { context: Arg, @@ -1024,7 +1031,8 @@ impl InstructionMetadata for Instruction { Self::LoadSuperAttr { .. } => (1 + (oparg & 1), 3), Self::LoadSuperAttrAttr => (1, 3), Self::LoadSuperAttrMethod => (2, 3), - Self::MakeCell { .. } => (0, 0), + Self::MakeCell { .. } => (1, 0), + Self::RestoreCell { .. } => (0, 1), Self::MakeFunction { .. } => (1, 1), Self::MapAdd { .. } => (1 + (oparg - 1), 3 + (oparg - 1)), Self::MatchClass { .. } => (1, 3), @@ -1320,6 +1328,7 @@ impl InstructionMetadata for Instruction { ) } Self::MakeCell { i } => w!(MAKE_CELL, cell_name = i), + Self::RestoreCell { i } => w!(RESTORE_CELL, cell_name = i), Self::MakeFunction => w!(MAKE_FUNCTION), Self::MapAdd { i } => w!(MAP_ADD, i), Self::MatchClass { count } => w!(MATCH_CLASS, count), diff --git a/crates/vm/src/frame.rs b/crates/vm/src/frame.rs index d8e3d35f9e3..1efe51462a2 100644 --- a/crates/vm/src/frame.rs +++ b/crates/vm/src/frame.rs @@ -906,7 +906,19 @@ impl Frame { } } if !code.cellvars.is_empty() || !code.freevars.is_empty() { + let fastlocals = unsafe { (*self.iframe.get()).localsplus.fastlocals() }; for (i, &k) in code.cellvars.iter().enumerate() { + // When a variable appears in both varnames and cellvars + // (inlined comprehension with scope tweak), the fastlocal + // value takes precedence, matching CPython FrameLocalsProxy. + let has_fastlocal = code + .varnames + .iter() + .position(|&v| v == k) + .is_some_and(|idx| fastlocals.get(idx).is_some_and(|v| v.is_some())); + if has_fastlocal { + continue; + } let cell_value = self.get_cell_contents(i); match locals_map.ass_subscript(k, cell_value, vm) { Ok(()) => {} @@ -2699,13 +2711,12 @@ impl ExecutingFrame<'_> { Ok(None) } Instruction::LoadFastAndClear { var_num } => { - // Load value and clear the slot (for inlined comprehensions) - // If slot is empty, push None (not an error - variable may not exist yet) + // Save current slot value and clear it (for inlined comprehensions). + // Pushes NULL (None at Option level) if slot was empty, so that + // StoreFast can restore the empty state after the comprehension. let idx = var_num.get(arg); - let x = self.localsplus.fastlocals_mut()[idx] - .take() - .unwrap_or_else(|| vm.ctx.none()); - self.push_value(x); + let x = self.localsplus.fastlocals_mut()[idx].take(); + self.push_value_opt(x); Ok(None) } Instruction::LoadFastCheck { var_num } => { @@ -2854,8 +2865,29 @@ impl ExecutingFrame<'_> { Ok(None) } Instruction::MakeFunction => self.execute_make_function(vm), - Instruction::MakeCell { .. } => { - // Cell creation is handled at frame creation time in RustPython + Instruction::MakeCell { i } => { + // PEP 709: Save the current cell object on the stack and + // create a fresh empty cell for the inlined comprehension. + // The old cell is restored afterwards via RestoreCell. + let cell_idx = i.get(arg) as usize; + let nlocals = self.code.varnames.len(); + let old_cell = self.localsplus.fastlocals_mut()[nlocals + cell_idx] + .take() + .expect("cell slot empty"); + let new_cell = PyCell::default().into_ref(&vm.ctx).into(); + self.localsplus.fastlocals_mut()[nlocals + cell_idx] = Some(new_cell); + // Push the old cell object itself + self.push_value(old_cell); + Ok(None) + } + Instruction::RestoreCell { i } => { + // PEP 709: Restore the saved cell object after an inlined + // comprehension. Pops the old cell from the stack and writes + // it back to the cell slot, replacing the temporary cell. + let cell_idx = i.get(arg) as usize; + let nlocals = self.code.varnames.len(); + let old_cell = self.pop_value(); + self.localsplus.fastlocals_mut()[nlocals + cell_idx] = Some(old_cell); Ok(None) } Instruction::MapAdd { i } => { @@ -3298,9 +3330,10 @@ impl ExecutingFrame<'_> { Ok(None) } Instruction::StoreFast { var_num } => { - let value = self.pop_value(); + // pop_value_opt: allows NULL from LoadFastAndClear restore path + let value = self.pop_value_opt(); let fastlocals = self.localsplus.fastlocals_mut(); - fastlocals[var_num.get(arg)] = Some(value); + fastlocals[var_num.get(arg)] = value; Ok(None) } Instruction::StoreFastLoadFast { var_nums } => { From 6723f2fa060ed2c95452241c5021e8a56e87f90f Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Tue, 17 Mar 2026 20:15:08 +0900 Subject: [PATCH 2/4] Add localspluskinds, unify DEREF to localsplus index - Add CO_FAST_LOCAL/CELL/FREE/HIDDEN constants and localspluskinds field to CodeObject for per-slot metadata - Change DEREF instruction opargs from cell-relative indices (NameIdx) to localsplus absolute indices (oparg::VarNum) - Add fixup_deref_opargs pass in ir.rs to convert cell-relative indices to localsplus indices after finalization - Replace get_cell_name with get_localsplus_name in InstrDisplayContext trait - Update VM cell_ref/get_cell_contents/set_cell_contents to use localsplus indices directly (no nlocals offset) - Update function.rs cell2arg, super.rs __class__ lookup with explicit nlocals offsets --- Lib/_opcode_metadata.py | 1 - crates/codegen/src/compile.rs | 54 ++-- crates/codegen/src/ir.rs | 155 +++++++--- crates/compiler-core/src/bytecode.rs | 34 +- .../compiler-core/src/bytecode/instruction.rs | 23 +- crates/compiler-core/src/marshal.rs | 60 +++- crates/vm/src/builtins/code.rs | 35 ++- crates/vm/src/builtins/function.rs | 23 +- crates/vm/src/builtins/super.rs | 22 +- crates/vm/src/frame.rs | 290 ++++++++++-------- crates/vm/src/version.rs | 3 +- crates/vm/src/vm/context.rs | 2 +- 12 files changed, 433 insertions(+), 269 deletions(-) diff --git a/Lib/_opcode_metadata.py b/Lib/_opcode_metadata.py index f1656a35cb1..bb55ee423cf 100644 --- a/Lib/_opcode_metadata.py +++ b/Lib/_opcode_metadata.py @@ -128,7 +128,6 @@ 'UNPACK_EX': 118, 'UNPACK_SEQUENCE': 119, 'YIELD_VALUE': 120, - 'RESTORE_CELL': 121, 'RESUME': 128, 'BINARY_OP_ADD_FLOAT': 129, 'BINARY_OP_ADD_INT': 130, diff --git a/crates/codegen/src/compile.rs b/crates/codegen/src/compile.rs index 4e95d1b576a..14329e34d43 100644 --- a/crates/codegen/src/compile.rs +++ b/crates/codegen/src/compile.rs @@ -692,26 +692,28 @@ impl Compiler { .expect("symbol_table_stack is empty! This is a compiler bug.") } - /// Get the index of a free variable. - fn get_free_var_index(&mut self, name: &str) -> CompileResult { + /// Get the cell-relative index of a free variable. + /// Returns ncells + freevar_idx. Fixed up to localsplus index during finalize. + fn get_free_var_index(&mut self, name: &str) -> CompileResult { let info = self.code_stack.last_mut().unwrap(); let idx = info .metadata .freevars .get_index_of(name) .unwrap_or_else(|| info.metadata.freevars.insert_full(name.to_owned()).0); - Ok((idx + info.metadata.cellvars.len()).to_u32()) + Ok((idx + info.metadata.cellvars.len()).to_u32().into()) } - /// Get the index of a cell variable. - fn get_cell_var_index(&mut self, name: &str) -> CompileResult { + /// Get the cell-relative index of a cell variable. + /// Returns cellvar_idx. Fixed up to localsplus index during finalize. + fn get_cell_var_index(&mut self, name: &str) -> CompileResult { let info = self.code_stack.last_mut().unwrap(); let idx = info .metadata .cellvars .get_index_of(name) .unwrap_or_else(|| info.metadata.cellvars.insert_full(name.to_owned()).0); - Ok(idx.to_u32()) + Ok(idx.to_u32().into()) } /// Get the index of a local variable. @@ -1149,6 +1151,19 @@ impl Compiler { self.set_qualname(); } + // Emit COPY_FREE_VARS and MAKE_CELL prolog before RESUME + { + let nfrees = self.code_stack.last().unwrap().metadata.freevars.len(); + if nfrees > 0 { + emit!(self, Instruction::CopyFreeVars { n: nfrees as u32 }); + } + let ncells = self.code_stack.last().unwrap().metadata.cellvars.len(); + for i in 0..ncells { + let i_varnum: oparg::VarNum = (i as u32).into(); + emit!(self, Instruction::MakeCell { i: i_varnum }); + } + } + // Emit RESUME (handles async preamble and module lineno 0) // CPython: LOCATION(lineno, lineno, 0, 0), then loc.lineno = 0 for module self.emit_resume_for_scope(scope_type, lineno); @@ -8035,16 +8050,15 @@ impl Compiler { // Step 2: Save local variables that will be shadowed by the comprehension. // For each variable, we push the fast local value via LoadFastAndClear. - // For CELL variables, we additionally push the cell content via MakeCell - // (which saves the old cell value and clears the cell for the comprehension). - // Track cell indices to restore them later. - let mut cell_indices: Vec> = Vec::new(); + // For merged CELL variables, LoadFastAndClear saves the cell object from + // the merged slot, and MAKE_CELL creates a new empty cell in-place. + // MAKE_CELL has no stack effect (operates only on fastlocals). let mut total_stack_items: usize = 0; for name in &pushed_locals { let var_num = self.varname(name)?; emit!(self, Instruction::LoadFastAndClear { var_num }); total_stack_items += 1; - // If the comp symbol is CELL, emit MAKE_CELL to save cell value + // If the comp symbol is CELL, emit MAKE_CELL to create fresh cell if let Some(comp_sym) = comp_table.symbols.get(name) { if comp_sym.scope == SymbolScope::Cell { let i = if self @@ -8058,13 +8072,7 @@ impl Compiler { self.get_cell_var_index(name)? }; emit!(self, Instruction::MakeCell { i }); - cell_indices.push(Some(i)); - total_stack_items += 1; - } else { - cell_indices.push(None); } - } else { - cell_indices.push(None); } } @@ -8192,10 +8200,7 @@ impl Compiler { i: u32::try_from(total_stack_items + 1).unwrap() } ); - for (name, cell_idx) in pushed_locals.iter().rev().zip(cell_indices.iter().rev()) { - if let Some(i) = cell_idx { - emit!(self, Instruction::RestoreCell { i: *i }); - } + for name in pushed_locals.iter().rev() { let var_num = self.varname(name)?; emit!(self, Instruction::StoreFast { var_num }); } @@ -8216,11 +8221,8 @@ impl Compiler { ); } - // Restore saved locals and cell values - for (name, cell_idx) in pushed_locals.iter().rev().zip(cell_indices.iter().rev()) { - if let Some(i) = cell_idx { - emit!(self, Instruction::RestoreCell { i: *i }); - } + // Restore saved locals (StoreFast restores the saved cell object for merged cells) + for name in pushed_locals.iter().rev() { let var_num = self.varname(name)?; emit!(self, Instruction::StoreFast { var_num }); } diff --git a/crates/codegen/src/ir.rs b/crates/codegen/src/ir.rs index 8a34fced545..5ba56eb0ec5 100644 --- a/crates/codegen/src/ir.rs +++ b/crates/codegen/src/ir.rs @@ -8,9 +8,10 @@ use num_traits::ToPrimitive; use rustpython_compiler_core::{ OneIndexed, SourceLocation, bytecode::{ - AnyInstruction, Arg, CodeFlags, CodeObject, CodeUnit, CodeUnits, ConstantData, - ExceptionTableEntry, InstrDisplayContext, Instruction, InstructionMetadata, Label, OpArg, - PseudoInstruction, PyCodeLocationInfoKind, encode_exception_table, oparg, + AnyInstruction, Arg, CO_FAST_CELL, CO_FAST_FREE, CO_FAST_LOCAL, CodeFlags, CodeObject, + CodeUnit, CodeUnits, ConstantData, ExceptionTableEntry, InstrDisplayContext, Instruction, + InstructionMetadata, Label, OpArg, PseudoInstruction, PyCodeLocationInfoKind, + encode_exception_table, oparg, }, varint::{write_signed_varint, write_varint}, }; @@ -210,7 +211,6 @@ impl CodeInfo { self.optimize_load_global_push_null(); let max_stackdepth = self.max_stackdepth()?; - let cell2arg = self.cell2arg(); let Self { flags, @@ -247,8 +247,12 @@ impl CodeInfo { let mut locations = Vec::new(); let mut linetable_locations: Vec = Vec::new(); - // Convert pseudo ops and remove resulting NOPs (keep line-marker NOPs) - convert_pseudo_ops(&mut blocks, varname_cache.len() as u32); + // Build cellfixedoffsets for cell-local merging + let cellfixedoffsets = + build_cellfixedoffsets(&varname_cache, &cellvar_cache, &freevar_cache); + // Convert pseudo ops (LoadClosure uses cellfixedoffsets) and fixup DEREF opargs + convert_pseudo_ops(&mut blocks, &cellfixedoffsets); + fixup_deref_opargs(&mut blocks, &cellfixedoffsets); // Remove redundant NOPs, keeping line-marker NOPs only when // they are needed to preserve tracing. let mut block_order = Vec::new(); @@ -482,6 +486,35 @@ impl CodeInfo { // Generate exception table before moving source_path let exceptiontable = generate_exception_table(&blocks, &block_to_index); + // Build localspluskinds with cell-local merging + let nlocals = varname_cache.len(); + let ncells = cellvar_cache.len(); + let nfrees = freevar_cache.len(); + let numdropped = cellvar_cache + .iter() + .filter(|cv| varname_cache.contains(cv.as_str())) + .count(); + let nlocalsplus = nlocals + ncells - numdropped + nfrees; + let mut localspluskinds = vec![0u8; nlocalsplus]; + // Mark locals + for kind in localspluskinds.iter_mut().take(nlocals) { + *kind = CO_FAST_LOCAL; + } + // Mark cells (merged and non-merged) + for (i, cellvar) in cellvar_cache.iter().enumerate() { + let idx = cellfixedoffsets[i] as usize; + if varname_cache.contains(cellvar.as_str()) { + localspluskinds[idx] |= CO_FAST_CELL; // merged: LOCAL | CELL + } else { + localspluskinds[idx] = CO_FAST_CELL; + } + } + // Mark frees + for i in 0..nfrees { + let idx = cellfixedoffsets[ncells + i] as usize; + localspluskinds[idx] = CO_FAST_FREE; + } + Ok(CodeObject { flags, posonlyarg_count, @@ -500,43 +533,12 @@ impl CodeInfo { varnames: varname_cache.into_iter().collect(), cellvars: cellvar_cache.into_iter().collect(), freevars: freevar_cache.into_iter().collect(), - cell2arg, + localspluskinds: localspluskinds.into_boxed_slice(), linetable, exceptiontable, }) } - fn cell2arg(&self) -> Option> { - if self.metadata.cellvars.is_empty() { - return None; - } - - let total_args = self.metadata.argcount - + self.metadata.kwonlyargcount - + self.flags.contains(CodeFlags::VARARGS) as u32 - + self.flags.contains(CodeFlags::VARKEYWORDS) as u32; - - let mut found_cellarg = false; - let cell2arg = self - .metadata - .cellvars - .iter() - .map(|var| { - self.metadata - .varnames - .get_index_of(var) - // check that it's actually an arg - .filter(|i| *i < total_args as usize) - .map_or(-1, |i| { - found_cellarg = true; - i as i32 - }) - }) - .collect::>(); - - if found_cellarg { Some(cell2arg) } else { None } - } - fn dce(&mut self) { for block in &mut self.blocks { let mut last_instr = None; @@ -1107,12 +1109,19 @@ impl InstrDisplayContext for CodeInfo { self.metadata.varnames[var_num.as_usize()].as_ref() } - fn get_cell_name(&self, i: usize) -> &str { - self.metadata - .cellvars - .get_index(i) - .unwrap_or_else(|| &self.metadata.freevars[i - self.metadata.cellvars.len()]) - .as_ref() + fn get_localsplus_name(&self, var_num: oparg::VarNum) -> &str { + let idx = var_num.as_usize(); + let nlocals = self.metadata.varnames.len(); + if idx < nlocals { + self.metadata.varnames[idx].as_ref() + } else { + let cell_idx = idx - nlocals; + self.metadata + .cellvars + .get_index(cell_idx) + .unwrap_or_else(|| &self.metadata.freevars[cell_idx - self.metadata.cellvars.len()]) + .as_ref() + } } } @@ -1768,7 +1777,7 @@ pub(crate) fn label_exception_targets(blocks: &mut [Block]) { /// Convert remaining pseudo ops to real instructions or NOP. /// flowgraph.c convert_pseudo_ops -pub(crate) fn convert_pseudo_ops(blocks: &mut [Block], varnames_len: u32) { +pub(crate) fn convert_pseudo_ops(blocks: &mut [Block], cellfixedoffsets: &[u32]) { for block in blocks.iter_mut() { for info in &mut block.instructions { let Some(pseudo) = info.instr.pseudo() else { @@ -1786,9 +1795,10 @@ pub(crate) fn convert_pseudo_ops(blocks: &mut [Block], varnames_len: u32) { PseudoInstruction::PopBlock => { info.instr = Instruction::Nop.into(); } - // LOAD_CLOSURE → LOAD_FAST (with varnames offset) + // LOAD_CLOSURE → LOAD_FAST (using cellfixedoffsets for merged layout) PseudoInstruction::LoadClosure { i } => { - let new_idx = varnames_len + i.get(info.arg); + let cell_relative = i.get(info.arg) as usize; + let new_idx = cellfixedoffsets[cell_relative]; info.arg = OpArg::new(new_idx); info.instr = Instruction::LoadFast { var_num: Arg::marker(), @@ -1808,3 +1818,54 @@ pub(crate) fn convert_pseudo_ops(blocks: &mut [Block], varnames_len: u32) { } } } + +/// Build cellfixedoffsets mapping: cell/free index -> localsplus index. +/// Merged cells (cellvar also in varnames) get the local slot index. +/// Non-merged cells get slots after nlocals. Free vars follow. +pub(crate) fn build_cellfixedoffsets( + varnames: &IndexSet, + cellvars: &IndexSet, + freevars: &IndexSet, +) -> Vec { + let nlocals = varnames.len(); + let ncells = cellvars.len(); + let nfrees = freevars.len(); + let mut fixed = Vec::with_capacity(ncells + nfrees); + let mut numdropped = 0usize; + for (i, cellvar) in cellvars.iter().enumerate() { + if let Some(local_idx) = varnames.get_index_of(cellvar) { + fixed.push(local_idx as u32); + numdropped += 1; + } else { + fixed.push((nlocals + i - numdropped) as u32); + } + } + for i in 0..nfrees { + fixed.push((nlocals + ncells - numdropped + i) as u32); + } + fixed +} + +/// Convert DEREF instruction opargs from cell-relative indices to localsplus indices +/// using the cellfixedoffsets mapping. +pub(crate) fn fixup_deref_opargs(blocks: &mut [Block], cellfixedoffsets: &[u32]) { + for block in blocks.iter_mut() { + for info in &mut block.instructions { + let Some(instr) = info.instr.real() else { + continue; + }; + let needs_fixup = matches!( + instr, + Instruction::LoadDeref { .. } + | Instruction::StoreDeref { .. } + | Instruction::DeleteDeref { .. } + | Instruction::LoadFromDictOrDeref { .. } + | Instruction::MakeCell { .. } + ); + if needs_fixup { + let cell_relative = u32::from(info.arg) as usize; + info.arg = OpArg::new(cellfixedoffsets[cell_relative]); + } + } + } +} diff --git a/crates/compiler-core/src/bytecode.rs b/crates/compiler-core/src/bytecode.rs index fff56de5f50..15940b68d1b 100644 --- a/crates/compiler-core/src/bytecode.rs +++ b/crates/compiler-core/src/bytecode.rs @@ -334,6 +334,12 @@ impl IndexMut for [T] { } } +/// Per-slot kind flags for localsplus (co_localspluskinds). +pub const CO_FAST_HIDDEN: u8 = 0x10; +pub const CO_FAST_LOCAL: u8 = 0x20; +pub const CO_FAST_CELL: u8 = 0x40; +pub const CO_FAST_FREE: u8 = 0x80; + /// Primary container of a single code object. Each python function has /// a code object. Also a module has a code object. #[derive(Clone)] @@ -352,12 +358,14 @@ pub struct CodeObject { pub obj_name: C::Name, /// Qualified name of the object (like CPython's co_qualname) pub qualname: C::Name, - pub cell2arg: Option>, pub constants: Constants, pub names: Box<[C::Name]>, pub varnames: Box<[C::Name]>, pub cellvars: Box<[C::Name]>, pub freevars: Box<[C::Name]>, + /// Per-slot kind flags: CO_FAST_LOCAL, CO_FAST_CELL, CO_FAST_FREE, CO_FAST_HIDDEN. + /// Length = nlocalsplus (nlocals + ncells + nfrees). + pub localspluskinds: Box<[u8]>, /// Line number table (CPython 3.11+ format) pub linetable: Box<[u8]>, /// Exception handling table @@ -1080,7 +1088,7 @@ impl CodeObject { kwonlyarg_count: self.kwonlyarg_count, first_line_number: self.first_line_number, max_stackdepth: self.max_stackdepth, - cell2arg: self.cell2arg, + localspluskinds: self.localspluskinds, linetable: self.linetable, exceptiontable: self.exceptiontable, } @@ -1112,7 +1120,7 @@ impl CodeObject { kwonlyarg_count: self.kwonlyarg_count, first_line_number: self.first_line_number, max_stackdepth: self.max_stackdepth, - cell2arg: self.cell2arg.clone(), + localspluskinds: self.localspluskinds.clone(), linetable: self.linetable.clone(), exceptiontable: self.exceptiontable.clone(), } @@ -1141,7 +1149,8 @@ pub trait InstrDisplayContext { fn get_varname(&self, var_num: oparg::VarNum) -> &str; - fn get_cell_name(&self, i: usize) -> &str; + /// Get name for a localsplus index (used by DEREF instructions). + fn get_localsplus_name(&self, var_num: oparg::VarNum) -> &str; } impl InstrDisplayContext for CodeObject { @@ -1159,11 +1168,18 @@ impl InstrDisplayContext for CodeObject { self.varnames[var_num].as_ref() } - fn get_cell_name(&self, i: usize) -> &str { - self.cellvars - .get(i) - .unwrap_or_else(|| &self.freevars[i - self.cellvars.len()]) - .as_ref() + fn get_localsplus_name(&self, var_num: oparg::VarNum) -> &str { + let idx = var_num.as_usize(); + let nlocals = self.varnames.len(); + if idx < nlocals { + self.varnames[idx].as_ref() + } else { + let cell_idx = idx - nlocals; + self.cellvars + .get(cell_idx) + .unwrap_or_else(|| &self.freevars[cell_idx - self.cellvars.len()]) + .as_ref() + } } } diff --git a/crates/compiler-core/src/bytecode/instruction.rs b/crates/compiler-core/src/bytecode/instruction.rs index e36c3b9cb56..6544c675c22 100644 --- a/crates/compiler-core/src/bytecode/instruction.rs +++ b/crates/compiler-core/src/bytecode/instruction.rs @@ -130,7 +130,7 @@ pub enum Instruction { namei: Arg, } = 61, DeleteDeref { - i: Arg, + i: Arg, } = 62, DeleteFast { var_num: Arg, @@ -189,7 +189,7 @@ pub enum Instruction { consti: Arg, } = 82, LoadDeref { - i: Arg, + i: Arg, } = 83, LoadFast { var_num: Arg, @@ -210,7 +210,7 @@ pub enum Instruction { var_nums: Arg, } = 89, LoadFromDictOrDeref { - i: Arg, + i: Arg, } = 90, LoadFromDictOrGlobals { i: Arg, @@ -231,7 +231,7 @@ pub enum Instruction { namei: Arg, } = 96, MakeCell { - i: Arg, + i: Arg, } = 97, MapAdd { i: Arg, @@ -273,7 +273,7 @@ pub enum Instruction { namei: Arg, } = 110, StoreDeref { - i: Arg, + i: Arg, } = 111, StoreFast { var_num: Arg, @@ -302,13 +302,6 @@ pub enum Instruction { YieldValue { arg: Arg, } = 120, - /// PEP 709: Restore a saved cell object after an inlined comprehension. - /// Pops the old cell object from the stack and writes it back to the - /// cell slot at `fastlocals[nlocals + i]`, replacing the temporary cell - /// created by `MakeCell`. - RestoreCell { - i: Arg, - } = 121, // CPython 3.14 RESUME (128) Resume { context: Arg, @@ -1031,8 +1024,7 @@ impl InstructionMetadata for Instruction { Self::LoadSuperAttr { .. } => (1 + (oparg & 1), 3), Self::LoadSuperAttrAttr => (1, 3), Self::LoadSuperAttrMethod => (2, 3), - Self::MakeCell { .. } => (1, 0), - Self::RestoreCell { .. } => (0, 1), + Self::MakeCell { .. } => (0, 0), Self::MakeFunction { .. } => (1, 1), Self::MapAdd { .. } => (1 + (oparg - 1), 3 + (oparg - 1)), Self::MatchClass { .. } => (1, 3), @@ -1136,7 +1128,7 @@ impl InstructionMetadata for Instruction { let varname = |var_num: oparg::VarNum| ctx.get_varname(var_num); let name = |i: u32| ctx.get_name(i as usize); - let cell_name = |i: u32| ctx.get_cell_name(i as usize); + let cell_name = |i: oparg::VarNum| ctx.get_localsplus_name(i); let fmt_const = |op: &str, arg: OpArg, @@ -1328,7 +1320,6 @@ impl InstructionMetadata for Instruction { ) } Self::MakeCell { i } => w!(MAKE_CELL, cell_name = i), - Self::RestoreCell { i } => w!(RESTORE_CELL, cell_name = i), Self::MakeFunction => w!(MAKE_FUNCTION), Self::MapAdd { i } => w!(MAP_ADD, i), Self::MatchClass { count } => w!(MATCH_CLASS, count), diff --git a/crates/compiler-core/src/marshal.rs b/crates/compiler-core/src/marshal.rs index ba3cf7a35c3..770b8531824 100644 --- a/crates/compiler-core/src/marshal.rs +++ b/crates/compiler-core/src/marshal.rs @@ -1,5 +1,5 @@ use crate::{OneIndexed, SourceLocation, bytecode::*}; -use alloc::{boxed::Box, vec::Vec}; +use alloc::{boxed::Box, vec, vec::Vec}; use core::convert::Infallible; use malachite_bigint::{BigInt, Sign}; use num_complex::Complex64; @@ -228,14 +228,11 @@ pub fn deserialize_code( let len = rdr.read_u32()?; let qualname = bag.make_name(rdr.read_str(len)?); - let len = rdr.read_u32()?; - let cell2arg = (len != 0) - .then(|| { - (0..len) - .map(|_| Ok(rdr.read_u32()? as i32)) - .collect::>>() - }) - .transpose()?; + // Read and discard legacy cell2arg data for backwards compatibility + let cell2arg_len = rdr.read_u32()?; + for _ in 0..cell2arg_len { + let _ = rdr.read_u32()?; + } let len = rdr.read_u32()?; let constants = (0..len) @@ -267,6 +264,42 @@ pub fn deserialize_code( .to_vec() .into_boxed_slice(); + // Build localspluskinds with cell-local merging + let localspluskinds = { + use crate::bytecode::*; + let nlocals = varnames.len(); + let ncells = cellvars.len(); + let nfrees = freevars.len(); + // Count merged cells (cellvar also in varnames) + let numdropped = cellvars.iter().filter(|cv| { + varnames.iter().any(|v| v.as_ref() == cv.as_ref()) + }).count(); + let nlocalsplus = nlocals + ncells - numdropped + nfrees; + let mut kinds = vec![0u8; nlocalsplus]; + // Mark locals + for kind in kinds.iter_mut().take(nlocals) { + *kind = CO_FAST_LOCAL; + } + // Build cellfixedoffsets and mark cells + let mut cell_numdropped = 0usize; + for (i, cv) in cellvars.iter().enumerate() { + let merged_idx = varnames.iter().position(|v| v.as_ref() == cv.as_ref()); + if let Some(local_idx) = merged_idx { + kinds[local_idx] |= CO_FAST_CELL; // merged: LOCAL | CELL + cell_numdropped += 1; + } else { + let idx = nlocals + i - cell_numdropped; + kinds[idx] = CO_FAST_CELL; + } + } + // Mark frees + let free_start = nlocals + ncells - numdropped; + for i in 0..nfrees { + kinds[free_start + i] = CO_FAST_FREE; + } + kinds.into_boxed_slice() + }; + Ok(CodeObject { instructions, locations, @@ -279,12 +312,12 @@ pub fn deserialize_code( max_stackdepth, obj_name, qualname, - cell2arg, constants, names, varnames, cellvars, freevars, + localspluskinds, linetable, exceptiontable, }) @@ -687,11 +720,8 @@ pub fn serialize_code(buf: &mut W, code: &CodeObject) write_vec(buf, code.obj_name.as_ref().as_bytes()); write_vec(buf, code.qualname.as_ref().as_bytes()); - let cell2arg = code.cell2arg.as_deref().unwrap_or(&[]); - write_len(buf, cell2arg.len()); - for &i in cell2arg { - buf.write_u32(i as u32) - } + // Write empty cell2arg for backwards compatibility + write_len(buf, 0); write_len(buf, code.constants.len()); for constant in &*code.constants { diff --git a/crates/vm/src/builtins/code.rs b/crates/vm/src/builtins/code.rs index 4ab4c7fefd3..65bab0d3a9b 100644 --- a/crates/vm/src/builtins/code.rs +++ b/crates/vm/src/builtins/code.rs @@ -633,6 +633,37 @@ impl Constructor for PyCode { )], > = vec![(loc, loc); instructions.len()].into_boxed_slice(); + // Build localspluskinds with cell-local merging + let localspluskinds = { + use rustpython_compiler_core::bytecode::*; + let nlocals = varnames.len(); + let ncells = cellvars.len(); + let nfrees = freevars.len(); + let numdropped = cellvars.iter().filter(|cv| { + varnames.iter().any(|v| *v == **cv) + }).count(); + let nlocalsplus = nlocals + ncells - numdropped + nfrees; + let mut kinds = vec![0u8; nlocalsplus]; + for kind in kinds.iter_mut().take(nlocals) { + *kind = CO_FAST_LOCAL; + } + let mut cell_numdropped = 0usize; + for (i, cv) in cellvars.iter().enumerate() { + let merged_idx = varnames.iter().position(|v| **v == **cv); + if let Some(local_idx) = merged_idx { + kinds[local_idx] |= CO_FAST_CELL; + cell_numdropped += 1; + } else { + kinds[nlocals + i - cell_numdropped] = CO_FAST_CELL; + } + } + let free_start = nlocals + ncells - numdropped; + for i in 0..nfrees { + kinds[free_start + i] = CO_FAST_FREE; + } + kinds.into_boxed_slice() + }; + // Build the CodeObject let code = CodeObject { instructions, @@ -650,12 +681,12 @@ impl Constructor for PyCode { max_stackdepth: args.stacksize, obj_name: vm.ctx.intern_str(args.name.as_wtf8()), qualname: vm.ctx.intern_str(args.qualname.as_wtf8()), - cell2arg: None, // TODO: reuse `fn cell2arg` constants, names, varnames, cellvars, freevars, + localspluskinds, linetable: args.linetable.as_bytes().to_vec().into_boxed_slice(), exceptiontable: args.exceptiontable.as_bytes().to_vec().into_boxed_slice(), }; @@ -1237,7 +1268,7 @@ impl PyCode { .collect(), cellvars, freevars, - cell2arg: self.code.cell2arg.clone(), + localspluskinds: self.code.localspluskinds.clone(), linetable, exceptiontable, }; diff --git a/crates/vm/src/builtins/function.rs b/crates/vm/src/builtins/function.rs index 0003720c669..0072ca4b724 100644 --- a/crates/vm/src/builtins/function.rs +++ b/crates/vm/src/builtins/function.rs @@ -64,7 +64,7 @@ pub struct PyFunction { code: PyAtomicRef, globals: PyDictRef, builtins: PyObjectRef, - closure: Option>>, + pub(crate) closure: Option>>, defaults_and_kwdefaults: PyMutex<(Option, Option)>, name: PyMutex, qualname: PyMutex, @@ -443,13 +443,6 @@ impl PyFunction { } } - if let Some(cell2arg) = code.cell2arg.as_deref() { - for (cell_idx, arg_idx) in cell2arg.iter().enumerate().filter(|(_, i)| **i != -1) { - let x = fastlocals[*arg_idx as usize].take(); - frame.set_cell_contents(cell_idx, x); - } - } - Ok(()) } @@ -725,14 +718,6 @@ impl Py { } } - if let Some(cell2arg) = code.cell2arg.as_deref() { - let fastlocals = unsafe { frame.fastlocals_mut() }; - for (cell_idx, arg_idx) in cell2arg.iter().enumerate().filter(|(_, i)| **i != -1) { - let x = fastlocals[*arg_idx as usize].take(); - frame.set_cell_contents(cell_idx, x); - } - } - frame } @@ -780,11 +765,7 @@ pub(crate) fn datastack_frame_size_bytes_for_code(code: &Py) -> Option()) } diff --git a/crates/vm/src/builtins/super.rs b/crates/vm/src/builtins/super.rs index 01bdfa6749e..2a3b83314d4 100644 --- a/crates/vm/src/builtins/super.rs +++ b/crates/vm/src/builtins/super.rs @@ -7,6 +7,7 @@ See also [CPython source code.](https://github.com/python/cpython/blob/50b48572d use super::{PyStr, PyType, PyTypeRef}; use crate::{ AsObject, Context, Py, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine, + builtins::function::PyCell, class::PyClassImpl, common::lock::PyRwLock, function::{FuncArgs, IntoFuncArgs, OptionalArg}, @@ -86,27 +87,28 @@ impl Initializer for PySuper { return Err(vm.new_runtime_error("super(): no arguments")); } // SAFETY: Frame is current and not concurrently mutated. + use rustpython_compiler_core::bytecode::CO_FAST_CELL; let obj = unsafe { frame.fastlocals() }[0] .clone() - .or_else(|| { - if let Some(cell2arg) = frame.code.cell2arg.as_deref() { - cell2arg[..frame.code.cellvars.len()] - .iter() - .enumerate() - .find(|(_, arg_idx)| **arg_idx == 0) - .and_then(|(cell_idx, _)| frame.get_cell_contents(cell_idx)) + .and_then(|val| { + // If slot 0 is a merged cell (LOCAL|CELL), extract value from cell + if frame.code.localspluskinds.first().is_some_and(|&k| k & CO_FAST_CELL != 0) { + val.downcast_ref::().and_then(|c| c.get()) } else { - None + Some(val) } }) .ok_or_else(|| vm.new_runtime_error("super(): arg[0] deleted"))?; let mut typ = None; + // Search for __class__ in freevars using localspluskinds + let nlocalsplus = frame.code.localspluskinds.len(); + let nfrees = frame.code.freevars.len(); + let free_start = nlocalsplus - nfrees; for (i, var) in frame.code.freevars.iter().enumerate() { if var.as_bytes() == b"__class__" { - let i = frame.code.cellvars.len() + i; let class = frame - .get_cell_contents(i) + .get_cell_contents(free_start + i) .ok_or_else(|| vm.new_runtime_error("super(): empty __class__ cell"))?; typ = Some(class.downcast().map_err(|o| { vm.new_type_error(format!( diff --git a/crates/vm/src/frame.rs b/crates/vm/src/frame.rs index 1efe51462a2..18085ec4978 100644 --- a/crates/vm/src/frame.rs +++ b/crates/vm/src/frame.rs @@ -41,7 +41,6 @@ use crate::{ use alloc::fmt; use bstr::ByteSlice; use core::cell::UnsafeCell; -use core::iter::zip; use core::sync::atomic; use core::sync::atomic::AtomicPtr; use core::sync::atomic::Ordering::{Acquire, Relaxed}; @@ -684,29 +683,17 @@ impl Frame { use_datastack: bool, vm: &VirtualMachine, ) -> Self { - let nlocals = code.varnames.len(); - let num_cells = code.cellvars.len(); - let nfrees = closure.len(); - - let nlocalsplus = nlocals - .checked_add(num_cells) - .and_then(|v| v.checked_add(nfrees)) - .expect("Frame::new: nlocalsplus overflow"); + let nlocalsplus = code.localspluskinds.len(); let max_stackdepth = code.max_stackdepth as usize; - let mut localsplus = if use_datastack { + let localsplus = if use_datastack { LocalsPlus::new_on_datastack(nlocalsplus, max_stackdepth, vm) } else { LocalsPlus::new(nlocalsplus, max_stackdepth) }; - // Store cell/free variable objects directly in localsplus - let fastlocals = localsplus.fastlocals_mut(); - for i in 0..num_cells { - fastlocals[nlocals + i] = Some(PyCell::default().into_ref(&vm.ctx).into()); - } - for (i, cell) in closure.iter().enumerate() { - fastlocals[nlocals + num_cells + i] = Some(cell.clone().into()); - } + // Free vars and cells are now set up by COPY_FREE_VARS and MAKE_CELL + // instructions emitted at function entry. No pre-creation needed. + let _ = closure; let iframe = InterpreterFrame { localsplus, @@ -791,30 +778,17 @@ impl Frame { } } - /// Get cell contents by cell index. Reads through fastlocals (no state lock needed). - pub(crate) fn get_cell_contents(&self, cell_idx: usize) -> Option { - let nlocals = self.code.varnames.len(); + /// Get cell contents by localsplus index. + pub(crate) fn get_cell_contents(&self, localsplus_idx: usize) -> Option { // SAFETY: Frame not executing; no concurrent mutation. let fastlocals = unsafe { (*self.iframe.get()).localsplus.fastlocals() }; fastlocals - .get(nlocals + cell_idx) + .get(localsplus_idx) .and_then(|slot| slot.as_ref()) .and_then(|obj| obj.downcast_ref::()) .and_then(|cell| cell.get()) } - /// Set cell contents by cell index. Only safe to call before frame execution starts. - pub(crate) fn set_cell_contents(&self, cell_idx: usize, value: Option) { - let nlocals = self.code.varnames.len(); - // SAFETY: Called before frame execution starts. - let fastlocals = unsafe { (*self.iframe.get()).localsplus.fastlocals() }; - fastlocals[nlocals + cell_idx] - .as_ref() - .and_then(|obj| obj.downcast_ref::()) - .expect("cell slot empty or not a PyCell") - .set(value); - } - /// Store a borrowed back-reference to the owning generator/coroutine. /// The caller must ensure the generator outlives the frame. pub fn set_generator(&self, generator: &PyObject) { @@ -888,53 +862,78 @@ impl Frame { } pub fn locals(&self, vm: &VirtualMachine) -> PyResult { + use rustpython_compiler_core::bytecode::{CO_FAST_CELL, CO_FAST_FREE, CO_FAST_HIDDEN, CO_FAST_LOCAL}; // SAFETY: Either the frame is not executing (caller checked owner), // or we're in a trace callback on the same thread that's executing. let locals = &self.locals; let code = &**self.code; - let map = &code.varnames; - let j = core::cmp::min(map.len(), code.varnames.len()); let locals_map = locals.mapping(vm); - if !code.varnames.is_empty() { - let fastlocals = unsafe { (*self.iframe.get()).localsplus.fastlocals() }; - for (&k, v) in zip(&map[..j], fastlocals) { - match locals_map.ass_subscript(k, v.clone(), vm) { - Ok(()) => {} - Err(e) if e.fast_isinstance(vm.ctx.exceptions.key_error) => {} - Err(e) => return Err(e), - } + let fastlocals = unsafe { (*self.iframe.get()).localsplus.fastlocals() }; + + // Iterate through all localsplus slots using localspluskinds + let nlocalsplus = code.localspluskinds.len(); + let nfrees = code.freevars.len(); + let free_start = nlocalsplus - nfrees; + let is_optimized = code.flags.contains(bytecode::CodeFlags::OPTIMIZED); + + // Track which non-merged cellvar index we're at + let mut nonmerged_cell_idx = 0; + + for (i, &kind) in code.localspluskinds.iter().enumerate() { + if kind & CO_FAST_HIDDEN != 0 { + continue; } - } - if !code.cellvars.is_empty() || !code.freevars.is_empty() { - let fastlocals = unsafe { (*self.iframe.get()).localsplus.fastlocals() }; - for (i, &k) in code.cellvars.iter().enumerate() { - // When a variable appears in both varnames and cellvars - // (inlined comprehension with scope tweak), the fastlocal - // value takes precedence, matching CPython FrameLocalsProxy. - let has_fastlocal = code - .varnames - .iter() - .position(|&v| v == k) - .is_some_and(|idx| fastlocals.get(idx).is_some_and(|v| v.is_some())); - if has_fastlocal { - continue; - } - let cell_value = self.get_cell_contents(i); - match locals_map.ass_subscript(k, cell_value, vm) { - Ok(()) => {} - Err(e) if e.fast_isinstance(vm.ctx.exceptions.key_error) => {} - Err(e) => return Err(e), - } + + // Free variables only included for optimized (function-like) scopes. + // Class/module scopes should not expose free vars in locals(). + if kind == CO_FAST_FREE && !is_optimized { + continue; } - if code.flags.contains(bytecode::CodeFlags::OPTIMIZED) { - for (i, &k) in code.freevars.iter().enumerate() { - let cell_value = self.get_cell_contents(code.cellvars.len() + i); - match locals_map.ass_subscript(k, cell_value, vm) { - Ok(()) => {} - Err(e) if e.fast_isinstance(vm.ctx.exceptions.key_error) => {} - Err(e) => return Err(e), + + // Get the name for this slot + let name = if kind & CO_FAST_LOCAL != 0 { + code.varnames[i] + } else if kind & CO_FAST_FREE != 0 { + code.freevars[i - free_start] + } else if kind & CO_FAST_CELL != 0 { + // Non-merged cell: find the name by skipping merged cellvars + let mut found_name = None; + let mut skip = nonmerged_cell_idx; + for cv in code.cellvars.iter() { + let is_merged = code.varnames.iter().any(|&v| v == *cv); + if !is_merged { + if skip == 0 { + found_name = Some(*cv); + break; + } + skip -= 1; } } + nonmerged_cell_idx += 1; + match found_name { + Some(n) => n, + None => continue, + } + } else { + continue; + }; + + // Get the value + let value = if kind & (CO_FAST_CELL | CO_FAST_FREE) != 0 { + // Cell or free var: extract value from PyCell + fastlocals[i] + .as_ref() + .and_then(|obj| obj.downcast_ref::()) + .and_then(|cell| cell.get()) + } else { + // Regular local + fastlocals[i].clone() + }; + + match locals_map.ass_subscript(name, value, vm) { + Ok(()) => {} + Err(e) if e.fast_isinstance(vm.ctx.exceptions.key_error) => {} + Err(e) => return Err(e), } } Ok(locals.clone_mapping(vm)) @@ -1337,13 +1336,12 @@ impl ExecutingFrame<'_> { self.lasti.load(Relaxed) } - /// Access the PyCellRef at the given cell/free variable index. - /// `cell_idx` is 0-based: 0..ncells for cellvars, ncells.. for freevars. + /// Access the PyCellRef at the given localsplus index. #[inline(always)] - fn cell_ref(&self, cell_idx: usize) -> &PyCell { - let nlocals = self.code.varnames.len(); - self.localsplus.fastlocals()[nlocals + cell_idx] - .as_ref() + fn cell_ref(&self, localsplus_idx: usize) -> &PyCell { + let fastlocals = self.localsplus.fastlocals(); + let slot = &fastlocals[localsplus_idx]; + slot.as_ref() .expect("cell slot empty") .downcast_ref::() .expect("cell slot is not a PyCell") @@ -1883,18 +1881,72 @@ impl ExecutingFrame<'_> { } } - fn unbound_cell_exception(&self, i: usize, vm: &VirtualMachine) -> PyBaseExceptionRef { - if let Some(&name) = self.code.cellvars.get(i) { + fn unbound_cell_exception( + &self, + localsplus_idx: usize, + vm: &VirtualMachine, + ) -> PyBaseExceptionRef { + use rustpython_compiler_core::bytecode::CO_FAST_FREE; + let kind = self + .code + .localspluskinds + .get(localsplus_idx) + .copied() + .unwrap_or(0); + if kind & CO_FAST_FREE != 0 { + let name = self.localsplus_name(localsplus_idx); + vm.new_name_error( + format!("cannot access free variable '{name}' where it is not associated with a value in enclosing scope"), + name.to_owned(), + ) + } else { + // Both merged cells (LOCAL|CELL) and non-merged cells get unbound local error + let name = self.localsplus_name(localsplus_idx); vm.new_exception_msg( vm.ctx.exceptions.unbound_local_error.to_owned(), format!("local variable '{name}' referenced before assignment").into(), ) + } + } + + /// Get the variable name for a localsplus index. + fn localsplus_name(&self, idx: usize) -> &'static PyStrInterned { + use rustpython_compiler_core::bytecode::{CO_FAST_CELL, CO_FAST_FREE, CO_FAST_LOCAL}; + let nlocals = self.code.varnames.len(); + let kind = self.code.localspluskinds.get(idx).copied().unwrap_or(0); + if kind & CO_FAST_LOCAL != 0 { + // Merged cell or regular local: name is in varnames + self.code.varnames[idx] + } else if kind & CO_FAST_FREE != 0 { + // Free var: slots are at the end of localsplus + let nlocalsplus = self.code.localspluskinds.len(); + let nfrees = self.code.freevars.len(); + let free_start = nlocalsplus - nfrees; + self.code.freevars[idx - free_start] + } else if kind & CO_FAST_CELL != 0 { + // Non-merged cell: count how many non-merged cell slots are before + // this index to find the corresponding cellvars entry. + // Non-merged cellvars appear in their original order (skipping merged ones). + let nonmerged_pos = self.code.localspluskinds[nlocals..idx] + .iter() + .filter(|&&k| k == CO_FAST_CELL) + .count(); + // Skip merged cellvars to find the right one + let mut cv_idx = 0; + let mut nonmerged_count = 0; + for (i, name) in self.code.cellvars.iter().enumerate() { + let is_merged = self.code.varnames.iter().any(|v| *v == *name); + if !is_merged { + if nonmerged_count == nonmerged_pos { + cv_idx = i; + break; + } + nonmerged_count += 1; + } + } + self.code.cellvars[cv_idx] } else { - let name = self.code.freevars[i - self.code.cellvars.len()]; - vm.new_name_error( - format!("cannot access free variable '{name}' where it is not associated with a value in enclosing scope"), - name.to_owned(), - ) + self.code.varnames[idx] } } @@ -2165,13 +2217,29 @@ impl ExecutingFrame<'_> { self.push_stackref_opt(value); Ok(None) } - Instruction::CopyFreeVars { .. } => { - // Free vars are already set up at frame creation time in RustPython + Instruction::CopyFreeVars { n } => { + let n = n.get(arg) as usize; + if n > 0 { + let closure = self + .object + .func_obj + .as_ref() + .and_then(|f| f.downcast_ref::()) + .and_then(|f| f.closure.as_ref()); + let nlocalsplus = self.code.localspluskinds.len(); + let freevar_start = nlocalsplus - n; + let fastlocals = self.localsplus.fastlocals_mut(); + if let Some(closure) = closure { + for i in 0..n { + fastlocals[freevar_start + i] = Some(closure[i].clone().into()); + } + } + } Ok(None) } Instruction::DeleteAttr { namei: idx } => self.delete_attr(vm, idx.get(arg)), Instruction::DeleteDeref { i } => { - self.cell_ref(i.get(arg) as usize).set(None); + self.cell_ref(i.get(arg).as_usize()).set(None); Ok(None) } Instruction::DeleteFast { var_num } => { @@ -2597,12 +2665,8 @@ impl ExecutingFrame<'_> { Instruction::LoadFromDictOrDeref { i } => { // Pop dict from stack (locals or classdict depending on context) let class_dict = self.pop_value(); - let i = i.get(arg) as usize; - let name = if i < self.code.cellvars.len() { - self.code.cellvars[i] - } else { - self.code.freevars[i - self.code.cellvars.len()] - }; + let idx = i.get(arg).as_usize(); + let name = self.localsplus_name(idx); // Only treat KeyError as "not found", propagate other exceptions let value = if let Some(dict_obj) = class_dict.downcast_ref::() { dict_obj.get_item_opt(name, vm)? @@ -2616,9 +2680,9 @@ impl ExecutingFrame<'_> { self.push_value(match value { Some(v) => v, None => self - .cell_ref(i) + .cell_ref(idx) .get() - .ok_or_else(|| self.unbound_cell_exception(i, vm))?, + .ok_or_else(|| self.unbound_cell_exception(idx, vm))?, }); Ok(None) } @@ -2684,7 +2748,7 @@ impl ExecutingFrame<'_> { Ok(None) } Instruction::LoadDeref { i } => { - let idx = i.get(arg) as usize; + let idx = i.get(arg).as_usize(); let x = self .cell_ref(idx) .get() @@ -2866,28 +2930,14 @@ impl ExecutingFrame<'_> { } Instruction::MakeFunction => self.execute_make_function(vm), Instruction::MakeCell { i } => { - // PEP 709: Save the current cell object on the stack and - // create a fresh empty cell for the inlined comprehension. - // The old cell is restored afterwards via RestoreCell. - let cell_idx = i.get(arg) as usize; - let nlocals = self.code.varnames.len(); - let old_cell = self.localsplus.fastlocals_mut()[nlocals + cell_idx] - .take() - .expect("cell slot empty"); - let new_cell = PyCell::default().into_ref(&vm.ctx).into(); - self.localsplus.fastlocals_mut()[nlocals + cell_idx] = Some(new_cell); - // Push the old cell object itself - self.push_value(old_cell); - Ok(None) - } - Instruction::RestoreCell { i } => { - // PEP 709: Restore the saved cell object after an inlined - // comprehension. Pops the old cell from the stack and writes - // it back to the cell slot, replacing the temporary cell. - let cell_idx = i.get(arg) as usize; - let nlocals = self.code.varnames.len(); - let old_cell = self.pop_value(); - self.localsplus.fastlocals_mut()[nlocals + cell_idx] = Some(old_cell); + // Wrap the current slot value (if any) in a new PyCell. + // For merged cells (LOCAL|CELL), this wraps the argument value. + // For non-merged cells, this creates an empty cell. + let idx = i.get(arg).as_usize(); + let fastlocals = self.localsplus.fastlocals_mut(); + let initial = fastlocals[idx].take(); + let cell = PyCell::new(initial).into_ref(&vm.ctx).into(); + fastlocals[idx] = Some(cell); Ok(None) } Instruction::MapAdd { i } => { @@ -3326,7 +3376,7 @@ impl ExecutingFrame<'_> { } Instruction::StoreDeref { i } => { let value = self.pop_value(); - self.cell_ref(i.get(arg) as usize).set(Some(value)); + self.cell_ref(i.get(arg).as_usize()).set(Some(value)); Ok(None) } Instruction::StoreFast { var_num } => { diff --git a/crates/vm/src/version.rs b/crates/vm/src/version.rs index a75a6f47de6..21efecd6c5a 100644 --- a/crates/vm/src/version.rs +++ b/crates/vm/src/version.rs @@ -129,7 +129,8 @@ pub fn get_git_datetime() -> String { } // Must be aligned to Lib/importlib/_bootstrap_external.py -pub const PYC_MAGIC_NUMBER: u16 = 2996; +// Bumped to 2997 for MAKE_CELL/COPY_FREE_VARS prolog and cell-local merging +pub const PYC_MAGIC_NUMBER: u16 = 2997; // CPython format: magic_number | ('\r' << 16) | ('\n' << 24) // This protects against text-mode file reads diff --git a/crates/vm/src/vm/context.rs b/crates/vm/src/vm/context.rs index dc0af9386fe..46500ab3c2a 100644 --- a/crates/vm/src/vm/context.rs +++ b/crates/vm/src/vm/context.rs @@ -427,12 +427,12 @@ impl Context { max_stackdepth: 2, obj_name: names.__init__, qualname: names.__init__, - cell2arg: None, constants: core::iter::empty().collect(), names: Vec::new().into_boxed_slice(), varnames: Vec::new().into_boxed_slice(), cellvars: Vec::new().into_boxed_slice(), freevars: Vec::new().into_boxed_slice(), + localspluskinds: Vec::new().into_boxed_slice(), linetable: Vec::new().into_boxed_slice(), exceptiontable: Vec::new().into_boxed_slice(), }; From c18b4647a82ce3b1faa27f2a7e1a6284333fa7b7 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Tue, 17 Mar 2026 23:19:28 +0900 Subject: [PATCH 3/4] Fix clippy warnings, formatting, restore _opcode_metadata.py Fix cast_possible_truncation, nonminimal_bool, collapsible_if, manual_contains clippy lints. Restore _opcode_metadata.py to upstream/main version (3.14 aligned). Pre-copy closure cells in Frame::new for coroutine locals(). Handle raw values in merged cell slots during inlined comps. Exclude async comprehensions from inlining path. --- .cspell.dict/cpython.txt | 1 + .cspell.json | 1 + Lib/_opcode_metadata.py | 307 +++++++++++++++++++--------- crates/codegen/src/compile.rs | 72 +++++-- crates/codegen/src/ir.rs | 14 +- crates/compiler-core/src/marshal.rs | 7 +- crates/vm/src/builtins/code.rs | 7 +- crates/vm/src/builtins/super.rs | 7 +- crates/vm/src/frame.rs | 60 ++++-- 9 files changed, 341 insertions(+), 135 deletions(-) diff --git a/.cspell.dict/cpython.txt b/.cspell.dict/cpython.txt index a9fbc8f4318..38e8656c3b2 100644 --- a/.cspell.dict/cpython.txt +++ b/.cspell.dict/cpython.txt @@ -109,6 +109,7 @@ lineiterator linetable loadfast localsplus +localspluskinds Lshift lsprof MAXBLOCKS diff --git a/.cspell.json b/.cspell.json index f6887a7bb35..b7efb2e2311 100644 --- a/.cspell.json +++ b/.cspell.json @@ -70,6 +70,7 @@ "lossily", "mcache", "oparg", + "opargs", "pyc", "significand", "summands", diff --git a/Lib/_opcode_metadata.py b/Lib/_opcode_metadata.py index bb55ee423cf..4da6e507736 100644 --- a/Lib/_opcode_metadata.py +++ b/Lib/_opcode_metadata.py @@ -1,16 +1,223 @@ # This file is generated by scripts/generate_opcode_metadata.py -# for RustPython bytecode format (CPython 3.13 compatible opcode numbers). +# for RustPython bytecode format (CPython 3.14 compatible opcode numbers). # Do not edit! -_specializations = {} +_specializations = { + "RESUME": [ + "RESUME_CHECK", + ], + "LOAD_CONST": [ + "LOAD_CONST_MORTAL", + "LOAD_CONST_IMMORTAL", + ], + "TO_BOOL": [ + "TO_BOOL_ALWAYS_TRUE", + "TO_BOOL_BOOL", + "TO_BOOL_INT", + "TO_BOOL_LIST", + "TO_BOOL_NONE", + "TO_BOOL_STR", + ], + "BINARY_OP": [ + "BINARY_OP_MULTIPLY_INT", + "BINARY_OP_ADD_INT", + "BINARY_OP_SUBTRACT_INT", + "BINARY_OP_MULTIPLY_FLOAT", + "BINARY_OP_ADD_FLOAT", + "BINARY_OP_SUBTRACT_FLOAT", + "BINARY_OP_ADD_UNICODE", + "BINARY_OP_SUBSCR_LIST_INT", + "BINARY_OP_SUBSCR_LIST_SLICE", + "BINARY_OP_SUBSCR_TUPLE_INT", + "BINARY_OP_SUBSCR_STR_INT", + "BINARY_OP_SUBSCR_DICT", + "BINARY_OP_SUBSCR_GETITEM", + "BINARY_OP_EXTEND", + "BINARY_OP_INPLACE_ADD_UNICODE", + ], + "STORE_SUBSCR": [ + "STORE_SUBSCR_DICT", + "STORE_SUBSCR_LIST_INT", + ], + "SEND": [ + "SEND_GEN", + ], + "UNPACK_SEQUENCE": [ + "UNPACK_SEQUENCE_TWO_TUPLE", + "UNPACK_SEQUENCE_TUPLE", + "UNPACK_SEQUENCE_LIST", + ], + "STORE_ATTR": [ + "STORE_ATTR_INSTANCE_VALUE", + "STORE_ATTR_SLOT", + "STORE_ATTR_WITH_HINT", + ], + "LOAD_GLOBAL": [ + "LOAD_GLOBAL_MODULE", + "LOAD_GLOBAL_BUILTIN", + ], + "LOAD_SUPER_ATTR": [ + "LOAD_SUPER_ATTR_ATTR", + "LOAD_SUPER_ATTR_METHOD", + ], + "LOAD_ATTR": [ + "LOAD_ATTR_INSTANCE_VALUE", + "LOAD_ATTR_MODULE", + "LOAD_ATTR_WITH_HINT", + "LOAD_ATTR_SLOT", + "LOAD_ATTR_CLASS", + "LOAD_ATTR_CLASS_WITH_METACLASS_CHECK", + "LOAD_ATTR_PROPERTY", + "LOAD_ATTR_GETATTRIBUTE_OVERRIDDEN", + "LOAD_ATTR_METHOD_WITH_VALUES", + "LOAD_ATTR_METHOD_NO_DICT", + "LOAD_ATTR_METHOD_LAZY_DICT", + "LOAD_ATTR_NONDESCRIPTOR_WITH_VALUES", + "LOAD_ATTR_NONDESCRIPTOR_NO_DICT", + ], + "COMPARE_OP": [ + "COMPARE_OP_FLOAT", + "COMPARE_OP_INT", + "COMPARE_OP_STR", + ], + "CONTAINS_OP": [ + "CONTAINS_OP_SET", + "CONTAINS_OP_DICT", + ], + "JUMP_BACKWARD": [ + "JUMP_BACKWARD_NO_JIT", + "JUMP_BACKWARD_JIT", + ], + "FOR_ITER": [ + "FOR_ITER_LIST", + "FOR_ITER_TUPLE", + "FOR_ITER_RANGE", + "FOR_ITER_GEN", + ], + "CALL": [ + "CALL_BOUND_METHOD_EXACT_ARGS", + "CALL_PY_EXACT_ARGS", + "CALL_TYPE_1", + "CALL_STR_1", + "CALL_TUPLE_1", + "CALL_BUILTIN_CLASS", + "CALL_BUILTIN_O", + "CALL_BUILTIN_FAST", + "CALL_BUILTIN_FAST_WITH_KEYWORDS", + "CALL_LEN", + "CALL_ISINSTANCE", + "CALL_LIST_APPEND", + "CALL_METHOD_DESCRIPTOR_O", + "CALL_METHOD_DESCRIPTOR_FAST_WITH_KEYWORDS", + "CALL_METHOD_DESCRIPTOR_NOARGS", + "CALL_METHOD_DESCRIPTOR_FAST", + "CALL_ALLOC_AND_ENTER_INIT", + "CALL_PY_GENERAL", + "CALL_BOUND_METHOD_GENERAL", + "CALL_NON_PY_GENERAL", + ], + "CALL_KW": [ + "CALL_KW_BOUND_METHOD", + "CALL_KW_PY", + "CALL_KW_NON_PY", + ], +} -_specialized_opmap = {} +_specialized_opmap = { + 'BINARY_OP_ADD_FLOAT': 129, + 'BINARY_OP_ADD_INT': 130, + 'BINARY_OP_ADD_UNICODE': 131, + 'BINARY_OP_EXTEND': 132, + 'BINARY_OP_INPLACE_ADD_UNICODE': 3, + 'BINARY_OP_MULTIPLY_FLOAT': 133, + 'BINARY_OP_MULTIPLY_INT': 134, + 'BINARY_OP_SUBSCR_DICT': 135, + 'BINARY_OP_SUBSCR_GETITEM': 136, + 'BINARY_OP_SUBSCR_LIST_INT': 137, + 'BINARY_OP_SUBSCR_LIST_SLICE': 138, + 'BINARY_OP_SUBSCR_STR_INT': 139, + 'BINARY_OP_SUBSCR_TUPLE_INT': 140, + 'BINARY_OP_SUBTRACT_FLOAT': 141, + 'BINARY_OP_SUBTRACT_INT': 142, + 'CALL_ALLOC_AND_ENTER_INIT': 143, + 'CALL_BOUND_METHOD_EXACT_ARGS': 144, + 'CALL_BOUND_METHOD_GENERAL': 145, + 'CALL_BUILTIN_CLASS': 146, + 'CALL_BUILTIN_FAST': 147, + 'CALL_BUILTIN_FAST_WITH_KEYWORDS': 148, + 'CALL_BUILTIN_O': 149, + 'CALL_ISINSTANCE': 150, + 'CALL_KW_BOUND_METHOD': 151, + 'CALL_KW_NON_PY': 152, + 'CALL_KW_PY': 153, + 'CALL_LEN': 154, + 'CALL_LIST_APPEND': 155, + 'CALL_METHOD_DESCRIPTOR_FAST': 156, + 'CALL_METHOD_DESCRIPTOR_FAST_WITH_KEYWORDS': 157, + 'CALL_METHOD_DESCRIPTOR_NOARGS': 158, + 'CALL_METHOD_DESCRIPTOR_O': 159, + 'CALL_NON_PY_GENERAL': 160, + 'CALL_PY_EXACT_ARGS': 161, + 'CALL_PY_GENERAL': 162, + 'CALL_STR_1': 163, + 'CALL_TUPLE_1': 164, + 'CALL_TYPE_1': 165, + 'COMPARE_OP_FLOAT': 166, + 'COMPARE_OP_INT': 167, + 'COMPARE_OP_STR': 168, + 'CONTAINS_OP_DICT': 169, + 'CONTAINS_OP_SET': 170, + 'FOR_ITER_GEN': 171, + 'FOR_ITER_LIST': 172, + 'FOR_ITER_RANGE': 173, + 'FOR_ITER_TUPLE': 174, + 'JUMP_BACKWARD_JIT': 175, + 'JUMP_BACKWARD_NO_JIT': 176, + 'LOAD_ATTR_CLASS': 177, + 'LOAD_ATTR_CLASS_WITH_METACLASS_CHECK': 178, + 'LOAD_ATTR_GETATTRIBUTE_OVERRIDDEN': 179, + 'LOAD_ATTR_INSTANCE_VALUE': 180, + 'LOAD_ATTR_METHOD_LAZY_DICT': 181, + 'LOAD_ATTR_METHOD_NO_DICT': 182, + 'LOAD_ATTR_METHOD_WITH_VALUES': 183, + 'LOAD_ATTR_MODULE': 184, + 'LOAD_ATTR_NONDESCRIPTOR_NO_DICT': 185, + 'LOAD_ATTR_NONDESCRIPTOR_WITH_VALUES': 186, + 'LOAD_ATTR_PROPERTY': 187, + 'LOAD_ATTR_SLOT': 188, + 'LOAD_ATTR_WITH_HINT': 189, + 'LOAD_CONST_IMMORTAL': 190, + 'LOAD_CONST_MORTAL': 191, + 'LOAD_GLOBAL_BUILTIN': 192, + 'LOAD_GLOBAL_MODULE': 193, + 'LOAD_SUPER_ATTR_ATTR': 194, + 'LOAD_SUPER_ATTR_METHOD': 195, + 'RESUME_CHECK': 196, + 'SEND_GEN': 197, + 'STORE_ATTR_INSTANCE_VALUE': 198, + 'STORE_ATTR_SLOT': 199, + 'STORE_ATTR_WITH_HINT': 200, + 'STORE_SUBSCR_DICT': 201, + 'STORE_SUBSCR_LIST_INT': 202, + 'TO_BOOL_ALWAYS_TRUE': 203, + 'TO_BOOL_BOOL': 204, + 'TO_BOOL_INT': 205, + 'TO_BOOL_LIST': 206, + 'TO_BOOL_NONE': 207, + 'TO_BOOL_STR': 208, + 'UNPACK_SEQUENCE_LIST': 209, + 'UNPACK_SEQUENCE_TUPLE': 210, + 'UNPACK_SEQUENCE_TWO_TUPLE': 211, +} opmap = { 'CACHE': 0, + 'RESERVED': 17, + 'RESUME': 128, + 'INSTRUMENTED_LINE': 254, + 'ENTER_EXECUTOR': 255, 'BINARY_SLICE': 1, 'BUILD_TEMPLATE': 2, - 'BINARY_OP_INPLACE_ADD_UNICODE': 3, 'CALL_FUNCTION_EX': 4, 'CHECK_EG_MATCH': 5, 'CHECK_EXC_MATCH': 6, @@ -24,7 +231,6 @@ 'GET_AITER': 14, 'GET_ANEXT': 15, 'GET_ITER': 16, - 'RESERVED': 17, 'GET_LEN': 18, 'GET_YIELD_FROM_ITER': 19, 'INTERPRETER_EXIT': 20, @@ -128,90 +334,6 @@ 'UNPACK_EX': 118, 'UNPACK_SEQUENCE': 119, 'YIELD_VALUE': 120, - 'RESUME': 128, - 'BINARY_OP_ADD_FLOAT': 129, - 'BINARY_OP_ADD_INT': 130, - 'BINARY_OP_ADD_UNICODE': 131, - 'BINARY_OP_EXTEND': 132, - 'BINARY_OP_MULTIPLY_FLOAT': 133, - 'BINARY_OP_MULTIPLY_INT': 134, - 'BINARY_OP_SUBSCR_DICT': 135, - 'BINARY_OP_SUBSCR_GETITEM': 136, - 'BINARY_OP_SUBSCR_LIST_INT': 137, - 'BINARY_OP_SUBSCR_LIST_SLICE': 138, - 'BINARY_OP_SUBSCR_STR_INT': 139, - 'BINARY_OP_SUBSCR_TUPLE_INT': 140, - 'BINARY_OP_SUBTRACT_FLOAT': 141, - 'BINARY_OP_SUBTRACT_INT': 142, - 'CALL_ALLOC_AND_ENTER_INIT': 143, - 'CALL_BOUND_METHOD_EXACT_ARGS': 144, - 'CALL_BOUND_METHOD_GENERAL': 145, - 'CALL_BUILTIN_CLASS': 146, - 'CALL_BUILTIN_FAST': 147, - 'CALL_BUILTIN_FAST_WITH_KEYWORDS': 148, - 'CALL_BUILTIN_O': 149, - 'CALL_ISINSTANCE': 150, - 'CALL_KW_BOUND_METHOD': 151, - 'CALL_KW_NON_PY': 152, - 'CALL_KW_PY': 153, - 'CALL_LEN': 154, - 'CALL_LIST_APPEND': 155, - 'CALL_METHOD_DESCRIPTOR_FAST': 156, - 'CALL_METHOD_DESCRIPTOR_FAST_WITH_KEYWORDS': 157, - 'CALL_METHOD_DESCRIPTOR_NOARGS': 158, - 'CALL_METHOD_DESCRIPTOR_O': 159, - 'CALL_NON_PY_GENERAL': 160, - 'CALL_PY_EXACT_ARGS': 161, - 'CALL_PY_GENERAL': 162, - 'CALL_STR_1': 163, - 'CALL_TUPLE_1': 164, - 'CALL_TYPE_1': 165, - 'COMPARE_OP_FLOAT': 166, - 'COMPARE_OP_INT': 167, - 'COMPARE_OP_STR': 168, - 'CONTAINS_OP_DICT': 169, - 'CONTAINS_OP_SET': 170, - 'FOR_ITER_GEN': 171, - 'FOR_ITER_LIST': 172, - 'FOR_ITER_RANGE': 173, - 'FOR_ITER_TUPLE': 174, - 'JUMP_BACKWARD_JIT': 175, - 'JUMP_BACKWARD_NO_JIT': 176, - 'LOAD_ATTR_CLASS': 177, - 'LOAD_ATTR_CLASS_WITH_METACLASS_CHECK': 178, - 'LOAD_ATTR_GETATTRIBUTE_OVERRIDDEN': 179, - 'LOAD_ATTR_INSTANCE_VALUE': 180, - 'LOAD_ATTR_METHOD_LAZY_DICT': 181, - 'LOAD_ATTR_METHOD_NO_DICT': 182, - 'LOAD_ATTR_METHOD_WITH_VALUES': 183, - 'LOAD_ATTR_MODULE': 184, - 'LOAD_ATTR_NONDESCRIPTOR_NO_DICT': 185, - 'LOAD_ATTR_NONDESCRIPTOR_WITH_VALUES': 186, - 'LOAD_ATTR_PROPERTY': 187, - 'LOAD_ATTR_SLOT': 188, - 'LOAD_ATTR_WITH_HINT': 189, - 'LOAD_CONST_IMMORTAL': 190, - 'LOAD_CONST_MORTAL': 191, - 'LOAD_GLOBAL_BUILTIN': 192, - 'LOAD_GLOBAL_MODULE': 193, - 'LOAD_SUPER_ATTR_ATTR': 194, - 'LOAD_SUPER_ATTR_METHOD': 195, - 'RESUME_CHECK': 196, - 'SEND_GEN': 197, - 'STORE_ATTR_INSTANCE_VALUE': 198, - 'STORE_ATTR_SLOT': 199, - 'STORE_ATTR_WITH_HINT': 200, - 'STORE_SUBSCR_DICT': 201, - 'STORE_SUBSCR_LIST_INT': 202, - 'TO_BOOL_ALWAYS_TRUE': 203, - 'TO_BOOL_BOOL': 204, - 'TO_BOOL_INT': 205, - 'TO_BOOL_LIST': 206, - 'TO_BOOL_NONE': 207, - 'TO_BOOL_STR': 208, - 'UNPACK_SEQUENCE_LIST': 209, - 'UNPACK_SEQUENCE_TUPLE': 210, - 'UNPACK_SEQUENCE_TWO_TUPLE': 211, 'INSTRUMENTED_END_FOR': 234, 'INSTRUMENTED_POP_ITER': 235, 'INSTRUMENTED_END_SEND': 236, @@ -232,8 +354,6 @@ 'INSTRUMENTED_CALL_KW': 251, 'INSTRUMENTED_CALL_FUNCTION_EX': 252, 'INSTRUMENTED_JUMP_BACKWARD': 253, - 'INSTRUMENTED_LINE': 254, - 'ENTER_EXECUTOR': 255, 'ANNOTATIONS_PLACEHOLDER': 256, 'JUMP': 257, 'JUMP_IF_FALSE': 258, @@ -247,6 +367,5 @@ 'STORE_FAST_MAYBE_NULL': 266, } -# CPython 3.13 compatible: opcodes < 44 have no argument -HAVE_ARGUMENT = 44 -MIN_INSTRUMENTED_OPCODE = 236 +HAVE_ARGUMENT = 43 +MIN_INSTRUMENTED_OPCODE = 234 diff --git a/crates/codegen/src/compile.rs b/crates/codegen/src/compile.rs index 14329e34d43..fc81320fa91 100644 --- a/crates/codegen/src/compile.rs +++ b/crates/codegen/src/compile.rs @@ -1155,11 +1155,16 @@ impl Compiler { { let nfrees = self.code_stack.last().unwrap().metadata.freevars.len(); if nfrees > 0 { - emit!(self, Instruction::CopyFreeVars { n: nfrees as u32 }); + 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 = (i as u32).into(); + let i_varnum: oparg::VarNum = u32::try_from(i).expect("too many cellvars").into(); emit!(self, Instruction::MakeCell { i: i_varnum }); } } @@ -2022,6 +2027,19 @@ impl Compiler { _ => unreachable!("Invalid scope for Deref operation"), }; + // Mark cell variables accessed inside inlined comprehensions as hidden + if self.current_code_info().in_inlined_comp { + let info = self.code_stack.last_mut().unwrap(); + if info + .metadata + .fast_hidden + .get(name.as_ref()) + .is_none_or(|&v| v) + { + info.metadata.fast_hidden.insert(name.to_string(), true); + } + } + match usage { NameUsage::Load => { // ClassBlock (not inlined comp): LOAD_LOCALS first, then LOAD_FROM_DICT_OR_DEREF @@ -2043,6 +2061,18 @@ impl Compiler { } NameOp::Fast => { let var_num = self.get_local_var_index(&name)?; + // Mark variables accessed inside inlined comprehensions as hidden + if self.current_code_info().in_inlined_comp { + let info = self.code_stack.last_mut().unwrap(); + if info + .metadata + .fast_hidden + .get(name.as_ref()) + .is_none_or(|&v| v) + { + info.metadata.fast_hidden.insert(name.to_string(), true); + } + } match usage { NameUsage::Load => emit!(self, Instruction::LoadFast { var_num }), NameUsage::Store => emit!(self, Instruction::StoreFast { var_num }), @@ -7349,6 +7379,14 @@ impl Compiler { node_index: _, range: _, }) => { + // Walrus targets in inlined comps should NOT be hidden from locals() + if self.current_code_info().in_inlined_comp + && let ast::Expr::Name(ast::ExprName { id, .. }) = target.as_ref() + { + let name = self.mangle(id.as_str()); + let info = self.code_stack.last_mut().unwrap(); + info.metadata.fast_hidden.insert(name.to_string(), false); + } self.compile_expression(value)?; emit!(self, Instruction::Copy { i: 1 }); self.compile_store(target)?; @@ -7801,7 +7839,7 @@ impl Compiler { // We must have at least one generator: assert!(!generators.is_empty()); - if is_inlined { + 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; @@ -8059,20 +8097,20 @@ impl Compiler { emit!(self, Instruction::LoadFastAndClear { var_num }); total_stack_items += 1; // If the comp symbol is CELL, emit MAKE_CELL to create fresh cell - if let Some(comp_sym) = comp_table.symbols.get(name) { - if comp_sym.scope == SymbolScope::Cell { - let i = if self - .current_symbol_table() - .symbols - .get(name) - .is_some_and(|s| s.scope == SymbolScope::Free) - { - self.get_free_var_index(name)? - } else { - self.get_cell_var_index(name)? - }; - emit!(self, Instruction::MakeCell { i }); - } + if let Some(comp_sym) = comp_table.symbols.get(name) + && comp_sym.scope == SymbolScope::Cell + { + let i = if self + .current_symbol_table() + .symbols + .get(name) + .is_some_and(|s| s.scope == SymbolScope::Free) + { + self.get_free_var_index(name)? + } else { + self.get_cell_var_index(name)? + }; + emit!(self, Instruction::MakeCell { i }); } } diff --git a/crates/codegen/src/ir.rs b/crates/codegen/src/ir.rs index 5ba56eb0ec5..6be851e15b5 100644 --- a/crates/codegen/src/ir.rs +++ b/crates/codegen/src/ir.rs @@ -8,9 +8,9 @@ use num_traits::ToPrimitive; use rustpython_compiler_core::{ OneIndexed, SourceLocation, bytecode::{ - AnyInstruction, Arg, CO_FAST_CELL, CO_FAST_FREE, CO_FAST_LOCAL, CodeFlags, CodeObject, - CodeUnit, CodeUnits, ConstantData, ExceptionTableEntry, InstrDisplayContext, Instruction, - InstructionMetadata, Label, OpArg, PseudoInstruction, PyCodeLocationInfoKind, + AnyInstruction, Arg, CO_FAST_CELL, CO_FAST_FREE, CO_FAST_HIDDEN, CO_FAST_LOCAL, CodeFlags, + CodeObject, CodeUnit, CodeUnits, ConstantData, ExceptionTableEntry, InstrDisplayContext, + Instruction, InstructionMetadata, Label, OpArg, PseudoInstruction, PyCodeLocationInfoKind, encode_exception_table, oparg, }, varint::{write_signed_varint, write_varint}, @@ -236,7 +236,7 @@ impl CodeInfo { varnames: varname_cache, cellvars: cellvar_cache, freevars: freevar_cache, - fast_hidden: _, + fast_hidden, argcount: arg_count, posonlyargcount: posonlyarg_count, kwonlyargcount: kwonlyarg_count, @@ -514,6 +514,12 @@ impl CodeInfo { let idx = cellfixedoffsets[ncells + i] as usize; localspluskinds[idx] = CO_FAST_FREE; } + // Apply CO_FAST_HIDDEN for inlined comprehension variables + for (name, &hidden) in &fast_hidden { + if hidden && let Some(idx) = varname_cache.get_index_of(name.as_str()) { + localspluskinds[idx] |= CO_FAST_HIDDEN; + } + } Ok(CodeObject { flags, diff --git a/crates/compiler-core/src/marshal.rs b/crates/compiler-core/src/marshal.rs index 770b8531824..c47d41f1233 100644 --- a/crates/compiler-core/src/marshal.rs +++ b/crates/compiler-core/src/marshal.rs @@ -271,9 +271,10 @@ pub fn deserialize_code( let ncells = cellvars.len(); let nfrees = freevars.len(); // Count merged cells (cellvar also in varnames) - let numdropped = cellvars.iter().filter(|cv| { - varnames.iter().any(|v| v.as_ref() == cv.as_ref()) - }).count(); + let numdropped = cellvars + .iter() + .filter(|cv| varnames.iter().any(|v| v.as_ref() == cv.as_ref())) + .count(); let nlocalsplus = nlocals + ncells - numdropped + nfrees; let mut kinds = vec![0u8; nlocalsplus]; // Mark locals diff --git a/crates/vm/src/builtins/code.rs b/crates/vm/src/builtins/code.rs index 65bab0d3a9b..0bf193914c2 100644 --- a/crates/vm/src/builtins/code.rs +++ b/crates/vm/src/builtins/code.rs @@ -639,9 +639,10 @@ impl Constructor for PyCode { let nlocals = varnames.len(); let ncells = cellvars.len(); let nfrees = freevars.len(); - let numdropped = cellvars.iter().filter(|cv| { - varnames.iter().any(|v| *v == **cv) - }).count(); + let numdropped = cellvars + .iter() + .filter(|cv| varnames.iter().any(|v| *v == **cv)) + .count(); let nlocalsplus = nlocals + ncells - numdropped + nfrees; let mut kinds = vec![0u8; nlocalsplus]; for kind in kinds.iter_mut().take(nlocals) { diff --git a/crates/vm/src/builtins/super.rs b/crates/vm/src/builtins/super.rs index 2a3b83314d4..88998dc8c8a 100644 --- a/crates/vm/src/builtins/super.rs +++ b/crates/vm/src/builtins/super.rs @@ -92,7 +92,12 @@ impl Initializer for PySuper { .clone() .and_then(|val| { // If slot 0 is a merged cell (LOCAL|CELL), extract value from cell - if frame.code.localspluskinds.first().is_some_and(|&k| k & CO_FAST_CELL != 0) { + if frame + .code + .localspluskinds + .first() + .is_some_and(|&k| k & CO_FAST_CELL != 0) + { val.downcast_ref::().and_then(|c| c.get()) } else { Some(val) diff --git a/crates/vm/src/frame.rs b/crates/vm/src/frame.rs index 18085ec4978..fef682ff686 100644 --- a/crates/vm/src/frame.rs +++ b/crates/vm/src/frame.rs @@ -685,15 +685,25 @@ impl Frame { ) -> Self { let nlocalsplus = code.localspluskinds.len(); let max_stackdepth = code.max_stackdepth as usize; - let localsplus = if use_datastack { + let mut localsplus = if use_datastack { LocalsPlus::new_on_datastack(nlocalsplus, max_stackdepth, vm) } else { LocalsPlus::new(nlocalsplus, max_stackdepth) }; - // Free vars and cells are now set up by COPY_FREE_VARS and MAKE_CELL - // instructions emitted at function entry. No pre-creation needed. - let _ = closure; + // Pre-copy closure cells into free var slots so that locals() works + // even before COPY_FREE_VARS runs (e.g. coroutine before first send). + // COPY_FREE_VARS will overwrite these on first execution. + { + let nfrees = code.freevars.len(); + if nfrees > 0 { + let freevar_start = nlocalsplus - nfrees; + let fastlocals = localsplus.fastlocals_mut(); + for (i, cell) in closure.iter().enumerate() { + fastlocals[freevar_start + i] = Some(cell.clone().into()); + } + } + } let iframe = InterpreterFrame { localsplus, @@ -862,7 +872,9 @@ impl Frame { } pub fn locals(&self, vm: &VirtualMachine) -> PyResult { - use rustpython_compiler_core::bytecode::{CO_FAST_CELL, CO_FAST_FREE, CO_FAST_HIDDEN, CO_FAST_LOCAL}; + use rustpython_compiler_core::bytecode::{ + CO_FAST_CELL, CO_FAST_FREE, CO_FAST_HIDDEN, CO_FAST_LOCAL, + }; // SAFETY: Either the frame is not executing (caller checked owner), // or we're in a trace callback on the same thread that's executing. let locals = &self.locals; @@ -881,7 +893,24 @@ impl Frame { for (i, &kind) in code.localspluskinds.iter().enumerate() { if kind & CO_FAST_HIDDEN != 0 { - continue; + // Hidden variables are only skipped when their slot is empty. + // After a comprehension restores values, they should appear in locals(). + let slot_empty = match fastlocals[i].as_ref() { + None => true, + Some(obj) => { + if kind & (CO_FAST_CELL | CO_FAST_FREE) != 0 { + // If it's a PyCell, check if the cell is empty. + // If it's a raw value (merged cell during inlined comp), not empty. + obj.downcast_ref::() + .is_some_and(|cell| cell.get().is_none()) + } else { + false + } + } + }; + if slot_empty { + continue; + } } // Free variables only included for optimized (function-like) scopes. @@ -900,7 +929,7 @@ impl Frame { let mut found_name = None; let mut skip = nonmerged_cell_idx; for cv in code.cellvars.iter() { - let is_merged = code.varnames.iter().any(|&v| v == *cv); + let is_merged = code.varnames.contains(cv); if !is_merged { if skip == 0 { found_name = Some(*cv); @@ -920,11 +949,16 @@ impl Frame { // Get the value let value = if kind & (CO_FAST_CELL | CO_FAST_FREE) != 0 { - // Cell or free var: extract value from PyCell - fastlocals[i] - .as_ref() - .and_then(|obj| obj.downcast_ref::()) - .and_then(|cell| cell.get()) + // Cell or free var: extract value from PyCell. + // During inlined comprehensions, a merged cell slot may hold a raw + // value (not a PyCell) after LOAD_FAST_AND_CLEAR + STORE_FAST. + fastlocals[i].as_ref().and_then(|obj| { + if let Some(cell) = obj.downcast_ref::() { + cell.get() + } else { + Some(obj.clone()) + } + }) } else { // Regular local fastlocals[i].clone() @@ -1935,7 +1969,7 @@ impl ExecutingFrame<'_> { let mut cv_idx = 0; let mut nonmerged_count = 0; for (i, name) in self.code.cellvars.iter().enumerate() { - let is_merged = self.code.varnames.iter().any(|v| *v == *name); + let is_merged = self.code.varnames.contains(name); if !is_merged { if nonmerged_count == nonmerged_pos { cv_idx = i; From b25e850b2e0ac0ab4295e1050e8cf894e958665d Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Fri, 20 Mar 2026 05:23:18 +0900 Subject: [PATCH 4/4] Exclude async/await comprehensions from PEP 709 inlining in symboltable Async comprehensions and comprehensions with await in the element expression need their own coroutine scope and cannot be inlined. The symboltable builder was not checking these conditions, causing incorrect symbol scope resolution when an async comprehension was nested inside an inlined comprehension (e.g. [[x async for x in g] for j in items]). --- crates/codegen/src/symboltable.rs | 25 ++++++++++++++++++++++++- 1 file changed, 24 insertions(+), 1 deletion(-) diff --git a/crates/codegen/src/symboltable.rs b/crates/codegen/src/symboltable.rs index b86e0f87b21..b7ef3e7eaba 100644 --- a/crates/codegen/src/symboltable.rs +++ b/crates/codegen/src/symboltable.rs @@ -299,6 +299,26 @@ fn drop_class_free(symbol_table: &mut SymbolTable, newfree: &mut IndexSet bool { + use ast::visitor::Visitor; + struct AwaitFinder(bool); + impl ast::visitor::Visitor<'_> for AwaitFinder { + fn visit_expr(&mut self, expr: &ast::Expr) { + if !self.0 { + if matches!(expr, ast::Expr::Await(_)) { + self.0 = true; + } else { + ast::visitor::walk_expr(self, expr); + } + } + } + } + let mut finder = AwaitFinder(false); + finder.visit_expr(expr); + finder.0 +} + /// PEP 709: Merge symbols from an inlined comprehension into the parent scope. /// Matches symtable.c inline_comprehension(). fn inline_comprehension( @@ -2102,9 +2122,12 @@ impl SymbolTableBuilder { // but only inside function-like scopes (fastlocals). // Module/class scope uses STORE_NAME which is incompatible // with LOAD_FAST_AND_CLEAR / STORE_FAST save/restore. + // Async comprehensions cannot be inlined because they need + // their own coroutine scope. // Note: tables.last() is the comprehension scope we just pushed, // so we check the second-to-last for the parent scope. - if !is_generator { + let element_has_await = expr_contains_await(elt1) || elt2.is_some_and(expr_contains_await); + if !is_generator && !has_async_gen && !element_has_await { let parent = self.tables.iter().rev().nth(1); let parent_is_func = parent.is_some_and(|t| { matches!(