Skip to content

Commit e58f796

Browse files
youknowoneyouknowone
committed
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
1 parent d62481f commit e58f796

File tree

3 files changed

+116
-53
lines changed

3 files changed

+116
-53
lines changed

‎crates/codegen/src/compile.rs‎

Lines changed: 33 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -968,12 +968,21 @@ impl Compiler {
968968
Ok(())
969969
}
970970

971-
/// Check if this is an inlined comprehension context (PEP 709)
972-
/// Currently disabled - always returns false to avoid stack issues
973-
fn is_inlined_comprehension_context(&self, _comprehension_type: ComprehensionType) -> bool {
974-
// TODO: Implement PEP 709 inlined comprehensions properly
975-
// For now, disabled to avoid stack underflow issues
976-
false
971+
/// Check if this is an inlined comprehension context (PEP 709).
972+
/// Only inline in function-like scopes (fastlocals) — module/class
973+
/// level uses STORE_NAME which is incompatible with LOAD_FAST_AND_CLEAR.
974+
/// Generator expressions are never inlined.
975+
fn is_inlined_comprehension_context(&self, comprehension_type: ComprehensionType) -> bool {
976+
if comprehension_type == ComprehensionType::Generator {
977+
return false;
978+
}
979+
if !self.ctx.in_func() {
980+
return false;
981+
}
982+
self.symbol_table_stack
983+
.last()
984+
.and_then(|t| t.sub_tables.get(t.next_sub_table))
985+
.is_some_and(|st| st.comp_inlined)
977986
}
978987

979988
/// Enter a new scope
@@ -7640,12 +7649,15 @@ impl Compiler {
76407649

76417650
if is_inlined {
76427651
// PEP 709: Inlined comprehension - compile inline without new scope
7643-
return self.compile_inlined_comprehension(
7652+
self.current_code_info().in_inlined_comp = true;
7653+
let result = self.compile_inlined_comprehension(
76447654
init_collection,
76457655
generators,
76467656
compile_element,
76477657
has_an_async_gen,
76487658
);
7659+
self.current_code_info().in_inlined_comp = false;
7660+
return result;
76497661
}
76507662

76517663
// Non-inlined path: create a new code object (generator expressions, etc.)
@@ -7672,9 +7684,6 @@ impl Compiler {
76727684
// Create magnificent function <listcomp>:
76737685
self.push_output(flags, 1, 1, 0, name.to_owned())?;
76747686

7675-
// Mark that we're in an inlined comprehension
7676-
self.current_code_info().in_inlined_comp = true;
7677-
76787687
// Set qualname for comprehension
76797688
self.set_qualname();
76807689

@@ -7837,15 +7846,23 @@ impl Compiler {
78377846
compile_element: &dyn Fn(&mut Self) -> CompileResult<()>,
78387847
_has_an_async_gen: bool,
78397848
) -> CompileResult<()> {
7840-
// PEP 709: Consume the comprehension's sub_table (but we won't use it as a separate scope)
7841-
// We need to consume it to keep sub_tables in sync with AST traversal order.
7849+
// PEP 709: Consume the comprehension's sub_table (but we won't use it as a separate scope).
78427850
// The symbols are already merged into parent scope by analyze_symbol_table.
7843-
let _comp_table = self
7851+
// Splice the comprehension's sub_tables into the parent so nested scopes
7852+
// (e.g. inner comprehensions, lambdas) can still find their sub_tables.
7853+
let current_table = self
78447854
.symbol_table_stack
78457855
.last_mut()
7846-
.expect("no current symbol table")
7847-
.sub_tables
7848-
.remove(0);
7856+
.expect("no current symbol table");
7857+
let comp_table = current_table.sub_tables[current_table.next_sub_table].clone();
7858+
current_table.next_sub_table += 1;
7859+
// Insert the comprehension's sub_tables right after the consumed entry
7860+
if !comp_table.sub_tables.is_empty() {
7861+
let insert_pos = current_table.next_sub_table;
7862+
for (i, st) in comp_table.sub_tables.iter().enumerate() {
7863+
current_table.sub_tables.insert(insert_pos + i, st.clone());
7864+
}
7865+
}
78497866

78507867
// Collect local variables that need to be saved/restored
78517868
// These are variables bound in the comprehension (iteration vars from targets)

‎crates/codegen/src/symboltable.rs‎

Lines changed: 75 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -459,34 +459,62 @@ impl SymbolTableAnalyzer {
459459

460460
symbol_table.symbols = info.0;
461461

462-
// PEP 709: Merge symbols from inlined comprehensions into parent scope
463-
// Only merge symbols that are actually bound in the comprehension,
464-
// not references to outer scope variables (Free symbols).
462+
// PEP 709: Merge symbols from inlined comprehensions into parent scope.
463+
// If a comprehension-bound name conflicts with an existing parent symbol,
464+
// disable inlining for that comprehension (the save/restore mechanism
465+
// via LOAD_FAST_AND_CLEAR / STORE_FAST can't handle scope mismatches).
465466
const BOUND_FLAGS: SymbolFlags = SymbolFlags::ASSIGNED
466467
.union(SymbolFlags::PARAMETER)
467468
.union(SymbolFlags::ITER)
468469
.union(SymbolFlags::ASSIGNED_IN_COMPREHENSION);
469470

470-
for sub_table in sub_tables.iter() {
471-
if sub_table.comp_inlined {
472-
for (name, sub_symbol) in &sub_table.symbols {
473-
// Skip the .0 parameter - it's internal to the comprehension
474-
if name == ".0" {
475-
continue;
476-
}
477-
// Only merge symbols that are bound in the comprehension
478-
// Skip Free references to outer scope variables
479-
if !sub_symbol.flags.intersects(BOUND_FLAGS) {
480-
continue;
481-
}
482-
// If the symbol doesn't exist in parent, add it
483-
if !symbol_table.symbols.contains_key(name) {
484-
let mut symbol = sub_symbol.clone();
485-
// Mark as local in parent scope
486-
symbol.scope = SymbolScope::Local;
487-
symbol_table.symbols.insert(name.clone(), symbol);
488-
}
471+
// Track symbols added by inlined comprehensions, so later comps
472+
// can detect when a reference would resolve to a comp-local.
473+
let mut comp_added_symbols: IndexSet<String> = IndexSet::default();
474+
475+
for sub_table in sub_tables.iter_mut() {
476+
if !sub_table.comp_inlined {
477+
continue;
478+
}
479+
// Don't inline if the comprehension contains nested scopes
480+
// (lambdas, inner comprehensions, nested functions) — these need
481+
// Cell/Free variable handling that inlining doesn't support yet.
482+
if !sub_table.sub_tables.is_empty() {
483+
sub_table.comp_inlined = false;
484+
continue;
485+
}
486+
// Don't inline if a bound comprehension name conflicts with parent
487+
let has_bound_conflict = sub_table.symbols.iter().any(|(name, sym)| {
488+
name != ".0"
489+
&& sym.flags.intersects(BOUND_FLAGS)
490+
&& symbol_table.symbols.contains_key(name)
491+
});
492+
// Don't inline if a non-bound reference would resolve to a
493+
// symbol added by a previous inlined comprehension
494+
let has_ref_conflict = sub_table.symbols.iter().any(|(name, sym)| {
495+
name != ".0"
496+
&& !sym.flags.intersects(BOUND_FLAGS)
497+
&& comp_added_symbols.contains(name)
498+
});
499+
if has_bound_conflict || has_ref_conflict {
500+
sub_table.comp_inlined = false;
501+
continue;
502+
}
503+
for (name, sub_symbol) in &sub_table.symbols {
504+
if name == ".0" {
505+
continue;
506+
}
507+
if symbol_table.symbols.contains_key(name) {
508+
continue;
509+
}
510+
let mut symbol = sub_symbol.clone();
511+
if sub_symbol.flags.intersects(BOUND_FLAGS) {
512+
symbol.scope = SymbolScope::Local;
513+
comp_added_symbols.insert(name.clone());
489514
}
515+
// Non-bound symbols keep their analyzed scope from the
516+
// comprehension sub_table (e.g., GlobalImplicit, Free).
517+
symbol_table.symbols.insert(name.clone(), symbol);
490518
}
491519
}
492520

@@ -2037,13 +2065,31 @@ impl SymbolTableBuilder {
20372065
self.line_index_start(range),
20382066
);
20392067

2040-
// PEP 709: inlined comprehensions are not yet implemented in the
2041-
// compiler (is_inlined_comprehension_context always returns false),
2042-
// so do NOT mark comp_inlined here. Setting it would cause the
2043-
// symbol-table analyzer to merge comprehension-local symbols into
2044-
// the parent scope, while the compiler still emits a separate code
2045-
// object — leading to the merged symbols being missing from the
2046-
// comprehension's own symbol table lookup.
2068+
// PEP 709: Mark non-generator comprehensions for inlining,
2069+
// but only inside function-like scopes (fastlocals).
2070+
// Module/class scope uses STORE_NAME which is incompatible
2071+
// with LOAD_FAST_AND_CLEAR / STORE_FAST save/restore.
2072+
// Note: tables.last() is the comprehension scope we just pushed,
2073+
// so we check the second-to-last for the parent scope.
2074+
if !is_generator {
2075+
let parent_is_func = self
2076+
.tables
2077+
.iter()
2078+
.rev()
2079+
.nth(1)
2080+
.is_some_and(|t| {
2081+
matches!(
2082+
t.typ,
2083+
CompilerScope::Function
2084+
| CompilerScope::AsyncFunction
2085+
| CompilerScope::Lambda
2086+
| CompilerScope::Comprehension
2087+
)
2088+
});
2089+
if parent_is_func {
2090+
self.tables.last_mut().unwrap().comp_inlined = true;
2091+
}
2092+
}
20472093

20482094
// Register the passed argument to the generator function as the name ".0"
20492095
self.register_name(".0", SymbolUsage::Parameter, range)?;

‎crates/vm/src/frame.rs‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2699,13 +2699,12 @@ impl ExecutingFrame<'_> {
26992699
Ok(None)
27002700
}
27012701
Instruction::LoadFastAndClear { var_num: idx } => {
2702-
// Load value and clear the slot (for inlined comprehensions)
2703-
// If slot is empty, push None (not an error - variable may not exist yet)
2702+
// Save current slot value and clear it (for inlined comprehensions).
2703+
// Pushes NULL (None at Option level) if slot was empty, so that
2704+
// StoreFast can restore the empty state after the comprehension.
27042705
let idx = idx.get(arg) as usize;
2705-
let x = self.localsplus.fastlocals_mut()[idx]
2706-
.take()
2707-
.unwrap_or_else(|| vm.ctx.none());
2708-
self.push_value(x);
2706+
let x = self.localsplus.fastlocals_mut()[idx].take();
2707+
self.push_value_opt(x);
27092708
Ok(None)
27102709
}
27112710
Instruction::LoadFastCheck { var_num: idx } => {
@@ -3300,9 +3299,10 @@ impl ExecutingFrame<'_> {
33003299
Ok(None)
33013300
}
33023301
Instruction::StoreFast { var_num: idx } => {
3303-
let value = self.pop_value();
3302+
// pop_value_opt: allows NULL from LoadFastAndClear restore path
3303+
let value = self.pop_value_opt();
33043304
let fastlocals = self.localsplus.fastlocals_mut();
3305-
fastlocals[idx.get(arg) as usize] = Some(value);
3305+
fastlocals[idx.get(arg) as usize] = value;
33063306
Ok(None)
33073307
}
33083308
Instruction::StoreFastLoadFast { var_nums } => {

0 commit comments

Comments
 (0)