Keep the budget in typed-array length loops that assign their array

A loop to a typed array's length reads the length again on every
iteration, so a body that assigns the array's binding changes the limit
for the next iteration. The strip-mined form took the length once per
block of 64 iterations, and after the body assigned a shorter array it
ran on to the block's end, reading past the new length.

The loop to a typed array's length is now strip-mined only when the
array is a local of the function that the body never assigns
(stableReceiver on the match); others keep the budget and test the
length on every iteration.

New corpus program: native-worker-length-loop-reassign.ts.

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

@@ -1216,14 +1216,15 @@ export function emitStmtBody(emitter: LlEmitter, s: IrStmt): void {
       // length, polls between blocks of WORKER_POLL_BLOCK iterations instead
       // of counting a budget in its body: the inner loop over a block then
       // has no exit besides its own bound, which LLVM can vectorize. A typed
-      // array's length is fixed, so each block can take it as its limit.
+      // array's length is fixed, so each block can take it as its limit when
+      // the loop does not assign the array's binding.
       const stripMined =
         capturedLets.length === 0 &&
         integerSlot !== null &&
         ((countedLoop !== null &&
           countedLimit !== null &&
           Math.abs(countedLoop.step) * WORKER_POLL_BLOCK <= 2 ** 52) ||
-          (integerLoop !== null && integerArrayLoop === null)) &&
+          (integerLoop !== null && integerArrayLoop === null && integerLoop.stableReceiver)) &&
         emitter.workerLoopPolls(s) &&
         loopAccessesTypedArrays(s);
       const workerBudget = stripMined ? null : emitter.workerLoopBudget(s);

packages/compiler/src/ir/integer-loops.ts

