Keep a callback's captured scalars in registers in loops that call it (#1818)

packages/compiler/src/backend/llvm/emit-function.ts

@@ -83,6 +83,8 @@ export function emitFunction(emitter: LlEmitter, fn: IrFunction): string {
   }
   emitter.currentConstantCallbacks =
     emitter.debug === null ? (emitter.constantCallbacks.get(fn.name) ?? new Map()) : new Map();
+  emitter.knownCallbacks.clear();
+  emitter.callbackVersionsEnabled = emitter.debug === null && !fn.async && !fn.generator;
   // Preserve concrete source bindings for debugger inspection. Suspended
   // functions keep the established array lifetime across continuations.
   emitter.streamingSplitsEnabled = emitter.debug === null && !fn.async && !fn.generator;

packages/compiler/src/backend/llvm/emit-statements.ts

@@ -55,6 +55,7 @@ import { emitBytesLoopVersionGuard } from "./bytes-loop-versions.js";
 import { mangleLocal } from "../mangle.js";
 import { LlvmUnsupportedError } from "./unsupported.js";
 import type { LlValue } from "./expr-context.js";
+import { everyExprChild, everyStmtChild } from "../../ir/traverse.js";
 import { widenInteger } from "./integer-values.js";
 import { SCALAR_UNION_RESULT } from "./int32-abi.js";
 import { emitLiteralSwitch } from "./switch-dispatch.js";
@@ -76,6 +77,9 @@ import {
 } from "./string-construction.js";
 import { emitStringInputs } from "./string-lifetimes.js";
 
