Skip to content

lowering: accept a destructuring element parameter in the fold-loop pass #90

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

Merged
simontreanor merged 1 commit into from
Aug 30, 2026
Merged

lowering: accept a destructuring element parameter in the fold-loop pass #90

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
12 changes: 9 additions & 3 deletions INTERNALS.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -77,15 +77,21 @@ accumulator). The pass (`src/lowering/fold_loop.rs`, hooked in `lower_applicatio
`PYFUN_NO_FOLD_OPT` is the kill switch) recognizes the common case and rewrites
`functools.reduce(f, xs, acc)` into a `for`-loop over a **mutable** accumulator, turning the
copy-returning ops into in-place mutations (`Map.add`→`m[k]=v`, `List.concat`→`.append`/`.extend`,
`Set.add`→`.add`), collapsing the build to linear. This adds two Python-IR nodes: `PyStmt::For` and
`PyStmt::SubscriptAssign`.
`Set.add`→`.add`), collapsing the build to linear. This adds two Python-IR nodes: `PyStmt::For`
(whose target is a `PyForTarget`: a name or a nested tuple of names) and `PyStmt::SubscriptAssign`.

**Soundness is what the pass must protect.** The rewrite is observable only through a *retained reference* to a
mutated container, so the pass is a set of conservative **syntactic** proof obligations on the AST,
checked with no side effects on the lowerer (a rejected fold falls through to the byte-identical
`_pf_fold` lowering). A fold qualifies only when: it is a fully-applied `Seq.fold`/`List.fold`
(exactly 3 args); the folder is a 2-ary lambda literal or a **top-level** named `let` (an inlinable
body, not a `mut`/extern/imported member/parameter); the accumulator is a fresh literal collection or
body, not a `mut`/extern/imported member/parameter); the accumulator parameter is a plain name (the
pass substitutes it by name), while the element parameter may be any irrefutable pattern, since it is
only ever the loop target: a name, a wildcard, or a tuple of those becomes Python's own
`for (p, l) in steps:` header (`PyForTarget`), and any other shape binds a temp and unpacks on the
first line of the body, so `fun m (p, l) -> Map.add p l m` stays linear like its `fst`/`snd`
spelling and the pattern's bound names count as body binders for the P8 collision check; the
accumulator is a fresh literal collection or
a flat tuple of them (a `Var` init is rejected — it may be read after the fold); every slot is
threaded **position-preservingly** (no swap, no duplication, no cross-slot storage, no closure
capture, no escape to a user function — retention is unknowable, so reject); reads of a slot use a
Expand Down
2 changes: 1 addition & 1 deletion src/lowering/decode_spec.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -392,7 +392,7 @@ impl Lowerer {
args: vec![decoded],
}));
out.push(PyStmt::For {
target: item,
target: item.into(),
iter: in_name(),
body,
});
Expand Down
119 changes: 99 additions & 20 deletions src/lowering/fold_loop.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,9 +27,9 @@
use std::collections::HashSet;

use crate::parser::ast::{BlockStmt, CeItem, Expr, ExprKind, InterpPart, Param, Pattern};
use crate::python_emitter::{PyCase, PyExpr, PyStmt};
use crate::python_emitter::{PyCase, PyExpr, PyForTarget, PyStmt};

use super::{LowerError, Lowered, Lowerer, py_value_name};
use super::{LowerError, Lowered, Lowerer, py_value_name, unpack_into_as};

/// One recognized in-place update at a fold tail leaf, with the (still-unlowered)
/// argument expressions the mutation needs.
Expand Down Expand Up @@ -80,8 +80,9 @@ enum InitKind {
pub(super) struct FoldPlan<'a> {
/// A single-slot accumulator (vs a flat tuple of slots).
single: bool,
/// The folder's element parameter — becomes the loop variable.
elem_param: String,
/// The folder's element parameter (an irrefutable pattern) — becomes the
/// loop target: a name, or a tuple Python unpacks itself (`for (p, l) in …`).
elem: &'a Pattern,
/// The accumulator slot local names (length 1 for a single accumulator).
slots: Vec<String>,
/// The folder's accumulator parameter (the destructure scrutinee / single slot).
Expand Down Expand Up @@ -167,19 +168,23 @@ impl Lowerer {
#[allow(clippy::too_many_arguments)]
fn plan_fold<'a>(
&self,
params: &[Param],
params: &'a [Param],
body: &'a Expr,
init: &'a Expr,
xs: &'a Expr,
same_frame: bool,
locals: &HashSet<String>,
enclosing: &HashSet<String>,
) -> Option<FoldPlan<'a>> {
// A destructuring folder parameter (`fun (a, b) c -> …`) has no single
// name to substitute, and the pass rewrites by name — reject and fall
// through to the byte-identical `_pf_fold` lowering.
// A destructuring ACCUMULATOR parameter (`fun (a, b) c -> …`) has no
// single name to substitute, and the pass rewrites the accumulator by
// name — reject and fall through to the byte-identical `_pf_fold`
// lowering. The ELEMENT parameter is only ever the loop target, so any
// irrefutable pattern is fine there: its bound names are body binders
// for the P8 check, exactly where the non-loop lowering puts them.
let acc_param = params[0].name()?.to_string();
let elem_param = params[1].name()?.to_string();
let elem = &params[1].pattern;
let elem_names: Vec<String> = elem.bound_names();

// P3 + P4: classify the accumulator shape. Each slot's init must be a
// fresh literal (`Map.empty`/`Set.empty`/a list literal) or a bare `Var`
Expand Down Expand Up @@ -218,7 +223,7 @@ impl Lowerer {
// enclosing local (inlining would clobber it).
let mut introduced = binders;
introduced.extend(slots.iter().cloned());
introduced.insert(elem_param.clone());
introduced.extend(elem_names.iter().cloned());
if !introduced.is_disjoint(enclosing) {
return None;
}
Expand All @@ -229,7 +234,7 @@ impl Lowerer {
// the call site's own frame by construction, so the check is skipped.
if !same_frame {
let mut bound: HashSet<String> = sensitive.clone();
bound.insert(elem_param.clone());
bound.extend(elem_names.iter().cloned());
let mut free = HashSet::new();
collect_free(body, &bound, &mut free);
if !free.is_disjoint(enclosing) {
Expand Down Expand Up @@ -271,14 +276,14 @@ impl Lowerer {
// `def`, so `lower_var` rerouting is identical); a lambda or local folder
// additionally sees the call site's locals (where its free vars resolve).
let mut base_locals: HashSet<String> = sensitive.clone();
base_locals.insert(elem_param.clone());
base_locals.extend(elem_names.iter().cloned());
if same_frame {
base_locals.extend(locals.iter().cloned());
}

Some(FoldPlan {
single,
elem_param,
elem,
slots,
acc_param,
inits,
Expand Down Expand Up @@ -371,19 +376,27 @@ impl Lowerer {
// The collection, evaluated once after the inits (P10), under site locals.
let (xs_stmts, iter) = self.lower_value(plan.xs, &plan.site_locals)?;
stmts.extend(xs_stmts);
// The loop target: a name or a tuple of names Python unpacks itself;
// any other irrefutable shape binds a temp and unpacks on the first
// line of the body, as a destructuring parameter does in a `def`.
let (target, mut body) = match loop_target(plan.elem) {
Some(target) => (target, Vec::new()),
None => {
let tmp = self.fresh_tmp();
let mut unpack = Vec::new();
unpack_into_as(plan.elem, &tmp, &tmp, &|n| py_value_name(n), &mut unpack);
(PyForTarget::Name(tmp), unpack)
}
};
// The inlined folder body as the loop body.
let body = self.lower_fold_step(
body.extend(self.lower_fold_step(
plan.step_body,
&plan.base_locals,
&plan.slots,
&plan.acc_param,
plan.single,
)?;
stmts.push(PyStmt::For {
target: py_value_name(&plan.elem_param),
iter,
body,
});
)?);
stmts.push(PyStmt::For { target, iter, body });
// The result: the mutated accumulator(s) — a fresh container graph (P11).
let value = if plan.single {
PyExpr::Name(py_value_name(&plan.slots[0]))
Expand Down Expand Up @@ -1298,3 +1311,69 @@ pub(super) fn collect_free(e: &Expr, bound: &HashSet<String>, out: &mut HashSet<
| ExprKind::OpFunc(_) => {}
}
}

/// The `for`-loop target for an element pattern, when Python can unpack it
/// itself: a name, a wildcard (`_`), or a tuple of those, nested. Anything
/// else needs a temp and an explicit unpack.
fn loop_target(pattern: &Pattern) -> Option<PyForTarget> {
match pattern {
Pattern::Var { name, .. } => Some(PyForTarget::Name(py_value_name(name))),
Pattern::Wildcard => Some(PyForTarget::Name("_".to_string())),
Pattern::Tuple { elems } => elems
.iter()
.map(loop_target)
.collect::<Option<Vec<_>>>()
.map(PyForTarget::Tuple),
_ => None,
}
}

#[cfg(test)]
mod tests {
use super::*;
use crate::parser::ast::NodeSpan;

fn var(name: &str) -> Pattern {
Pattern::Var {
name: name.to_string(),
span: NodeSpan::new(crate::lexer::Span::new(0, 0)),
}
}

#[test]
fn loop_target_covers_names_wildcards_and_nested_tuples() {
assert_eq!(loop_target(&var("x")), Some(PyForTarget::Name("x".into())));
assert_eq!(
loop_target(&Pattern::Wildcard),
Some(PyForTarget::Name("_".into()))
);
let nested = Pattern::Tuple {
elems: vec![
var("a"),
Pattern::Tuple {
elems: vec![var("b"), Pattern::Wildcard],
},
],
};
assert_eq!(
loop_target(&nested),
Some(PyForTarget::Tuple(vec![
PyForTarget::Name("a".into()),
PyForTarget::Tuple(vec![
PyForTarget::Name("b".into()),
PyForTarget::Name("_".into()),
]),
]))
);
}

#[test]
fn loop_target_refuses_a_shape_python_cannot_unpack_in_the_header() {
// A literal element (not a parameter shape the parser admits today, but
// the emit path must still fall back to a temp + explicit unpack).
let shape = Pattern::Tuple {
elems: vec![var("a"), Pattern::Int(0)],
};
assert_eq!(loop_target(&shape), None);
}
}
34 changes: 17 additions & 17 deletions src/lowering/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4795,7 +4795,7 @@ fn list_prelude(used: &BTreeSet<&'static str>) -> Vec<PyStmt> {
value,
};
let for_ = |target: &str, iter: PyExpr, body: Vec<PyStmt>| PyStmt::For {
target: target.to_string(),
target: target.into(),
iter,
body,
};
Expand Down Expand Up @@ -4894,7 +4894,7 @@ fn list_prelude(used: &BTreeSet<&'static str>) -> Vec<PyStmt> {
value: PyExpr::List(vec![]),
},
PyStmt::For {
target: "x".to_string(),
target: "x".into(),
iter: name("xs"),
body: vec![PyStmt::Expr(PyExpr::Call {
func: Box::new(PyExpr::Attribute {
Expand Down Expand Up @@ -5742,7 +5742,7 @@ fn collection_prelude(used: &BTreeSet<&'static str>) -> Vec<PyStmt> {
value: PyExpr::List(vec![]),
},
PyStmt::For {
target: "x".to_string(),
target: "x".into(),
iter: name("xs"),
body: vec![
PyStmt::Assign {
Expand Down Expand Up @@ -6024,7 +6024,7 @@ fn collection_prelude(used: &BTreeSet<&'static str>) -> Vec<PyStmt> {
&["f", "xs"],
vec![
PyStmt::For {
target: "p".to_string(),
target: "p".into(),
iter: call("enumerate", vec![name("xs")]),
body: vec![PyStmt::If {
test: PyExpr::Call {
Expand Down Expand Up @@ -6241,7 +6241,7 @@ fn collection_prelude(used: &BTreeSet<&'static str>) -> Vec<PyStmt> {
value: call("set", vec![]),
},
PyStmt::For {
target: "x".to_string(),
target: "x".into(),
iter: name("xs"),
body: vec![PyStmt::If {
test: PyExpr::Not(Box::new(binop(
Expand Down Expand Up @@ -6373,7 +6373,7 @@ fn collection_prelude(used: &BTreeSet<&'static str>) -> Vec<PyStmt> {
&["f", "xs"],
vec![
PyStmt::For {
target: "x".to_string(),
target: "x".into(),
iter: name("xs"),
body: vec![PyStmt::Expr(PyExpr::Call {
func: Box::new(name("f")),
Expand Down Expand Up @@ -6426,7 +6426,7 @@ fn collection_prelude(used: &BTreeSet<&'static str>) -> Vec<PyStmt> {
value: call("set", vec![]),
},
PyStmt::For {
target: "x".to_string(),
target: "x".into(),
iter: name("s"),
body: vec![PyStmt::If {
test: PyExpr::Call {
Expand Down Expand Up @@ -6489,7 +6489,7 @@ fn collection_prelude(used: &BTreeSet<&'static str>) -> Vec<PyStmt> {
value: call("dict", vec![]),
},
PyStmt::For {
target: "kv".to_string(),
target: "kv".into(),
iter: method(name("m"), "items", vec![]),
body: vec![PyStmt::SubscriptAssign {
obj: name("out"),
Expand Down Expand Up @@ -6525,7 +6525,7 @@ fn collection_prelude(used: &BTreeSet<&'static str>) -> Vec<PyStmt> {
value: call("dict", vec![]),
},
PyStmt::For {
target: "kv".to_string(),
target: "kv".into(),
iter: method(name("m"), "items", vec![]),
body: vec![PyStmt::If {
test: PyExpr::Call {
Expand Down Expand Up @@ -6564,7 +6564,7 @@ fn collection_prelude(used: &BTreeSet<&'static str>) -> Vec<PyStmt> {
&["f", "acc", "m"],
vec![
PyStmt::For {
target: "kv".to_string(),
target: "kv".into(),
iter: method(name("m"), "items", vec![]),
body: vec![PyStmt::Assign {
target: "acc".to_string(),
Expand Down Expand Up @@ -6593,7 +6593,7 @@ fn collection_prelude(used: &BTreeSet<&'static str>) -> Vec<PyStmt> {
&["f", "m"],
vec![
PyStmt::For {
target: "kv".to_string(),
target: "kv".into(),
iter: method(name("m"), "items", vec![]),
body: vec![PyStmt::If {
test: PyExpr::Call {
Expand Down Expand Up @@ -6621,7 +6621,7 @@ fn collection_prelude(used: &BTreeSet<&'static str>) -> Vec<PyStmt> {
&["f", "m"],
vec![
PyStmt::For {
target: "kv".to_string(),
target: "kv".into(),
iter: method(name("m"), "items", vec![]),
body: vec![PyStmt::If {
test: PyExpr::Not(Box::new(PyExpr::Call {
Expand Down Expand Up @@ -6658,7 +6658,7 @@ fn collection_prelude(used: &BTreeSet<&'static str>) -> Vec<PyStmt> {
value: call("dict", vec![]),
},
PyStmt::For {
target: "kv".to_string(),
target: "kv".into(),
iter: method(name("m"), "items", vec![]),
body: vec![PyStmt::If {
test: PyExpr::Call {
Expand Down Expand Up @@ -6937,7 +6937,7 @@ fn collection_prelude(used: &BTreeSet<&'static str>) -> Vec<PyStmt> {
value: call("set", vec![]),
},
PyStmt::For {
target: "x".to_string(),
target: "x".into(),
iter: name("xs"),
body: vec![
PyStmt::Assign {
Expand Down Expand Up @@ -7013,7 +7013,7 @@ fn collection_prelude(used: &BTreeSet<&'static str>) -> Vec<PyStmt> {
value: call("None_", vec![]),
},
PyStmt::For {
target: "x".to_string(),
target: "x".into(),
iter: name("xs"),
body: vec![PyStmt::Assign {
target: "out".to_string(),
Expand Down Expand Up @@ -7082,7 +7082,7 @@ fn collection_prelude(used: &BTreeSet<&'static str>) -> Vec<PyStmt> {
&["f", "s"],
vec![
PyStmt::For {
target: "x".to_string(),
target: "x".into(),
iter: name("s"),
body: vec![PyStmt::Expr(PyExpr::Call {
func: Box::new(name("f")),
Expand All @@ -7097,7 +7097,7 @@ fn collection_prelude(used: &BTreeSet<&'static str>) -> Vec<PyStmt> {
&["f", "m"],
vec![
PyStmt::For {
target: "kv".to_string(),
target: "kv".into(),
iter: method(name("m"), "items", vec![]),
body: vec![PyStmt::Expr(PyExpr::Call {
func: Box::new(name("f")),
Expand Down
4 changes: 2 additions & 2 deletions src/lowering/self_tail_call.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -191,7 +191,7 @@ fn binds(stmts: &[PyStmt], name: &str) -> bool {
stmts.iter().any(|s| match s {
PyStmt::Assign { target, .. } => target == name,
PyStmt::UnpackAssign { targets, .. } => targets.iter().any(|t| t == name),
PyStmt::For { target, body, .. } => target == name || binds(body, name),
PyStmt::For { target, body, .. } => target.binds(name) || binds(body, name),
PyStmt::FuncDef { name: n, .. } => n == name,
PyStmt::ClassDef { name: n, .. } => n == name,
PyStmt::If { body, orelse, .. } => binds(body, name) || binds(orelse, name),
Expand Down Expand Up @@ -252,7 +252,7 @@ fn collect_bound(stmts: &[PyStmt], out: &mut HashSet<String>) {
}
PyStmt::UnpackAssign { targets, .. } => out.extend(targets.iter().cloned()),
PyStmt::For { target, body, .. } => {
out.insert(target.clone());
out.extend(target.names());
collect_bound(body, out);
}
PyStmt::If { body, orelse, .. } => {
Expand Down
Loading