Version dense array loops bounded by an invariant limit (#1816)

packages/compiler/src/backend/llvm/dense-array-loops.test.ts

@@ -1,5 +1,6 @@
 import { expect, test } from "vitest";
 import { BOOL, F64, arrayOf, type IrExpr, type IrLocal, type IrStmt } from "../../ir/ir.js";
+import type { IntegerCountedForLoop } from "../../ir/integer-loops.js";
 import {
   countedDenseShape,
   matchDenseArrayLoop,
@@ -111,16 +112,23 @@ test("loops that may resize, rebind or move storage keep the ordinary copy only"
   expect([...rebound!.arrays.keys()]).toEqual(["a"]);
 });
 
+const counted = (
+  fields: Pick<IntegerCountedForLoop, "start" | "limit" | "step" | "inclusive" | "range">,
+): IntegerCountedForLoop => ({
+  localId: "i",
+  guarded: false,
+  limitGuarded: false,
+  startGuarded: false,
+  guardLimit: 0,
+  ...fields,
+});
+
 test("counted shapes need a non-negative counter strictly below the length", () => {
   const ascending = (limit: IrExpr, inclusive: boolean, min = 0) =>
-    countedDenseShape(loop([]), {
-      localId: "i",
-      start: num(min),
-      limit,
-      step: 1,
-      inclusive,
-      range: { min, max: 100 },
-    });
+    countedDenseShape(
+      loop([]),
+      counted({ start: num(min), limit, step: 1, inclusive, range: { min, max: 100 } }),
+    );
   const minus = (offset: number): IrExpr => ({
     kind: "bin",
     op: "-",
@@ -129,19 +137,58 @@ test("counted shapes need a non-negative counter strictly below the length", ()
     type: F64,
     loc,
   });
-  expect(ascending(method("a", "length"), false)?.length.localId).toBe("a");
-  expect(ascending(minus(1), true)?.length.localId).toBe("a");
+  expect(ascending(method("a", "length"), false)?.length?.localId).toBe("a");
+  expect(ascending(minus(1), true)?.length?.localId).toBe("a");
   expect(ascending(method("a", "length"), true)).toBeNull();
   expect(ascending(minus(1), false, -1)).toBeNull();
+  // Another invariant limit bounds the counter itself; the entry test
+  // compares it with every indexed array's hole-free prefix.
+  const limited = ascending(ref("n"), false);
+  expect(limited?.length).toBeNull();
+  expect(limited?.limit?.limit).toEqual(ref("n"));
+  expect(ascending(ref("n"), false, -1)).toBeNull();
   const descending = (offset: number) =>
-    countedDenseShape(loop([]), {
-      localId: "i",
-      start: minus(offset),
-      limit: num(0),
-      step: -1,
-      inclusive: true,
-      range: { min: 0, max: 100 },
-    });
-  expect(descending(1)?.length.localId).toBe("a");
+    countedDenseShape(
+      loop([]),
+      counted({
+        start: minus(offset),
+        limit: num(0),
+        step: -1,
+        inclusive: true,
+        range: { min: 0, max: 100 },
+      }),
+    );
+  expect(descending(1)?.length?.localId).toBe("a");
   expect(descending(0)).toBeNull();
 });
+
+test("a loop to another invariant limit covers each array whose length it keeps", () => {
+  const limited: DenseLoopShape = {
+    counter: "i",
+    length: null,
+    limit: counted({
+      start: num(0),
+      limit: ref("n"),
+      step: 1,
+      inclusive: false,
+      range: { min: 0, max: 100 },
+    }),
+    range: { min: 0, max: 100 },
+  };
+  const matchLimited = (body: IrStmt[], unaliased = new Set<string>()) =>
+    matchDenseArrayLoop(loop(body), limited, locals, unaliased, new Set(["a", "b"]));
+  const copy = write("b", read("a"));
+  const dense = matchLimited([copy]);
+  expect([...dense!.arrays.keys()].sort()).toEqual(["a", "b"]);
+  expect(dense!.length).toBeNull();
+  expect(dense!.limit?.limit).toEqual(ref("n"));
+  // An array the loop may resize is not covered; neither is one it may
+  // share storage with.
+  const resizing = [
+    { kind: "exprStmt" as const, expr: read("a"), loc },
+    { kind: "exprStmt" as const, expr: method("b", "getNumber", [ref("i")]), loc },
+    { kind: "exprStmt" as const, expr: method("b", "push", [num(1)]), loc },
+  ];
+  expect([...matchLimited(resizing, new Set(["b"]))!.arrays.keys()]).toEqual(["a"]);
+  expect(matchLimited(resizing)).toBeNull();
+});

packages/compiler/src/backend/llvm/dense-array-loops.ts

@@ -15,22 +15,33 @@
  * test bounds `a.length` by 2^31 - 1, so the fast copy's counter is an int32.
  * Arrays longer than that are sparse in practice (a dense one would need
  * 16 GB of slots) and take the ordinary copy, as does every array with a
- * hole, present undefined or sparse storage below `a.length`. */
+ * hole, present undefined or sparse storage below `a.length`.
+ *
+ * A counted loop with another invariant limit, such as `j < n` after
+ * `const n = a.length` or `j < 5`, is covered the same way: the test bounds
+ * the loop's exclusive limit by 2^31 - 1 and by each indexed array's
+ * `filled`, and the loop must provably never change any of their lengths. */
 
 import type { IrExpr, IrLocal, IrStmt } from "../../ir/ir.js";
 import type { IntegerRange } from "../../ir/integer-ranges.js";
 import { preservesArrayLength } from "../../ir/array-length-stability.js";
 import { everyExprChild, everyStmtChild } from "../../ir/traverse.js";
+import type { IntegerCountedForLoop } from "../../ir/integer-loops.js";
 import type { LlvmEmitterContext } from "./expr-context.js";
 import { loadArrCount, SCR_ARR_FIELD } from "./common.js";
+import { emitCountedLoopLimit } from "./counted-loops.js";
 
 export const INT32_MAX = 2147483647;
 
 export interface DenseArrayLoop {
   /** The loop's counter binding. */
   counter: string;
-  /** The array whose length bounds the counter. */
-  length: IrExpr & { kind: "varRef" };
+  /** The array whose length bounds the counter, or null for a counted loop
+   * bounded by `limit`. */
+  length: (IrExpr & { kind: "varRef" }) | null;
+  /** The ascending counted loop whose invariant limit bounds the counter,
+   * when `length` is null. */
+  limit: IntegerCountedForLoop | null;
   /** Arrays whose accesses at the counter are proven to be inside the
    * hole-free prefix in the fast copy, with a read of each binding. */
   arrays: ReadonlyMap<string, IrExpr & { kind: "varRef" }>;
@@ -39,12 +50,15 @@ export interface DenseArrayLoop {
 }
 
 /** The fast copy duplicates the loop; keep the duplicated code small. */
-const MAX_LOOP_NODES = 120;
+const MAX_LOOP_NODES = 160;
 
 /** How a matched integer loop bounds its counter below `length.length`. */
 export interface DenseLoopShape {
   counter: string;
-  length: IrExpr & { kind: "varRef" };
+  /** The array whose length bounds the counter, or null for `limit`. */
+  length: (IrExpr & { kind: "varRef" }) | null;
+  /** An ascending counted loop bounded by another invariant limit. */
+  limit?: IntegerCountedForLoop;
   /** The counter's proven range on every pass of the ordinary loop; the
    * fast copy keeps its minimum, which must be at least 0. */
   range: IntegerRange;
@@ -73,22 +87,18 @@ function lengthOffset(e: IrExpr): { receiver: IrExpr & { kind: "varRef" }; offse
 
 /** The array bounding a counted loop's counter below its length: an
  * ascending loop compared below `a.length - c`, or a descending one starting
- * at `a.length - c` (c >= 1). */
+ * at `a.length - c` (c >= 1). Another ascending loop is bounded by its own
+ * invariant limit. */
 export function countedDenseShape(
   loop: IrStmt & { kind: "for" },
-  counted: {
-    localId: string;
-    start: IrExpr;
-    limit: IrExpr;
-    step: number;
-    inclusive: boolean;
-    range: IntegerRange;
-  },
+  counted: IntegerCountedForLoop,
 ): DenseLoopShape | null {
   if (counted.range.min < 0) return null;
   if (counted.step > 0) {
     const bound = lengthOffset(counted.limit);
-    if (!bound || bound.offset < (counted.inclusive ? 1 : 0)) return null;
+    if (!bound)
+      return { counter: counted.localId, length: null, limit: counted, range: counted.range };
+    if (bound.offset < (counted.inclusive ? 1 : 0)) return null;
     return { counter: counted.localId, length: bound.receiver, range: counted.range };
   }
   // The start is evaluated once, before the loop; the counter only falls.
@@ -143,9 +153,11 @@ export function matchDenseArrayLoop(
     return everyStmtChild(s, expr, stmt);
   };
   if (!loop.body.every(stmt) || !expr(loop.cond) || nested) return null;
+  const nodesOfLoop = [loop.cond, ...loop.body];
   if (
+    length !== null &&
     !preservesArrayLength(
-      [loop.cond, ...loop.body],
+      nodesOfLoop,
       { receiver: length.localId, inRange: counter },
       locals,
       unaliased,
@@ -161,7 +173,14 @@ export function matchDenseArrayLoop(
       local.boxed ||
       local.tdz ||
       disturbed.has(id) ||
-      (id !== length.localId && !initialized.has(id))
+      (id !== length?.localId && !initialized.has(id))
+    )
+      continue;
+    // Without a bounding array, the entry test reads each array's `filled`
+    // once: the loop must not change any of their lengths.
+    if (
+      length === null &&
+      !preservesArrayLength(nodesOfLoop, { receiver: id, inRange: counter }, locals, unaliased)
     )
       continue;
     arrays.set(id, ref);
@@ -170,13 +189,15 @@ export function matchDenseArrayLoop(
   return {
     counter,
     length,
+    limit: length === null ? (shape.limit ?? null) : null,
     arrays,
     range: { min: shape.range.min, max: Math.min(shape.range.max, INT32_MAX - 1) },
   };
 }
 
-/** The fast copy's entry test: the bounding length is at most 2^31 - 1 and
- * at most every indexed array's hole-free prefix bound. */
+/** The fast copy's entry test: the bounding length (or the counted loop's
+ * exclusive limit) is at most 2^31 - 1 and at most every indexed array's
+ * hole-free prefix bound. */
 export function emitDenseLoopGuard(host: LlvmEmitterContext, loop: DenseArrayLoop): string {
   const B = host.B;
   const size = host.sizeType;
@@ -186,7 +207,29 @@ export function emitDenseLoopGuard(host: LlvmEmitterContext, loop: DenseArrayLoo
     host.markMemoryPointer(ptr, "array:header");
     return ptr;
   };
-  const receiver = host.emitStableReceiver(loop.length, []);
+  if (loop.length === null) {
+    // A limit at or below the start runs no pass, and passes the test.
+    const limit = emitCountedLoopLimit(host, loop.limit!);
+    const small = B.tmp();
+    B.line(`${small} = icmp sle i64 ${limit}, ${INT32_MAX}`);
+    let valid = small;
+    for (const ref of loop.arrays.values()) {
+      const array = host.emitStableReceiver(ref, []);
+      const filledPtr = header(array.name, SCR_ARR_FIELD.filled),
+        filled = B.tmp(),
+        wide = B.tmp(),
+        covered = B.tmp(),
+        both = B.tmp();
+      B.line(`${filled} = load i32, ptr ${filledPtr}${host.fieldAliasAttachment(filledPtr)}`);
+      B.line(`${wide} = zext i32 ${filled} to i64`);
+      B.line(`${covered} = icmp sle i64 ${limit}, ${wide}`);
+      B.line(`${both} = and i1 ${valid}, ${covered}`);
+      valid = both;
+    }
+    return valid;
+  }
+  const length = loop.length;
+  const receiver = host.emitStableReceiver(length, []);
   const lenPtr = header(receiver.name, SCR_ARR_FIELD.len),
     len = B.tmp(),
     small = B.tmp();
@@ -194,7 +237,7 @@ export function emitDenseLoopGuard(host: LlvmEmitterContext, loop: DenseArrayLoo
   B.line(`${small} = icmp ule ${size} ${len}, ${INT32_MAX}`);
   let valid = small;
   for (const [id, ref] of loop.arrays) {
-    const array = id === loop.length.localId ? receiver : host.emitStableReceiver(ref, []);
+    const array = id === length.localId ? receiver : host.emitStableReceiver(ref, []);
     const filledPtr = header(array.name, SCR_ARR_FIELD.filled),
       filled = B.tmp(),
       covered = B.tmp(),

packages/compiler/src/backend/llvm/local-array-reads.ts

@@ -586,6 +586,24 @@ export function emitBorrowedArrayRead(
   const B = host.B;
   const array = host.emitStableReceiver(read.arr, [read.index]);
   const loopIndex = host.emitIntegerLoopIndex(read.index);
+  if (loopIndex !== null && host.provenDenseIndex(read.arr, read.index)) {
+    // In a dense array loop's fast copy, the entry test proved the counter
+    // below the array's hole-free prefix: the slot holds a present value.
+    const dataPtr = B.tmp(),
+      data = B.tmp(),
+      valuePtr = B.tmp(),
+      value = B.tmp();
+    B.line(
+      `${dataPtr} = getelementptr inbounds %ScrArr, ptr ${array.name}, i32 0, i32 ${SCR_ARR_FIELD.data}`,
+    );
+    host.markMemoryPointer(dataPtr, "array:header");
+    B.line(`${data} = load ptr, ptr ${dataPtr}${host.fieldAliasAttachment(dataPtr)}`);
+    B.line(`${valuePtr} = getelementptr inbounds i64, ptr ${data}, ${host.sizeType} ${loopIndex}`);
+    host.markMemoryPointer(valuePtr, "array:elements");
+    B.line(`${value} = load ptr, ptr ${valuePtr}${host.fieldAliasAttachment(valuePtr)}`);
+    B.line(`store ptr ${value}, ptr ${slot}`);
+    return;
+  }
   const index = host.emitExpr(read.index);
   // A proven integer index (an int32 local, say) needs no double
   // round-trip; a negative one compares above the capacity as unsigned.

packages/compiler/src/ir/array-length-stability.test.ts

@@ -167,3 +167,29 @@ test("fresh local arrays read only as receivers have no alias", () => {
   const rebound = [...body, { kind: "assign", localId: "out", value: arr("a"), loc } as IrStmt];
   expect(findUnaliasedLocalArrays(fn(rebound)).has("out")).toBe(false);
 });
+
+test("number and boolean writes to class fields run no user code", () => {
+  const obj: IrExpr = {
+    kind: "varRef",
+    localId: "o",
+    type: { kind: "object", className: "P" },
+    loc,
+  };
+  const scope = new Map(locals).set("o", {
+    id: "o",
+    name: "o",
+    type: { kind: "object", className: "P" },
+    mutable: false,
+  });
+  const fieldSet = (value: IrExpr): IrStmt => ({
+    kind: "fieldSet",
+    obj,
+    className: "P",
+    field: "x",
+    value,
+    loc,
+  });
+  expect(preserves([fieldSet(read("a", ref("i")))], "i", new Set(), scope)).toBe(true);
+  // A reference stored into a field releases the old one.
+  expect(preserves([fieldSet(arr("b"))], "i", new Set(), scope)).toBe(false);
+});

packages/compiler/src/ir/array-length-stability.ts

@@ -264,6 +264,15 @@ export function preservesArrayLength(
         );
       case "return":
         return s.value === null || expr(s.value);
+      case "fieldSet":
+        // A number or boolean stored into a class field: accessors are
+        // calls, not field writes, and the old value holds no reference to
+        // release, so the write runs no user code.
+        return (
+          (s.value.type.kind === "f64" || s.value.type.kind === "bool") &&
+          expr(s.obj) &&
+          expr(s.value)
+        );
       default:
         if (!PLAIN_STATEMENTS.has(s.kind)) return false;
     }

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/7320-array-dense-loop-invariant-limits.ts.json

@@ -0,0 +1,6 @@
+{
+ "order": [
+  "<repo>/tests/corpus/7320-array-dense-loop-invariant-limits.ts"
+ ],
+ "diags": []
+}

tests/corpus/7320-array-dense-loop-invariant-limits.ts

@@ -0,0 +1,101 @@
+// Index loops bounded by an invariant limit other than the array's own length
+// (a local copy of a length, a literal, an inclusive bound), over arrays the
+// loop never resizes: class instances whose number and boolean fields the loop
+// writes, and number arrays that are dense, short, holey or partially filled.
+
+class Body {
+  x: number;
+  v = 0;
+  hit = false;
+  next: Body | null = null;
+
+  constructor(x: number) {
+    this.x = x;
+  }
+}
+
+function pairs(bodies: Body[]): string {
+  const n = bodies.length;
+  for (let i = 0; i < n; i++) {
+    const bi = bodies[i]!;
+    let vi = bi.v;
+    for (let j = i + 1; j < n; j++) {
+      const bj = bodies[j]!;
+      const d = bi.x - bj.x;
+      const f = 1 / (d * d + 1);
+      vi -= d * f;
+      bj.v += d * f;
+      bj.hit = f > 0.1;
+    }
+    bi.v = vi;
+  }
+  return bodies.map((b) => `${b.v.toFixed(6)}${b.hit ? "*" : ""}`).join(" ");
+}
+
+function firstThree(bodies: Body[]): number {
+  let s = 0;
+  for (let j = 0; j < 3; j++) s += bodies[j]!.x;
+  return s;
+}
+
+function inclusive(bodies: Body[]): number {
+  const last = bodies.length - 1;
+  let s = 0;
+  for (let j = 1; j <= last; j++) s += bodies[j]!.x * j;
+  return s;
+}
+
+function linked(bodies: Body[]): string {
+  const n = bodies.length;
+  // A reference field write: the loop keeps the ordinary copy.
+  for (let j = 0; j < n - 1; j++) bodies[j]!.next = bodies[j + 1]!;
+  let count = 0;
+  for (let b: Body | null = bodies[0]!; b !== null; b = b.next) count++;
+  return `${count}`;
+}
+
+function sumTo(a: number[], limit: number): string {
+  const n = limit;
+  let s = 0;
+  let missing = 0;
+  for (let i = 0; i < n; i++) {
+    const v = a[i];
+    if (v === undefined) missing++;
+    else s += v;
+  }
+  return `${s} ${missing}`;
+}
+
+function scaleTo(a: number[], out: number[]): string {
+  const n = a.length;
+  for (let i = 0; i < n; i++) out[i] = a[i] * 2;
+  return `${out.length}:${out.join(",")}`;
+}
+
+function window(a: number[]): number {
+  let s = 0;
+  for (let i = 2; i < 6; i++) s += a[i] * i;
+  return s;
+}
+
+const bodies = [new Body(0), new Body(1.5), new Body(-2), new Body(4), new Body(0.25)];
+console.log(pairs(bodies));
+console.log(pairs(bodies));
+console.log(firstThree(bodies), inclusive(bodies), inclusive([new Body(7)]));
+console.log(linked([new Body(1), new Body(2), new Body(3)]));
+
+const dense = [1, 2, 3, 4, 5, 6, 7, 8];
+console.log(sumTo(dense, 8), sumTo(dense, 5), sumTo(dense, 0), sumTo(dense, -3));
+// Limits past the length or the filled prefix take the checked loop.
+console.log(sumTo(dense, 10), sumTo([], 2));
+const holey = [1];
+holey[3] = 4;
+console.log(sumTo(holey, 4), sumTo(holey, 1));
+const shortened = [1, 2, 3, 4, 5, 6];
+shortened.length = 3;
+shortened.length = 6;
+console.log(sumTo(shortened, 6), sumTo(shortened, 3));
+console.log(scaleTo(dense, new Array<number>(8).fill(0)));
+console.log(scaleTo(dense, [9, 9]));
+console.log(scaleTo([1, 2, 3], []));
+console.log(window(dense), window([1, 2, 3, 4, 5, 6]));

tests/harness/corpus-costs/7320-array-dense-loop-invariant-limits.ts.json

@@ -0,0 +1 @@
+{ "plain": 2600, "san": 31600 }