Skip to content

Commit 3b7a16b

Browse files
youknowoneyouknowone
committed
Match CPython PEP 709 inlining strategy
Replace conservative bail-out approach with CPython-matching strategy: Symbol table (symboltable.rs): - Remove conflict-based bail-out checks (bound conflict, ref conflict, nested scopes) - Implement inline_comprehension() matching CPython's symtable.c: track inlined_cells, handle __class__ in ClassBlock, merge symbols with proper is_free_in_any_child check - Promote LOCAL to CELL for inlined_cells, set COMP_CELL flag - Splice inlined comp children into parent's sub_tables - Add can_see_class_scope check to inlining decision Compiler (compile.rs): - Implement TweakInlinedComprehensionScopes: temporarily swap symbol scopes during comprehension body compilation - Implement RevertInlinedComprehensionScopes: restore original scopes - Push ALL locally-bound names (not just iteration targets) for save/restore, excluding walrus operator targets - Emit MakeCell/RestoreCell for CELL variables in comprehensions - Re-enable async comprehension inlining with proper SetupFinally/EndAsyncFor handling - Include COMP_CELL symbols in cellvars VM (frame.rs, instruction.rs): - Add RestoreCell instruction for proper cell save/restore - MakeCell now saves old cell to stack and creates new empty cell
1 parent 711b83b commit 3b7a16b

File tree

4 files changed

+311
-153
lines changed

4 files changed

+311
-153
lines changed

‎crates/codegen/src/compile.rs‎

