diff --git a/crates/perry-transform/src/generator/lower.rs b/crates/perry-transform/src/generator/lower.rs index e910100e89..fcef02f432 100644 --- a/crates/perry-transform/src/generator/lower.rs +++ b/crates/perry-transform/src/generator/lower.rs @@ -139,6 +139,17 @@ pub fn transform_generator_function_with_extra_captures( await_async_generator_yield_operands(&mut func.body, next_local_id); } + // #6345: decide which loop bindings must NOT be hoisted into the + // activation-wide box frame, and snapshot the ones that outlive a suspend + // into per-state locals. Both run BEFORE `linearize_body` so the inserted + // `Let`s land in the same state as the closure that reads them, and before + // `local_id_before` below so the new ids are not swept into + // `extra_local_ids` (which is force-preallocated). + let per_iteration_ids = collect_per_iteration_ids(&func.body); + let mut no_hoist_ids = + snapshot_suspended_loop_captures(&mut func.body, next_local_id, &per_iteration_ids); + no_hoist_ids.extend(per_iteration_ids); + let state_id = alloc_local(next_local_id); let done_id = alloc_local(next_local_id); let sent_id = alloc_local(next_local_id); // value passed by caller via next(val) @@ -219,8 +230,23 @@ pub fn transform_generator_function_with_extra_captures( } } - // Collect hoisted var IDs first so we know which Lets to rewrite - let hoisted_for_rewrite = collect_hoisted_vars(&func.body); + // Collect hoisted var IDs first so we know which Lets to rewrite. + // + // #6345: NOT every body `Let` may be hoisted. A `let`/`const` declared in a + // loop gets a FRESH binding per iteration, and a closure made in iteration + // k must capture iteration k's binding. Hoisting moves the declaration into + // the activation-wide `PreallocateBoxes` frame (one box per call), which + // collapses every iteration onto a single cell — so all closures read the + // last value (`for (let i…) { const j = i; fns.push(() => j) }` printed the + // final `j` N times). `no_hoist_ids` (computed above) holds the bindings + // that keep their in-loop declaration — where codegen re-executes and + // re-boxes them every iteration, exactly as the non-async path does — plus + // the per-state snapshot locals. `var` and anything else live across an + // `await` is excluded there and keeps today's hoisting. + let hoisted_for_rewrite: Vec<(LocalId, String, Type)> = collect_hoisted_vars(&func.body) + .into_iter() + .filter(|(id, _, _)| !no_hoist_ids.contains(id)) + .collect(); let mut hoisted_ids: std::collections::HashSet = hoisted_for_rewrite.iter().map(|(id, _, _)| *id).collect(); // The lifted param prologue defines locals (destructured targets + temps) diff --git a/crates/perry-transform/src/generator/mod.rs b/crates/perry-transform/src/generator/mod.rs index 83c91ddb0c..912d2a2ded 100644 --- a/crates/perry-transform/src/generator/mod.rs +++ b/crates/perry-transform/src/generator/mod.rs @@ -18,6 +18,7 @@ mod id_scan; mod iter_result_rewrite; mod linearize; mod lower; +mod per_iteration; mod rewrite_returns; // Explicit named re-exports so siblings can reach each other via @@ -41,6 +42,7 @@ pub(crate) use linearize::{ pub(crate) use lower::{ transform_generator_function, transform_generator_function_with_extra_captures, }; +pub(crate) use per_iteration::{collect_per_iteration_ids, snapshot_suspended_loop_captures}; pub(crate) use rewrite_returns::{ body_contains_return, prepend_done_before_returns, rewrite_catch_returns_to_iter_result, rewrite_iter_results_in_stmts, rewrite_returns_as_done, rewrite_returns_to_labeled_break, diff --git a/crates/perry-transform/src/generator/per_iteration.rs b/crates/perry-transform/src/generator/per_iteration.rs new file mode 100644 index 0000000000..2ab55a5af3 --- /dev/null +++ b/crates/perry-transform/src/generator/per_iteration.rs @@ -0,0 +1,742 @@ +//! #6345: keep per-iteration `let`/`const` bindings per-iteration across the +//! async/generator state-machine transform. +//! +//! # The bug +//! +//! `generator::lower` hoists EVERY `Stmt::Let` in the function body into one +//! function-activation-wide frame: `collect_hoisted_vars` gathers them, +//! `Stmt::PreallocateBoxes` allocates one box per id in the entry block (once +//! per call), and `rewrite_hoisted_lets_in_stmts` demotes each `Stmt::Let` to +//! a bare `Expr(LocalSet(id, init))`. +//! +//! That is *required* for any binding whose live range crosses a suspend +//! point: the step closure returns at every `await`/`yield` and is re-entered +//! later, so its plain locals (allocas) do not survive — only the boxes it +//! captured do. +//! +//! But it is *wrong* for a `let`/`const` declared inside a loop. JS gives each +//! iteration a FRESH binding, and a closure created in iteration `k` captures +//! iteration `k`'s binding. Collapsing all iterations onto a single box makes +//! every closure observe the LAST value: +//! +//! ```ignore +//! async function main() { +//! const fns = []; +//! for (let i = 0; i < 9; i++) { +//! const j = i; +//! fns.push(() => console.log(j)); // node: 0..8 perry (pre-fix): 8 x9 +//! } +//! fns.forEach(f => f()); +//! await 0; // <- the await is the only trigger +//! } +//! ``` +//! +//! # Why the non-async path is correct +//! +//! It simply leaves the `Stmt::Let` inside the loop body. Codegen therefore +//! re-executes the declaration on every iteration — re-allocating its box when +//! the binding needs one (`boxed_vars.rs`) — and for a binding nothing ever +//! writes, `boxed_vars` declines to box it at all, so the closure +//! snapshot-captures the value at creation time. Both routes yield a distinct +//! per-iteration binding. The async path loses this the moment the `Let` is +//! hoisted out of the loop. +//! +//! # The fix +//! +//! Hoist only what actually has to be hoisted. This pass returns the ids whose +//! `Stmt::Let` must be LEFT IN PLACE inside the loop; `generator::lower` +//! subtracts them from the hoisted set, so they keep their declaration, keep +//! their per-iteration box/snapshot, and never reach `PreallocateBoxes`. +//! +//! An id qualifies only when all of these hold: +//! +//! 1. It is declared by a `Stmt::Let` inside a loop (`for` / `while` / +//! `do-while`; `for-of`/`for-in` desugar to `Stmt::For` before we run). +//! 2. It is genuinely block-scoped, NOT a `var`. Two independent guards: +//! HIR emits a `var` as a function-top pre-declaration `Let` *plus* an +//! in-place `Let` sharing the same id, so `decl_sites > 1` means `var`; +//! and a `var` read outside the loop shows up as `total_refs > +//! refs_inside_the_loop` (the #2308 test). A `var` must stay hoisted — one +//! function-scoped binding, closures see the last value, which is what node +//! does and what perry already does correctly. +//! 3. Its live range does not cross a suspend point, so it can never be read +//! after the step closure has been re-entered: +//! - a loop-body binding qualifies when no use of it appears at or after a +//! `yield` that follows its declaration (so `const data = await read(f); +//! cbs.push(() => data)` — the canonical async loop — still qualifies, +//! because the `await` precedes the declaration); +//! - a `for`-init binding is loop-carried (the condition and update re-read +//! it after the body), so it qualifies only when the loop contains no +//! suspend at all. +//! +//! Anything we cannot prove safe keeps today's hoisting. The conservative +//! direction is to hoist: un-hoisting a binding that IS live across a suspend +//! would lose its value on resume, which is a worse bug than the one being +//! fixed. + +use super::hoist_yields::expr_contains_yield; +use crate::unroll::escape_analysis::{ + count_local_refs_expr, count_local_refs_stmt, count_local_refs_stmts, +}; +use perry_hir::walker::walk_expr_children_mut; +use perry_hir::{Expr, Stmt}; +use perry_types::{LocalId, Type}; +use std::collections::{HashMap, HashSet}; + +/// Ids whose `Stmt::Let` the state-machine transform must leave in place so +/// each loop iteration gets a fresh binding. See the module docs. +pub(crate) fn collect_per_iteration_ids(body: &[Stmt]) -> HashSet { + let mut total: HashMap = HashMap::new(); + count_local_refs_stmts(body, &mut total); + + // `var` detection: HIR pre-declares a `var` at function top AND re-declares + // it in place, so its id owns two `Let` sites. A block-scoped `let`/`const` + // owns exactly one. + let mut decl_sites: HashMap = HashMap::new(); + count_decl_sites(body, &mut decl_sites); + + let mut out = HashSet::new(); + scan_stmts(body, &total, &decl_sites, &mut out); + out +} + +/// Walk every statement list looking for loops. Each loop is analyzed on its +/// own; nested loops are reached by the recursion and analyzed against the +/// same whole-body reference counts. +fn scan_stmts( + stmts: &[Stmt], + total: &HashMap, + decl_sites: &HashMap, + out: &mut HashSet, +) { + for s in stmts { + if let Some(l) = as_loop(s) { + analyze_loop(l, total, decl_sites, out); + } + each_child_stmt_list(s, &mut |list| scan_stmts(list, total, decl_sites, out)); + } +} + +fn analyze_loop( + loop_stmt: &Stmt, + total: &HashMap, + decl_sites: &HashMap, + out: &mut HashSet, +) { + // Reference counts confined to this loop (init + condition + update + body). + // An id whose every use is inside the loop cannot be a `var` that leaks out. + let mut inside: HashMap = HashMap::new(); + count_local_refs_stmt(loop_stmt, &mut inside); + + let suspends = loop_contains_suspend(loop_stmt); + let block_scoped = |id: LocalId| -> bool { + decl_sites.get(&id).copied().unwrap_or(0) == 1 + && total.get(&id).copied().unwrap_or(0) == inside.get(&id).copied().unwrap_or(0) + }; + + // `for (let i = 0; …; i++)` — the binding is loop-carried: the condition and + // the update re-read it *after* the body has run, so if the body suspends, + // `i` is live across that suspend and has to stay boxed. With no suspend + // anywhere in the loop the whole loop runs inside a single state and `i` can + // keep its per-iteration declaration. + if !suspends { + if let Stmt::For { + init: Some(init), .. + } = loop_stmt + { + if let Stmt::Let { + id, + init: init_expr, + .. + } = init.as_ref() + { + if !init_expr.as_ref().is_some_and(expr_contains_yield) && block_scoped(*id) { + out.insert(*id); + } + } + } + } + + classify_block(loop_body(loop_stmt), &block_scoped, out); +} + +/// Classify the `Stmt::Let`s of one block inside a loop. +/// +/// Descends through non-loop nesting (`if` / `try` / `switch` / labeled) — +/// `for-in` in particular parks its key binding in an `if` arm inside the +/// desugared `for` body — but stops at nested loops, which `scan_stmts` +/// analyzes in their own right. +fn classify_block( + block: &[Stmt], + block_scoped: &dyn Fn(LocalId) -> bool, + out: &mut HashSet, +) { + for (i, stmt) in block.iter().enumerate() { + if let Stmt::Let { id, init, .. } = stmt { + // A `let __tmp = yield …;` IS the state split: the linearizer + // assigns it in the resumed state, so it must stay a boxed + // cross-state local. + let splits_state = init.as_ref().is_some_and(expr_contains_yield); + if !splits_state && block_scoped(*id) && !used_after_suspend(*id, &block[i + 1..]) { + out.insert(*id); + } + } + if !is_loop(stmt) { + each_child_stmt_list(stmt, &mut |list| classify_block(list, block_scoped, out)); + } + } +} + +/// Is `id` read at or after a suspend point in `rest` (the statements that +/// follow its declaration in the same block)? +/// +/// A statement that contains BOTH a yield and a use of `id` is treated as a +/// use-after-suspend: we cannot order the two without a finer analysis, and +/// guessing wrong would drop the binding on resume. +fn used_after_suspend(id: LocalId, rest: &[Stmt]) -> bool { + let mut suspended = false; + for stmt in rest { + let mut refs: HashMap = HashMap::new(); + count_local_refs_stmt(stmt, &mut refs); + let uses = refs.get(&id).copied().unwrap_or(0) > 0; + let yields = stmt_contains_suspend(stmt); + if uses && (suspended || yields) { + return true; + } + suspended |= yields; + } + false +} + +/// The loop `stmt` is, looking through any `Stmt::Labeled` wrappers. +/// +/// `outer: for (…) { … continue outer; … }` lowers to `Labeled { body: For }`, +/// so a bare `matches!(stmt, Stmt::For { .. })` test silently skips every +/// labeled loop — and `each_child_stmt_list` unwraps `Labeled` straight to the +/// inner loop's BODY, so the loop statement itself would never be analyzed at +/// all. Its per-iteration bindings would then keep collapsing onto one box. +fn as_loop(stmt: &Stmt) -> Option<&Stmt> { + match stmt { + Stmt::For { .. } | Stmt::While { .. } | Stmt::DoWhile { .. } => Some(stmt), + Stmt::Labeled { body, .. } => as_loop(body), + _ => None, + } +} + +fn is_loop(stmt: &Stmt) -> bool { + as_loop(stmt).is_some() +} + +fn loop_body(stmt: &Stmt) -> &[Stmt] { + match stmt { + Stmt::For { body, .. } | Stmt::While { body, .. } | Stmt::DoWhile { body, .. } => body, + _ => &[], + } +} + +fn loop_contains_suspend(stmt: &Stmt) -> bool { + match stmt { + Stmt::For { + init, + condition, + update, + body, + } => { + init.as_ref().is_some_and(|i| stmt_contains_suspend(i)) + || condition.as_ref().is_some_and(expr_contains_yield) + || update.as_ref().is_some_and(expr_contains_yield) + || body.iter().any(stmt_contains_suspend) + } + Stmt::While { condition, body } | Stmt::DoWhile { body, condition } => { + expr_contains_yield(condition) || body.iter().any(stmt_contains_suspend) + } + _ => false, + } +} + +/// Conservative "does this statement suspend?" — scans every expression it +/// owns (not just the normalized `yield` statement shapes `body_contains_yield` +/// recognizes) and recurses through all nested blocks, including nested loops. +/// `expr_contains_yield` stops at `Expr::Closure`, so a nested async arrow's +/// own `await`s correctly do not count as a suspend of THIS function. +fn stmt_contains_suspend(stmt: &Stmt) -> bool { + let mut found = false; + each_expr(stmt, &mut |e| { + if !found && expr_contains_yield(e) { + found = true; + } + }); + if found { + return true; + } + let mut nested = false; + each_child_stmt_list(stmt, &mut |list| { + if !nested && list.iter().any(stmt_contains_suspend) { + nested = true; + } + }); + nested +} + +/// Invoke `f` on each expression directly owned by `stmt` (not those inside +/// nested statement lists — `each_child_stmt_list` covers those). +fn each_expr(stmt: &Stmt, f: &mut F) { + match stmt { + Stmt::Let { init, .. } => { + if let Some(e) = init { + f(e); + } + } + Stmt::Expr(e) | Stmt::Throw(e) => f(e), + Stmt::Return(opt) => { + if let Some(e) = opt { + f(e); + } + } + Stmt::If { condition, .. } => f(condition), + Stmt::While { condition, .. } | Stmt::DoWhile { condition, .. } => f(condition), + Stmt::For { + init, + condition, + update, + .. + } => { + if let Some(i) = init { + each_expr(i, f); + } + if let Some(c) = condition { + f(c); + } + if let Some(u) = update { + f(u); + } + } + Stmt::Switch { + discriminant, + cases, + } => { + f(discriminant); + for c in cases { + if let Some(t) = &c.test { + f(t); + } + } + } + Stmt::Labeled { body, .. } => each_expr(body, f), + _ => {} + } +} + +/// Invoke `f` on each nested `&[Stmt]` directly owned by `stmt`. +fn each_child_stmt_list(stmt: &Stmt, f: &mut F) { + match stmt { + Stmt::If { + then_branch, + else_branch, + .. + } => { + f(then_branch); + if let Some(eb) = else_branch { + f(eb); + } + } + Stmt::While { body, .. } | Stmt::DoWhile { body, .. } => f(body), + Stmt::For { init, body, .. } => { + if let Some(i) = init { + each_child_stmt_list(i, f); + } + f(body); + } + Stmt::Switch { cases, .. } => { + for c in cases { + f(&c.body); + } + } + Stmt::Try { + body, + catch, + finally, + } => { + f(body); + if let Some(c) = catch { + f(&c.body); + } + if let Some(fin) = finally { + f(fin); + } + } + Stmt::Labeled { body, .. } => each_child_stmt_list(body, f), + _ => {} + } +} + +// --------------------------------------------------------------------------- +// #6345 part 2: per-iteration bindings that ARE live across a suspend +// --------------------------------------------------------------------------- +// +// `collect_per_iteration_ids` can only un-hoist a binding whose live range +// stays inside one state. When the binding really is read after an `await`, +// it has to keep its cross-state box — and that box is shared by every +// iteration, so a closure capturing it still sees the last value: +// +// ```ignore +// for (let i = 0; i < 9; i++) { await step(); tasks.push(() => i); } +// for (let i = 0; i < 9; i++) { const j = i; await step(); fns.push(() => j); } +// ``` +// +// The step closure is re-entered on every resume and its capture slots are +// fixed for its lifetime, so the box pointer itself cannot be swapped per +// iteration — the only per-activation storage that survives a suspend IS that +// box. What we can do instead is give the CLOSURE the value rather than the +// cell: at the moment the closure is built, the shared box holds the current +// iteration's value, so copying it into a fresh local declared immediately +// before the closure — a local that lives and dies inside a single state — +// hands each closure its own binding. +// +// This is only sound when nothing may write the binding after the closure +// captures it, and HIR already answers that question: a captured id lands in +// `mutable_captures` when it is assigned anywhere (by a closure OR by the +// enclosing scope — verified for both), and stays in plain `captures` when it +// is only read. So we snapshot exactly `captures \ mutable_captures`. A +// `for`-init counter is deliberately kept out of `mutable_captures` by HIR +// despite `i++`, which is precisely the per-iteration semantics we want. +// Bindings a closure writes (`let acc = 0; const add = () => acc += 5`) keep +// their shared box and are untouched. + +/// Copy per-iteration bindings that outlive a suspend into a fresh local at +/// each closure-creation site. Returns the ids of the locals it introduced; +/// they must be kept out of the hoisted set (they are per-state by design). +/// +/// Runs BEFORE linearization so the inserted `Let`s land in the same state as +/// the closure that reads them. +pub(crate) fn snapshot_suspended_loop_captures( + body: &mut Vec, + next_local_id: &mut LocalId, + per_iteration: &HashSet, +) -> HashSet { + let mut total: HashMap = HashMap::new(); + count_local_refs_stmts(body, &mut total); + let mut decl_sites: HashMap = HashMap::new(); + count_decl_sites(body, &mut decl_sites); + + let ctx = SnapshotCtx { + total, + decl_sites, + per_iteration: per_iteration.clone(), + }; + let mut introduced = HashSet::new(); + walk_snapshot(body, &HashSet::new(), &ctx, next_local_id, &mut introduced); + introduced +} + +struct SnapshotCtx { + total: HashMap, + decl_sites: HashMap, + per_iteration: HashSet, +} + +/// Per-iteration bindings of `loop_stmt` that are STILL hoisted — i.e. block +/// scoped, declared by this loop, and not already un-hoisted by +/// `collect_per_iteration_ids` (those need no snapshot; they keep a real +/// per-iteration declaration, which also preserves write-sharing). +fn hoisted_loop_bindings(loop_stmt: &Stmt, ctx: &SnapshotCtx) -> HashSet { + let mut inside: HashMap = HashMap::new(); + count_local_refs_stmt(loop_stmt, &mut inside); + + let mut declared: HashSet = HashSet::new(); + if let Stmt::For { + init: Some(init), .. + } = loop_stmt + { + if let Stmt::Let { id, .. } = init.as_ref() { + declared.insert(*id); + } + } + collect_block_lets(loop_body(loop_stmt), &mut declared); + + declared + .into_iter() + .filter(|id| { + !ctx.per_iteration.contains(id) + && ctx.decl_sites.get(id).copied().unwrap_or(0) == 1 + && ctx.total.get(id).copied().unwrap_or(0) == inside.get(id).copied().unwrap_or(0) + }) + .collect() +} + +/// `Stmt::Let` ids in this block and its non-loop nesting. Nested loops own +/// their own bindings and are handled when the walk reaches them. +fn collect_block_lets(block: &[Stmt], out: &mut HashSet) { + for s in block { + if let Stmt::Let { id, .. } = s { + out.insert(*id); + } + if !is_loop(s) { + each_child_stmt_list(s, &mut |list| collect_block_lets(list, out)); + } + } +} + +fn walk_snapshot( + stmts: &mut Vec, + active: &HashSet, + ctx: &SnapshotCtx, + next_local_id: &mut LocalId, + introduced: &mut HashSet, +) { + // Descend first: a suspending loop adds its still-hoisted per-iteration + // bindings to the set visible to closures built anywhere inside it. + for stmt in stmts.iter_mut() { + // `as_loop` looks through `Stmt::Labeled` — `outer: for (…)` is a + // labeled wrapper around the loop, and its bindings are per-iteration + // just the same. + let nested = match as_loop(stmt) { + Some(l) if loop_contains_suspend(l) => { + let mut a = active.clone(); + a.extend(hoisted_loop_bindings(l, ctx)); + a + } + _ => active.clone(), + }; + each_child_stmt_list_mut(stmt, &mut |list| { + walk_snapshot(list, &nested, ctx, next_local_id, introduced) + }); + } + + // Then rewrite the closures created at THIS level, inserting each + // snapshot `Let` immediately before the statement that builds the closure + // (same block, no intervening `await` — so the same state). + let taken: Vec = stmts.drain(..).collect(); + let mut out: Vec = Vec::with_capacity(taken.len()); + for mut stmt in taken { + let mut snap_map: HashMap = HashMap::new(); + each_expr_mut(&mut stmt, &mut |e| { + snapshot_closures_in_expr(e, active, &mut snap_map, next_local_id) + }); + let mut pairs: Vec<(LocalId, LocalId)> = snap_map.into_iter().collect(); + pairs.sort(); + for (orig, snap) in pairs { + introduced.insert(snap); + out.push(Stmt::Let { + id: snap, + name: format!("__periter_{}", snap), + ty: Type::Any, + mutable: false, + init: Some(Expr::LocalGet(orig)), + }); + } + out.push(stmt); + } + *stmts = out; +} + +/// Find closures in `e` (not descending into a closure's own body — the whole +/// subtree of a rewritten closure is handled in one go) and redirect their +/// read-only captures of `active` bindings to per-iteration snapshots. +fn snapshot_closures_in_expr( + e: &mut Expr, + active: &HashSet, + snap_map: &mut HashMap, + next_local_id: &mut LocalId, +) { + if let Expr::Closure { + captures, + mutable_captures, + .. + } = e + { + let targets: Vec = captures + .iter() + .copied() + .filter(|id| active.contains(id) && !mutable_captures.contains(id)) + .collect(); + if !targets.is_empty() { + let mut local_map: HashMap = HashMap::new(); + for orig in &targets { + let snap = *snap_map.entry(*orig).or_insert_with(|| { + let id = *next_local_id; + *next_local_id = next_local_id.saturating_add(1); + id + }); + local_map.insert(*orig, snap); + } + // Substitute on a copy, then prove — with the same exhaustive + // counter the `var` analysis trusts — that no reference to the + // original id survives anywhere in the closure. If one does, this + // rewrite would leave the closure reading an id it no longer + // captures, so roll the whole closure back and keep the (merely + // stale) shared-box behaviour instead of inventing a new bug. + let original = e.clone(); + rename_in_expr(e, &local_map); + let mut leftover: HashMap = HashMap::new(); + count_local_refs_expr(e, &mut leftover); + if targets + .iter() + .any(|orig| leftover.get(orig).copied().unwrap_or(0) > 0) + { + *e = original; + for orig in &targets { + snap_map.remove(orig); + } + } + } + return; + } + walk_expr_children_mut(e, &mut |child| { + snapshot_closures_in_expr(child, active, snap_map, next_local_id) + }); +} + +/// Rewrite USE sites of the mapped ids throughout an expression subtree, +/// including capture lists and nested closure bodies. Declaration ids +/// (`Stmt::Let`) are never rewritten: every binding owns a unique LocalId, so +/// a mapped id can only ever appear as a use inside this subtree. +fn rename_in_expr(e: &mut Expr, map: &HashMap) { + let sub = |id: &mut LocalId| { + if let Some(new) = map.get(id) { + *id = *new; + } + }; + match e { + Expr::LocalGet(id) | Expr::Update { id, .. } => sub(id), + Expr::LocalSet(id, _) => sub(id), + Expr::ArrayPush { array_id, .. } + | Expr::ArrayPushSpread { array_id, .. } + | Expr::ArrayUnshift { array_id, .. } + | Expr::ArraySplice { array_id, .. } + | Expr::ArrayCopyWithin { array_id, .. } => sub(array_id), + Expr::ArrayPop(id) | Expr::ArrayShift(id) => sub(id), + Expr::SetAdd { set_id, .. } => sub(set_id), + Expr::Closure { + captures, + mutable_captures, + body, + .. + } => { + for c in captures.iter_mut() { + sub(c); + } + for c in mutable_captures.iter_mut() { + sub(c); + } + for s in body.iter_mut() { + rename_in_stmt(s, map); + } + } + _ => {} + } + walk_expr_children_mut(e, &mut |child| rename_in_expr(child, map)); +} + +fn rename_in_stmt(stmt: &mut Stmt, map: &HashMap) { + each_expr_mut(stmt, &mut |e| rename_in_expr(e, map)); + each_child_stmt_list_mut(stmt, &mut |list| { + for s in list.iter_mut() { + rename_in_stmt(s, map); + } + }); +} + +fn each_expr_mut(stmt: &mut Stmt, f: &mut F) { + match stmt { + Stmt::Let { init, .. } => { + if let Some(e) = init { + f(e); + } + } + Stmt::Expr(e) | Stmt::Throw(e) => f(e), + Stmt::Return(opt) => { + if let Some(e) = opt { + f(e); + } + } + Stmt::If { condition, .. } => f(condition), + Stmt::While { condition, .. } | Stmt::DoWhile { condition, .. } => f(condition), + Stmt::For { + init, + condition, + update, + .. + } => { + if let Some(i) = init { + each_expr_mut(i, f); + } + if let Some(c) = condition { + f(c); + } + if let Some(u) = update { + f(u); + } + } + Stmt::Switch { + discriminant, + cases, + } => { + f(discriminant); + for c in cases { + if let Some(t) = &mut c.test { + f(t); + } + } + } + Stmt::Labeled { body, .. } => each_expr_mut(body, f), + _ => {} + } +} + +fn each_child_stmt_list_mut)>(stmt: &mut Stmt, f: &mut F) { + match stmt { + Stmt::If { + then_branch, + else_branch, + .. + } => { + f(then_branch); + if let Some(eb) = else_branch { + f(eb); + } + } + Stmt::While { body, .. } | Stmt::DoWhile { body, .. } | Stmt::For { body, .. } => f(body), + Stmt::Switch { cases, .. } => { + for c in cases { + f(&mut c.body); + } + } + Stmt::Try { + body, + catch, + finally, + } => { + f(body); + if let Some(c) = catch { + f(&mut c.body); + } + if let Some(fin) = finally { + f(fin); + } + } + Stmt::Labeled { body, .. } => each_child_stmt_list_mut(body, f), + _ => {} + } +} + +/// Count `Stmt::Let` DECLARATION sites per id (a `var` owns two: the +/// function-top pre-declaration and the in-place one). Does not descend into +/// closure bodies — a closure's locals carry their own distinct ids. +fn count_decl_sites(stmts: &[Stmt], out: &mut HashMap) { + for s in stmts { + if let Stmt::Let { id, .. } = s { + *out.entry(*id).or_insert(0) += 1; + } + // A `for (let i = …; …)` init is a `Stmt::Let` that `each_child_stmt_list` + // does not surface (it owns no statement LIST). Count it explicitly, or + // every for-init binding would look like it has zero declaration sites + // and be rejected by the `decl_sites == 1` block-scoped test. + if let Stmt::For { init: Some(i), .. } = s { + if let Stmt::Let { id, .. } = i.as_ref() { + *out.entry(*id).or_insert(0) += 1; + } + } + each_child_stmt_list(s, &mut |list| count_decl_sites(list, out)); + } +} diff --git a/crates/perry-transform/src/unroll/escape_analysis.rs b/crates/perry-transform/src/unroll/escape_analysis.rs index e9be9b2c49..09d6cc125d 100644 --- a/crates/perry-transform/src/unroll/escape_analysis.rs +++ b/crates/perry-transform/src/unroll/escape_analysis.rs @@ -121,13 +121,20 @@ fn each_child_stmt_list(stmt: &Stmt, f: &mut F) { /// `stmts`. Mirrors the id-bearing variants of `scan_expr_for_max_local` /// but accumulates per-id counts so the #2308 analysis can tell whether an /// id is used outside a given loop body. -fn count_local_refs_stmts(stmts: &[Stmt], counts: &mut HashMap) { +/// +/// Also reused by the generator's #6345 per-iteration analysis +/// (`generator::per_iteration`), which needs the same "is this id used +/// outside the loop that declares it?" question answered. Kept in one place +/// so a newly-added id-bearing `Expr` variant only has to be taught to +/// `count_local_refs_expr` once — missing one would silently under-count and +/// let a `var` be mistaken for a block-scoped binding. +pub(crate) fn count_local_refs_stmts(stmts: &[Stmt], counts: &mut HashMap) { for s in stmts { count_local_refs_stmt(s, counts); } } -fn count_local_refs_stmt(stmt: &Stmt, counts: &mut HashMap) { +pub(crate) fn count_local_refs_stmt(stmt: &Stmt, counts: &mut HashMap) { match stmt { // A `Stmt::Let` id is a DECLARATION, not a use — don't count it; // only walk its initializer for uses. @@ -209,7 +216,7 @@ fn count_local_refs_stmt(stmt: &Stmt, counts: &mut HashMap) { } } -fn count_local_refs_expr(expr: &Expr, counts: &mut HashMap) { +pub(crate) fn count_local_refs_expr(expr: &Expr, counts: &mut HashMap) { fn bump(counts: &mut HashMap, id: LocalId) { *counts.entry(id).or_insert(0) += 1; } diff --git a/crates/perry-transform/src/unroll/mod.rs b/crates/perry-transform/src/unroll/mod.rs index 87e12eaf77..faab87d7e2 100644 --- a/crates/perry-transform/src/unroll/mod.rs +++ b/crates/perry-transform/src/unroll/mod.rs @@ -73,7 +73,7 @@ use perry_hir::{CallArg, CompareOp, Expr, Module, Stmt, UpdateOp}; use perry_types::{FuncId, LocalId}; use std::collections::HashMap; -mod escape_analysis; +pub(crate) mod escape_analysis; use escape_analysis::compute_loop_escaping_ids; // #5293: the max-LocalId / max-FuncId scans were copy-pasted here; route through diff --git a/test-files/test_gap_6345_async_loop_per_iteration_binding.ts b/test-files/test_gap_6345_async_loop_per_iteration_binding.ts new file mode 100644 index 0000000000..e93889b1b2 --- /dev/null +++ b/test-files/test_gap_6345_async_loop_per_iteration_binding.ts @@ -0,0 +1,325 @@ +// #6345: a `let`/`const` declared in a loop body inside an ASYNC function must +// keep its per-iteration binding. The async-to-generator transform used to +// hoist every body local into one activation-wide box, so every closure made in +// the loop observed the LAST iteration's value (silent wrong answer, exit 0). +// +// Trip counts are 9 on purpose: the static-loop unroller (MAX_TRIP_COUNT = 8) +// mints fresh ids per unrolled copy and would mask the bug entirely at <= 8. + +const out: string[] = []; +const log = (...a: unknown[]) => out.push(a.join(" ")); +const tick = () => new Promise((r) => r()); + +// --- the reported repro: `const j = i` captured, await AFTER the loop -------- +async function constCapture() { + const fns: (() => void)[] = []; + for (let i = 0; i < 9; i++) { + const j = i; + fns.push(() => log("constCapture", j)); + } + fns.forEach((f) => f()); + await 0; +} + +// --- capturing the loop variable directly ------------------------------------ +async function loopVarCapture() { + const fns: (() => void)[] = []; + for (let i = 0; i < 9; i++) { + fns.push(() => log("loopVarCapture", i)); + } + fns.forEach((f) => f()); + await 0; +} + +// --- await INSIDE the body, binding declared BEFORE the await ----------------- +// (the capture has to survive the suspend/resume) +async function awaitInBodyAfterCapture() { + const fns: (() => void)[] = []; + for (let i = 0; i < 9; i++) { + const j = i; + fns.push(() => log("awaitInBodyAfterCapture", j)); + await tick(); + } + fns.forEach((f) => f()); +} + +// --- await INSIDE the body, binding declared AFTER the await ------------------ +async function awaitInBodyBeforeCapture() { + const fns: (() => void)[] = []; + for (let i = 0; i < 9; i++) { + const j = i; + await tick(); + fns.push(() => log("awaitInBodyBeforeCapture", j)); + } + fns.forEach((f) => f()); +} + +// --- loop variable captured, await inside the body --------------------------- +async function loopVarCaptureWithAwait() { + const fns: (() => void)[] = []; + for (let i = 0; i < 9; i++) { + fns.push(() => log("loopVarCaptureWithAwait", i)); + await tick(); + } + fns.forEach((f) => f()); +} + +// --- while / do-while --------------------------------------------------------- +async function whileLoop() { + const fns: (() => void)[] = []; + let i = 0; + while (i < 9) { + const j = i; + fns.push(() => log("whileLoop", j)); + i++; + } + fns.forEach((f) => f()); + await 0; +} + +async function doWhileLoop() { + const fns: (() => void)[] = []; + let i = 0; + do { + const j = i; + fns.push(() => log("doWhileLoop", j)); + i++; + } while (i < 9); + fns.forEach((f) => f()); + await 0; +} + +// --- for-of / for-in / for-await-of ------------------------------------------ +async function forOf() { + const fns: (() => void)[] = []; + for (const x of [10, 20, 30, 40, 50, 60, 70, 80, 90]) { + fns.push(() => log("forOf", x)); + } + fns.forEach((f) => f()); + await 0; +} + +async function forIn() { + const fns: (() => void)[] = []; + for (const k in { a: 1, b: 2, c: 3, d: 4, e: 5, f: 6, g: 7, h: 8, i: 9 }) { + fns.push(() => log("forIn", k)); + } + fns.forEach((f) => f()); + await 0; +} + +async function forAwaitOf() { + const fns: (() => void)[] = []; + for await (const x of [1, 2, 3, 4, 5, 6, 7, 8, 9]) { + const y = x * 10; + fns.push(() => log("forAwaitOf", y)); + } + fns.forEach((f) => f()); +} + +// --- nested loops, and one closure capturing BOTH levels ---------------------- +async function nested() { + const fns: (() => void)[] = []; + for (let a = 0; a < 3; a++) { + for (let b = 0; b < 3; b++) { + fns.push(() => log("nested", a, b)); + } + } + fns.forEach((f) => f()); + await 0; +} + +async function nestedAwaitInInner() { + const fns: (() => void)[] = []; + for (let a = 0; a < 3; a++) { + for (let b = 0; b < 3; b++) { + fns.push(() => log("nestedAwaitInInner", a, b)); + await tick(); + } + } + fns.forEach((f) => f()); +} + +// --- `var` must STAY function-scoped: ONE binding, closures see the last value +async function varStaysVar() { + const fns: (() => void)[] = []; + for (var i = 0; i < 9; i++) { + var v = i; + fns.push(() => log("varStaysVar", v)); + } + fns.forEach((f) => f()); + log("varStaysVar after", v, i); + await 0; +} + +// --- a binding the closure WRITES keeps its shared box (no snapshot) ---------- +async function closureWritesBinding() { + const fns: (() => void)[] = []; + for (let i = 0; i < 9; i++) { + let acc = 0; + const add = () => { + acc += 5; + }; + add(); + await tick(); + fns.push(() => log("closureWritesBinding", acc)); + } + fns.forEach((f) => f()); +} + +// --- outer scope mutates the binding AFTER the closure captured it ------------ +// (same binding, so the closure must observe the update) +async function outerWritesAfterCapture() { + const fns: (() => void)[] = []; + for (let i = 0; i < 9; i++) { + let z = i; + fns.push(() => log("outerWritesAfterCapture", z)); + z += 100; + } + fns.forEach((f) => f()); + await 0; +} + +// --- labeled loop + continue -------------------------------------------------- +async function labeledLoop() { + const fns: (() => void)[] = []; + outer: for (let i = 0; i < 9; i++) { + if (i === 3) continue outer; + const j = i; + await tick(); + fns.push(() => log("labeledLoop", j)); + } + fns.forEach((f) => f()); +} + +// --- try / catch / finally inside a suspending loop --------------------------- +async function tryCatchInLoop() { + const fns: (() => void)[] = []; + for (let i = 0; i < 5; i++) { + try { + const j = i; + await tick(); + if (i === 2) throw new Error("boom"); + fns.push(() => log("tryCatchInLoop ok", j)); + } catch { + const k = i; + fns.push(() => log("tryCatchInLoop err", k)); + } finally { + const f = i; + fns.push(() => log("tryCatchInLoop fin", f)); + } + } + fns.forEach((f) => f()); +} + +// --- switch inside a suspending loop ------------------------------------------ +async function switchInLoop() { + const fns: (() => void)[] = []; + for (let i = 0; i < 9; i++) { + switch (i % 2) { + case 0: { + const e = i; + await tick(); + fns.push(() => log("switchInLoop even", e)); + break; + } + default: { + const o = i; + fns.push(() => log("switchInLoop odd", o)); + } + } + } + fns.forEach((f) => f()); +} + +// --- a captured ARRAY binding mutated from inside the closure ------------------ +// (exercises the id rewrite on array fast-path expressions) +async function arrayBindingCapture() { + const fns: (() => void)[] = []; + for (let i = 0; i < 9; i++) { + const row: number[] = []; + fns.push(() => { + row.push(i); + log("arrayBindingCapture", i, row.length, row[0]); + }); + await tick(); + } + fns.forEach((f) => f()); +} + +// --- destructured binding in a suspending loop --------------------------------- +async function destructuredBinding() { + const fns: (() => void)[] = []; + for (const o of [{ v: 1 }, { v: 2 }, { v: 3 }, { v: 4 }, { v: 5 }]) { + const { v } = o; + await tick(); + fns.push(() => log("destructuredBinding", v)); + } + fns.forEach((f) => f()); +} + +// --- a closure nested inside a closure, capturing two loop levels --------------- +async function nestedClosureCapture() { + const fns: (() => void)[] = []; + for (let a = 0; a < 3; a++) { + for (let b = 0; b < 3; b++) { + fns.push(() => { + const inner = () => log("nestedClosureCapture", a, b); + inner(); + }); + await tick(); + } + } + fns.forEach((f) => f()); +} + +// --- SYNC generator: same state-machine lowering, no async ---------------------- +function* syncGenerator(): Generator { + const fns: (() => number)[] = []; + for (let i = 0; i < 9; i++) { + const j = i; + fns.push(() => j); + yield i; + } + for (const f of fns) yield f(); +} + +// --- non-async control: must be unchanged -------------------------------------- +function plainFunction() { + const fns: (() => void)[] = []; + for (let i = 0; i < 9; i++) { + const j = i; + fns.push(() => log("plainFunction", j)); + } + fns.forEach((f) => f()); +} + +async function main() { + await constCapture(); + await loopVarCapture(); + await awaitInBodyAfterCapture(); + await awaitInBodyBeforeCapture(); + await loopVarCaptureWithAwait(); + await whileLoop(); + await doWhileLoop(); + await forOf(); + await forIn(); + await forAwaitOf(); + await nested(); + await nestedAwaitInInner(); + await varStaysVar(); + await closureWritesBinding(); + await outerWritesAfterCapture(); + await labeledLoop(); + await tryCatchInLoop(); + await switchInLoop(); + await arrayBindingCapture(); + await destructuredBinding(); + await nestedClosureCapture(); + for (const v of syncGenerator()) log("syncGenerator", v); + plainFunction(); + + for (const line of out) console.log(line); +} + +main();