@@ -12,6 +12,9 @@ import { preservesArrayLength } from "./array-length-stability.js";
 export interface IntegerBytesForLoop {
   localId: string;
   limitReceiver: IrExpr;
+  /** The receiver is a local of the function that the body never assigns
+   * (a typed array's length is then the same on every iteration). */
+  stableReceiver: boolean;
 }
 
 export interface IntegerCountedForLoop {
@@ -682,7 +685,12 @@ export function matchIntegerBytesForLoop(
   if (!isUnitIncrement(stmt.update, init.localId)) return null;
   if (writesLocal(stmt.body, init.localId)) return null;
 
-  return { localId: init.localId, limitReceiver: cond.right.receiver };
+  return {
+    localId: init.localId,
+    limitReceiver: cond.right.receiver,
+    stableReceiver:
+      receiverLocal !== undefined && !writesLocal(stmt.body, cond.right.receiver.localId),
+  };
 }
 
 /** Array lengths are uint32, so zero-based loops and nested `j = i + 1`
@@ -729,5 +737,11 @@ export function matchIntegerArrayForLoop(
     writesLocal(stmt.body, local.id)
   )
     return null;
-  return { localId: local.id, limitReceiver: cond.right.receiver };
+  return {
+    localId: local.id,
+    limitReceiver: cond.right.receiver,
+    stableReceiver:
+      locals.has(cond.right.receiver.localId) &&
+      !writesLocal(stmt.body, cond.right.receiver.localId),
+  };
 }

packages/compiler/test/counted-loop-emission.test.ts

@@ -58,6 +58,14 @@ test("worker programs poll counted loops over typed arrays between blocks of ite
         "export function codes(a: Uint8Array): void {",
         "  for (let i = 0; i < a.length; ++i) a[i] = a[i] === 97 ? 0 : 3;",
         "}",
+        "export function shrink(a: Uint8Array): number {",
+        "  let t = 0;",
+        "  for (let i = 0; i < a.length; ++i) {",
+        "    t += a[i]!;",
+        "    if (i === 2) a = a.subarray(0, 3);",
+        "  }",
+        "  return t;",
+        "}",
         "export function mix(n: number): number {",
         "  let t = 0;",
         "  for (let i = 0; i < n; i++) t = (t * 31 + i) % 1000003;",
@@ -68,7 +76,7 @@ test("worker programs poll counted loops over typed arrays between blocks of ite
         "scale(a, b, 8);",
         "const c = new Uint8Array(8).fill(97);",
         "codes(c);",
-        "console.log(b[7], mix(8), c[0]);",
+        "console.log(b[7], mix(8), c[0], shrink(new Uint8Array(8).fill(1)));",
         workers ? "if (!isMainThread) new Worker(new URL(import.meta.url));" : "",
       ].join("\n");
     const emit = async (workers: boolean): Promise<string> => {
@@ -113,6 +121,11 @@ test("worker programs poll counted loops over typed arrays between blocks of ite
     expect(codes).toMatch(/@llvm\.uadd\.sat\.i64\(i64 %\w+, i64 64\)/);
     expect(codes).toMatch(/@llvm\.umin\.i64\(i64 %\w+, i64 %\w+\)/);
     expect(codes).toMatch(/^loop\.poll\d+:$/m);
+    // A loop that assigns its array reads the new length on the next
+    // iteration: it keeps the budget.
+    const shrink = bodies(workerModule, "shrink");
+    expect(shrink).toContain("worker cancellation budget");
+    expect(shrink).not.toMatch(/^loop\.poll\d+:$/m);
     const mix = bodies(workerModule, "mix");
     expect(mix).toContain("worker cancellation budget");
     expect(mix).not.toContain("loop.poll");

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/native-worker-length-loop-reassign.ts.json

@@ -0,0 +1,6 @@
+{
+ "order": [
+  "<repo>/tests/corpus/native-worker-length-loop-reassign.ts"
+ ],
+ "diags": []
+}

tests/corpus/native-worker-length-loop-reassign.ts

@@ -0,0 +1,58 @@
+// Loops to a typed array's length in a program that starts workers, where
+// the body assigns the loop's array: the condition reads the new array's
+// length on the next iteration, shorter or longer.
+import { Worker, isMainThread, parentPort } from "node:worker_threads";
+
+function shrink(values: Float64Array): string {
+  let total = 0;
+  let steps = 0;
+  for (let i = 0; i < values.length; i++) {
+    total += values[i]!;
+    steps++;
+    if (i === 100) values = values.subarray(0, 50);
+  }
+  return `shrink ${steps} ${total}`;
+}
+
+function grow(values: Uint8Array, longer: Uint8Array): string {
+  let total = 0;
+  let steps = 0;
+  for (let i = 0; i < values.length; i++) {
+    total += values[i]!;
+    steps++;
+    if (i === 10) values = longer;
+  }
+  return `grow ${steps} ${total}`;
+}
+
+function swap(a: Int32Array, b: Int32Array): string {
+  let total = 0;
+  let steps = 0;
+  for (let i = 0; i < a.length; i++) {
+    total += a[i]!;
+    steps++;
+    if (i % 70 === 69) {
+      const t = a;
+      a = b;
+      b = t;
+    }
+  }
+  return `swap ${steps} ${total}`;
+}
+
+function run(): string[] {
+  const values = new Float64Array(300).map((_, i) => i * 0.5);
+  const short = new Uint8Array(20).fill(3);
+  const long = new Uint8Array(200).map((_, i) => i & 7);
+  const a = new Int32Array(100).map((_, i) => i);
+  const b = new Int32Array(160).map((_, i) => -i);
+  return [shrink(values), grow(short, long), swap(a, b)];
+}
+
+if (isMainThread) {
+  for (const line of run()) console.log(line);
+  const worker = new Worker(new URL(import.meta.url));
+  worker.on("message", (lines: string[]) => {
+    for (const line of lines) console.log("worker", line);
+  });
+} else parentPort!.postMessage(run());

tests/harness/corpus-costs/native-worker-length-loop-reassign.ts.json

@@ -0,0 +1 @@
+{ "plain": 3158, "san": 20404 }