Skip to content

Enable PEP 709 inlined comprehensions #7412

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
youknowone wants to merge 2 commits into
base: main
Choose a base branch
from
+117 −53
Open

Enable PEP 709 inlined comprehensions #7412

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
1219da9
Enable PEP 709 inlined comprehensions for function-like scopes
youknowone Mar 12, 2026
9121365
Save/restore in_inlined_comp flag for nested comprehensions
youknowone Mar 12, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 33 additions & 16 deletions crates/codegen/src/compile.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -968,12 +968,21 @@ impl Compiler {
Ok(())
}

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

/// Enter a new scope
Expand Down Expand Up @@ -7650,12 +7659,15 @@ impl Compiler {

if is_inlined {
// PEP 709: Inlined comprehension - compile inline without new scope
return self.compile_inlined_comprehension(
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 = false;
return result;
Comment on lines +7662 to +7671
Copy link
Contributor

@coderabbitai coderabbitai bot Mar 12, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🔴 Critical

Don't route async comprehensions through the inline path yet.

This branch now makes async list/set/dict comprehensions hit compile_inlined_comprehension(), but the inline async loop still lacks the SetupFinally { delta: after_block } / AsyncComprehensionGenerator scaffolding that the non-inlined path uses at Lines 7733-7745. In the inline version, after_block is never targeted for the async case, so END_ASYNC_FOR at Line 7989 will not run when the iterator terminates.

Possible short-term fix 
-        if is_inlined {
+        if is_inlined && !has_an_async_gen && !element_contains_await {
             // PEP 709: Inlined comprehension - compile inline without new scope
             let was_in_inlined_comp = self.current_code_info().in_inlined_comp;
             self.current_code_info().in_inlined_comp = true;
             let result = self.compile_inlined_comprehension(
                 init_collection,
🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@crates/codegen/src/compile.rs` around lines 7662 - 7671, The inlined path is
incorrectly used for async comprehensions; modify the branch around
current_code_info().in_inlined_comp and compile_inlined_comprehension so that if
has_an_async_gen is true you do NOT set in_inlined_comp or call
compile_inlined_comprehension (i.e., fall through to the non-inlined
comprehension code path that builds the AsyncComprehensionGenerator and emits
SetupFinally/END_ASYNC_FOR). Specifically, add a guard on has_an_async_gen
before calling compile_inlined_comprehension and only use the inlined path when
has_an_async_gen is false.

}

// Non-inlined path: create a new code object (generator expressions, etc.)
Expand All @@ -7682,9 +7694,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();

Expand Down Expand Up @@ -7847,15 +7856,23 @@ impl Compiler {
compile_element: &dyn Fn(&mut Self) -> CompileResult<()>,
_has_an_async_gen: 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 (but we won't use it as a separate scope).
// The symbols are already merged into parent scope by analyze_symbol_table.
let _comp_table = self
// Splice the comprehension's sub_tables into the parent so nested scopes
// (e.g. inner comprehensions, lambdas) can still find their sub_tables.
let current_table = self
.symbol_table_stack
.last_mut()
.expect("no current symbol table")
.sub_tables
.remove(0);
.expect("no current symbol table");
let comp_table = current_table.sub_tables[current_table.next_sub_table].clone();
current_table.next_sub_table += 1;
// Insert the comprehension's sub_tables right after the consumed entry
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());
}
}

// Collect local variables that need to be saved/restored
// These are variables bound in the comprehension (iteration vars from targets)
Expand Down
104 changes: 75 additions & 29 deletions crates/codegen/src/symboltable.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -459,34 +459,62 @@ impl SymbolTableAnalyzer {

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).
// PEP 709: Merge symbols from inlined comprehensions into parent scope.
// If a comprehension-bound name conflicts with an existing parent symbol,
// disable inlining for that comprehension (the save/restore mechanism
// via LOAD_FAST_AND_CLEAR / STORE_FAST can't handle scope mismatches).
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);
}
// Track symbols added by inlined comprehensions, so later comps
// can detect when a reference would resolve to a comp-local.
let mut comp_added_symbols: IndexSet<String> = IndexSet::default();

for sub_table in sub_tables.iter_mut() {
if !sub_table.comp_inlined {
continue;
}
// Don't inline if the comprehension contains nested scopes
// (lambdas, inner comprehensions, nested functions) — these need
// Cell/Free variable handling that inlining doesn't support yet.
if !sub_table.sub_tables.is_empty() {
sub_table.comp_inlined = false;
continue;
}
// Don't inline if a bound comprehension name conflicts with parent
let has_bound_conflict = sub_table.symbols.iter().any(|(name, sym)| {
name != ".0"
&& sym.flags.intersects(BOUND_FLAGS)
&& symbol_table.symbols.contains_key(name)
});
// Don't inline if a non-bound reference would resolve to a
// symbol added by a previous inlined comprehension
let has_ref_conflict = sub_table.symbols.iter().any(|(name, sym)| {
name != ".0"
&& !sym.flags.intersects(BOUND_FLAGS)
&& comp_added_symbols.contains(name)
});
if has_bound_conflict || has_ref_conflict {
sub_table.comp_inlined = false;
continue;
}
for (name, sub_symbol) in &sub_table.symbols {
if name == ".0" {
continue;
}
if symbol_table.symbols.contains_key(name) {
continue;
}
let mut symbol = sub_symbol.clone();
if sub_symbol.flags.intersects(BOUND_FLAGS) {
symbol.scope = SymbolScope::Local;
comp_added_symbols.insert(name.clone());
}
// Non-bound symbols keep their analyzed scope from the
// comprehension sub_table (e.g., GlobalImplicit, Free).
symbol_table.symbols.insert(name.clone(), symbol);
}
Comment on lines +475 to 518
Copy link
Contributor

@coderabbitai coderabbitai bot Mar 12, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🔴 Critical

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo "1) Merge path and parent analysis:"
sed -n '460,535p' crates/codegen/src/symboltable.rs

echo
echo "2) Inner-scope promotion logic:"
sed -n '718,740p' crates/codegen/src/symboltable.rs

echo
echo "Expected: the second snippet has no guard skipping st.comp_inlined before returning SymbolScope::Cell."

Repository: RustPython/RustPython

Length of output: 4707

🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo "=== All calls to found_in_inner_scope() ==="
rg 'found_in_inner_scope' crates/codegen/src/symboltable.rs -B 2 -A 2

echo
echo "=== Lines around parent analysis (line 523) ==="
sed -n '520,530p' crates/codegen/src/symboltable.rs

echo
echo "=== Check if comp_inlined is used elsewhere in symbol analysis ==="
rg 'comp_inlined' crates/codegen/src/symboltable.rs | head -20

Repository: RustPython/RustPython

Length of output: 1439

🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo "=== Confirm analyze_symbol definition and all params passed ==="
sed -n '635,680p' crates/codegen/src/symboltable.rs

echo
echo "=== Check comment context for comp_inlined behavior ==="
sed -n '462,475p' crates/codegen/src/symboltable.rs

echo
echo "=== Verify the field definition and its purpose ==="
rg 'comp_inlined:' crates/codegen/src/symboltable.rs -B 5 -A 1 | head -30

Repository: RustPython/RustPython

Length of output: 3373

Skip inlined comprehensions during closure promotion.

After merging inlined comprehension symbols into the parent scope (lines 475–518), the parent analysis at line 523 still consults the full sub_tables slice via analyze_symbol(). The found_in_inner_scope() method lacks a guard to skip comp_inlined tables, so an outer local referenced from an inlined comprehension can still be promoted to Cell. This undermines the PEP 709 fast-local model for cases like def f(x): return [x + y for y in z].

Modify found_in_inner_scope() to skip tables where st.comp_inlined is true:

fn found_in_inner_scope(
    &self,
    sub_tables: &[SymbolTable],
    name: &str,
    st_typ: CompilerScope,
) -> Option<SymbolScope> {
    sub_tables.iter().find_map(|st| {
+        if st.comp_inlined {
+            return None;
+        }
         let sym = st.symbols.get(name)?;
         if sym.scope == SymbolScope::Free || sym.flags.contains(SymbolFlags::FREE_CLASS) {
             if st_typ == CompilerScope::Class && name != "__class__" {
                 None
             } else {
🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@crates/codegen/src/symboltable.rs` around lines 475 - 518,
found_in_inner_scope() currently examines sub_tables when deciding to promote
outer locals to Cells but doesn’t skip comprehension tables that were inlined
(st.comp_inlined), so closures get wrongly promoted after we've already merged
those symbols; update found_in_inner_scope() to ignore any sub-table where
st.comp_inlined is true (e.g., add a guard at the start of the loop or
conditional to continue if st.comp_inlined) so inlined comprehension tables are
not considered during closure promotion after merging via analyze_symbol().

}

Expand Down Expand Up @@ -2037,13 +2065,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.
// Note: tables.last() is the comprehension scope we just pushed,
// so we check the second-to-last for the parent scope.
if !is_generator {
let parent_is_func = self
.tables
.iter()
.rev()
.nth(1)
.is_some_and(|t| {
matches!(
t.typ,
CompilerScope::Function
| CompilerScope::AsyncFunction
| CompilerScope::Lambda
| CompilerScope::Comprehension
)
});
if parent_is_func {
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)?;
Expand Down
16 changes: 8 additions & 8 deletions crates/vm/src/frame.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2699,13 +2699,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 } => {
Expand Down Expand Up @@ -3300,9 +3299,10 @@ impl ExecutingFrame<'_> {
Ok(None)
}
Instruction::StoreFast { var_num } => {
let value = self.pop_value();
// pop_value_opt: allows NULL from LoadFastAndClear restore path
let value = self.pop_value_opt();
let fastlocals = self.localsplus.fastlocals_mut();
fastlocals[var_num.get(arg)] = Some(value);
fastlocals[var_num.get(arg)] = value;
Ok(None)
}
Instruction::StoreFastLoadFast { var_nums } => {
Expand Down
Loading