Compile large closure environments in linear time

packages/compiler/src/backend/llvm/blocks.test.ts

@@ -1,5 +1,5 @@
 import { expect, test } from "vitest";
-import { BlockBuilder } from "./blocks.js";
+import { BlockBuilder, volatileSlotAccess } from "./blocks.js";
 import { f64Lit } from "./common.js";
 
 test("entry allocations precede instructions and temporaries stay unique", () => {
@@ -106,3 +106,34 @@ test("inclusive loops and nested loops keep labels separate", () => {
   expect(text).toContain("fcmp ole double");
   expect(text.match(/alloca double/g)).toHaveLength(2);
 });
+
+test("volatile slots render their loads and stores volatile and nothing else", () => {
+  const slots = new Set(["%sc_l_a_0"]);
+  expect(volatileSlotAccess("  %t1 = load ptr, ptr %sc_l_a_0 ; a", slots)).toBe(
+    "  %t1 = load volatile ptr, ptr %sc_l_a_0 ; a",
+  );
+  expect(volatileSlotAccess("  store ptr %t1, ptr %sc_l_a_0, !dbg !7", slots)).toBe(
+    "  store volatile ptr %t1, ptr %sc_l_a_0, !dbg !7",
+  );
+  expect(volatileSlotAccess("  store ptr null, ptr %sc_l_a_0", slots)).toBe(
+    "  store volatile ptr null, ptr %sc_l_a_0",
+  );
+  // Other addresses, a slot stored as a value, prefixes and other opcodes.
+  for (const line of [
+    "  %t2 = load ptr, ptr %sc_l_a_01",
+    "  store ptr %sc_l_a_0, ptr %sc_l_b_0",
+    "  %t3 = getelementptr inbounds ptr, ptr %sc_l_a_0, i64 1",
+    "  call void @scr_box_release(ptr %sc_l_a_0)",
+    "  %t4 = load volatile ptr, ptr %sc_l_a_0",
+  ])
+    expect(volatileSlotAccess(line, slots)).toBe(line);
+
+  const blocks = new BlockBuilder();
+  blocks.volatileSlots = slots;
+  blocks.entryAllocas.push("%sc_l_a_0 = alloca ptr");
+  blocks.line("store ptr null, ptr %sc_l_a_0");
+  blocks.terminate("ret void");
+  expect(blocks.render()).toBe(
+    "entry:\n  %sc_l_a_0 = alloca ptr\n  store volatile ptr null, ptr %sc_l_a_0\n  ret void",
+  );
+});

packages/compiler/src/backend/llvm/blocks.ts

@@ -31,6 +31,9 @@ export class BlockBuilder {
   /** Temps known to hold immortal values (temp names are only unique
    * within one builder). */
   readonly immortalTemps = new Set<string>();
+  /** Entry slots whose loads and stores render volatile, so no LLVM pass
+   * promotes them to SSA values (see volatileSlotAccess). */
+  volatileSlots: ReadonlySet<string> | null = null;
 
   constructor() {
     this.cur = { label: "entry", lines: [], term: null };
@@ -130,9 +133,14 @@ export class BlockBuilder {
   }
 
   render(): string {
+    const slots = this.volatileSlots;
     return this.blocks
       .map((b, i) => {
-        const lines = i === 0 ? [...this.entryAllocas.map((l) => `  ${l}`), ...b.lines] : b.lines;
+        const body =
+          slots === null || slots.size === 0
+            ? b.lines
+            : b.lines.map((line) => volatileSlotAccess(line, slots));
+        const lines = i === 0 ? [...this.entryAllocas.map((l) => `  ${l}`), ...body] : body;
         // A block left unterminated is a structurally unreachable join
         // (every predecessor jumped elsewhere); the verifier still wants a
         // terminator.
@@ -142,3 +150,36 @@ export class BlockBuilder {
       .join("\n");
   }
 }
+
+/** `line` with its load or store made volatile when the address operand is
+ * one of `slots`. Lines are this builder's own instruction text, where the
+ * address is the first `ptr %name` operand after the opcode (a stored value
+ * precedes it without a comma). */
+export function volatileSlotAccess(line: string, slots: ReadonlySet<string>): string {
+  let op = 0;
+  while (line.charCodeAt(op) === 32) op++;
+  if (line.charCodeAt(op) === 37) {
+    const assign = line.indexOf(" = ", op);
+    if (assign < 0) return line;
+    op = assign + 3;
+  }
+  const kind = line.startsWith("load ", op) ? 5 : line.startsWith("store ", op) ? 6 : 0;
+  if (kind === 0 || line.startsWith("volatile ", op + kind)) return line;
+  const address = line.indexOf(", ptr %", op + kind);
+  if (address < 0) return line;
+  const start = address + 6;
+  let end = start + 1;
+  while (end < line.length && isNameCode(line.charCodeAt(end))) end++;
+  if (!slots.has(line.slice(start, end))) return line;
+  return `${line.slice(0, op + kind)}volatile ${line.slice(op + kind)}`;
+}
+
+function isNameCode(code: number): boolean {
+  return (
+    (code >= 48 && code <= 57) ||
+    (code >= 65 && code <= 90) ||
+    (code >= 97 && code <= 122) ||
+    code === 95 ||
+    code === 46
+  );
+}

packages/compiler/src/backend/llvm/emitter.test.ts

@@ -126,6 +126,61 @@ test("the shared indexed-read helpers carry an inline hint", () => {
   expect(define("plain")).not.toContain("inlinehint");
 });
 
+test("bodies holding thousands of capture boxes compile unoptimized with boxes in memory", () => {
+  const number = (value: number): IrExpr => ({ kind: "numLit", value, type: F64, loc });
+  const read = (localId: string): IrExpr => ({ kind: "varRef", localId, type: F64, loc });
+  // A declaring body with `count` boxed bindings, and a closure body that
+  // receives `count` of them through its environment.
+  const declaring = (name: string, count: number): IrFunction => {
+    const ids = Array.from({ length: count }, (_, i) => `${name}${i}`);
+    return {
+      name,
+      params: [],
+      returnType: F64,
+      locals: ids.map((id) => ({ id, name: id, type: F64, mutable: true, boxed: true as const })),
+      body: [
+        ...ids.map((id, i): IrStmt => ({ kind: "varDecl", localId: id, init: number(i), loc })),
+        { kind: "return", value: read(ids[count - 1]!), loc },
+      ],
+      loc,
+    };
+  };
+  const receiving = (name: string, count: number): IrFunction => {
+    const ids = Array.from({ length: count }, (_, i) => `${name}${i}`);
+    return {
+      name,
+      params: [],
+      returnType: F64,
+      captures: ids.map((id) => ({ localId: id, name: id, type: F64 })),
+      locals: ids.map((id) => ({ id, name: id, type: F64, mutable: true, boxed: true as const })),
+      body: [{ kind: "return", value: read(ids[0]!), loc }],
+      loc,
+    };
+  };
+  const module = moduleFor([]);
+  const llvm = emitLlvmModule({
+    ...module,
+    functions: [
+      ...module.functions,
+      declaring("wide", 512),
+      declaring("narrow", 511),
+      receiving("env", 512),
+    ],
+  });
+  const body = (name: string): string => {
+    const start = llvm.split("\n").findIndex((line) => line.endsWith(`{ ; ${name}`));
+    const lines = llvm.split("\n").slice(start);
+    return lines.slice(0, lines.indexOf("}")).join("\n");
+  };
+  expect(body("wide")).toMatch(/\) noinline optnone #0 \{/);
+  expect(body("wide")).toMatch(/store volatile ptr %[st]\d+, ptr %sc_l_wide0\b/);
+  expect(body("wide")).toMatch(/= load volatile ptr, ptr %sc_l_wide511\b/);
+  expect(body("env")).toMatch(/\) noinline optnone #0 \{/);
+  expect(body("env")).toMatch(/store volatile ptr %t\d+, ptr %sc_l_env0 ; captured env0/);
+  expect(body("narrow")).not.toContain("optnone");
+  expect(body("narrow")).not.toContain("volatile");
+});
+
 test("throwing-call count does not multiply identical scope cleanup", () => {
   for (const boxed of [false, true]) {
     const small = emitLlvmModule(moduleFor(Array.from({ length: 4 }, call), { boxed }));

packages/compiler/src/backend/llvm/emitter.ts

@@ -548,10 +548,35 @@ const AWAIT_SUSPENSIONS = /@(?:scr_await_(?:void|f64|bool|str|ref|hop)\b|sc_a[wd
  * own fiber and promise. */
 const COROUTINE_SUSPENDS =
   /@(?:scr_(?:await_|module_await\b|promise_await_settled\b|test_sub\b|gen_yield_|async_gen_|fiber_(?!self\b|promise\b))|sc_a[wd]_)/;
+/** Captured bindings (boxes a body declares or receives) from which a
+ * function compiles without LLVM optimization (see boxSlotCount). */
+const UNOPTIMIZED_BOX_SLOTS = 512;
 /** ScrExcKind tags of a generator slot holding a number or a boolean. */
 const SCR_EXC_F64 = 1;
 const SCR_EXC_BOOL = 2;
 
+/** The capture boxes `fn` holds: its boxed bindings and the bindings its
+ * environment brings in. Each box pointer is a value from its binding (or
+ * function entry) to its last use, and every one of them crosses the calls
+ * in between. With thousands of them in one body, as in a factory whose
+ * returned methods share a large environment, LLVM's inliner, instruction
+ * selection and register allocation grow quadratically and take minutes.
+ * Past UNOPTIMIZED_BOX_SLOTS, which ordinary functions stay far below, the
+ * body is compiled `optnone` (nothing inlines into it and code generation
+ * takes its fast paths) and its box slots are accessed volatile, so no
+ * pass promotes them to values even where a pipeline runs IR passes over
+ * optnone bodies, as the native helper's does. Each use reloads its box
+ * pointer and code generation stays linear. Semantics, boxes and
+ * environments are unchanged; only that body's machine code is slower. */
+function boxSlotCount(fn: IrFunction): number {
+  const captured = new Set<string>();
+  for (const capture of fn.captures ?? []) captured.add(capture.localId);
+  for (const capture of fn.classCaptures ?? []) captured.add(capture.localId);
+  let count = 0;
+  for (const local of fn.locals) if (local.boxed === true || captured.has(local.id)) count++;
+  return count;
+}
+
 /** Whether any node of `e` names the local (reads, writes, declarations,
  * captures). */
 function mentionsLocal(e: IrExpr, localId: string): boolean {
@@ -6817,6 +6842,15 @@ export class LlEmitter {
         B.line(`store ptr null, ptr %${mangleLocal(local.id)}`);
       }
     }
+    // A body holding thousands of capture boxes compiles unoptimized, its
+    // box pointers kept in their slots (see boxSlotCount).
+    const unoptimized = this.currentCoro === null && boxSlotCount(fn) >= UNOPTIMIZED_BOX_SLOTS;
+    if (unoptimized) {
+      const slots = new Set<string>();
+      for (const local of this.currentLocals.values())
+        if (local.boxed || this.captureIds.has(local.id)) slots.add(`%${mangleLocal(local.id)}`);
+      B.volatileSlots = slots;
+    }
     // Captured bindings come in through the environment — borrowed for the
     // whole call (the closure owns them): bound here, never released here.
     (fn.captures ?? []).forEach((c, i) => {
@@ -7020,6 +7054,7 @@ export class LlEmitter {
     if (fn.name.startsWith("%dyn.class.protocol.fallback.")) attrs = `noinline ${attrs}`;
     else if (fn.name.startsWith("%dyn.class.protocol.")) attrs = `inlinehint ${attrs}`;
     if (fn.name.startsWith(RESULT_ENTRY_PREFIX)) attrs = `inlinehint ${attrs}`;
+    if (unoptimized && attrs === FN_ATTRS) attrs = `noinline optnone ${attrs}`;
     const debug = this.debugScope === null ? "" : ` !dbg ${this.debugScope}`;
     const symbol = borrowed ? mangleBorrowedFunction(fn.name) : mangleFunction(fn.name);
     this.privateGlobalSlots.clear();

packages/compiler/src/backend/llvm/lazy-captures.test.ts

@@ -47,7 +47,7 @@ function lifted(name: string, body: IrStmt[], captured = "options"): IrFunction
     loc,
   };
 }
-function lazy(...functions: IrFunction[]): Set<string> {
+function lazy(...functions: IrFunction[]): ReadonlySet<string> {
   const map = new Map(functions.map((fn) => [fn.name, fn]));
   return new LazyCaptures(map).parameters(functions[0]!);
 }

packages/compiler/src/backend/llvm/lazy-captures.ts

@@ -26,9 +26,15 @@ function rebound(node: IrStmt | IrExpr): string | null {
  * closure then allocate no box. The proof follows each capture into nested
  * function bodies; class environments keep the eager box. */
 export class LazyCaptures {
-  private readonly unchanged = new Map<string, boolean>();
   private readonly stable = new Map<string, Set<string>>();
   private readonly once = new Map<string, Set<string>>();
+  private readonly unchangedParameters = new Map<IrFunction, Set<string>>();
+  /** Per function, every local rebound there or through nested captures;
+   * the second map counts declarations separately (see rebindings). Each
+   * body is walked once per map, so the queries stay linear in the IR. */
+  private readonly rebound = new Map<IrFunction, Set<string>>();
+  private readonly reboundAfterDeclaration = new Map<IrFunction, Set<string>>();
+  private readonly pending = new Set<IrFunction>();
   private sites: Map<string, { owner: string; captures: string[] }[]> | null = null;
 
   constructor(private readonly functions: ReadonlyMap<string, IrFunction>) {}
@@ -112,8 +118,11 @@ export class LazyCaptures {
     return sites;
   }
 
-  parameters(fn: IrFunction): Set<string> {
+  parameters(fn: IrFunction): ReadonlySet<string> {
+    const known = this.unchangedParameters.get(fn);
+    if (known !== undefined) return known;
     const result = new Set<string>();
+    this.unchangedParameters.set(fn, result);
     if (fn.async || fn.generator) return result;
     const locals = new Map(fn.locals.map((local) => [local.id, local]));
     for (const param of fn.params) {
@@ -132,18 +141,33 @@ export class LazyCaptures {
     ids: ReadonlySet<string>,
     declarationsBind = false,
   ): Set<string> {
+    const all = this.reboundIn(fn, declarationsBind);
     const changed = new Set<string>();
+    for (const id of ids) if (all.has(id)) changed.add(id);
+    return changed;
+  }
+
+  /** Every local rebound in `fn` or, through the closures it creates, in
+   * any function that captures it, directly or through further nesting.
+   * With `declarationsBind`, declarations are counted separately. */
+  private reboundIn(fn: IrFunction, declarationsBind: boolean): ReadonlySet<string> {
+    const memo = declarationsBind ? this.reboundAfterDeclaration : this.rebound;
+    const known = memo.get(fn);
+    if (known !== undefined) return known;
+    const changed = new Set<string>();
+    // Capture chains follow lexical nesting and cannot cycle; on malformed
+    // input a function reached again while its walk is pending reports
+    // every capture as rebound (see isUnchanged).
+    if (!declarationsBind) this.pending.add(fn);
     const nested: { target: string; local: string; source: string }[] = [];
     const visit = (node: IrStmt | IrExpr): boolean => {
-      // With `declarationsBind`, declarations are counted separately.
       const target = declarationsBind && node.kind === "varDecl" ? null : rebound(node);
-      if (target !== null && ids.has(target)) changed.add(target);
+      if (target !== null) changed.add(target);
       if (node.kind === "classRef") {
-        for (const id of node.captures ?? []) if (ids.has(id)) changed.add(id);
+        for (const id of node.captures ?? []) changed.add(id);
       } else if (node.kind === "closure") {
         const callee = this.functions.get(node.fnName);
         node.captures.forEach((id, index) => {
-          if (!ids.has(id)) return;
           const capture = callee?.captures?.[index];
           if (!capture) changed.add(id);
           else nested.push({ target: node.fnName, local: capture.localId, source: id });
@@ -156,19 +180,14 @@ export class LazyCaptures {
       if (!changed.has(edge.source) && !this.isUnchanged(edge.target, edge.local))
         changed.add(edge.source);
     }
+    if (!declarationsBind) this.pending.delete(fn);
+    memo.set(fn, changed);
     return changed;
   }
 
   private isUnchanged(fnName: string, localId: string): boolean {
-    const key = `${fnName}\0${localId}`;
-    const known = this.unchanged.get(key);
-    if (known !== undefined) return known;
-    // Capture chains follow lexical nesting and cannot cycle; the
-    // provisional entry only bounds work on malformed input.
-    this.unchanged.set(key, false);
     const fn = this.functions.get(fnName);
-    const result = fn !== undefined && this.rebindings(fn, new Set([localId])).size === 0;
-    this.unchanged.set(key, result);
-    return result;
+    if (fn === undefined || this.pending.has(fn)) return false;
+    return !this.reboundIn(fn, false).has(localId);
   }
 }