Call typed-array callbacks directly and visit on an integer counter

The helpers of the typed-array callback methods (forEach, map, filter,
some, every, find, findIndex, findLast, findLastIndex, reduce,
reduceRight) called their callback through the closure's code pointer on
every element, and one helper served every call with the same signature,
so the callback was never inlined. A literal callback now gets a helper
of its own, as for the array methods, and the constant-callback census
accepts the typed-array helpers, so the helper calls that function
directly.

The visit is now a loop on its own counter, `for (let i = 0; i < a.length;
i++)` or down from the last index, instead of a loop on a counter
declared before it, so the counter is an integer and element reads need no
index test. reduce and reduceRight without an initial value read their
seed at the first visited index and start one past it.

Adds reduce, some and map callbacks to the typed-array-kernels workload.

benchmarks/runtime/README.md

@@ -26,7 +26,7 @@ Application-shaped workloads for measuring the speed of compiled executables. Ea
 | `region-builds` | cpu | a watch-mode build loop over `LifetimeRegion`: each build binds a program in a new region and releases the previous one, which a cache keeps referenced for two builds |
 | `template-render` | cpu | template literals, escaping, `map`/`join` |
 | `thread-slots`, `thread-slots-cache` | cpu | worker threads reading one published type graph from an `InternTable` and keeping per-thread state for each type they visit: a `ThreadSlots` family, or a direct-mapped cache by id in front of a `Map` by id (the `-cache` entry); under Node.js both are per thread |
-| `typed-array-kernels` | cpu | element-wise loops that run to one typed array's length and index others at the same position: `Float64Array` axpy, `Float32Array` blend, `Uint8Array` mask, wrapping `Int32Array` add, `Uint8ClampedArray` brighten, RGBA to gray |
+| `typed-array-kernels` | cpu | element-wise loops that run to one typed array's length and index others at the same position: `Float64Array` axpy, `Float32Array` blend, `Uint8Array` mask, wrapping `Int32Array` add, `Uint8ClampedArray` brighten, RGBA to gray, and `reduce`/`some`/`map` callbacks |
 | `typed-array-scan`, `typed-array-scan-workers` | cpu | element-by-element scan loops over `Uint8Array`, `Uint16Array`, `Int32Array` and `Float64Array`; the second entry is a worker-enabled executable (it can start a worker), so each access also checks for shared storage |
 | `unicode-lexer` | cpu | a hand-written lexer over a 2.7M-unit bundled module whose comments and strings are in several languages: `charCodeAt` scans, line starts |
 | `validate-errors` | cpu | throw/catch on a fraction of inputs, small call chains |

benchmarks/runtime/workloads/typed-array-kernels.ts

@@ -21,6 +21,14 @@ function brighten(src: Uint8ClampedArray, gain: number, out: Uint8ClampedArray):
   for (let i = 0; i < out.length; i++) out[i] = src[i]! * gain + 8;
 }
 
