diff --git a/crates/perry-transform/src/generator/lower.rs b/crates/perry-transform/src/generator/lower.rs index d5a5402ac2..ec6e5ba1f6 100644 --- a/crates/perry-transform/src/generator/lower.rs +++ b/crates/perry-transform/src/generator/lower.rs @@ -139,6 +139,16 @@ pub fn transform_generator_function_with_extra_captures( await_async_generator_yield_operands(&mut func.body, next_local_id); } + // #6354: a per-iteration binding a closure WRITES that also outlives a + // suspend is fixed by neither #6345 path (a value snapshot would drop the + // write, so it stays in `mutable_captures` and keeps its shared box). Back + // each such binding with a one-element heap cell BEFORE the #6345 passes + // run: the cell reference is then a read-only per-iteration capture the + // snapshot below handles, while writes go to the shared element. See + // `per_iteration.rs` part 3. + let cell_ids = collect_written_suspended_loop_captures(&func.body); + rewrite_written_captures_to_cells(&mut func.body, &cell_ids); + // #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 diff --git a/crates/perry-transform/src/generator/mod.rs b/crates/perry-transform/src/generator/mod.rs index 912d2a2ded..102784cd97 100644 --- a/crates/perry-transform/src/generator/mod.rs +++ b/crates/perry-transform/src/generator/mod.rs @@ -42,7 +42,10 @@ 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 per_iteration::{ + collect_per_iteration_ids, collect_written_suspended_loop_captures, + rewrite_written_captures_to_cells, 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 index 2ab55a5af3..8608593dc8 100644 --- a/crates/perry-transform/src/generator/per_iteration.rs +++ b/crates/perry-transform/src/generator/per_iteration.rs @@ -78,8 +78,8 @@ 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_hir::walker::{walk_expr_children, walk_expr_children_mut}; +use perry_hir::{BinaryOp, Expr, Stmt, UpdateOp, WithSetFallback}; use perry_types::{LocalId, Type}; use std::collections::{HashMap, HashSet}; @@ -740,3 +740,317 @@ fn count_decl_sites(stmts: &[Stmt], out: &mut HashMap) { each_child_stmt_list(s, &mut |list| count_decl_sites(list, out)); } } + +// --------------------------------------------------------------------------- +// #6354: per-iteration bindings a closure WRITES that also outlive a suspend +// --------------------------------------------------------------------------- +// +// `snapshot_suspended_loop_captures` (part 2 above) can only hand a closure a +// per-iteration binding when nobody writes it after capture — it copies the +// *value*, and a value snapshot would silently drop any later write. It is +// therefore gated to `captures \ mutable_captures`. A binding a closure assigns +// (`let acc = i; const bump = () => { acc += 100; }`) stays in `mutable_captures`, +// keeps its single activation-wide box, and every closure created in the loop +// then observes the LAST iteration's value: +// +// ```ignore +// for (let i = 0; i < 9; i++) { +// let acc = i; +// const bump = () => { acc += 100; }; // WRITES acc +// bump(); +// await tick(); // acc read after the suspend +// fns.push(() => acc); // node: 100..108 perry: 108 x9 +// } +// ``` +// +// The fix reduces this WRITE case to the already-solved READ-ONLY case by +// backing the binding with a one-element heap cell. The binding VARIABLE then +// holds a per-iteration array *reference* that is never reassigned (only its +// element is), so it lands in `captures \ mutable_captures` and part 2 snapshots +// the reference per iteration; writes go to the shared element and stay visible +// to every closure of the same iteration. Concretely `acc` above becomes: +// +// ```ignore +// let acc = [i]; +// const bump = () => { acc[0] += 100; }; +// bump(); +// await tick(); +// fns.push(() => acc[0]); +// ``` +// +// which perry already compiles correctly. The rewrite runs BEFORE part 1/part 2 +// so those passes see the cell form. + +/// Loop-body block-scoped `let` bindings that (a) some closure in the loop lists +/// in its `mutable_captures` and (b) are live across a suspend inside the loop. +/// These are exactly the bindings #6345 leaves collapsed: read-only captures are +/// snapshotted, and bindings whose live range stays in one state keep their +/// per-iteration declaration, so neither path covers a written binding that also +/// outlives an `await`. Each returned id should be rewritten to a heap cell by +/// [`rewrite_written_captures_to_cells`]. +pub(crate) fn collect_written_suspended_loop_captures(body: &[Stmt]) -> HashSet { + // The union of every closure's `mutable_captures` (including nested + // closures) — HIR marks an id here iff it is assigned somewhere it is + // captured, whether by a closure or by the enclosing scope after capture. + let mut mutably_captured: HashSet = HashSet::new(); + collect_mutable_captures_stmts(body, &mut mutably_captured); + if mutably_captured.is_empty() { + return HashSet::new(); + } + + 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 mut out = HashSet::new(); + scan_written_loops(body, &total, &decl_sites, &mutably_captured, &mut out); + + // Exclude anything referenced in a form the cell rewrite cannot express (a + // bare-`LocalId` array/set mutation intrinsic, or a `with` fallback). Such a + // binding is left as-is — the pre-existing collapse persists for that rare + // shape, but no new corruption is introduced. + if !out.is_empty() { + let mut unsafe_ids: HashSet = HashSet::new(); + collect_cell_unsafe_ids(body, &mut unsafe_ids); + out.retain(|id| !unsafe_ids.contains(id)); + } + out +} + +/// Ids that appear in a reference form [`rewrite_written_captures_to_cells`] +/// does NOT rewrite: array/set mutation intrinsics keyed on a bare `LocalId` +/// (`arr.push(x)` → `ArrayPush { array_id }`, `set.add(x)`, `arr.pop()`, …) and +/// `with`-statement fallbacks. Turning such a binding into a cell would leave +/// the intrinsic pointing at the one-element cell array instead of the value it +/// holds, so these ids are excluded from the candidate set. +fn collect_cell_unsafe_ids(stmts: &[Stmt], out: &mut HashSet) { + for s in stmts { + each_expr(s, &mut |e| collect_cell_unsafe_ids_expr(e, out)); + each_child_stmt_list(s, &mut |list| collect_cell_unsafe_ids(list, out)); + } +} + +fn collect_cell_unsafe_ids_expr(e: &Expr, out: &mut HashSet) { + match e { + Expr::ArrayPush { array_id, .. } + | Expr::ArrayPushSpread { array_id, .. } + | Expr::ArrayUnshift { array_id, .. } + | Expr::ArraySplice { array_id, .. } + | Expr::ArrayCopyWithin { array_id, .. } => { + out.insert(*array_id); + } + Expr::ArrayPop(id) | Expr::ArrayShift(id) => { + out.insert(*id); + } + Expr::SetAdd { set_id, .. } => { + out.insert(*set_id); + } + Expr::WithSet { fallback, .. } => { + if let WithSetFallback::Local(id) | WithSetFallback::SloppyImplicit(id) = fallback { + out.insert(*id); + } + } + // A closure body can mutate a captured binding the same way — descend. + Expr::Closure { body, .. } => collect_cell_unsafe_ids(body, out), + _ => {} + } + walk_expr_children(e, &mut |child| collect_cell_unsafe_ids_expr(child, out)); +} + +fn scan_written_loops( + stmts: &[Stmt], + total: &HashMap, + decl_sites: &HashMap, + mutably_captured: &HashSet, + out: &mut HashSet, +) { + for s in stmts { + if let Some(l) = as_loop(s) { + analyze_written_loop(l, total, decl_sites, mutably_captured, out); + } + each_child_stmt_list(s, &mut |list| { + scan_written_loops(list, total, decl_sites, mutably_captured, out) + }); + } +} + +fn analyze_written_loop( + loop_stmt: &Stmt, + total: &HashMap, + decl_sites: &HashMap, + mutably_captured: &HashSet, + out: &mut HashSet, +) { + let mut inside: HashMap = HashMap::new(); + count_local_refs_stmt(loop_stmt, &mut inside); + 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) + }; + classify_written_block(loop_body(loop_stmt), &block_scoped, mutably_captured, out); +} + +/// Mirror of `classify_block`, but selecting the bindings part 2 CANNOT fix: a +/// block-scoped loop-body `let` that some closure writes (`mutably_captured`) +/// and that is read at or after a suspend following its declaration. Descends +/// through non-loop nesting; nested loops are reached by `scan_written_loops`. +fn classify_written_block( + block: &[Stmt], + block_scoped: &dyn Fn(LocalId) -> bool, + mutably_captured: &HashSet, + out: &mut HashSet, +) { + for (i, stmt) in block.iter().enumerate() { + if let Stmt::Let { id, init, .. } = stmt { + let splits_state = init.as_ref().is_some_and(expr_contains_yield); + if !splits_state + && mutably_captured.contains(id) + && 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_written_block(list, block_scoped, mutably_captured, out) + }); + } + } +} + +/// Union every closure's `mutable_captures` in these statements, descending into +/// nested closure bodies (a deeper closure's write still makes the binding +/// "assigned after capture"). +fn collect_mutable_captures_stmts(stmts: &[Stmt], out: &mut HashSet) { + for s in stmts { + each_expr(s, &mut |e| collect_mutable_captures_expr(e, out)); + each_child_stmt_list(s, &mut |list| collect_mutable_captures_stmts(list, out)); + } +} + +fn collect_mutable_captures_expr(e: &Expr, out: &mut HashSet) { + if let Expr::Closure { + mutable_captures, + body, + .. + } = e + { + out.extend(mutable_captures.iter().copied()); + collect_mutable_captures_stmts(body, out); + } + walk_expr_children(e, &mut |child| collect_mutable_captures_expr(child, out)); +} + +/// Rewrite each id in `cells` from a scalar binding into a one-element heap +/// cell: its declaration becomes `let c = [init]`, every read/write/update of +/// `c` becomes an indexed access on `c[0]`, and every closure that captured `c` +/// keeps it as a now read-only capture (dropped from `mutable_captures`). See +/// the module docs — after this, `c`'s VALUE (the array reference) is never +/// reassigned, so part 2 hands each iteration its own cell. +pub(crate) fn rewrite_written_captures_to_cells(body: &mut [Stmt], cells: &HashSet) { + if cells.is_empty() { + return; + } + for s in body.iter_mut() { + rewrite_cells_in_stmt(s, cells); + } +} + +fn rewrite_cells_in_stmt(stmt: &mut Stmt, cells: &HashSet) { + match stmt { + // The candidate's declaration: wrap its initializer in a one-element + // array so the slot holds a fresh cell each iteration. `let c;` (no + // init) seeds `[undefined]`. + Stmt::Let { id, init, ty, .. } if cells.contains(id) => { + let mut inner = init.take().unwrap_or(Expr::Undefined); + rewrite_cells_in_expr(&mut inner, cells); + *init = Some(Expr::Array(vec![inner])); + *ty = Type::Array(Box::new(Type::Any)); + } + _ => { + each_expr_mut(stmt, &mut |e| rewrite_cells_in_expr(e, cells)); + } + } + each_child_stmt_list_mut(stmt, &mut |list| { + for s in list.iter_mut() { + rewrite_cells_in_stmt(s, cells); + } + }); +} + +fn rewrite_cells_in_expr(e: &mut Expr, cells: &HashSet) { + match e { + Expr::LocalGet(id) if cells.contains(id) => { + let id = *id; + *e = Expr::IndexGet { + object: Box::new(Expr::LocalGet(id)), + index: Box::new(Expr::Integer(0)), + }; + // Do NOT descend — the freshly built `LocalGet(id)` is the cell + // reference itself, not another read of the (now indexed) binding. + return; + } + Expr::LocalSet(id, value) if cells.contains(id) => { + let id = *id; + // Rewrite reads of other cells (and of this cell) inside the value + // BEFORE lifting it out, e.g. `c = c + 1` → `c[0] = c[0] + 1`. + rewrite_cells_in_expr(value, cells); + let value = std::mem::replace(value.as_mut(), Expr::Undefined); + *e = Expr::IndexSet { + object: Box::new(Expr::LocalGet(id)), + index: Box::new(Expr::Integer(0)), + value: Box::new(value), + }; + return; + } + Expr::Update { id, op, prefix } if cells.contains(id) => { + let id = *id; + let op = match op { + UpdateOp::Increment => BinaryOp::Add, + UpdateOp::Decrement => BinaryOp::Sub, + }; + let prefix = *prefix; + *e = Expr::IndexUpdate { + object: Box::new(Expr::LocalGet(id)), + index: Box::new(Expr::Integer(0)), + op, + prefix, + }; + return; + } + Expr::Closure { + captures, + mutable_captures, + body, + .. + } => { + // The cell is now mutated only through `IndexSet`/`IndexUpdate` on + // its element — the VARIABLE is never reassigned — so it is a + // read-only capture. Drop it from `mutable_captures` (keeping it in + // `captures` so the body can still read the array reference) so part + // 2's `captures \ mutable_captures` gate snapshots it per iteration. + let demoted: Vec = mutable_captures + .iter() + .copied() + .filter(|c| cells.contains(c)) + .collect(); + mutable_captures.retain(|c| !cells.contains(c)); + for c in demoted { + if !captures.contains(&c) { + captures.push(c); + } + } + for s in body.iter_mut() { + rewrite_cells_in_stmt(s, cells); + } + // Fall through to `walk_expr_children_mut`, which for a closure + // visits ONLY its param defaults (not the body — handled above) — + // so a cell referenced in a default (`(x = acc) => …`) is rewritten + // too. It does not re-descend into the body, so no double-rewrite. + } + _ => {} + } + walk_expr_children_mut(e, &mut |child| rewrite_cells_in_expr(child, cells)); +} diff --git a/test-files/test_gap_6354_async_written_binding_across_await.ts b/test-files/test_gap_6354_async_written_binding_across_await.ts new file mode 100644 index 0000000000..44d15fb87f --- /dev/null +++ b/test-files/test_gap_6354_async_written_binding_across_await.ts @@ -0,0 +1,269 @@ +// #6354: a per-iteration `let` that a closure WRITES *and* that is still read +// after an `await` in the same loop body used to collapse onto one binding — +// every closure observed the last iteration's value (silent wrong answer, +// exit 0). +// +// This is the residual left by #6345: its snapshot only copies READ-ONLY +// captures (a value snapshot would drop a later write), so a written binding +// stayed in `mutable_captures`, kept its single activation-wide box, and +// collapsed. The fix backs such a binding with a one-element heap cell so the +// binding VARIABLE is a per-iteration read-only reference (snapshotted per +// iteration) while writes go to the shared element. +// +// 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. +// Every binding here is iteration-DEPENDENT (`= i`, not `= 0`) so a collapse to +// the last value is distinguishable from the correct per-iteration answer. + +const out: string[] = []; +const log = (...a: unknown[]) => out.push(a.join(" ")); +const tick = () => new Promise((r) => r()); + +// --- the reported repro: closure writes `acc`, read after the suspend --------- +async function closureWrite() { + const fns: (() => void)[] = []; + for (let i = 0; i < 9; i++) { + let acc = i; + const bump = () => { + acc += 100; + }; + bump(); + await tick(); + fns.push(() => log("closureWrite", acc)); + } + fns.forEach((f) => f()); +} + +// --- closure writes via ++/-- (Update, not a compound LocalSet) --------------- +async function closureUpdate() { + const fns: (() => void)[] = []; + for (let i = 0; i < 9; i++) { + let n = i * 10; + const inc = () => { + n++; + }; + inc(); + inc(); + await tick(); + fns.push(() => log("closureUpdate", n)); + } + fns.forEach((f) => f()); +} + +// --- the enclosing scope writes the binding after a closure captured it, ------ +// with the await INSIDE the loop body (so the binding is live across the +// suspend and cannot be un-hoisted). +async function outerWriteAcrossAwait() { + const fns: (() => void)[] = []; + for (let i = 0; i < 9; i++) { + let z = i; + fns.push(() => log("outerWriteAcrossAwait", z)); + z += 100; + await tick(); + } + fns.forEach((f) => f()); +} + +// --- write-sharing must survive: two closures over the same binding, and a ---- +// write that happens AFTER the suspend must be visible to a reader closure +// created BEFORE it (a value snapshot would break this). +async function sharedWriteAfterSuspend() { + const readers: (() => void)[] = []; + for (let i = 0; i < 9; i++) { + let acc = i; + const bump = () => { + acc += 100; + }; + const reader = () => log("sharedWriteAfterSuspend", acc); + bump(); + await tick(); + bump(); // second write, AFTER the suspend + readers.push(reader); + } + readers.forEach((f) => f()); +} + +// --- two independent written bindings in one loop body ------------------------ +async function twoBindings() { + const fns: (() => void)[] = []; + for (let i = 0; i < 9; i++) { + let a = i; + let b = i * 100; + const bump = () => { + a += 1; + b += 1; + }; + bump(); + await tick(); + fns.push(() => log("twoBindings", a, b)); + } + fns.forEach((f) => f()); +} + +// --- while loop --------------------------------------------------------------- +async function whileLoop() { + const fns: (() => void)[] = []; + let i = 0; + while (i < 9) { + let acc = i; + const bump = () => { + acc += 100; + }; + bump(); + await tick(); + fns.push(() => log("whileLoop", acc)); + i++; + } + fns.forEach((f) => f()); +} + +// --- do/while loop ------------------------------------------------------------ +async function doWhileLoop() { + const fns: (() => void)[] = []; + let i = 0; + do { + let acc = i; + const bump = () => { + acc += 100; + }; + bump(); + await tick(); + fns.push(() => log("doWhileLoop", acc)); + i++; + } while (i < 9); + fns.forEach((f) => f()); +} + +// --- for-of loop -------------------------------------------------------------- +async function forOfLoop() { + const fns: (() => void)[] = []; + for (const x of [0, 1, 2, 3, 4, 5, 6, 7, 8]) { + let acc = x; + const bump = () => { + acc += 100; + }; + bump(); + await tick(); + fns.push(() => log("forOfLoop", acc)); + } + fns.forEach((f) => f()); +} + +// --- nested loops: the inner binding is per (i, j) ---------------------------- +async function nestedLoops() { + const fns: (() => void)[] = []; + for (let i = 0; i < 3; i++) { + for (let j = 0; j < 3; j++) { + let s = i * 10 + j; + const bump = () => { + s += 100; + }; + bump(); + await tick(); + fns.push(() => log("nestedLoops", s)); + } + } + fns.forEach((f) => f()); +} + +// --- a `var` (function-scoped) must NOT be turned per-iteration: node reports -- +// the last value, and so must perry. This guards against over-application. +async function varStaysCollapsed() { + const fns: (() => void)[] = []; + for (let i = 0; i < 9; i++) { + var acc = i; + const bump = () => { + acc += 100; + }; + bump(); + await tick(); + fns.push(() => log("varStaysCollapsed", acc)); + } + fns.forEach((f) => f()); +} + +// --- the written binding is read only inside a closure PARAM DEFAULT ---------- +// (a closure default is evaluated in the enclosing scope; the rewrite must +// reach it too, not just the closure body). +async function paramDefaultCapture() { + const fns: (() => number)[] = []; + for (let i = 0; i < 9; i++) { + let acc = i; + const bump = () => { + acc += 100; + }; + bump(); + await tick(); + const reader = (x = acc) => x; // `acc` referenced in a param DEFAULT + fns.push(() => reader()); + } + for (const f of fns) log("paramDefaultCapture", f()); +} + +// --- the enclosing scope writes the binding AFTER the suspend ----------------- +async function outerWriteAfterAwait() { + const fns: (() => void)[] = []; + for (let i = 0; i < 9; i++) { + let acc = i; + fns.push(() => log("outerWriteAfterAwait", acc)); + await tick(); + acc += 100; // written AFTER the suspend, in the enclosing scope + } + fns.forEach((f) => f()); +} + +// --- sync generator variant: a written binding live across a `yield` ---------- +function* syncGen() { + const fns: (() => void)[] = []; + for (let i = 0; i < 9; i++) { + let acc = i; + const bump = () => { + acc += 100; + }; + bump(); + yield i; + fns.push(() => log("syncGen", acc)); + } + fns.forEach((f) => f()); +} + +// --- async generator variant: same residual, driven by the same machinery ----- +async function* asyncGen() { + const fns: (() => void)[] = []; + for (let i = 0; i < 9; i++) { + let acc = i; + const bump = () => { + acc += 100; + }; + bump(); + await tick(); + yield i; + fns.push(() => log("asyncGen", acc)); + } + fns.forEach((f) => f()); +} + +async function main() { + await closureWrite(); + await closureUpdate(); + await outerWriteAcrossAwait(); + await sharedWriteAfterSuspend(); + await twoBindings(); + await whileLoop(); + await doWhileLoop(); + await forOfLoop(); + await nestedLoops(); + await paramDefaultCapture(); + await outerWriteAfterAwait(); + await varStaysCollapsed(); + for (const _ of syncGen()) { + // drain + } + for await (const _ of asyncGen()) { + // drain + } + + for (const line of out) console.log(line); +} + +main();