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

A folder that names its element by destructuring (`fun m (p, l) -> Map.add
p l m`) was rejected by the in-place fold pass and fell back to `_pf_fold`,
which copies the accumulator every step and goes quadratic. Only the
accumulator parameter is substituted by name; the element parameter is the
loop target, so any irrefutable pattern works there.

`PyStmt::For` now takes a `PyForTarget` (a name or a nested tuple of names),
so a tuple element emits Python's own `for (p, l) in steps:` header. Any
other irrefutable shape binds a temp and unpacks on the first body line via
the existing `unpack_into_as`. The pattern's bound names feed the P8
collision check as body binders. The accumulator parameter keeps its rules.

Author: simontreanor

INTERNALS.md

@@ -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

src/lowering/decode_spec.rs

@@ -392,7 +392,7 @@ impl Lowerer {
                     args: vec![decoded],
                 }));
                 out.push(PyStmt::For {
-                    target: item,
+                    target: item.into(),
                     iter: in_name(),
                     body,
                 });

src/lowering/fold_loop.rs

@@ -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.
@@ -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).
@@ -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`
@@ -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;
         }
@@ -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) {
@@ -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,
@@ -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]))
@@ -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);
+    }
+}

src/lowering/mod.rs

@@ -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,
     };
@@ -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 {
@@ -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 {
@@ -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 {
@@ -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(
@@ -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")),
@@ -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 {
@@ -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"),
@@ -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 {
@@ -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(),
@@ -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 {
@@ -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 {
@@ -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 {
@@ -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 {
@@ -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(),
@@ -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")),
@@ -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")),

src/lowering/self_tail_call.rs

@@ -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),
@@ -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, .. } => {