diff --git a/packages/compiler/src/backend/llvm/emitter.ts b/packages/compiler/src/backend/llvm/emitter.ts index 2c3a714574..126ff36e3a 100644 --- a/packages/compiler/src/backend/llvm/emitter.ts +++ b/packages/compiler/src/backend/llvm/emitter.ts @@ -100,6 +100,7 @@ import { } from "../../ir/integer-ranges.js"; import { unshareNodes } from "../../ir/unshare.js"; import { copyStableCaptures } from "../../ir/capture-values.js"; +import { resolveSelfCaptures } from "../../ir/self-captures.js"; import { analyzeInt32Slots, dynamicFieldStores, type Int32Slots } from "../../ir/int32-slots.js"; import { analyzeFinalFields, constructorClass, type FinalFields } from "../../ir/final-fields.js"; import { foldConstantGlobals } from "../../ir/constant-globals.js"; @@ -486,7 +487,9 @@ export function emitLlvmModule(input: IrModule, options: LlvmTargetOptions = {}) specializeNumericCalls( foldConstantGlobals( copyStableCaptures( - hoistFinalFieldReads(freezeInitOnlyGlobals(dropUnwrittenPropertyBags(mod))), + resolveSelfCaptures( + hoistFinalFieldReads(freezeInitOnlyGlobals(dropUnwrittenPropertyBags(mod))), + ), ), ), mangleFunction, @@ -514,7 +517,9 @@ export function emitLlvmModuleSource( specializeNumericCalls( foldConstantGlobals( copyStableCaptures( - hoistFinalFieldReads(freezeInitOnlyGlobals(dropUnwrittenPropertyBags(mod))), + resolveSelfCaptures( + hoistFinalFieldReads(freezeInitOnlyGlobals(dropUnwrittenPropertyBags(mod))), + ), ), ), mangleFunction, diff --git a/packages/compiler/src/ir/self-captures.test.ts b/packages/compiler/src/ir/self-captures.test.ts new file mode 100644 index 0000000000..bd03725aa7 --- /dev/null +++ b/packages/compiler/src/ir/self-captures.test.ts @@ -0,0 +1,206 @@ +import { expect, test } from "vitest"; +import { + F64, + funcOf, + type IrExpr, + type IrFunction, + type IrLocal, + type IrModule, + type IrStmt, +} from "./ir.js"; +import { resolveSelfCaptures } from "./self-captures.js"; + +const loc = { file: "self.ts", start: 0, end: 0 }; +const FN = funcOf([], F64); +const ref = (localId: string, type = FN): IrExpr => ({ kind: "varRef", localId, type, loc }); +const call = (localId: string): IrExpr => ({ + kind: "callValue", + callee: ref(localId), + args: [], + type: F64, + loc, +}); +const self: IrExpr = { kind: "selfRef", type: FN, loc }; +const closure = (fnName: string, captures: string[], extra: object = {}): IrExpr => ({ + kind: "closure", + fnName, + captures, + type: FN, + loc, + ...extra, +}); +const decl = (localId: string, init: IrExpr | null): IrStmt => ({ + kind: "varDecl", + localId, + init, + loc, +}); +const assign = (localId: string, value: IrExpr, initializes = false): IrStmt => ({ + kind: "assign", + localId, + value, + ...(initializes ? { initializes: true as const } : {}), + loc, +}); +const ret = (e: IrExpr): IrStmt => ({ kind: "return", value: e, loc }); +const use = (e: IrExpr): IrStmt => ({ kind: "exprStmt", expr: e, loc }); +const local = (id: string, extra: Partial = {}): IrLocal => ({ + id, + name: id.split(".")[0]!, + type: FN, + mutable: false, + boxed: true, + ...extra, +}); + +function outer(body: IrStmt[], locals: IrLocal[]): IrFunction { + return { name: "outer", params: [], returnType: F64, locals, body, loc }; +} +function lifted(name: string, body: IrStmt[], captured: IrLocal[]): IrFunction { + return { + name, + params: [], + returnType: F64, + captures: captured.map((l) => ({ localId: l.id, name: l.name, type: l.type })), + locals: captured, + body, + loc, + }; +} +function run(...functions: IrFunction[]): Map { + const mod = { functions } as unknown as IrModule; + return new Map(resolveSelfCaptures(mod).functions.map((fn) => [fn.name, fn])); +} + +test("a closure stored in the dead-zone binding it captures reads itself", () => { + const fns = run( + outer( + [decl("f.0", null), assign("f.0", closure("walk", ["f.0"]), true), ret(call("f.0"))], + [local("f.0", { mutable: true, tdz: true })], + ), + lifted("walk", [ret(call("f.0"))], [local("f.0", { tdz: true })]), + ); + const walk = fns.get("walk")!; + expect(walk.captures).toEqual([]); + expect(walk.locals).toEqual([]); + expect(walk.body).toEqual([ret({ kind: "callValue", callee: self, args: [], type: F64, loc })]); + const origin = fns.get("outer")!; + expect(origin.body[1]).toEqual(assign("f.0", closure("walk", []), true)); + // The dead-zone box still represents the binding's uninitialized state. + expect(origin.locals[0]).toMatchObject({ boxed: true, tdz: true }); +}); + +test("a binding that only held the closure loses its box", () => { + const n = local("n.0", { type: F64 }); + const fns = run( + outer( + [ + decl("n.0", { kind: "numLit", value: 1, type: F64, loc }), + decl("f.0", closure("walk", ["n.0", "f.0"])), + ret(call("f.0")), + ], + [n, local("f.0")], + ), + lifted("walk", [use(ref("n.0", F64)), ret(call("f.0"))], [n, local("f.0")]), + ); + expect(fns.get("walk")!.captures!.map((c) => c.localId)).toEqual(["n.0"]); + const origin = fns.get("outer")!; + expect(origin.body[1]).toEqual(decl("f.0", closure("walk", ["n.0"]))); + expect(origin.locals.find((l) => l.id === "f.0")!.boxed).toBeUndefined(); + expect(origin.locals.find((l) => l.id === "n.0")!.boxed).toBe(true); +}); + +test("hoisted declarations qualify; later writes and second sites do not", () => { + const body = (extra: IrStmt[]): IrStmt[] => [ + decl("f.0", null), + assign("f.0", closure("walk", ["f.0"])), + ...extra, + ret(call("f.0")), + ]; + const walk = lifted("walk", [ret(call("f.0"))], [local("f.0")]); + const hoisted = run(outer(body([]), [local("f.0", { mutable: true })]), walk); + expect(hoisted.get("walk")!.captures).toEqual([]); + expect(hoisted.get("outer")!.locals[0]!.boxed).toBeUndefined(); + + const rewritten = run( + outer(body([assign("f.0", closure("other", []))]), [local("f.0", { mutable: true })]), + walk, + lifted("other", [ret({ kind: "numLit", value: 0, type: F64, loc })], []), + ); + expect(rewritten.get("walk")).toBe(walk); + + // The same body stored somewhere else: its environment must keep the box. + const elsewhere = run( + outer( + [...body([]).slice(0, 2), use(closure("walk", ["f.0"])), ret(call("f.0"))], + [local("f.0", { mutable: true })], + ), + walk, + ); + expect(elsewhere.get("walk")).toBe(walk); +}); + +test("an assignment in another statement list, an adapter or a nested capture keeps the box", () => { + const walk = lifted("walk", [ret(call("f.0"))], [local("f.0")]); + const branch = run( + outer( + [ + decl("f.0", null), + { kind: "block", body: [assign("f.0", closure("walk", ["f.0"]))], loc }, + ret(call("f.0")), + ], + [local("f.0", { mutable: true })], + ), + walk, + ); + expect(branch.get("walk")).toBe(walk); + + const adapter = run( + outer( + [decl("f.0", closure("walk", ["f.0"], { adapts: true })), ret(call("f.0"))], + [local("f.0")], + ), + walk, + ); + expect(adapter.get("walk")).toBe(walk); + + const nested = lifted("walk", [use(closure("inner", ["f.0"])), ret(call("f.0"))], [local("f.0")]); + const recaptured = run( + outer([decl("f.0", closure("walk", ["f.0"])), ret(call("f.0"))], [local("f.0")]), + nested, + lifted("inner", [ret(call("f.1"))], [local("f.1")]), + ); + expect(recaptured.get("walk")).toBe(nested); +}); + +test("mutually recursive closures only drop their own capture", () => { + const fns = run( + outer( + [ + decl("a.0", null), + decl("b.0", null), + assign("a.0", closure("fa", ["a.0", "b.0"]), true), + assign("b.0", closure("fb", ["a.0", "b.0"]), true), + ret(call("a.0")), + ], + [local("a.0", { mutable: true, tdz: true }), local("b.0", { mutable: true, tdz: true })], + ), + lifted( + "fa", + [use(call("b.0")), ret(call("a.0"))], + [local("a.0", { tdz: true }), local("b.0", { tdz: true })], + ), + lifted( + "fb", + [use(call("a.0")), ret(call("b.0"))], + [local("a.0", { tdz: true }), local("b.0", { tdz: true })], + ), + ); + expect(fns.get("fa")!.captures!.map((c) => c.localId)).toEqual(["b.0"]); + expect(fns.get("fb")!.captures!.map((c) => c.localId)).toEqual(["a.0"]); + expect(fns.get("outer")!.body.slice(2)).toEqual([ + assign("a.0", closure("fa", ["b.0"]), true), + assign("b.0", closure("fb", ["a.0"]), true), + ret(call("a.0")), + ]); +}); diff --git a/packages/compiler/src/ir/self-captures.ts b/packages/compiler/src/ir/self-captures.ts new file mode 100644 index 0000000000..e009e9a116 --- /dev/null +++ b/packages/compiler/src/ir/self-captures.ts @@ -0,0 +1,229 @@ +import { CaptureFlow } from "./capture-values.js"; +import { + typeEquals, + type IrExpr, + type IrFunction, + type IrLocal, + type IrModule, + type IrStmt, +} from "./ir.js"; +import { everyExprChild, everyStmtChild, mapExprChildren, mapStmtChildren } from "./traverse.js"; + +/** Closures that capture the binding they are stored in read themselves + * through `selfRef` instead of through that binding's box. + * + * `const visit = (node) => { ... visit(child) ... }` lowers to a closure + * over the box of `visit`, and the box then holds the closure: a two-object + * reference cycle that reference counting cannot free. Each closure made + * this way is buffered as a cycle candidate and reclaimed only when a + * collector pass walks it. Nested function declarations already read their + * own name through `selfRef`; this pass gives arrows and function + * expressions bound to a binding the same form once the IR proves that the + * binding always holds the closure itself while the closure can run: + * + * - every creation site of the closure body stores the new closure, as is, + * into a boxed local of the creating function: either the initializer of + * that local's only declaration, or the only assignment of a local whose + * only declaration is empty and precedes the assignment in the same + * statement list (so each box is filled once, and a closure can only run + * after the statement that stores it); + * - nothing else writes the binding, in the creating function or in any + * body that shares its box (CaptureFlow.unchanged), and no local class + * captures it; + * - the closure body is synchronous, is not a generator, and does not pass + * the binding on to a nested closure (its environment would need a box of + * its own, allocated on every call). + * + * The capture then disappears from the body and from every creation site, + * and each read of it becomes `selfRef` (the running closure, retained), + * which is the same pointer the box would have produced. The creating + * function keeps its binding; it loses its box when no other closure + * captures it, except for temporal-dead-zone bindings, whose box still + * carries the empty state its own reads test. */ +export function resolveSelfCaptures(mod: IrModule): IrModule { + const functions = new Map(mod.functions.map((fn) => [fn.name, fn])); + const sites = new Map(); + /** Closure nodes stored into a local by a qualifying definition, with that + * local and the number of writes the definition itself makes. */ + const definitions = new Map(); + const declarations = new Map>(); + for (const fn of mod.functions) { + const counts = new Map(); + declarations.set(fn.name, counts); + const expr = (e: IrExpr): boolean => { + if (e.kind === "closure") { + let list = sites.get(e.fnName); + if (list === undefined) sites.set(e.fnName, (list = [])); + list.push({ owner: fn.name, closure: e }); + } + return everyExprChild(e, expr, stmt); + }; + const stmt = (s: IrStmt): boolean => { + if (s.kind === "varDecl") counts.set(s.localId, (counts.get(s.localId) ?? 0) + 1); + return everyStmtChild(s, expr, stmt); + }; + fn.body.every(stmt); + for (const list of statementLists(fn.body)) { + list.forEach((s, i) => { + if (s.kind !== "varDecl") return; + if (s.init !== null) { + if (s.init.kind === "closure") definitions.set(s.init, { local: s.localId, writes: 0 }); + return; + } + for (let j = i + 1; j < list.length; j++) { + const next = list[j]!; + if (next.kind === "assign" && next.localId === s.localId) { + if (next.value.kind === "closure") + definitions.set(next.value, { local: s.localId, writes: 1 }); + return; + } + } + }); + } + } + + let flow: CaptureFlow | null = null; + /** Per closure body, the capture index that names its own closure. */ + const resolved = new Map(); + for (const fn of mod.functions) { + const captures = fn.captures; + if (captures === undefined || captures.length === 0) continue; + if (fn.async === true || fn.generator !== undefined || fn.classCaptures !== undefined) continue; + const list = sites.get(fn.name); + if (list === undefined || list.length === 0) continue; + const recaptured = capturedIds(fn.body); + for (let index = 0; index < captures.length; index++) { + const capture = captures[index]!; + if (capture.type.kind !== "func" || recaptured.has(capture.localId)) continue; + const qualifies = list.every(({ owner, closure }) => { + if (closure.adapts === true || !typeEquals(closure.type, capture.type)) return false; + const definition = definitions.get(closure); + const source = closure.captures[index]; + if (definition === undefined || source !== definition.local) return false; + const creator = functions.get(owner)!; + const local = creator.locals.find((l) => l.id === source); + if (local?.boxed !== true || !typeEquals(local.type, closure.type)) return false; + if (creator.params.some((p) => p.localId === source)) return false; + if (creator.captures?.some((c) => c.localId === source)) return false; + if (creator.classCaptures?.some((c) => c.localId === source)) return false; + if (declarations.get(owner)!.get(source) !== 1) return false; + flow ??= new CaptureFlow(mod.functions); + return flow.unchanged(owner, source, definition.writes); + }); + if (qualifies) { + resolved.set(fn.name, index); + break; + } + } + } + if (resolved.size === 0) return mod; + + const rewritten = mod.functions.map((fn) => { + const own = resolved.get(fn.name); + const self = own === undefined ? null : fn.captures![own]!; + let changed = false; + const expr = (e: IrExpr): IrExpr => { + if (self !== null && e.kind === "varRef" && e.localId === self.localId) { + changed = true; + return { kind: "selfRef", type: self.type, loc: e.loc }; + } + if (e.kind === "closure") { + const drop = resolved.get(e.fnName); + if (drop !== undefined) { + changed = true; + return { ...e, captures: e.captures.filter((_, i) => i !== drop) }; + } + } + return mapExprChildren(e, expr, stmt); + }; + const stmt = (s: IrStmt): IrStmt => mapStmtChildren(s, expr, stmt); + const body = fn.body.map(stmt); + if (!changed && self === null) return fn; + // A binding no environment captures any more needs no box. A dead-zone + // binding keeps the box that represents its uninitialized state. + const stillCaptured = capturedIds(body); + // The callback returns IrLocal explicitly: inferred from both returns, + // its type would be the unboxed copy's, and the self-hosted compiler + // would narrow every kept local to that shape, dropping `boxed`. + const locals = fn.locals + .filter((local) => self === null || local.id !== self.localId) + .map((local): IrLocal => { + if (local.boxed !== true || local.tdz === true || stillCaptured.has(local.id)) return local; + if (fn.captures?.some((c) => c.localId === local.id)) return local; + if (fn.classCaptures?.some((c) => c.localId === local.id)) return local; + if (!wasSource(fn, local.id, resolved)) return local; + const { boxed: _boxed, ...plain } = local; + return plain; + }); + const next: IrFunction = { ...fn, locals, body }; + if (self !== null) next.captures = fn.captures!.filter((_, i) => i !== own); + return next; + }); + return { ...mod, functions: rewritten }; +} + +/** Whether `fn` stored a resolved closure into `id`, the only case in which + * this pass may unbox a binding (others stay exactly as lowered). */ +function wasSource(fn: IrFunction, id: string, resolved: ReadonlyMap): boolean { + let found = false; + const expr = (e: IrExpr): boolean => { + if (e.kind === "closure") { + const index = resolved.get(e.fnName); + if (index !== undefined && e.captures[index] === id) found = true; + } + return !found && everyExprChild(e, expr, stmt); + }; + const stmt = (s: IrStmt): boolean => everyStmtChild(s, expr, stmt); + fn.body.every(stmt); + return found; +} + +/** Locals some closure or class created in `body` captures. */ +function capturedIds(body: readonly IrStmt[]): Set { + const ids = new Set(); + const expr = (e: IrExpr): boolean => { + if (e.kind === "closure" || e.kind === "classRef") + for (const id of e.captures ?? []) ids.add(id); + return everyExprChild(e, expr, stmt); + }; + const stmt = (s: IrStmt): boolean => everyStmtChild(s, expr, stmt); + body.every(stmt); + return ids; +} + +/** `body` and every statement list nested in its statements (not the ones + * inside expressions). Statements of one list run in order, each at most + * once per execution of the list. */ +function statementLists(body: readonly IrStmt[]): (readonly IrStmt[])[] { + const lists: (readonly IrStmt[])[] = []; + const visit = (list: readonly IrStmt[]): void => { + lists.push(list); + for (const s of list) { + switch (s.kind) { + case "block": + case "while": + case "doWhile": + case "for": + case "forOf": + visit(s.body); + break; + case "if": + visit(s.then); + if (s.else_ !== null) visit(s.else_); + break; + case "switch": + for (const c of s.cases) visit(c.body); + break; + case "tryCatch": + visit(s.tryBody); + if (s.catchBody !== null) visit(s.catchBody); + if (s.finallyBody !== null) visit(s.finallyBody); + break; + default: + break; + } + } + }; + visit(body); + return lists; +} diff --git a/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/5990-self-capturing-closures.ts.json b/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/5990-self-capturing-closures.ts.json new file mode 100644 index 0000000000..5f707aaf25 --- /dev/null +++ b/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/5990-self-capturing-closures.ts.json @@ -0,0 +1,6 @@ +{ + "order": [ + "/tests/corpus/5990-self-capturing-closures.ts" + ], + "diags": [] +} diff --git a/tests/corpus/5990-self-capturing-closures.ts b/tests/corpus/5990-self-capturing-closures.ts new file mode 100644 index 0000000000..274e7b8d73 --- /dev/null +++ b/tests/corpus/5990-self-capturing-closures.ts @@ -0,0 +1,124 @@ +// Closures that call themselves through the binding they are stored in: +// const arrows, function expressions, closures made in loops, closures that +// also hand themselves to nested callbacks, and a binding that is later +// rebound (its closure must see the new value). Each read of the binding +// inside the body is the same function object the binding holds. The +// sanitized lane's RC audit fails if any of these closures leaks or is +// freed while reachable. + +const fact = (n: number): number => (n <= 1 ? 1 : n * fact(n - 1)); +console.log(fact(10)); + +const pick = (n: number): (() => number) | null => (n === 0 ? null : () => n); +const identity = (depth: number): unknown => (depth === 0 ? identity : identity(depth - 1)); +console.log(identity(3) === identity, pick(0) === null, pick(2)!()); + +const fib = function (n: number): number { + return n < 2 ? n : fib(n - 1) + fib(n - 2); +}; +console.log(fib(15)); + +const outerName = function innerName(n: number): number { + return n <= 0 ? 0 : 1 + outerName(n - 1) + innerName(0); +}; +console.log(outerName(5)); + +// One closure per iteration, each recursing into itself and reading its +// own iteration's value. +const perIteration: ((n: number) => number)[] = []; +for (let i = 0; i < 4; i++) { + const step = (n: number): number => (n === 0 ? i : step(n - 1) + 10); + perIteration.push(step); +} +console.log(perIteration.map((f) => f(2)).join(" ")); +console.log(perIteration[0] === perIteration[1], perIteration[2] === perIteration[2]); + +// The closure also passes itself to nested callbacks. +interface Tree { + value: number; + kids: Tree[]; +} +const leaf = (value: number): Tree => ({ value, kids: [] }); +const tree: Tree = { + value: 1, + kids: [{ value: 2, kids: [leaf(3), leaf(4)] }, leaf(5)], +}; +const sum = (node: Tree): number => node.value + node.kids.map((kid) => sum(kid)).reduce((a, b) => a + b, 0); +console.log(sum(tree)); +const visited: number[] = []; +const walk = (node: Tree): void => { + visited.push(node.value); + node.kids.forEach(walk); +}; +walk(tree); +console.log(visited.join(",")); + +// Captured state next to the self reference. +function counterTo(limit: number): () => number { + let calls = 0; + const tick = (): number => { + calls++; + return calls < limit ? tick() : calls; + }; + return tick; +} +const tick = counterTo(5); +console.log(tick(), tick()); + +// A rebound binding: the first closure calls whatever the binding holds now. +let rebound = (n: number): number => (n <= 0 ? 0 : 1 + rebound(n - 1)); +const first = rebound; +console.log(first(3)); +rebound = (n: number): number => n * 100; +console.log(first(3), rebound(3)); + +// Mutual recursion keeps both bindings live. +const isEven = (n: number): boolean => (n === 0 ? true : isOdd(n - 1)); +const isOdd = (n: number): boolean => (n === 0 ? false : isEven(n - 1)); +console.log(isEven(10), isOdd(7), isEven(3)); + +// Identity as a map key from inside the body. +const seen = new Map(); +const keyed = (n: number): number => { + seen.set(keyed, (seen.get(keyed) ?? 0) + 1); + return n === 0 ? seen.size : keyed(n - 1); +}; +console.log(keyed(4), seen.get(keyed)); + +// An object that holds a closure capturing the object (a real cycle). +class Widget { + label: string; + onChange: (() => string) | null = null; + constructor(label: string) { + this.label = label; + } +} +function widgets(count: number): number { + let total = 0; + for (let i = 0; i < count; i++) { + const widget = new Widget(`w${i}`); + const describe = (depth: number): string => (depth === 0 ? widget.label : describe(depth - 1)); + widget.onChange = () => describe(2); + total += widget.onChange().length; + } + return total; +} +console.log(widgets(1000)); + +// Many short-lived recursive closures. +let checksum = 0; +for (let round = 0; round < 20000; round++) { + const down = (n: number): number => (n === 0 ? round % 7 : down(n - 1)); + checksum += down(3); +} +console.log(checksum); + +// Generators and async functions that name themselves. +const countdown = function* (n: number): Generator { + if (n < 0) return; + yield n; + yield* countdown(n - 1); +}; +console.log([...countdown(3)].join("")); +const later = async (n: number): Promise => (n === 0 ? 0 : 1 + (await later(n - 1))); +later(3).then((value) => console.log("async", value));