Recover nested generator comprehensions through statement parsing

docs/LIMITATIONS.md

@@ -143,57 +143,28 @@ rendering bugs:
   (300+ regression-driven special cases), so it is left documented
   rather than patched blind.
 
-### Nested comprehension inside a genexp (3 modules, deliberately not patched)
-
-`({c: v for c, v in zip(cols, row) if c} for row in rows)` — a genexp whose
-*element* (or a call argument inside it) is an inlined (PEP 709) dict
-comprehension. The genexp keeps its own code object and is decompiled by
-the comprehension walker, which has no notion of a nested level: the
-inner `BUILD_MAP 0` accumulator lands on the walker's stack, `SWAP` is
-not modeled, and the following `FOR_ITER` treats it as a second `for`
-clause (`... for c, v in {}`), so the output does not parse.
-
-A frame-stack prototype (one frame per comprehension level, `SWAP`
-modeled, the accumulator placeholder replaced by Rc identity when a level
-ends, `pending_or_filter` flushed per level) was implemented **twice** and
-reverted both times. What it achieved, verified by runtime-equivalence
-checks against the original sources: the whole-element shape and the
-call-argument shapes *as minimal repros* became semantically correct
-(`({c: v for c, v in zip(cols, t) if c} for t in rows)`,
-`(f(**d) for x in xs)`, `(f(x, y=1) for x in xs)`; the last two were
-fixed for real and are committed).
-
-It still produced *plausible but semantically wrong* output for the three
-real modules (`pandas/core/methods/to_dict.py`,
-`sklearn/linear_model/_logistic.py`, `metrics/pairwise.py`), e.g.
-
-    fd(func, ret, s, X, Y[s, ...], **{k: (v[s] if k == "Y" else v) for k, v in kwds.items()})
-    ->  (fd if k == 'Y' else (func, ret, s, X, Y[s, ...])(*{}, **{**k}) for s in ...)
-
-The blocker is not the nesting itself: the walker's stack model is a
-heuristic that happens to agree with the VM on the shapes the corpus
-exercises, and a nested level exposes several concrete divergences at
-once (all observed with probes during the third attempt, 2026-09-15):
-
-  * `GET_ITER` is a no-op, so the 3.12 prologue's `SWAP`s reorder the
-    wrong slots (the VM's is `[saved…, iter]` -> `BUILD_* 0` -> `SWAP 2`;
-    the walker has no `iter` on the stack there);
-  * the `LOAD_FAST_AND_CLEAR` save slots are pushed but never popped —
-    the epilogue's `STORE_FAST`s restore them in the VM, while the walker
-    treats stores as target assignments, so they linger on the stack;
-  * `PUSH_NULL` markers then sit between the level's value and the call
-    slot, and the `CALL_FUNCTION_EX` pops pick up the wrong values
-    (`fd(...)` came out as `fd if k == 'Y' else (func, ret, X, Y)(*{}, **{**k})`).
-
-Fixing one of these shifts the divergence to the next (truncating the
-level's stack to its base fixed the kwargs slot but stripped the outer
-operands, and excluding the save slots then leaked the `PUSH_NULL`
-marker into the output), which is exactly why this needs the call model
-reworked rather than another local patch. A correct fix
-therefore starts with making that call model faithful (or by delegating
-the enclosing expression to the main engine's `exec`), and only then
-re-adding the nesting. Emitting a visible syntax error stays preferable
-to silently wrong code, so the modules remain in the failing set.
+### Nested comprehensions inside generator expressions
+
+Python 3.12+ generators containing an inlined list, set or dictionary
+comprehension now reuse the statement parser to recover the nested element.
+This fixes whole-element and call-argument shapes, including the former
+minimal reproduction:
+
+```python
+({c: v for c, v in zip(cols, row) if c} for row in rows)
+```
+
+The extractor accepts a completely recovered For/If/yield tree, preserves
+its yielded expression, and restores the actual outer iterable at the call
+site. It rejects unsupported statements and multiple sibling yield paths.
+The legacy comprehension walker remains the fallback when extraction does
+not prove that shape, and for generators without an inlined comprehension.
+
+The full-module cases previously listed here (`pandas/core/methods/to_dict.py`,
+`sklearn/linear_model/_logistic.py`, and `metrics/pairwise.py`) also exercise
+complex call, closure and control-flow reconstruction. They remain on the
+full-module verification backlog; repairing the nested-comprehension shape
+alone does not establish that every construct in those files is recovered.
 
 ### ~~Multi-link ternary on the value path~~ (fixed 2026-09-15)
 

src/decompiler.rs

@@ -50939,6 +50939,49 @@ impl<'a> Ctx<'a> {
 
 // =====================  comprehensions  =====================
 
+// The main statement parser already understands PEP 709's nested inline
+// comprehensions. Reuse it for generator elements instead of flattening an
+// inner dict/list/set comprehension into the generator's own clauses.
+fn generator_parts_from_body(
+    body: &[Stmt], generators: &mut Vec<Comprehension>,
+) -> Option<ExprRef> {
+    let mut result = None;
+    for statement in body {
+        match statement {
+            Stmt::For { target, iter, body, orelse, is_async } if orelse.is_empty() => {
+                if result.is_some() { return None; }
+                generators.push(Comprehension {
+                    target: target.clone(), iter: iter.clone(), ifs: Vec::new(),
+                    is_async: *is_async, if_line_break: false,
+                });
+                result = generator_parts_from_body(body, generators);
+                if result.is_none() { return None; }
+            }
+            Stmt::If { cond, body, orelse } if orelse.is_empty() => {
+                if result.is_some() { return None; }
+                generators.last_mut()?.ifs.push(cond.clone());
+                result = generator_parts_from_body(body, generators);
+                if result.is_none() { return None; }
+            }
+            Stmt::Expr(value) => {
+                match &**value {
+                    Expr::Yield(Some(value)) => {
+                        if result.is_some() { return None; }
+                        result = Some(value.clone());
+                    }
+                    Expr::Const(_) => {}, // Docstring.
+                    _ => return None,
+                }
+            }
+            Stmt::Pass | Stmt::Return(None) => {},
+            Stmt::Return(Some(value))
+                if matches!(&**value, Expr::Const(object) if matches!(&**object, PyObject::None)) => {},
+            _ => return None,
+        }
+    }
+    result
+}
+
 impl<'a> Ctx<'a> {
     /// Try to turn a completed call into a comprehension. The callee is a
     /// Function whose code object is named `<listcomp>` etc. The iterable is
@@ -51186,6 +51229,20 @@ fn genexpr_ternary_merge(
         let table = table_for(self.version).ok()?;
         let instrs = crate::bytecode::decode_instructions(code, table, self.version);
 
+        if self.version.at_least(3, 12) && matches!(_kind, CompKind::Generator)
+            && instrs.iter().any(|instruction| instruction.op == Op::LOAD_FAST_AND_CLEAR)
+        {
+            if let Some(statements) = decompile(code, self.version).ok().filter(|result| result.clean) {
+                let mut generators = Vec::new();
+                if let Some(element) = generator_parts_from_body(&statements.body, &mut generators) {
+                    if !generators.is_empty() {
+                        generators[0].iter = outer_iter.clone();
+                        return Some((element, None, generators));
+                    }
+                }
+            }
+        }
+
         struct PartialGen {
             target: Option<ExprRef>,
             iter: ExprRef,
@@ -53096,3 +53153,66 @@ impl<'a> Ctx<'a> {
         }
     }
 }
+
+#[cfg(test)]
+mod generator_parts_tests {
+    use super::*;
+
+    fn yielded_loop() -> Stmt {
+        Stmt::For {
+            target: Rc::new(Expr::Name("row".into())),
+            iter: Rc::new(Expr::Name("rows".into())),
+            body: vec![Stmt::Expr(Rc::new(Expr::Yield(Some(Rc::new(
+                Expr::Name("row".into()),
+            )))))],
+            orelse: vec![],
+            is_async: false,
+        }
+    }
+
+    #[test]
+    fn extraction_preserves_one_yield_path_and_rejects_siblings() {
+        let body = vec![yielded_loop()];
+        let mut generators = Vec::new();
+        let element = generator_parts_from_body(&body, &mut generators).unwrap();
+        assert!(matches!(&*element, Expr::Name(name) if name == "row"));
+        assert_eq!(generators.len(), 1);
+
+        let mut generators = Vec::new();
+        assert!(generator_parts_from_body(
+            &[yielded_loop(), yielded_loop()], &mut generators,
+        ).is_none());
+
+        let branch = Stmt::If {
+            cond: Rc::new(Expr::Name("ready".into())),
+            body: vec![Stmt::Expr(Rc::new(Expr::Yield(Some(Rc::new(
+                Expr::Name("row".into()),
+            )))))],
+            orelse: vec![],
+        };
+        let mut body = yielded_loop();
+        if let Stmt::For { body, .. } = &mut body {
+            *body = vec![branch.clone(), branch];
+        }
+        assert!(generator_parts_from_body(&[body], &mut Vec::new()).is_none());
+    }
+
+    #[test]
+    fn extraction_does_not_discard_assignments_or_loop_else() {
+        for name in ["saved", "__assert_armored__"] {
+            let assignment = Stmt::Assign {
+                targets: vec![Rc::new(Expr::Name(name.into()))],
+                value: Rc::new(Expr::Const(Rc::new(PyObject::None))),
+                wrap_before: vec![],
+            };
+            assert!(generator_parts_from_body(
+                &[assignment, yielded_loop()], &mut Vec::new(),
+            ).is_none());
+        }
+        let mut body = yielded_loop();
+        if let Stmt::For { orelse, .. } = &mut body {
+            *orelse = vec![Stmt::Pass];
+        }
+        assert!(generator_parts_from_body(&[body], &mut Vec::new()).is_none());
+    }
+}

tests/behavior/cases/b_nested_generator_comprehensions.py

@@ -0,0 +1,26 @@
+# MIN_VERSION: 3.12
+
+
+def dictionary(rows):
+    return tuple({key: row[key] for key in row} for row in rows)
+
+
+def call_arguments(rows, collect):
+    return tuple(collect([value * 2 for value in row]) for row in rows)
+
+
+def filtered(rows):
+    return tuple({key: row[key] for key in row if key} for row in rows if row)
+
+
+def unpacked(rows, columns):
+    return tuple({key: value for key, value in zip(columns, row) if key} for row in rows)
+
+
+def collect(values):
+    print("collect", values)
+    return sum(values)
+print(dictionary([{"a": 1, "b": 2}, {}, {"c": 3}]))
+print(call_arguments([[1, 2], [], [3]], collect))
+print(filtered([{}, {"": 0, "a": 1}]))
+print(unpacked([[1, 2], [3, 4]], ["a", "b"]))

tests/fixtures/nested_generator_comprehensions.py

@@ -0,0 +1,22 @@
+def dictionary(rows):
+    return tuple({key: row[key] for key in row} for row in rows)
+
+
+def list_elements(rows):
+    return tuple([value * 2 for value in row] for row in rows)
+
+
+def set_elements(rows):
+    return tuple({value * 2 for value in row} for row in rows)
+
+
+def call_arguments(rows, collect):
+    return tuple(collect([value * 2 for value in row]) for row in rows)
+
+
+def filtered(rows):
+    return tuple({key: row[key] for key in row if key} for row in rows if row)
+
+
+def unpacked(rows, columns):
+    return tuple({key: value for key, value in zip(columns, row) if key} for row in rows)

tests/integration.rs

@@ -169,3 +169,24 @@ fn py314_hello() {
         "py3.14 hello",
     );
 }
+
+#[test]
+fn nested_comprehensions_remain_generator_elements() {
+    for version in ["3.11", "3.12", "3.13", "3.14"] {
+        let fixture = format!("nested_generator_comprehensions.{version}.pyc");
+        let (src, clean) = decompile_fixture(&fixture);
+        assert!(clean, "{fixture}: output marked incomplete:\n{src}");
+        assert_contains(
+            &src,
+            &[
+                "return tuple(({key: row[key] for key in row} for row in rows))",
+                "return tuple(([value * 2 for value in row] for row in rows))",
+                "return tuple(({value * 2 for value in row} for row in rows))",
+                "return tuple((collect([value * 2 for value in row]) for row in rows))",
+                "return tuple(({key: row[key] for key in row if key} for row in rows if row))",
+                "return tuple(({key: value for key, value in zip(columns, row) if key} for row in rows))",
+            ],
+            &fixture,
+        );
+    }
+}