mirrors/RustPython
0
0 Fork 0
mirror of https://github.com/RustPython/RustPython.git synced 2026-10-28 10:44:36 +08:00
Code Activity

Enable PEP 709 inlined comprehensions (#7412)

Browse Source 
* Enable PEP 709 inlined comprehensions for function-like scopes

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

* 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

* 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.

* 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]).
This commit is contained in:
Jeong, YunWon authored and GitHub committed 2026-03-21 22:48:35 +09:00
1 parent 009efbb48a
commit fb9cafa8c9
14 files changed
+814 -344
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

.cspell.dict/cpython.txt
+1
View File
@@ -109,6 +109,7 @@ lineiterator
linetable
loadfast
localsplus
localspluskinds
Lshift
lsprof
MAXBLOCKS
.cspell.json
+1
View File
@@ -70,6 +70,7 @@
"lossily",
"mcache",
"oparg",
"opargs",
"pyc",
"significand",
"summands",
crates/codegen/src/compile.rs
+204 -87
View File
@@ -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<u32> {
/// 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<oparg::VarNum> {
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<u32> {
/// 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<oparg::VarNum> {
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<String> = 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 <listcomp>:
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<String>) {
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<AnyInstruction>,
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<String> = 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<String> = 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<String, Symbol> = 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(())
}
crates/codegen/src/ir.rs
+115 -48
View File
@@ -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<LineTableLocation> = 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<Box<[i32]>> {
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::<Box<[_]>>();
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<String>,
cellvars: &IndexSet<String>,
freevars: &IndexSet<String>,
) -> Vec<u32> {
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]);
}
}
}
}
crates/codegen/src/symboltable.rs
+149 -50
View File
@@ -299,6 +299,88 @@ fn drop_class_free(symbol_table: &mut SymbolTable, newfree: &mut IndexSet<String
}
}
/// Check if an expression contains an `await` node (shallow, not into nested scopes).
fn expr_contains_await(expr: &ast::Expr) -> 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<String>,
inlined_cells: &mut IndexSet<String>,
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<String, Symbol>;
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<String>, bool)> = Vec::new();
let mut annotation_free: Option<IndexSet<String>> = 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<String> = 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<SymbolScope> {
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)?;
crates/compiler-core/src/bytecode.rs
+25 -9
View File
@@ -334,6 +334,12 @@ impl<T> IndexMut<oparg::VarNum> 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<C: Constant = ConstantData> {
pub obj_name: C::Name,
/// Qualified name of the object (like CPython's co_qualname)
pub qualname: C::Name,
pub cell2arg: Option<Box<[i32]>>,
pub constants: Constants<C>,
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<C: Constant> CodeObject<C> {
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<C: Constant> CodeObject<C> {
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<C: Constant> InstrDisplayContext for CodeObject<C> {
@@ -1159,11 +1168,18 @@ impl<C: Constant> InstrDisplayContext for CodeObject<C> {
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()
}
}
}
crates/compiler-core/src/bytecode/instruction.rs
+6 -6
View File
@@ -130,7 +130,7 @@ pub enum Instruction {
namei: Arg<NameIdx>,
} = 61,
DeleteDeref {
i: Arg<NameIdx>,
i: Arg<oparg::VarNum>,
} = 62,
DeleteFast {
var_num: Arg<oparg::VarNum>,
@@ -189,7 +189,7 @@ pub enum Instruction {
consti: Arg<oparg::ConstIdx>,
} = 82,
LoadDeref {
i: Arg<NameIdx>,
i: Arg<oparg::VarNum>,
} = 83,
LoadFast {
var_num: Arg<oparg::VarNum>,
@@ -210,7 +210,7 @@ pub enum Instruction {
var_nums: Arg<oparg::VarNums>,
} = 89,
LoadFromDictOrDeref {
i: Arg<NameIdx>,
i: Arg<oparg::VarNum>,
} = 90,
LoadFromDictOrGlobals {
i: Arg<NameIdx>,
@@ -231,7 +231,7 @@ pub enum Instruction {
namei: Arg<LoadSuperAttr>,
} = 96,
MakeCell {
i: Arg<NameIdx>,
i: Arg<oparg::VarNum>,
} = 97,
MapAdd {
i: Arg<u32>,
@@ -273,7 +273,7 @@ pub enum Instruction {
namei: Arg<NameIdx>,
} = 110,
StoreDeref {
i: Arg<NameIdx>,
i: Arg<oparg::VarNum>,
} = 111,
StoreFast {
var_num: Arg<oparg::VarNum>,
@@ -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,
crates/compiler-core/src/marshal.rs
+46 -15
View File
@@ -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<R: Read, Bag: ConstantBag>(
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::<Result<Box<[i32]>>>()
})
.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<R: Read, Bag: ConstantBag>(
.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<R: Read, Bag: ConstantBag>(
max_stackdepth,
obj_name,
qualname,
cell2arg,
constants,
names,
varnames,
cellvars,
freevars,
localspluskinds,
linetable,
exceptiontable,
})
@@ -687,11 +721,8 @@ pub fn serialize_code<W: Write, C: Constant>(buf: &mut W, code: &CodeObject<C>)
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 {
crates/vm/src/builtins/code.rs
+34 -2
View File
@@ -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,
};
crates/vm/src/builtins/function.rs
+2 -21
View File
@@ -64,7 +64,7 @@ pub struct PyFunction {
code: PyAtomicRef<PyCode>,
globals: PyDictRef,
builtins: PyObjectRef,
closure: Option<PyRef<PyTuple<PyCellRef>>>,
pub(crate) closure: Option<PyRef<PyTuple<PyCellRef>>>,
defaults_and_kwdefaults: PyMutex<(Option<PyTupleRef>, Option<PyDictRef>)>,
name: PyMutex<PyStrRef>,
qualname: PyMutex<PyStrRef>,
@@ -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<PyFunction> {
}
}
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<PyCode>) -> Option<u
{
return None;
}
let nlocalsplus = code
.varnames
.len()
.checked_add(code.cellvars.len())?
.checked_add(code.freevars.len())?;
let nlocalsplus = code.localspluskinds.len();
let capacity = nlocalsplus.checked_add(code.max_stackdepth as usize)?;
capacity.checked_mul(core::mem::size_of::<usize>())
}
crates/vm/src/builtins/super.rs
+17 -10
View File
@@ -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::<PyCell>().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!(
crates/vm/src/frame.rs
+211 -94
View File
@@ -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<PyObjectRef> {
let nlocals = self.code.varnames.len();
/// Get cell contents by localsplus index.
pub(crate) fn get_cell_contents(&self, localsplus_idx: usize) -> Option<PyObjectRef> {
// 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::<PyCell>())
.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<PyObjectRef>) {
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::<PyCell>())
.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<ArgMapping> {
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::<PyCell>()
.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::<PyCell>() {
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::<PyCell>()
.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) {
vm.new_exception_msg(
vm.ctx.exceptions.unbound_local_error.to_owned(),
format!("local variable '{name}' referenced before assignment").into(),
)
} else {
let name = self.code.freevars[i - self.code.cellvars.len()];
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 {
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::<PyFunction>())
.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::<PyDict>() {
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 } => {
crates/vm/src/version.rs
+2 -1
View File
@@ -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
crates/vm/src/vm/context.rs
+1 -1
View File
@@ -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(),
};