Lines changed: 168 additions & 75 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ use crate::{
1313
IndexMap, IndexSet, ToPythonName,
1414
error::{CodegenError, CodegenErrorType, InternalError, PatternUnreachableReason},
1515
ir::{self, BlockIdx},
16-
symboltable::{self, CompilerScope, SymbolFlags, SymbolScope, SymbolTable},
16+
symboltable::{self, CompilerScope, Symbol, SymbolFlags, SymbolScope, SymbolTable},
1717
unparse::UnparseExpr,
1818
};
1919
use alloc::borrow::Cow;
@@ -1012,12 +1012,14 @@ impl Compiler {
10121012
// Use varnames from symbol table (already collected in definition order)
10131013
let varname_cache: IndexSet<String> = ste.varnames.iter().cloned().collect();
10141014

1015-
// Build cellvars using dictbytype (CELL scope, sorted)
1015+
// Build cellvars using dictbytype (CELL scope or COMP_CELL flag, sorted)
10161016
let mut cellvar_cache = IndexSet::default();
10171017
let mut cell_names: Vec<_> = ste
10181018
.symbols
10191019
.iter()
1020-
.filter(|(_, s)| s.scope == SymbolScope::Cell)
1020+
.filter(|(_, s)| {
1021+
s.scope == SymbolScope::Cell || s.flags.contains(SymbolFlags::COMP_CELL)
1022+
})
10211023
.map(|(name, _)| name.clone())
10221024
.collect();
10231025
cell_names.sort();
@@ -7656,13 +7658,16 @@ impl Compiler {
76567658
// We must have at least one generator:
76577659
assert!(!generators.is_empty());
76587660

7659-
if is_inlined && !has_an_async_gen && !element_contains_await {
7661+
if is_inlined {
76607662
// PEP 709: Inlined comprehension - compile inline without new scope
7661-
// Async comprehensions are not inlined due to complex exception handling
76627663
let was_in_inlined_comp = self.current_code_info().in_inlined_comp;
76637664
self.current_code_info().in_inlined_comp = true;
7664-
let result =
7665-
self.compile_inlined_comprehension(init_collection, generators, compile_element);
7665+
let result = self.compile_inlined_comprehension(
7666+
init_collection,
7667+
generators,
7668+
compile_element,
7669+
has_an_async_gen,
7670+
);
76667671
self.current_code_info().in_inlined_comp = was_in_inlined_comp;
76677672
return result;
76687673
}
@@ -7816,90 +7821,136 @@ impl Compiler {
78167821
Ok(())
78177822
}
78187823

7819-
/// Collect variable names from an assignment target expression
7820-
fn collect_target_names(&self, target: &ast::Expr, names: &mut Vec<String>) {
7821-
match target {
7822-
ast::Expr::Name(name) => {
7823-
let name_str = name.id.to_string();
7824-
if !names.contains(&name_str) {
7825-
names.push(name_str);
7826-
}
7827-
}
7828-
ast::Expr::Tuple(tuple) => {
7829-
for elt in &tuple.elts {
7830-
self.collect_target_names(elt, names);
7831-
}
7832-
}
7833-
ast::Expr::List(list) => {
7834-
for elt in &list.elts {
7835-
self.collect_target_names(elt, names);
7836-
}
7837-
}
7838-
ast::Expr::Starred(starred) => {
7839-
self.collect_target_names(&starred.value, names);
7840-
}
7841-
_ => {
7842-
// Other targets (attribute, subscript) don't bind local names
7843-
}
7844-
}
7845-
}
7846-
78477824
/// Compile an inlined comprehension (PEP 709)
78487825
/// This generates bytecode inline without creating a new code object
78497826
fn compile_inlined_comprehension(
78507827
&mut self,
78517828
init_collection: Option<AnyInstruction>,
78527829
generators: &[ast::Comprehension],
78537830
compile_element: &dyn Fn(&mut Self) -> CompileResult<()>,
7831+
has_async: bool,
78547832
) -> CompileResult<()> {
7855-
// PEP 709: Consume the comprehension's sub_table (but we won't use it as a separate scope).
7833+
// PEP 709: Consume the comprehension's sub_table.
78567834
// The symbols are already merged into parent scope by analyze_symbol_table.
7857-
// Splice the comprehension's sub_tables into the parent so nested scopes
7858-
// (e.g. inner comprehensions, lambdas) can still find their sub_tables.
7835+
// Splice the comprehension's children into the parent so nested scopes
7836+
// (e.g. inner comprehensions, lambdas) can be found by the compiler.
78597837
let current_table = self
78607838
.symbol_table_stack
78617839
.last_mut()
78627840
.expect("no current symbol table");
78637841
let comp_table = current_table.sub_tables[current_table.next_sub_table].clone();
78647842
current_table.next_sub_table += 1;
7865-
// Insert the comprehension's sub_tables right after the consumed entry
78667843
if !comp_table.sub_tables.is_empty() {
78677844
let insert_pos = current_table.next_sub_table;
78687845
for (i, st) in comp_table.sub_tables.iter().enumerate() {
78697846
current_table.sub_tables.insert(insert_pos + i, st.clone());
78707847
}
78717848
}
78727849

7873-
// Collect local variables that need to be saved/restored
7874-
// These are variables bound in the comprehension (iteration vars from targets)
7875-
let mut pushed_locals: Vec<String> = Vec::new();
7876-
for generator in generators {
7877-
self.collect_target_names(&generator.target, &mut pushed_locals);
7850+
// Step 1: Compile the outermost iterator BEFORE tweaking scopes
7851+
self.compile_expression(&generators[0].iter)?;
7852+
if has_async && generators[0].is_async {
7853+
emit!(self, Instruction::GetAIter);
7854+
} else {
7855+
emit!(self, Instruction::GetIter);
78787856
}
78797857

7880-
// Step 1: Compile the outermost iterator
7881-
// Async comprehensions are never inlined, so only sync iteration here
7882-
self.compile_expression(&generators[0].iter)?;
7883-
emit!(self, Instruction::GetIter);
7858+
// Collect local variables that need to be saved/restored.
7859+
// All DEF_LOCAL && !DEF_NONLOCAL names from the comp table, plus class block names.
7860+
let in_class_block = {
7861+
let ct = self.current_symbol_table();
7862+
ct.typ == CompilerScope::Class && !self.current_code_info().in_inlined_comp
7863+
};
7864+
let mut pushed_locals: Vec<String> = Vec::new();
7865+
for (name, sym) in &comp_table.symbols {
7866+
if sym.flags.contains(SymbolFlags::PARAMETER) {
7867+
continue; // skip .0
7868+
}
7869+
// Walrus operator targets (ASSIGNED_IN_COMPREHENSION without ITER)
7870+
// are not local to the comprehension; they leak to the outer scope.
7871+
let is_walrus = sym.flags.contains(SymbolFlags::ASSIGNED_IN_COMPREHENSION)
7872+
&& !sym.flags.contains(SymbolFlags::ITER);
7873+
let is_local = sym
7874+
.flags
7875+
.intersects(SymbolFlags::ASSIGNED | SymbolFlags::ITER)
7876+
&& !sym.flags.contains(SymbolFlags::NONLOCAL)
7877+
&& !is_walrus;
7878+
if is_local || in_class_block {
7879+
pushed_locals.push(name.clone());
7880+
}
7881+
}
7882+
7883+
// TweakInlinedComprehensionScopes: temporarily override parent symbols
7884+
// with comp scopes where they differ.
7885+
let mut temp_symbols: IndexMap<String, Symbol> = IndexMap::default();
7886+
for (name, comp_sym) in &comp_table.symbols {
7887+
if comp_sym.flags.contains(SymbolFlags::PARAMETER) {
7888+
continue; // skip .0
7889+
}
7890+
let comp_scope = comp_sym.scope;
7891+
7892+
let current_table = self.symbol_table_stack.last().expect("no symbol table");
7893+
if let Some(outer_sym) = current_table.symbols.get(name) {
7894+
let outer_scope = outer_sym.scope;
7895+
if (comp_scope != outer_scope
7896+
&& comp_scope != SymbolScope::Free
7897+
&& !(comp_scope == SymbolScope::Cell && outer_scope == SymbolScope::Free))
7898+
|| in_class_block
7899+
{
7900+
temp_symbols.insert(name.clone(), outer_sym.clone());
7901+
let current_table =
7902+
self.symbol_table_stack.last_mut().expect("no symbol table");
7903+
current_table.symbols.insert(name.clone(), comp_sym.clone());
7904+
}
7905+
}
7906+
}
78847907

7885-
// Step 2: Save local variables that will be shadowed by the comprehension
7908+
// Step 2: Save local variables that will be shadowed by the comprehension.
7909+
// For each variable, we push the fast local value via LoadFastAndClear.
7910+
// For CELL variables, we additionally push the cell content via MakeCell
7911+
// (which saves the old cell value and clears the cell for the comprehension).
7912+
// Track cell indices to restore them later.
7913+
let mut cell_indices: Vec<Option<u32>> = Vec::new();
7914+
let mut total_stack_items: usize = 0;
78867915
for name in &pushed_locals {
78877916
let var_num = self.varname(name)?;
78887917
emit!(self, Instruction::LoadFastAndClear { var_num });
7918+
total_stack_items += 1;
7919+
// If the comp symbol is CELL, emit MAKE_CELL to save cell value
7920+
if let Some(comp_sym) = comp_table.symbols.get(name) {
7921+
if comp_sym.scope == SymbolScope::Cell {
7922+
let i = if self
7923+
.current_symbol_table()
7924+
.symbols
7925+
.get(name)
7926+
.is_some_and(|s| s.scope == SymbolScope::Free)
7927+
{
7928+
self.get_free_var_index(name)?
7929+
} else {
7930+
self.get_cell_var_index(name)?
7931+
};
7932+
emit!(self, Instruction::MakeCell { i });
7933+
cell_indices.push(Some(i));
7934+
total_stack_items += 1;
7935+
} else {
7936+
cell_indices.push(None);
7937+
}
7938+
} else {
7939+
cell_indices.push(None);
7940+
}
78897941
}
78907942

7891-
// Step 3: SWAP iterator to TOS (above saved locals)
7892-
if !pushed_locals.is_empty() {
7943+
// Step 3: SWAP iterator to TOS (above saved locals + cell values)
7944+
if total_stack_items > 0 {
78937945
emit!(
78947946
self,
78957947
Instruction::Swap {
7896-
i: u32::try_from(pushed_locals.len() + 1).unwrap()
7948+
i: u32::try_from(total_stack_items + 1).unwrap()
78977949
}
78987950
);
78997951
}
79007952

79017953
// Step 4: Create the collection (list/set/dict)
7902-
// For generator expressions, init_collection is None
79037954
if let Some(init_collection) = init_collection {
79047955
self._emit(init_collection, OpArg::new(0), BlockIdx::NULL);
79057956
// SWAP to get iterator on top
@@ -7921,24 +7972,49 @@ impl Compiler {
79217972
}
79227973

79237974
// Step 5: Compile the comprehension loop(s)
7924-
let mut loop_labels = vec![];
7975+
let mut loop_labels: Vec<(BlockIdx, BlockIdx, BlockIdx, bool, BlockIdx)> = vec![];
79257976
for (i, generator) in generators.iter().enumerate() {
79267977
let loop_block = self.new_block();
79277978
let if_cleanup_block = self.new_block();
79287979
let after_block = self.new_block();
79297980

79307981
if i > 0 {
7931-
// For nested loops, compile the iterator expression
79327982
self.compile_expression(&generator.iter)?;
7933-
emit!(self, Instruction::GetIter);
7983+
if generator.is_async {
7984+
emit!(self, Instruction::GetAIter);
7985+
} else {
7986+
emit!(self, Instruction::GetIter);
7987+
}
79347988
}
79357989

79367990
self.switch_to_block(loop_block);
79377991

7938-
emit!(self, Instruction::ForIter { delta: after_block });
7939-
self.compile_store(&generator.target)?;
7992+
let mut end_async_for_target = BlockIdx::NULL;
7993+
if generator.is_async {
7994+
emit!(self, PseudoInstruction::SetupFinally { delta: after_block });
7995+
emit!(self, Instruction::GetANext);
7996+
self.push_fblock(
7997+
FBlockType::AsyncComprehensionGenerator,
7998+
loop_block,
7999+
after_block,
8000+
)?;
8001+
self.emit_load_const(ConstantData::None);
8002+
end_async_for_target = self.compile_yield_from_sequence(true)?;
8003+
emit!(self, PseudoInstruction::PopBlock);
8004+
self.pop_fblock(FBlockType::AsyncComprehensionGenerator);
8005+
self.compile_store(&generator.target)?;
8006+
} else {
8007+
emit!(self, Instruction::ForIter { delta: after_block });
8008+
self.compile_store(&generator.target)?;
8009+
}
79408010

7941-
loop_labels.push((loop_block, if_cleanup_block, after_block));
8011+
loop_labels.push((
8012+
loop_block,
8013+
if_cleanup_block,
8014+
after_block,
8015+
generator.is_async,
8016+
end_async_for_target,
8017+
));
79428018

79438019
// Evaluate the if conditions
79448020
for if_condition in &generator.ifs {
@@ -7949,20 +8025,26 @@ impl Compiler {
79498025
// Step 6: Compile the element expression and append to collection
79508026
compile_element(self)?;
79518027

7952-
// Step 7: Close all loops (sync only - async comprehensions are never inlined)
7953-
for (loop_block, if_cleanup_block, after_block) in loop_labels.iter().rev().copied() {
8028+
// Step 7: Close all loops
8029+
for &(loop_block, if_cleanup_block, after_block, is_async, end_async_for_target) in
8030+
loop_labels.iter().rev()
8031+
{
79548032
emit!(self, PseudoInstruction::Jump { delta: loop_block });
79558033

79568034
self.switch_to_block(if_cleanup_block);
79578035
emit!(self, PseudoInstruction::Jump { delta: loop_block });
79588036

79598037
self.switch_to_block(after_block);
7960-
emit!(self, Instruction::EndFor);
7961-
emit!(self, Instruction::PopIter);
8038+
if is_async {
8039+
self.emit_end_async_for(end_async_for_target);
8040+
} else {
8041+
emit!(self, Instruction::EndFor);
8042+
emit!(self, Instruction::PopIter);
8043+
}
79628044
}
79638045

7964-
// Step 8: Clean up - restore saved locals
7965-
if !pushed_locals.is_empty() {
8046+
// Step 8: Clean up - restore saved locals (and cell values)
8047+
if total_stack_items > 0 {
79668048
emit!(self, PseudoInstruction::PopBlock);
79678049
self.pop_fblock(FBlockType::TryExcept);
79688050

@@ -7971,19 +8053,21 @@ impl Compiler {
79718053

79728054
// Exception cleanup path
79738055
self.switch_to_block(cleanup_block);
7974-
// Stack: [saved_locals..., collection, exception]
7975-
// Swap to get collection out from under exception
8056+
// Stack: [saved_values..., collection, exception]
79768057
emit!(self, Instruction::Swap { i: 2 });
79778058
emit!(self, Instruction::PopTop); // Pop incomplete collection
79788059

7979-
// Restore locals
8060+
// Restore locals and cell values
79808061
emit!(
79818062
self,
79828063
Instruction::Swap {
7983-
i: u32::try_from(pushed_locals.len() + 1).unwrap()
8064+
i: u32::try_from(total_stack_items + 1).unwrap()
79848065
}
79858066
);
7986-
for name in pushed_locals.iter().rev() {
8067+
for (name, cell_idx) in pushed_locals.iter().rev().zip(cell_indices.iter().rev()) {
8068+
if let Some(i) = cell_idx {
8069+
emit!(self, Instruction::RestoreCell { i: *i });
8070+
}
79878071
let var_num = self.varname(name)?;
79888072
emit!(self, Instruction::StoreFast { var_num });
79898073
}
@@ -7994,22 +8078,31 @@ impl Compiler {
79948078
self.switch_to_block(end_block);
79958079
}
79968080

7997-
// SWAP result to TOS (above saved locals)
7998-
if !pushed_locals.is_empty() {
8081+
// SWAP result to TOS (above saved values)
8082+
if total_stack_items > 0 {
79998083
emit!(
80008084
self,
80018085
Instruction::Swap {
8002-
i: u32::try_from(pushed_locals.len() + 1).unwrap()
8086+
i: u32::try_from(total_stack_items + 1).unwrap()
80038087
}
80048088
);
80058089
}
80068090

8007-
// Restore saved locals
8008-
for name in pushed_locals.iter().rev() {
8091+
// Restore saved locals and cell values
8092+
for (name, cell_idx) in pushed_locals.iter().rev().zip(cell_indices.iter().rev()) {
8093+
if let Some(i) = cell_idx {
8094+
emit!(self, Instruction::RestoreCell { i: *i });
8095+
}
80098096
let var_num = self.varname(name)?;
80108097
emit!(self, Instruction::StoreFast { var_num });
80118098
}
80128099

8100+
// RevertInlinedComprehensionScopes: restore original symbols
8101+
let current_table = self.symbol_table_stack.last_mut().expect("no symbol table");
8102+
for (name, original_sym) in temp_symbols {
8103+
current_table.symbols.insert(name, original_sym);
8104+
}
8105+
80138106
Ok(())
80148107
}
80158108

0 commit comments

Comments
 (0)