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/crates/codegen/src/compile.rs b/crates/codegen/src/compile.rs index b0cbbf501f8..fc81320fa91 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; @@ -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. @@ -968,12 +970,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 +1014,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(); @@ -1138,6 +1151,24 @@ 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: u32::try_from(nfrees).expect("too many freevars"), + } + ); + } + let ncells = self.code_stack.last().unwrap().metadata.cellvars.len(); + for i in 0..ncells { + let i_varnum: oparg::VarNum = u32::try_from(i).expect("too many cellvars").into(); + emit!(self, Instruction::MakeCell { i: i_varnum }); + } + } + // 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); @@ -1996,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 @@ -2017,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 }), @@ -7323,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)?; @@ -7775,14 +7839,18 @@ 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 - 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 +7877,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 +8002,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 +8009,122 @@ 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 merged CELL variables, LoadFastAndClear saves the cell object from + // the merged slot, and MAKE_CELL creates a new empty cell in-place. + // MAKE_CELL has no stack effect (operates only on fastlocals). + let mut total_stack_items: usize = 0; for name in &pushed_locals { let var_num = self.varname(name)?; emit!(self, Instruction::LoadFastAndClear { var_num }); + total_stack_items += 1; + // If the comp symbol is CELL, emit MAKE_CELL to create fresh cell + if let Some(comp_sym) = comp_table.symbols.get(name) + && comp_sym.scope == SymbolScope::Cell + { + let i = if self + .current_symbol_table() + .symbols + .get(name) + .is_some_and(|s| s.scope == SymbolScope::Free) + { + self.get_free_var_index(name)? + } else { + self.get_cell_var_index(name)? + }; + emit!(self, Instruction::MakeCell { i }); + } } - // Step 3: SWAP iterator to TOS (above saved locals) - 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 +8146,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 +8162,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 +8200,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 +8211,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,16 +8227,15 @@ 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() { @@ -8138,22 +8249,28 @@ 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 + // 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 }); } + // 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/ir.rs b/crates/codegen/src/ir.rs index 8a34fced545..6be851e15b5 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_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}, }; @@ -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, @@ -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, @@ -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,41 @@ 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; + } + // 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, posonlyarg_count, @@ -500,43 +539,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 +1115,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 +1783,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 +1801,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 +1824,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/codegen/src/symboltable.rs b/crates/codegen/src/symboltable.rs index fdbdac2b2a7..b7ef3e7eaba 100644 --- a/crates/codegen/src/symboltable.rs +++ b/crates/codegen/src/symboltable.rs @@ -299,6 +299,88 @@ 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( + parent_symbols: &mut SymbolMap, + comp: &SymbolTable, + comp_free: &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 +474,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 +495,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 +515,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 +530,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 +770,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 +2118,31 @@ 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. + // 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. + 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!( + 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.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 cd85e2458d4..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, @@ -1128,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, diff --git a/crates/compiler-core/src/marshal.rs b/crates/compiler-core/src/marshal.rs index ba3cf7a35c3..c47d41f1233 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,43 @@ 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 +313,12 @@ pub fn deserialize_code( max_stackdepth, obj_name, qualname, - cell2arg, constants, names, varnames, cellvars, freevars, + localspluskinds, linetable, exceptiontable, }) @@ -687,11 +721,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..0bf193914c2 100644 --- a/crates/vm/src/builtins/code.rs +++ b/crates/vm/src/builtins/code.rs @@ -633,6 +633,38 @@ 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 +682,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 +1269,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..88998dc8c8a 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,33 @@ 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 d8e3d35f9e3..fef682ff686 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,14 +683,7 @@ 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 { LocalsPlus::new_on_datastack(nlocalsplus, max_stackdepth, vm) @@ -699,13 +691,18 @@ impl Frame { 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()); + // 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 { @@ -791,30 +788,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,41 +872,102 @@ 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 { + // 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; } } - } - if !code.cellvars.is_empty() || !code.freevars.is_empty() { - for (i, &k) in code.cellvars.iter().enumerate() { - 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.contains(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. + // 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() + }; + + 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)) @@ -1325,13 +1370,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") @@ -1871,18 +1915,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.contains(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] } } @@ -2153,13 +2251,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 } => { @@ -2585,12 +2699,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)? @@ -2604,9 +2714,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) } @@ -2672,7 +2782,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() @@ -2699,13 +2809,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 +2963,15 @@ 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 } => { + // 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 } => { @@ -3294,13 +3410,14 @@ 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 } => { - 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 } => { 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(), };