+/** A callback loop version copies the loop; keep the copies small. */
+const MAX_CALLBACK_VERSION_COST = 400;
+
 export function emitStmtBody(emitter: LlEmitter, s: IrStmt): void {
   const B = emitter.B;
   emitter.frames.push([]);
@@ -795,6 +799,11 @@ export function emitStmtBody(emitter: LlEmitter, s: IrStmt): void {
       break;
     }
     case "while": {
+      const callbacks = callbackVersionBindings(emitter, s);
+      if (callbacks) {
+        emitCallbackVersionedLoop(emitter, s, callbacks);
+        break;
+      }
       const workerBudget = emitter.workerLoopBudget(s);
       const le = B.newLabel("loop.e");
       emitter.emitLoopCopies(s, true, le, (exit) => {
@@ -828,6 +837,11 @@ export function emitStmtBody(emitter: LlEmitter, s: IrStmt): void {
       break;
     }
     case "doWhile": {
+      const callbacks = callbackVersionBindings(emitter, s);
+      if (callbacks) {
+        emitCallbackVersionedLoop(emitter, s, callbacks);
+        break;
+      }
       const workerBudget = emitter.workerLoopBudget(s);
       // Body first (runs at least once); continue jumps to the CONDITION.
       const le = B.newLabel("loop.e");
@@ -857,6 +871,11 @@ export function emitStmtBody(emitter: LlEmitter, s: IrStmt): void {
       break;
     }
     case "for": {
+      const callbacks = callbackVersionBindings(emitter, s);
+      if (callbacks) {
+        emitCallbackVersionedLoop(emitter, s, callbacks);
+        break;
+      }
       // The init's scope wraps the whole loop (break/continue must NOT
       // release it — scopeDepth captured after the push, C parity).
       const integerArrayLoop =
@@ -1888,6 +1907,101 @@ export function emitSwitch(emitter: LlEmitter, s: IrStmt & { kind: "switch" }):
   B.startBlock(end);
 }
 
+/** The loop copies a callback loop version needs: the bindings the loop
+ * calls that the constant-callback census resolves to one function, each
+ * with that function's direct entry. Only small loops of optimized
+ * synchronous functions, and only bindings declared before the loop,
+ * qualify: the guard reads each binding before the loop runs. Versions do
+ * not nest: a loop inside either copy of a version (including the
+ * versioned loop itself, emitted again for each copy) keeps its calls as
+ * they are, so a loop nest is emitted at most twice. */
+function callbackVersionBindings(
+  emitter: LlEmitter,
+  s: IrStmt & { kind: "for" | "while" | "doWhile" },
+): { ref: IrExpr & { kind: "varRef" }; direct: string }[] | null {
+  if (
+    !emitter.callbackVersionsEnabled ||
+    emitter.currentConstantCallbacks.size === 0 ||
+    emitter.callbackVersioning.size > 0
+  )
+    return null;
+  const found = new Map<string, { ref: IrExpr & { kind: "varRef" }; direct: string }>();
+  // Bindings the loop declares or writes (a callback declared in its body).
+  const bound = new Set<string>();
+  let cost = 0;
+  const expr = (e: IrExpr): boolean => {
+    if (++cost > MAX_CALLBACK_VERSION_COST) return false;
+    if (e.kind === "assignExpr") bound.add(e.localId);
+    if (e.kind === "callValue" && e.callee.kind === "varRef" && !found.has(e.callee.localId)) {
+      const local = emitter.currentLocals.get(e.callee.localId);
+      const target =
+        local && local.tdz !== true ? emitter.knownCallbackTarget(e.callee) : undefined;
+      if (target) found.set(e.callee.localId, { ref: e.callee, direct: target.direct });
+    }
+    return everyExprChild(e, expr, stmt);
+  };
+  const stmt = (x: IrStmt): boolean => {
+    if (++cost > MAX_CALLBACK_VERSION_COST) return false;
+    if (x.kind === "varDecl" || x.kind === "assign" || x.kind === "forOf") bound.add(x.localId);
+    else if (x.kind === "tryCatch" && x.catchLocalId !== null) bound.add(x.catchLocalId);
+    return everyStmtChild(x, expr, stmt);
+  };
+  if (!stmt(s)) return null;
+  const bindings = [...found.values()].filter(({ ref }) => !bound.has(ref.localId));
+  return bindings.length > 0 ? bindings : null;
+}
+
+/** A callback loop version. The guard tests, once before the loop, that
+ * each binding's closure holds its known function: a closure's code
+ * pointer is written only when the closure is built, and the census
+ * resolves only bindings the function never rebinds. The copy where the
+ * guard holds calls those functions directly, with no indirect fallback
+ * that could rewrite the closure or what it captures, so LLVM can keep a
+ * captured scalar the callback updates in a register. The other copy is the
+ * ordinary loop. */
+function emitCallbackVersionedLoop(
+  emitter: LlEmitter,
+  s: IrStmt & { kind: "for" | "while" | "doWhile" },
+  bindings: { ref: IrExpr & { kind: "varRef" }; direct: string }[],
+): void {
+  const B = emitter.B;
+  let ok = "true";
+  for (const { ref, direct } of bindings) {
+    // A borrowed read: the guard only compares the code pointer.
+    const closure = emitter.emitReadReceiver(ref).name;
+    const fnp = B.tmp(),
+      fn = B.tmp(),
+      same = B.tmp();
+    B.line(`${fnp} = getelementptr inbounds %ScrClosure, ptr ${closure}, i64 0, i32 1`);
+    B.line(`${fn} = load ptr, ptr ${fnp}`);
+    B.line(`${same} = icmp eq ptr ${fn}, ${direct}`);
+    if (ok === "true") ok = same;
+    else {
+      const both = B.tmp();
+      B.line(`${both} = and i1 ${ok}, ${same}`);
+      ok = both;
+    }
+  }
+  const fast = B.newLabel("callback.version.fast"),
+    slow = B.newLabel("callback.version.slow"),
+    done = B.newLabel("callback.version.done");
+  B.condBr(ok, fast, slow);
+  emitter.callbackVersioning.add(s);
+  try {
+    B.startBlock(slow);
+    emitter.emitStmt(s);
+    B.br(done);
+    B.startBlock(fast);
+    for (const { ref } of bindings) emitter.knownCallbacks.add(ref.localId);
+    emitter.emitStmt(s);
+    B.br(done);
+  } finally {
+    for (const { ref } of bindings) emitter.knownCallbacks.delete(ref.localId);
+    emitter.callbackVersioning.delete(s);
+  }
+  B.startBlock(done);
+}
+
 /** Entries in a row whose fast copy bailed before a speculative loop site
  * stops trying it. */
 const SPECULATION_BAILS = 8;

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

@@ -831,6 +831,12 @@ export class LlEmitter {
   readonly callLifetimes: CallLifetimes;
   readonly constantCallbacks: Map<string, Map<string, string>>;
   currentConstantCallbacks: ReadonlyMap<string, string> = new Map();
+  /** Callback bindings whose closures an enclosing callback loop version
+   * proved to hold their known function: calls take the direct entry with
+   * no fallback (emitCallbackVersionedLoop). */
+  knownCallbacks = new Set<string>();
+  readonly callbackVersioning = new Set<IrStmt>();
+  callbackVersionsEnabled = false;
   readonly stackCallbacks: StackCallbacks;
   readonly receiverReaders: AmbientReceiverReaders;
   readonly stackCaptures: StackCaptures;
@@ -3912,6 +3918,25 @@ export class LlEmitter {
     return b;
   }
 
+  /** The function a call through `callee` is known to reach (a literal
+   * closure, or a binding the constant-callback census resolves) and the
+   * entry it calls directly. Async and generator functions stay indirect. */
+  knownCallbackTarget(callee: IrExpr): { known: IrFunction; direct: string } | undefined {
+    const knownName =
+      callee.kind === "closure"
+        ? callee.fnName
+        : callee.kind === "varRef"
+          ? this.currentConstantCallbacks.get(callee.localId)
+          : undefined;
+    const known = knownName ? this.fnByName.get(knownName) : undefined;
+    if (!known || known.async || known.generator) return undefined;
+    if (known.captures === undefined) {
+      this.fnValues.add(known.name);
+      return { known, direct: `@${mangleWrapper(known.name)}` };
+    }
+    return { known, direct: `@${this.callTarget(known.name)}` };
+  }
+
   /** The payload slot of a capture box (ScrBox.slot, field 5). */
   boxSlot(box: string): string {
     const slot = this.B.tmp();
@@ -3925,12 +3950,18 @@ export class LlEmitter {
     const r = B.tmp();
     // Scalar payloads are read in place (scr_box_get_f64/bool are a
     // memcpy and a nonzero test of the same slot), so LLVM can keep a
-    // captured number in a register across a loop of calls.
+    // captured number in a register across a loop of calls. They have an
+    // alias category of their own: every other access to a box's payload
+    // is untagged or a runtime call.
     if (acc === "f64") {
-      B.line(`${r} = load double, ptr ${this.boxSlot(box)}`);
+      B.line(
+        `${r} = load double, ptr ${this.boxSlot(box)}${this.memoryAliasAttachment("box:scalar")}`,
+      );
     } else if (acc === "bool") {
       const raw = B.tmp();
-      B.line(`${raw} = load i64, ptr ${this.boxSlot(box)}`);
+      B.line(
+        `${raw} = load i64, ptr ${this.boxSlot(box)}${this.memoryAliasAttachment("box:scalar")}`,
+      );
       B.line(`${r} = icmp ne i64 ${raw}, 0`);
     } else {
       this.declare(`declare ptr @scr_box_get_ref(ptr)`);
@@ -3944,11 +3975,15 @@ export class LlEmitter {
     const B = this.B;
     const acc = boxAccess(t);
     if (acc === "f64") {
-      B.line(`store double ${value}, ptr ${this.boxSlot(box)}`);
+      B.line(
+        `store double ${value}, ptr ${this.boxSlot(box)}${this.memoryAliasAttachment("box:scalar")}`,
+      );
     } else if (acc === "bool") {
       const raw = B.tmp();
       B.line(`${raw} = zext i1 ${value} to i64`);
-      B.line(`store i64 ${raw}, ptr ${this.boxSlot(box)}`);
+      B.line(
+        `store i64 ${raw}, ptr ${this.boxSlot(box)}${this.memoryAliasAttachment("box:scalar")}`,
+      );
     } else {
       this.declare(`declare void @scr_box_set_ref(ptr, ptr)`);
       B.line(`call void @scr_box_set_ref(ptr ${box}, ptr ${value})`);

packages/compiler/src/backend/llvm/expr-calls.ts

@@ -29,7 +29,6 @@ import {
   mangleClassStruct,
   mangleFnClosure,
   mangleFunction,
-  mangleWrapper,
   mangleVtStruct,
   mangleAsyncDirectSpawn,
 } from "../mangle.js";
@@ -839,21 +838,9 @@ export function emitCallExpr(
         host.declare("declare void @scr_dyn_this_push_dyn(ptr)");
         B.line(`call void @scr_dyn_this_push_dyn(ptr ${receiver?.name ?? "null"})`);
       };
-      const knownName =
-        e.callee.kind === "closure"
-          ? e.callee.fnName
-          : e.callee.kind === "varRef"
-            ? host.currentConstantCallbacks.get(e.callee.localId)
-            : undefined;
-      const candidate = knownName ? host.fnByName.get(knownName) : undefined;
-      const known = candidate && !candidate.async && !candidate.generator ? candidate : undefined;
-      let direct: string | undefined;
-      if (known) {
-        if (known.captures === undefined) {
-          host.fnValues.add(known.name);
-          direct = `@${mangleWrapper(known.name)}`;
-        } else direct = `@${host.callTarget(known.name)}`;
-      }
+      const target = host.knownCallbackTarget(e.callee);
+      const known = target?.known;
+      const direct = target?.direct;
       // A known closure body that borrows reference parameters is entered
       // directly with lent arguments, as a direct call to a borrowing body
       // is: an argument is read without a reference when neither a later
@@ -939,7 +926,13 @@ export function emitCallExpr(
       let pending = host.indirectMayThrow;
       let pendingGuard: string | undefined;
       let t: string;
-      if (direct && e.callee.kind === "closure") {
+      // A literal closure, or a binding an enclosing callback loop version
+      // proved to hold its known function, enters it directly.
+      if (
+        direct &&
+        (e.callee.kind === "closure" ||
+          (e.callee.kind === "varRef" && host.knownCallbacks.has(e.callee.localId)))
+      ) {
         t = invoke(lending ?? direct, !host.receiverReaders.ignores(known!));
         if (quiet) pending = false;
       } else {

packages/compiler/test/callback-loop-versions.test.ts

@@ -0,0 +1,114 @@
+import { mkdtemp, readFile, rm, writeFile } from "node:fs/promises";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import { expect, test } from "vitest";
+import { compile } from "../src/index.js";
+
+async function emitLlvm(source: string): Promise<string> {
+  const dir = await mkdtemp(join(tmpdir(), "scriptc-callback-versions-"));
+  try {
+    const entry = join(dir, "main.ts");
+    const outPath = join(dir, "main.ll");
+    await writeFile(entry, source);
+    const result = await compile(entry, { outDir: dir, outPath, outputKind: "llvm" });
+    if (!result.ok) throw new Error(result.diagnostics.map((d) => d.message).join("\n"));
+    return await readFile(outPath, "utf8");
+  } finally {
+    await rm(dir, { recursive: true, force: true });
+  }
+}
+
+const body = (llvm: string, marker: string): string => {
+  const match = new RegExp(`^define internal [^\\n]*${marker}[^]*?^}`, "m").exec(llvm);
+  if (!match) throw new Error(`no function ${marker}`);
+  return match[0];
+};
+
+test("loops that call a known callback test its closure once and call it directly", async () => {
+  const llvm = await emitLlvm(`
+    function each(values: number[], visit: (v: number) => void): void {
+      for (let i = 0; i < values.length; i++) visit(values[i]!);
+    }
+    let sum = 0;
+    each([1, 2, 3], (v) => { sum += v; });
+    console.log(sum);
+  `);
+  const fn = body(llvm, "@sc_bf_each");
+  const fast = /^callback\.version\.fast\d+:$[^]*?^callback\.version\.done\d+:$/m.exec(fn)![0];
+  // The guard compares the closure's code pointer once, before the loop.
+  expect(fn.slice(0, fn.indexOf(fast))).toMatch(/icmp eq ptr %\S+, @/);
+  expect(fast).not.toContain("callback.indirect");
+  expect(fast).not.toContain("scr_dyn_this_push_dyn");
+  // The ordinary copy keeps the fallback through the closure's pointer.
+  expect(fn).toContain("callback.indirect");
+});
+
+test("captured number payloads carry their own alias tag", async () => {
+  const llvm = await emitLlvm(`
+    function total(values: number[]): number {
+      let sum = 0;
+      values.forEach((v) => { sum += v; });
+      return sum;
+    }
+    console.log(total([1, 2, 3]));
+  `);
+  // Every read and write through a box's payload slot (field 5) is tagged.
+  const slots = [
+    ...llvm.matchAll(/(%\S+) = getelementptr inbounds %ScrBox, ptr %\S+, i32 0, i32 5/g),
+  ];
+  expect(slots.length).toBeGreaterThan(0);
+  for (const [, slot] of slots) {
+    for (const access of llvm.matchAll(
+      new RegExp(`(load double, ptr|store double %\\S+, ptr) \\${slot}\\b[^\\n]*`, "g"),
+    ))
+      expect(access[0]).toMatch(/!tbaa !\d+$/);
+  }
+  expect(llvm).toMatch(/load double, ptr %\S+, !tbaa !\d+/);
+});
+
+test("a callback declared in the loop body is not read before the loop", async () => {
+  const llvm = await emitLlvm(`
+    function strides(): number {
+      let calls = 0;
+      for (let i = 0; i < 10; i++) {
+        const bump = (): void => { i += 2; };
+        bump();
+        calls++;
+      }
+      return calls;
+    }
+    console.log(strides());
+  `);
+  expect(llvm).not.toContain("callback.version");
+});
+
+test("loops inside a callback version are not versioned again", async () => {
+  // Each inner loop calls a callback a factory builds just before it; without
+  // the limit, every copy of a version would version the next loop again.
+  const llvm = await emitLlvm(`
+    function adder(k: number): (v: number) => number {
+      return (v: number) => v + k;
+    }
+    function nest(n: number, f: (v: number) => number): number {
+      let s = 0;
+      for (let i = 0; i < n; i++) {
+        s += f(i);
+        const g = adder(i);
+        for (let j = 0; j < n; j++) {
+          s += g(j);
+          const h = adder(j);
+          for (let k = 0; k < n; k++) s += h(k);
+        }
+      }
+      return s;
+    }
+    console.log(nest(3, (v) => v - 1), nest(2, (v) => v * 3));
+  `);
+  // One version for the outer loop of each compiled `nest`; its two copies
+  // keep the inner loops whole.
+  const versions = [...llvm.matchAll(/^define [^\n]*@sc_bf_nest[^]*?^}/gm)].map(
+    ([body]) => body.match(/^callback\.version\.fast\d+:$/gm)?.length ?? 0,
+  );
+  expect(versions.every((count) => count <= 1)).toBe(true);
+  expect(versions.some((count) => count === 1)).toBe(true);
+});

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/7300-callback-captured-scalars.ts.json

@@ -0,0 +1,6 @@
+{
+ "order": [
+  "<repo>/tests/corpus/7300-callback-captured-scalars.ts"
+ ],
+ "diags": []
+}

tests/corpus/7300-callback-captured-scalars.ts

@@ -0,0 +1,147 @@
+// Callbacks that update captured `let` numbers and booleans from loops: typed
+// array and array methods, an authored higher-order function called with one
+// function at several calls and with two functions at one, closures of one
+// function over different bindings, a callback that throws part way, nested
+// callbacks, callbacks declared in the loop body, and reads after the loop.
+
+function each(values: number[], visit: (v: number, i: number) => void): void {
+  for (let i = 0; i < values.length; i++) visit(values[i]!, i);
+}
+
+function until(limit: number, step: (i: number) => boolean): number {
+  let i = 0;
+  while (i < limit) {
+    if (!step(i)) break;
+    i++;
+  }
+  return i;
+}
+
+function twice(values: Int16Array, f: (v: number) => void): void {
+  let i = 0;
+  do {
+    for (let j = 0; j < values.length; j++) f(values[j]!);
+    i++;
+  } while (i < 2);
+}
+
+const ints = new Int32Array([5, -3, 8, 0, 12, -7, 4]);
+let total = 0;
+let negatives = 0;
+let sawZero = false;
+ints.forEach((v) => {
+  total = (total + v) | 0;
+  if (v < 0) negatives++;
+  if (v === 0) sawZero = true;
+});
+console.log("forEach", total, negatives, sawZero);
+
+let max = -Infinity;
+let maxAt = -1;
+const floats = new Float64Array([0.5, 9.25, -2, 9.25, 3]);
+floats.forEach((v, i) => {
+  if (v > max) {
+    max = v;
+    maxAt = i;
+  }
+});
+console.log("max", max, maxAt);
+
+let calls = 0;
+const doubled = ints.map((v) => {
+  calls++;
+  return v * 2 + calls;
+});
+console.log("map", doubled.join(","), calls);
+
+let seen = 0;
+const big = [3, 14, 15, 92, 65].filter((v) => {
+  seen += v;
+  return v > 10;
+});
+console.log("filter", big.join(","), seen);
+
+let sum = 0;
+let count = 0;
+const summing = (v: number): void => {
+  sum += v;
+  count++;
+};
+each([1, 2, 3], summing);
+each([10, 20], summing);
+console.log("each shared", sum, count);
+
+let a = 0;
+let b = 0;
+each([1, 2], (v) => {
+  a += v;
+});
+each([1, 2], (v, i) => {
+  b += v * (i + 1);
+});
+console.log("each two", a, b);
+
+const counters: (() => number)[] = [];
+for (let k = 0; k < 3; k++) {
+  let local = k * 100;
+  each([1, 1, 1], () => {
+    local++;
+  });
+  counters.push(() => local);
+}
+console.log("closures", counters.map((c) => c()).join(","));
+
+let steps = 0;
+console.log("until", until(10, (i) => (steps += i, steps < 12)), steps);
+
+let shorts = 0;
+twice(new Int16Array([7, -8, 9]), (v) => {
+  shorts = shorts * 2 + v;
+});
+console.log("twice", shorts);
+
+let partial = 0;
+try {
+  floats.forEach((v) => {
+    if (v < 0) throw new RangeError(`at ${partial}`);
+    partial += v;
+  });
+} catch (error) {
+  console.log("throw", (error as Error).message, partial);
+}
+
+let outer = 0;
+let inner = 0;
+[1, 2, 3].forEach((x) => {
+  outer += x;
+  ints.forEach((y) => {
+    if (y > x) inner++;
+  });
+});
+console.log("nested", outer, inner);
+
+let flag = false;
+let flips = 0;
+new Uint8Array([1, 0, 1, 1]).forEach((v) => {
+  if (v === 1) flag = !flag;
+  if (flag) flips++;
+});
+console.log("bool", flag, flips);
+
+// Callbacks declared in the loop body: each iteration binds its own.
+let strides = 0;
+for (let i = 0; i < 10; i++) {
+  const bump = (): void => {
+    i += 2;
+    strides++;
+  };
+  bump();
+}
+let countdown = 0;
+let k = 3;
+while (k > 0) {
+  const tick = (n: number): number => (n > 0 ? tick(n - 1) + 1 : 0);
+  countdown += tick(k);
+  k--;
+}
+console.log("declared in loop", strides, countdown);

tests/harness/corpus-costs/7300-callback-captured-scalars.ts.json

@@ -0,0 +1 @@
+{ "plain": 6000, "san": 7900 }