Keep while-loop counters bounded by a typed array length in integers (#2124)

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

@@ -2437,7 +2437,7 @@ function planProvenLoop(
   loop: Extract<IrStmt, { kind: "while" | "doWhile" | "for" }>,
 ): ProvenLoopPlan | null {
   if (
-    loop.kind !== "for" ||
+    loop.kind === "doWhile" ||
     !emitter.speculativeLoopsEnabled ||
     emitter.speculation !== null ||
     emitter.speculationResume !== null ||

packages/compiler/src/ir/integer-loop-facts.ts

@@ -152,7 +152,7 @@ export function boundedIntegerLoopFacts(
       e.receiver.kind === "varRef" &&
       stable(e.receiver.localId)
     )
-      return { min: 0, max: Number.MAX_SAFE_INTEGER };
+      return BYTES_LENGTH_RANGE;
     // Every evaluation of an array length lies in its range, so the
     // condition bounds each pass even when the body resizes the array.
     if (isArrayLength(e)) return ARRAY_LENGTH_RANGE;
@@ -225,6 +225,12 @@ export function boundedIntegerLoopFacts(
   return result;
 }
 
+/** Every length, byteLength and byteOffset of a typed array, DataView or
+ * Buffer: each counts elements or bytes of one allocation, and every
+ * supported target addresses less than 2^48 bytes. Counters and offsets
+ * derived from them then stay exact integers far below 2^53. */
+export const BYTES_LENGTH_RANGE: IntegerRange = { min: 0, max: 2 ** 48 };
+
 /** Every length an array can have: SCR_ARR_MAX_LENGTH (scr_array.c). */
 export const ARRAY_LENGTH_RANGE: IntegerRange = { min: 0, max: 4294967295 };
 
@@ -353,6 +359,13 @@ export function guardedLoopCounterFacts(
       stable(e.receiver.localId)
     )
       return { min: 0, max: STRING_MAX_UNITS };
+    if (
+      e.kind === "bytesIntrinsic" &&
+      (e.method === "length" || e.method === "byteLength") &&
+      e.receiver.kind === "varRef" &&
+      stable(e.receiver.localId)
+    )
+      return BYTES_LENGTH_RANGE;
     // An array length can change on any pass, but never leaves its range,
     // which is all a ceiling needs.
     if (isArrayLength(e)) return ARRAY_LENGTH_RANGE;
@@ -402,9 +415,84 @@ export function guardedLoopCounterFacts(
     if (Number.isSafeInteger(start.min) && Number.isSafeInteger(max))
       result.set(id, { min: start.min, max });
   }
+  // Companion counters: a binding the loop changes only by positive literal
+  // increments that no condition bounds (`out[k++] = ...` beside
+  // `while (i < n)`). A bounded counter that every pass advances before it
+  // can continue limits the number of passes, and each pass adds at most
+  // the companion's increments in the body. Increments inside nested loops
+  // can repeat within a pass, so those companions are not bounded.
+  let passes = Infinity;
+  for (const [id, range] of result) {
+    const hi = ceiling.get(id)!;
+    const step = guaranteedStep(loop.body, id);
+    if (step <= 0) continue;
+    passes = Math.min(passes, Math.floor((hi - range.min) / step) + 1);
+  }
+  if (!Number.isSafeInteger(passes) || passes < 0) return result;
+  const repeated = new Set<string>();
+  for (const inner of nested) for (const id of writtenLocals(inner.body)) repeated.add(id);
+  for (const [id, total] of budget) {
+    const local = locals.get(id),
+      start = initial.get(id);
+    if (result.has(id) || refused.has(id) || repeated.has(id) || ceiling.has(id)) continue;
+    if (!local || local.boxed || local.tdz || !start) continue;
+    const max = start.max + passes * total;
+    if (Number.isSafeInteger(start.min) && Number.isSafeInteger(max))
+      result.set(id, { min: start.min, max });
+  }
   return result;
 }
 
+/** The increments of counter `id` that every pass of `body` makes before it
+ * can start another pass: those of the leading statements up to the first
+ * one that can continue the loop or holds a nested region (0 when there is
+ * none). A pass that leaves the loop earlier starts no further pass, so the
+ * count bounds the passes that continue. */
+function guaranteedStep(body: readonly IrStmt[], id: string): number {
+  let step = 0;
+  for (const s of body) {
+    if (s.kind === "assign" && s.localId === id) {
+      const delta = increment(id, s.value);
+      if (delta === null || delta <= 0) return 0;
+      step += delta;
+      continue;
+    }
+    if (s.kind === "exprStmt" && s.expr.kind === "incDec" && s.expr.localId === id) {
+      if (s.expr.op !== "+") return 0;
+      step += 1;
+      continue;
+    }
+    if (mayContinue(s)) break;
+  }
+  return step;
+}
+
+/** Whether `s` can start another pass of the enclosing loop before its
+ * following statements run: it holds a `continue`, a loop or another
+ * region whose control flow is not walked here. */
+function mayContinue(s: IrStmt): boolean {
+  let found = false;
+  const stmt = (x: IrStmt): boolean => {
+    if (
+      x.kind === "continue" ||
+      x.kind === "for" ||
+      x.kind === "while" ||
+      x.kind === "doWhile" ||
+      x.kind === "forOf" ||
+      x.kind === "tryCatch" ||
+      x.kind === "switch" ||
+      (x.kind === "block" && !!x.labels?.length)
+    ) {
+      found = true;
+      return false;
+    }
+    return everyStmtChild(x, expr, stmt);
+  };
+  const expr = (e: IrExpr): boolean => everyExprChild(e, expr, stmt);
+  stmt(s);
+  return found;
+}
+
 /** The largest value counter `id` takes in one pass of `body`, which starts
  * at most at `entry`. The caller has checked that the body changes the
  * counter only by positive literal increments outside switch, exception and

packages/compiler/src/ir/integer-ranges.test.ts

@@ -828,6 +828,72 @@ test("a while condition on a string's length bounds a counter, and guarded incre
   expect(run((inside) => [observe(inside), bump, rebound]).after).toBeNull();
 });
 
+test("a while counter every pass advances bounds the companions the loop increments", () => {
+  // pos = 0; k = 0; while (pos < text.length) { pos++; observe(k); if (unknown === 1) k++; k++; } observe(k)
+  const text: IrExpr = { loc, kind: "varRef", localId: "text", type: STRING };
+  const length: IrExpr = {
+    loc,
+    kind: "strIntrinsic",
+    method: "length",
+    receiver: text,
+    args: [],
+    type: F64,
+  };
+  const set = (localId: string, value: IrExpr): IrStmt => ({ loc, kind: "assign", localId, value });
+  const bump = (localId: string): IrStmt => ({
+    loc,
+    kind: "exprStmt",
+    expr: { loc, kind: "incDec", op: "+", prefix: false, localId, type: F64 },
+  });
+  const maybe = (then: IrStmt[]): IrStmt => ({
+    loc,
+    kind: "if",
+    cond: condition("===", ref("unknown"), num(1)),
+    then,
+    else_: null,
+  });
+  const run = (body: (inside: IrExpr) => IrStmt[]) => {
+    const inside = ref("k"),
+      after = ref("k");
+    const loop: IrStmt = {
+      loc,
+      kind: "while",
+      cond: condition("<", ref("pos"), length),
+      body: body(inside),
+    };
+    const f = fn([set("pos", num(0)), set("k", num(0)), loop, observe(after)]);
+    f.locals.push(
+      { id: "pos", name: "pos", type: F64, mutable: true },
+      { id: "k", name: "k", type: F64, mutable: true },
+      { id: "text", name: "text", type: STRING, mutable: false },
+    );
+    const facts = analyzeIntegerRangeFacts(f);
+    return {
+      inside: facts.ranges.get(inside),
+      after: facts.ranges.get(after),
+      loop: facts.loopRanges.get(loop)?.get("k"),
+    };
+  };
+  // At most 2^31 - 1 passes, each adding at most 2.
+  const bounded = { min: 0, max: 4294967294 };
+  expect(run((inside) => [bump("pos"), observe(inside), maybe([bump("k")]), bump("k")])).toEqual({
+    inside: bounded,
+    after: bounded,
+    loop: bounded,
+  });
+  // A pass that can continue before advancing the counter, or an increment a
+  // nested loop repeats, bounds nothing.
+  const skip: IrStmt = maybe([{ loc, kind: "continue" }]);
+  expect(run((inside) => [skip, bump("pos"), observe(inside), bump("k")]).after).toBeNull();
+  const repeat: IrStmt = {
+    loc,
+    kind: "while",
+    cond: condition("===", ref("unknown"), num(1)),
+    body: [bump("k")],
+  };
+  expect(run((inside) => [bump("pos"), observe(inside), repeat]).after).toBeNull();
+});
+
 test("a nested while loop may advance a counter its own condition bounds", () => {
   // while (pos < text.length) { observe(pos); while (pos < text.length && unknown === 1) pos++; pos += 2; }
   const text: IrExpr = { loc, kind: "varRef", localId: "text", type: STRING };
@@ -1323,7 +1389,7 @@ test("a for loop bounds cursors it resets and copies, and their writes keep the
   }
 });
 
-test("a cursor of a loop over a typed array keeps its writes exact up to 2^53 - 1", () => {
+test("a cursor of a loop over a typed array keeps its writes within the largest length", () => {
   // for (x = 0; x < bytes.length; x++) changes = changes + 1;
   const bytes: IrExpr = {
     loc,
@@ -1355,7 +1421,7 @@ test("a cursor of a loop over a typed array keeps its writes exact up to 2^53 -
     { id: "changes", name: "changes", type: F64, mutable: true },
     { id: "bytes", name: "bytes", type: { kind: "bytes", elem: "u8" }, mutable: false },
   );
-  // The sum of the header range passes 2^53 - 1; the write takes the loop
-  // range instead.
-  expect(analyzeIntegerRanges(f).get(counted)).toEqual({ min: 0, max: Number.MAX_SAFE_INTEGER });
+  // The sum of the header range passes the largest typed array length; the
+  // write takes the loop range instead.
+  expect(analyzeIntegerRanges(f).get(counted)).toEqual({ min: 1, max: 2 ** 48 });
 });

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

@@ -2,6 +2,7 @@ import type { IrExpr, IrFunction, IrStmt } from "./ir.js";
 import { everyExprChild, everyStmtChild, everyStmtList } from "./traverse.js";
 import { byteNumberAccess, byteNumberRange } from "./byte-numbers.js";
 import {
+  BYTES_LENGTH_RANGE,
   boundedIntegerLoopFacts,
   guardedLoopCounterFacts,
   integerPathFallsThrough,
@@ -821,7 +822,7 @@ export function analyzeIntegerRangeFacts(
         const numeric = byteNumberAccess(e);
         if (numeric && !numeric.write) range = byteNumberRange(numeric);
         else if (e.method === "length" || e.method === "byteLength" || e.method === "byteOffset")
-          range = { min: 0, max: Number.MAX_SAFE_INTEGER };
+          range = BYTES_LENGTH_RANGE;
         else if (e.method === "get" && e.invalidNaN !== true && e.receiver.type.kind === "bytes") {
           const elem = e.receiver.type.elem;
           if (elem === "u8" || elem === "u8c") range = { min: 0, max: 255 };
@@ -1075,6 +1076,19 @@ export function analyzeIntegerRangeFacts(
           // inside and after the loop.
           const guarded =
             s.kind === "while" ? guardedLoopCounterFacts(s, loopLocals, facts) : noInvariants;
+          // The same bounds hold on entry and at every write in the loop,
+          // which is what a proven integer loop needs (loopRanges). A loop
+          // analyzed again keeps the bindings every analysis bounds.
+          if (s.kind === "while") {
+            const previous = loopRanges.get(s);
+            if (!previous) loopRanges.set(s, new Map(guarded));
+            else
+              for (const [id, range] of previous) {
+                const other = guarded.get(id);
+                if (other) previous.set(id, union(range, other));
+                else previous.delete(id);
+              }
+          }
           const loop = new Map(facts);
           havocWith(loop, [s], loop, guarded);
           for (const [id, range] of guarded) set(loop, id, range);

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/8002-while-byte-length-companions.ts.json

@@ -0,0 +1,6 @@
+{
+ "order": [
+  "<repo>/tests/corpus/8002-while-byte-length-companions.ts"
+ ],
+ "diags": []
+}

packages/compiler/test/while-induction-emission.test.ts

@@ -0,0 +1,94 @@
+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-while-induction-"));
+  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",
+    });
+    expect(result.ok, JSON.stringify(result.diagnostics)).toBe(true);
+    return await readFile(outPath, "utf8");
+  } finally {
+    await rm(dir, { recursive: true, force: true });
+  }
+}
+
+/** The definition holding the function's statements: the borrowed-argument
+ * body when there is one, otherwise the function itself. */
+function body(ll: string, name: string): string {
+  const definition = (prefix: string): number =>
+    ll.search(new RegExp(`define [^\\n]*@${prefix}_${name}\\(`));
+  const start = definition("sc_bf") >= 0 ? definition("sc_bf") : definition("sc_f");
+  expect(start).toBeGreaterThanOrEqual(0);
+  return ll.slice(start, ll.indexOf("\n}\n", start));
+}
+
+test("a while loop bounded by a byte length keeps its counter and companions in integers", async () => {
+  const ll = await emitLlvm(`
+function reversed(line: Uint8Array, target: Uint8Array): number {
+  const len = line.length;
+  let read = 0;
+  let written = 0;
+  let lines = 0;
+  while (read < len) {
+    const c = line[len - 1 - read]!;
+    read++;
+    if (c === 10) {
+      if (read === len) break;
+      read++;
+    }
+    target[written] = c;
+    written++;
+    if (written % 60 === 0) {
+      target[written] = 10;
+      written++;
+      lines++;
+    }
+  }
+  return written + lines;
+}
+console.log(reversed(new Uint8Array([65, 10, 66]), new Uint8Array(8)));
+`);
+  const fn = body(ll, "reversed");
+  expect(fn).toContain("alloca i64 ; integer read");
+  expect(fn).toContain("alloca i64 ; integer written");
+  expect(fn).toContain("alloca i64 ; integer lines");
+});
+
+test("a companion a nested loop advances, or a counter a continue can skip, stays a double", async () => {
+  const ll = await emitLlvm(`
+function nested(line: Uint8Array): number {
+  let read = 0;
+  let total = 0;
+  while (read < line.length) {
+    read++;
+    for (let j = 0; j < 2; j++) total++;
+  }
+  return total;
+}
+function skipped(line: Uint8Array): number {
+  let read = 0;
+  let total = 0;
+  while (read < line.length) {
+    if (line[read] === 0) continue;
+    read++;
+    total++;
+  }
+  return total;
+}
+console.log(nested(new Uint8Array(3)), skipped(new Uint8Array([1, 2])));
+`);
+  expect(body(ll, "nested")).not.toContain("; integer total");
+  // The counter compared with the length directly is an integer.
+  expect(body(ll, "skipped")).toContain("; integer read");
+  expect(body(ll, "skipped")).not.toContain("; integer total");
+});

tests/corpus/8002-while-byte-length-companions.ts

@@ -0,0 +1,201 @@
+// While loops bounded by a typed array or Buffer length, with companion
+// counters that only grow: reversed copies with line breaks, steps of two,
+// early exits, skipped passes, nested loops and loops inside loops.
+
+function reverseLines(
+  line: Uint8Array,
+  target: Uint8Array,
+  start: number,
+): number {
+  const len = line.length;
+  const right = len - 1;
+  let read = 0;
+  let written = 0;
+  let count = 0;
+  let lines = 0;
+  while (read < len) {
+    let c = line[right - read]!;
+    read++;
+    if (c === 10) {
+      if (read === len) break;
+      c = line[right - read]!;
+      read++;
+    }
+    target[written] = c + 1;
+    written++;
+    count++;
+    if ((count + start - lines) % 5 === 0) {
+      target[written] = 10;
+      written++;
+      count++;
+      lines++;
+    }
+  }
+  return written * 1000 + lines * 10 + count;
+}
+
+const text = Buffer.from("ABCDEFGHIJ\nKLMNOPQRST\nUVWXYZ\n");
+const out = new Uint8Array(64);
+for (const start of [0, 1, 3, 7]) {
+  const summary = reverseLines(text, out, start);
+  console.log(
+    start,
+    summary,
+    Buffer.from(out.subarray(0, Math.floor(summary / 1000))).toString("latin1"),
+  );
+}
+
+function pairs(values: Uint16Array): number[] {
+  let i = 0;
+  let sums = 0;
+  let odd = 0;
+  const result: number[] = [];
+  while (i < values.length) {
+    i += 2;
+    sums++;
+    if ((values[i - 2]! & 1) === 1) odd += 3;
+    result.push(values[i - 2]! + (values[i - 1] ?? -1));
+  }
+  result.push(i, sums, odd);
+  return result;
+}
+console.log(pairs(new Uint16Array([1, 2, 3, 4, 5])));
+console.log(pairs(new Uint16Array(0)));
+console.log(pairs(new Uint16Array([65535, 65535])));
+
+function firstZero(bytes: Uint8Array): string {
+  let i = 0;
+  let seen = 0;
+  let total = -5;
+  while (i < bytes.length) {
+    i++;
+    if (bytes[i - 1] === 0) break;
+    seen++;
+    total += bytes[i - 1]!;
+  }
+  return `${i} ${seen} ${total}`;
+}
+console.log(
+  firstZero(new Uint8Array([3, 4, 0, 9])),
+  firstZero(new Uint8Array([1, 2])),
+  firstZero(new Uint8Array(0)),
+);
+
+function skipping(bytes: Uint8Array): string {
+  let i = 0;
+  let kept = 0;
+  let skipped = 0;
+  while (i < bytes.length) {
+    const b = bytes[i]!;
+    i++;
+    if (b > 100) {
+      skipped++;
+      continue;
+    }
+    kept++;
+  }
+  return `${i} ${kept} ${skipped}`;
+}
+console.log(skipping(new Uint8Array([1, 200, 3, 250, 255, 7])));
+
+function repeated(bytes: Uint8Array): string {
+  let i = 0;
+  let inner = 0;
+  let outer = 0;
+  while (i < bytes.length) {
+    i++;
+    outer++;
+    let j = 0;
+    while (j < bytes[i - 1]!) {
+      j++;
+      inner++;
+    }
+  }
+  return `${i} ${outer} ${inner}`;
+}
+console.log(repeated(new Uint8Array([2, 0, 5, 1])));
+
+function grid(rows: Float64Array[]): string {
+  let total = 0;
+  let cells = 0;
+  for (const row of rows) {
+    let k = 0;
+    let doubled = 0;
+    while (k < row.length) {
+      k++;
+      cells++;
+      doubled += 2;
+      total += row[k - 1]! * doubled;
+    }
+    total += k * 0.5;
+  }
+  return `${cells} ${total}`;
+}
+console.log(
+  grid([
+    new Float64Array([0.5, 1.5]),
+    new Float64Array(0),
+    new Float64Array([-2, 3, 4.25]),
+  ]),
+);
+
+function views(buffer: ArrayBuffer): string {
+  const view = new DataView(buffer, 2);
+  let offset = 0;
+  let words = 0;
+  const parts: string[] = [];
+  while (offset < view.byteLength) {
+    offset += 2;
+    words++;
+    if (offset <= view.byteLength)
+      parts.push(view.getUint16(offset - 2, true).toString(16));
+  }
+  return `${offset} ${words} ${parts.join(",")}`;
+}
+const backing = new Uint8Array([9, 9, 1, 2, 3, 4, 5]).buffer;
+console.log(views(backing));
+
+function captured(bytes: Uint8Array): number {
+  let i = 0;
+  let n = 0;
+  const read = (): number => n;
+  while (i < bytes.length) {
+    i++;
+    n++;
+  }
+  return read() + i;
+}
+console.log(captured(new Uint8Array(4)));
+
+function fromNegative(bytes: Int8Array): string {
+  let i = 0;
+  let level = -3;
+  let low = -0;
+  while (i < bytes.length) {
+    i++;
+    level++;
+    if (bytes[i - 1]! < 0) low++;
+  }
+  return `${level} ${low} ${Object.is(low, -0)}`;
+}
+console.log(
+  fromNegative(new Int8Array([-1, 2, -3])),
+  fromNegative(new Int8Array(0)),
+);
+
+function throwing(bytes: Uint8Array): number {
+  let i = 0;
+  let n = 0;
+  while (i < bytes.length) {
+    i++;
+    n++;
+    if (bytes[i - 1] === 7) throw new Error(`stopped at ${i} after ${n}`);
+  }
+  return n;
+}
+try {
+  console.log(throwing(new Uint8Array([1, 2])));
+  console.log(throwing(new Uint8Array([1, 7, 3])));
+} catch (error) {
+  console.log((error as Error).message);
+}

tests/harness/corpus-costs/8002-while-byte-length-companions.ts.json

@@ -0,0 +1 @@
+{ "plain": 3948, "san": 3569 }