+// Callback methods over a whole typed array.
+function callbacks(samples: Float32Array, limit: number, data: Uint8Array): number {
+  const energy = samples.reduce((sum, v) => sum + v * v, 0);
+  const loud = samples.some((v) => v > limit) ? 1 : 0;
+  const halved = data.map((v) => v >> 1);
+  return energy / samples.length + loud + halved[halved.length >> 1]!;
+}
+
 function gray(rgba: Uint8Array, out: Uint8Array): void {
   for (let i = 0; i < out.length; i++) {
     const j = i * 4;
@@ -73,6 +81,7 @@ let byteSum = 0;
 let intSum = 0;
 let brightSum = 0;
 let lumaSum = 0;
+let callbackSum = 0;
 for (let r = 0; r < rounds; r++) {
   // Change one input element per round so no round repeats the last.
   const at = (r * 7919) & (n - 1);
@@ -88,6 +97,7 @@ for (let r = 0; r < rounds; r++) {
   addWrap(ia, ib, io);
   brighten(pixels, 0.75 + (r & 7) / 8, bright);
   gray(rgba, luma);
+  if ((r & 3) === 0) callbackSum += callbacks(fo, 140 + (r & 12), masked);
   const probe = (r * 104729) & (n - 1);
   f64Sum += d[probe]!;
   f32Sum += fo[probe]!;
@@ -102,3 +112,4 @@ console.log(`uint8 ${byteSum}`);
 console.log(`int32 ${intSum}`);
 console.log(`clamped ${brightSum}`);
 console.log(`gray ${lumaSum}`);
+console.log(`callbacks ${callbackSum.toFixed(3)}`);

packages/compiler/src/backend/llvm/constant-callbacks.test.ts

@@ -314,16 +314,19 @@ test("virtual calls admit every same-named method", () => {
   expect(new AmbientReceiverReaders(functions).ignores(caller)).toBe(false);
 });
 
-test("synthetic array HOF helpers are constant-callback consumers", () => {
-  const mod = fixture();
-  const consumer = mod.functions[1]!;
-  const entry = mod.functions[2]!;
-  delete consumer.ownsPrototype;
-  consumer.name = "%arr.some.0";
-  const call = entry.body[0]!;
-  if (call.kind === "exprStmt" && call.expr.kind === "call") call.expr.callee = "%arr.some.0";
-  expect(facts(mod).get("%arr.some.0")?.get("test")).toBe("predicate");
-});
+test.each(["%arr.some.0", "%bytes.some.0"])(
+  "synthetic array and typed-array HOF helpers are constant-callback consumers (%s)",
+  (name) => {
+    const mod = fixture();
+    const consumer = mod.functions[1]!;
+    const entry = mod.functions[2]!;
+    delete consumer.ownsPrototype;
+    consumer.name = name;
+    const call = entry.body[0]!;
+    if (call.kind === "exprStmt" && call.expr.kind === "call") call.expr.callee = name;
+    expect(facts(mod).get(name)?.get("test")).toBe("predicate");
+  },
+);
 
 function factory(name: string, returns: IrExpr[]): IrFunction {
   return {

packages/compiler/src/backend/llvm/constant-callbacks.ts

@@ -29,12 +29,16 @@ export function findConstantCallbacks(
     mod.functions
       .filter(
         (fn) =>
-          // Authored declarations, plus the frontend's synthetic array
-          // HOF loops (`%arr.<method>.<n>`): those are plain direct-call
-          // functions, and a literal callback gets a helper of its own.
+          // Authored declarations, plus the frontend's synthetic array and
+          // typed-array HOF loops (`%arr.<method>.<n>`, `%bytes.<method>.<n>`):
+          // those are plain direct-call functions, and a literal callback
+          // gets a helper of its own.
           // Generator declarations are entered through their spawn wrapper,
           // which moves the same arguments into the generator's frame.
-          (fn.ownsPrototype || fn.name.startsWith("%arr.") || fn.generator !== undefined) &&
+          (fn.ownsPrototype ||
+            fn.name.startsWith("%arr.") ||
+            fn.name.startsWith("%bytes.") ||
+            fn.generator !== undefined) &&
           (!fn.async || fn.generator !== undefined) &&
           !fn.captures &&
           !fn.classCaptures &&

packages/compiler/src/frontend/lowering/containers/array-callbacks.ts

@@ -452,7 +452,7 @@ function withFastCallback(fn: IrFunction, fastT: CallbackType | null, elem: IrTy
  * constant-callback census then calls it directly and LLVM can inline the
  * callback into the loop (TurboFan's builtin+callback inlining). Callback
  * values keep the shared, type-interned helper. */
-function callbackSite(fnArg: IrExpr): string {
+export function callbackSite(fnArg: IrExpr): string {
   return fnArg.kind === "closure" ? `@${fnArg.fnName}` : "";
 }
 

packages/compiler/src/frontend/lowering/containers/bytes.ts

@@ -43,7 +43,7 @@ import { lowerArraySpreadItems, lowerCheckedArrayFrom } from "./array-constructi
 import { lowerToNumberArgument } from "../lower-exprs.js";
 import { lowerBytesIteratorValue } from "../iterator-adapters.js";
 import { lowerArrayCallback } from "./callback-arguments.js";
-import { findPredicateOk, predicateToBoolean } from "./array-callbacks.js";
+import { callbackSite, findPredicateOk, predicateToBoolean } from "./array-callbacks.js";
 import { typeKey } from "../../type-mapper.js";
 
 /** TypedArray.prototype.sort and toSorted. The receiver and comparator
@@ -251,13 +251,28 @@ function lowerBytesCallbackCall(
           : method === "findIndex" || method === "findLastIndex"
             ? F64
             : lowerer.irTypeOf(call));
-  const helper = bytesCallbackHelper(lowerer, method, bytesT, resultT, fnRet, arity, hasInit, loc);
+  const helper = bytesCallbackHelper(
+    lowerer,
+    method,
+    bytesT,
+    resultT,
+    fnRet,
+    arity,
+    hasInit,
+    loc,
+    callbackSite(fnArg),
+  );
   const args: IrExpr[] = [receiver, fnArg];
   if (hasInit) args.push(lowerer.lowerExprExpecting(call.arguments.nodes[1]!, resultT));
   return { kind: "call", callee: helper, args, type: resultT, loc };
 }
 
-/** The interned helper of one typed-array callback method shape. */
+/** The interned helper of one typed-array callback method shape. A literal
+ * callback gets a helper of its own (`site`), which the backend then calls
+ * directly, as for the array methods. The visit is a loop on its own
+ * counter (`for (let i = 0; i < a.length; i++)`, or down from the last
+ * index), so the counter is an integer and the element reads need no
+ * index test. */
 function bytesCallbackHelper(
   lowerer: Lowerer,
   method: BytesCallbackMethod,
@@ -267,8 +282,9 @@ function bytesCallbackHelper(
   arity: number,
   hasInit: boolean,
   loc: SrcLoc,
+  site: string,
 ): string {
-  const key = `bytes.${method}:${bytesT.elem}:${typeKey(resultT)}:${typeKey(fnRet)}:${arity}:${hasInit ? "init" : ""}`;
+  const key = `bytes.${method}:${bytesT.elem}:${typeKey(resultT)}:${typeKey(fnRet)}:${arity}:${hasInit ? "init" : ""}${site}`;
   const existing = lowerer.arrHofHelpers.get(key);
   if (existing) return existing;
   const name = keyedHelperName(lowerer, `%bytes.${method}`, key);
@@ -300,10 +316,20 @@ function bytesCallbackHelper(
     type: fnRet,
     loc,
   });
+  // The forward bound is the receiver's length itself (fixed for a typed
+  // array, and equal to `n`), which the emitter's typed-array loops use.
+  const length = (): IrExpr => ({
+    kind: "bytesIntrinsic",
+    method: "length",
+    receiver: a(),
+    args: [],
+    type: F64,
+    loc,
+  });
   const inRange = (): IrExpr =>
     backward
       ? { kind: "bin", op: ">=", left: i(), right: numLit(0, loc), type: BOOL, loc }
-      : { kind: "bin", op: "<", left: i(), right: n(), type: BOOL, loc };
+      : { kind: "bin", op: "<", left: i(), right: length(), type: BOOL, loc };
   const advance = (): IrStmt => ({
     kind: "assign",
     localId: "i.0",
@@ -344,6 +370,7 @@ function bytesCallbackHelper(
     ...(hasInit ? [{ localId: "z.0", name: "z", type: resultT }] : []),
   ];
   const visit: IrStmt[] = [{ kind: "assign", localId: "v.0", value: at(i()), loc }];
+  let seeded = false;
   const prologue: IrStmt[] = [];
   let epilogue: IrStmt[];
   switch (method) {
@@ -459,10 +486,12 @@ function bytesCallbackHelper(
           loc,
         });
       } else {
+        // The seed is the first visited element; the loop starts past it.
+        seeded = true;
         prologue.push(
           {
             kind: "if",
-            cond: { kind: "unary", op: "!", operand: inRange(), type: BOOL, loc },
+            cond: { kind: "bin", op: "===", left: n(), right: numLit(0, loc), type: BOOL, loc },
             then: [
               {
                 kind: "throw",
@@ -486,8 +515,16 @@ function bytesCallbackHelper(
             else_: null,
             loc,
           },
-          { kind: "varDecl", localId: "acc.0", init: at(i()), loc },
-          advance(),
+          {
+            kind: "varDecl",
+            localId: "acc.0",
+            init: at(
+              backward
+                ? { kind: "bin", op: "-", left: n(), right: numLit(1, loc), type: F64, loc }
+                : numLit(0, loc),
+            ),
+            loc,
+          },
         );
       }
       visit.push({
@@ -506,17 +543,30 @@ function bytesCallbackHelper(
       init: { kind: "bytesIntrinsic", method: "length", receiver: a(), args: [], type: F64, loc },
       loc,
     },
-    {
-      kind: "varDecl",
-      localId: "i.0",
-      init: backward
-        ? { kind: "bin", op: "-", left: n(), right: numLit(1, loc), type: F64, loc }
-        : numLit(0, loc),
-      loc,
-    },
     { kind: "varDecl", localId: "v.0", init: numLit(0, loc), loc },
     ...prologue,
-    { kind: "for", init: null, cond: inRange(), update: advance(), body: visit, loc },
+    {
+      kind: "for",
+      init: {
+        kind: "varDecl",
+        localId: "i.0",
+        init: backward
+          ? {
+              kind: "bin",
+              op: "-",
+              left: n(),
+              right: numLit(seeded ? 2 : 1, loc),
+              type: F64,
+              loc,
+            }
+          : numLit(seeded ? 1 : 0, loc),
+        loc,
+      },
+      cond: inRange(),
+      update: advance(),
+      body: visit,
+      loc,
+    },
     ...epilogue,
   ];
   lowerer.liftedFns.push({ name, params, returnType: resultT, locals, body, loc });

packages/compiler/test/bytes-callback-helpers.test.ts

@@ -0,0 +1,68 @@
+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, deserializeModule, validateModule } from "../src/index.js";
+import type { IrFunction, IrModule, IrStmt } from "../src/ir/ir.js";
+
+async function lower(source: string): Promise<IrModule> {
+  const dir = await mkdtemp(join(tmpdir(), "scriptc-bytes-callbacks-"));
+  try {
+    const entry = join(dir, "main.ts");
+    const outPath = join(dir, "main.ir.json");
+    await writeFile(entry, source);
+    const result = await compile(entry, { outDir: dir, outPath, outputKind: "ir" });
+    if (!result.ok)
+      throw new Error(result.diagnostics.map((d) => `${d.code}: ${d.message}`).join("\n"));
+    const mod = deserializeModule(await readFile(outPath, "utf8"));
+    expect(validateModule(mod)).toEqual([]);
+    return mod;
+  } finally {
+    await rm(dir, { recursive: true, force: true });
+  }
+}
+
+const helpers = (mod: IrModule, method: string): IrFunction[] =>
+  mod.functions.filter((f) => f.name.startsWith(`%bytes.${method}.`));
+
+const visitLoop = (fn: IrFunction): IrStmt & { kind: "for" } => {
+  const loop = fn.body.find((s): s is IrStmt & { kind: "for" } => s.kind === "for");
+  if (!loop) throw new Error(`no loop in ${fn.name}`);
+  return loop;
+};
+
+test("literal callbacks get helpers of their own; function values share one", async () => {
+  const mod = await lower(`
+    const a = new Int32Array([1, 2, 3]);
+    const double = (v: number): number => v * 2;
+    console.log(a.map((v) => v + 1).join(","), a.map((v) => v - 1).join(","));
+    console.log(a.map(double).join(","), new Int32Array([4]).map(double).join(","));
+  `);
+  expect(helpers(mod, "map")).toHaveLength(3);
+});
+
+test("typed-array callback loops run on their own counter", async () => {
+  const mod = await lower(`
+    const a = new Float64Array([1.5, 2.5, 3.5]);
+    console.log(a.some((v) => v > 2), a.reduce((s, v) => s + v));
+    console.log(a.reduceRight((s, v) => s - v), a.findLastIndex((v) => v < 3));
+  `);
+  for (const method of ["some", "reduce"]) {
+    const loop = visitLoop(helpers(mod, method)[0]!);
+    // for (let i = <0 or 1 past reduce's seed>; i < a.length; i = i + 1)
+    expect(loop.init).toMatchObject({ kind: "varDecl", localId: "i.0", init: { kind: "numLit" } });
+    expect(loop.cond).toMatchObject({
+      kind: "bin",
+      op: "<",
+      right: { kind: "bytesIntrinsic", method: "length" },
+    });
+  }
+  expect(visitLoop(helpers(mod, "reduce")[0]!).init).toMatchObject({ init: { value: 1 } });
+  // Backward: from n - 2 past reduceRight's seed, and from n - 1.
+  expect(visitLoop(helpers(mod, "reduceRight")[0]!).init).toMatchObject({
+    init: { kind: "bin", op: "-", right: { value: 2 } },
+  });
+  expect(visitLoop(helpers(mod, "findLastIndex")[0]!).init).toMatchObject({
+    init: { kind: "bin", op: "-", right: { value: 1 } },
+  });
+});

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/6999-typed-array-callback-loops.ts.json

@@ -0,0 +1,6 @@
+{
+ "order": [
+  "<repo>/tests/corpus/6999-typed-array-callback-loops.ts"
+ ],
+ "diags": []
+}

tests/corpus/6999-typed-array-callback-loops.ts

@@ -0,0 +1,93 @@
+// Typed-array callback methods with literal callbacks, function values shared
+// by several calls, and callbacks chosen at run time: visit order and
+// arguments, writes during the visit, early exits, reduce seeds and the
+// empty-array TypeError, map conversions, Buffer results, throwing
+// callbacks, empty arrays and subarray receivers.
+
+const ints = new Int32Array([5, -3, 8, 0, 12, -7, 4]);
+const floats = new Float32Array([0.1, 2.5, -1.25, 3.75, 9.5]);
+
+const visits: string[] = [];
+ints.forEach((v, i, a) => {
+  visits.push(`${i}:${v}:${a === ints}`);
+});
+console.log("forEach", visits.join(" "));
+
+const writes = new Uint8Array([1, 2, 3, 4, 5, 6]);
+writes.forEach((v, i, a) => {
+  if (i + 1 < a.length) a[i + 1] = a[i + 1]! + v;
+});
+console.log("forEach writes ahead", writes.join(","));
+
+let calls = 0;
+const counted = (v: number): boolean => {
+  calls++;
+  return v > 7;
+};
+console.log("some", ints.some(counted), calls);
+calls = 0;
+console.log("every", ints.every((v) => (calls++, v < 10)), calls);
+calls = 0;
+console.log("find", ints.find(counted), calls, ints.find((v) => v > 100));
+calls = 0;
+console.log("findIndex", ints.findIndex(counted), calls, ints.findIndex((v) => v > 100));
+
+const order: number[] = [];
+console.log(
+  "findLast",
+  ints.findLast((v, i) => (order.push(i), v < 0)),
+  ints.findLastIndex((v) => v > 100),
+  order.join(","),
+);
+
+console.log("reduce", ints.reduce((acc, v) => acc + v), ints.reduce((acc, v, i) => acc + v * i, 100));
+const right: number[] = [];
+console.log(
+  "reduceRight",
+  floats.reduceRight((acc, v, i) => (right.push(i), acc - v)),
+  right.join(","),
+  floats.reduceRight((acc, v) => acc + v, -1),
+);
+console.log("single", new Int8Array([42]).reduce(() => -1), new Int8Array([42]).reduceRight(() => -1));
+function tryReduce(label: string, run: () => number): void {
+  try {
+    console.log(label, run());
+  } catch (error) {
+    console.log(label, (error as Error).name, (error as Error).message);
+  }
+}
+const emptyU16 = new Uint16Array(0);
+const emptyF64 = new Float64Array(0);
+tryReduce("empty u16 reduce", () => emptyU16.reduce((a, b) => a + b));
+tryReduce("empty f64 reduceRight", () => emptyF64.reduceRight((a, b) => a + b));
+tryReduce("empty with init", () => emptyU16.reduce((a, b) => a + b, 7) + emptyF64.reduceRight((a, b) => a + b, 1));
+console.log("empty", emptyU16.some(() => true), emptyF64.every(() => false), emptyU16.map((v) => v).length);
+
+console.log("map u8", new Uint8Array([250, 3, 128]).map((v) => v * 3).join(","));
+console.log("map clamped", new Uint8ClampedArray([250, 3, 128]).map((v) => v * 1.5 - 10).join(","));
+console.log("map i16", new Int16Array([30000, -30000, 7]).map((v, i) => v + i * 10000).join(","));
+console.log("map f32", floats.map((v) => v / 3).join(","));
+console.log("filter", ints.filter((v, i) => v > 0 && i % 2 === 0).join(","), floats.filter((v) => v > 1).join(","));
+
+const buf = Buffer.from("scriptc");
+const upper = buf.map((c) => c - 32);
+console.log("buffer map", Buffer.isBuffer(upper), upper.join(","), Buffer.isBuffer(buf.filter((c) => c > 112)));
+
+// One function value at several calls, and callbacks chosen at run time.
+const double = (v: number): number => v * 2;
+console.log("shared", new Int32Array([1, 2, 3]).map(double).join(","), new Int32Array([4, 5]).map(double).join(","));
+const pick = (k: number): ((v: number) => boolean) => (k > 0 ? (v) => v > k : (v) => v < k);
+console.log("chosen", ints.filter(pick(4)).join(","), ints.filter(pick(-4)).join(","));
+
+try {
+  let seen = 0;
+  floats.forEach((v) => {
+    seen++;
+    if (v < 0) throw new RangeError(`negative ${v} after ${seen}`);
+  });
+} catch (error) {
+  console.log("throw", (error as Error).name, (error as Error).message);
+}
+
+const view = ints.subarray(2, 5);
+console.log("subarray", view.map((v) => v + 1).join(","), view.reduce((a, v) => a * 10 + v, 0), view.findLastIndex((v) => v === 8));

tests/harness/corpus-costs/6999-typed-array-callback-loops.ts.json

@@ -0,0 +1 @@
+{ "plain": 4400, "san": 5000 }