Version while loops over typed arrays on their counters' bounds

Typed-array loop versions applied only to for loops, so a scanner written
with while loops, such as `while (pos < end && bytes[pos] === 32) pos++`,
kept one bounds test per access. A while loop counter is now a number
local that the loop changes only by positive integer increments (`c = c +
k`, `c++`, `c = (c + k) | 0`) written as statements of its body, blocks,
branches and nested while loops. An access at `counter + d` after a
condition compared the counter below a loop-invariant bound, with no
increment since, is in bounds for the whole loop when the guard finds the
counter's entry value an exact integer of at least -d and the bound at
most the length minus the access's extent. A counter with a `| 0`
increment also needs its loops' bounds at most 2^31 minus its increments,
so it never wraps.

A fixed check (a statically ranged index) now needs an end below 2^31.
The counter of a loop to a number[]'s length whose length the loop may
change reaches 2^32 - 2, and its typed-array accesses got a guard no
typed array passes and a second copy of the loop.

packages/compiler/src/backend/llvm/byte-windows.ts

@@ -43,6 +43,7 @@ export function findInitializedByteLoopBindings(
         if (stmt.init) list([stmt.init], inner);
         list(stmt.body, inner);
       }
+      if (stmt.kind === "while") entries.set(stmt, new Set(known));
       if (stmt.kind === "while" || stmt.kind === "doWhile") list(stmt.body, new Set(known));
     }
   }

packages/compiler/src/backend/llvm/bytes-loop-versions.test.ts

@@ -177,6 +177,11 @@ test("statically ranged indices compare one constant with the length", () => {
   ]);
   expect(version.affine).toEqual([]);
   expect(version.counted).toBeNull();
+  // A uint32 index (an array counter's range) gets no guard: no typed array
+  // in practice has 2^31 elements.
+  const wide = bin(">>>", bin("^", ref("i"), ref("k")), num(0));
+  const unbounded = counted([set(ref("out", f64s), num(0), get(ref("table", i32s), wide))]);
+  expect(match(fn([unbounded]), unbounded)!.fixed).toEqual([{ receiver: 0, end: 0 }]);
 });
 
 test("scaled counters and body temps resolve to their receiver and index", () => {
@@ -358,6 +363,63 @@ test("index temps of a version's counter are integers in the fast copy", () => {
   expect(fast).not.toContain("bytes.index");
 });
 
+test("while loop counters below a bound version their accesses", () => {
+  // while (i < n) { out[i] = src[i + 1]; i = i + 1 }
+  const step: IrStmt = { kind: "assign", localId: "i", value: bin("+", ref("i"), num(1)), loc };
+  const whileLoop = (body: IrStmt[]): IrStmt & { kind: "while" } => ({
+    kind: "while",
+    cond: { kind: "bin", op: "<", left: ref("i"), right: ref("n"), type: BOOL, loc },
+    body,
+    loc,
+  });
+  const copy = whileLoop([
+    set(ref("out", f64s), ref("i"), get(ref("src", f64s), bin("+", ref("i"), num(1)))),
+    step,
+  ]);
+  const declared = (loop: IrStmt): IrFunction =>
+    fn([{ kind: "varDecl", localId: "i", init: ref("k"), loc }, loop]);
+  const versionOf = (loop: IrStmt & { kind: "while" }): BytesLoopVersion | null => {
+    const f = declared(loop);
+    const ranges = analyzeIntegerRanges(f);
+    return matchBytesLoopVersion(loop, null, {
+      locals: new Map(f.locals.map((l) => [l.id, l])),
+      invariantGlobals: new Set(),
+      captured: new Set(),
+      ranges,
+      initialized: findInitializedByteLoopBindings(f).get(loop) ?? new Set(),
+      proven: new Set(),
+      workers: false,
+      owned: new Set(),
+    });
+  };
+  const version = versionOf(copy)!;
+  expect(version.bounds.size).toBe(2);
+  expect(version.bounded.map((c) => [version.receivers[c.receiver]!.ref, c.slack])).toEqual([
+    [ref("out", f64s), 0],
+    [ref("src", f64s), 1],
+  ]);
+  // `out[i]` needs `i` to start at 0 or above (`src[i + 1]` only at -1).
+  expect(version.counters.map((c) => [c.ref.kind === "varRef" && c.ref.localId, c.min])).toEqual([
+    ["i", 0],
+  ]);
+  // An access after the counter advances in the same pass stays checked.
+  const advanced = whileLoop([step, set(ref("out", f64s), ref("i"), num(0))]);
+  expect(versionOf(advanced)).toBeNull();
+  const mod: IrModule = {
+    irVersion: 17,
+    sourceFile: loc.file,
+    entry: "__main",
+    functions: [
+      { name: "__main", params: [], returnType: VOID, locals: [], body: [], loc },
+      declared(copy),
+    ],
+  };
+  const body = /^define internal [^\n]*@sc_bf_kernel[^]*?^}/m.exec(emitLlvmModule(mod))![0];
+  const fast = /^bytes\.version\.fast\d+:$[^]*?^bytes\.version\.done\d+:$/m.exec(body)![0];
+  expect(fast).not.toContain("bytes.index");
+  expect(body.slice(0, body.indexOf(fast))).toContain("bytes.index.invalid");
+});
+
 test("Float16Array accesses version in elements and share the element tag", () => {
   const f16s: IrType = { kind: "bytes", elem: "f16" };
   const withHalves = (body: IrStmt[]): IrFunction => {

packages/compiler/src/backend/llvm/bytes-loop-versions.ts

@@ -1,4 +1,4 @@
-import type { IrExpr, IrLocal, IrStmt } from "../../ir/ir.js";
+import { F64, type IrExpr, type IrLocal, type IrStmt } from "../../ir/ir.js";
 import type { IntegerRange, IntegerRanges } from "../../ir/integer-ranges.js";
 import type { IntegerCountedForLoop } from "../../ir/integer-loops.js";
 import type { LlvmEmitterContext } from "./expr-context.js";
@@ -23,6 +23,24 @@ interface KeyedTerm extends Term {
   key: string;
 }
 
+/** Every access of `receiver` at a while loop counter `c + d` happens while
+ * `c < bound` (`c <= bound` when `inclusive`) holds and before the counter
+ * advances, so it is in bounds when `bound + slack <= length`. */
+interface BoundedCheck {
+  receiver: number;
+  bound: IrExpr;
+  slack: number;
+}
+
+/** A while loop counter whose entry value the guard tests: an exact
+ * integer of at least `min`. A counter that `| 0` wraps also needs every
+ * bound of its loops at most `wrapLimit`, so that it never reaches 2^31. */
+interface CounterCheck {
+  ref: VarRef;
+  min: number;
+  wrapBounds: { bound: IrExpr; limit: number }[];
+}
+
 /** Every access of `receiver` at a static index range ends below `end + 1`. */
 interface FixedCheck {
   receiver: number;
@@ -51,6 +69,9 @@ export interface BytesLoopVersion {
   ownedReceivers: number[];
   fixed: FixedCheck[];
   affine: AffineCheck[];
+  /** While loops: accesses at a counter below a loop-invariant bound. */
+  bounded: BoundedCheck[];
+  counters: CounterCheck[];
   /** Keys of affine terms the guard tests at run time: exact integers
    * within GUARDED_TERM. */
   guarded: ReadonlySet<string>;
@@ -83,7 +104,7 @@ export interface BytesLoopFacts {
 }
 
 const MAX_COST = 400;
-const MAX_FIXED_END = 2 ** 32;
+const MAX_FIXED_END = 2 ** 31;
 const MAX_CHECKS = 8;
 const MAX_OFFSET = 2 ** 31;
 const SAFE = Number.MAX_SAFE_INTEGER;
@@ -116,7 +137,7 @@ const INVARIANT_OPS = new Set(["+", "-", "*", "&", "|", "^", "<<", ">>", ">>>"])
  * whenever the guard fails, and keeps every out-of-bounds read, ignored
  * write and NaN-index behavior. */
 export function matchBytesLoopVersion(
-  loop: IrStmt & { kind: "for" },
+  loop: IrStmt & { kind: "for" | "while" },
   counted: IntegerCountedForLoop | null,
   facts: BytesLoopFacts,
 ): BytesLoopVersion | null {
@@ -155,10 +176,10 @@ export function matchBytesLoopVersion(
   // writes is not invariant either: `for (let i = (a = b, 0); ...)` must
   // not keep the bounds the guard proved for the old `a`.
   if (
-    (loop.init !== null && !scanStmt(loop.init)) ||
+    (loop.kind === "for" && loop.init !== null && !scanStmt(loop.init)) ||
     !loop.body.every(scanStmt) ||
     (loop.cond !== null && !scanExpr(loop.cond)) ||
-    (loop.update !== null && !scanStmt(loop.update))
+    (loop.kind === "for" && loop.update !== null && !scanStmt(loop.update))
   )
     return null;
   // A body temp stands for its initializer only when it is a constant
@@ -337,6 +358,44 @@ export function matchBytesLoopVersion(
       return sum + Math.max(Math.abs(r.min), Math.abs(r.max));
     }, 0);
 
+  const whileAccesses = loop.kind === "while" ? whileCounterAccesses(loop) : null;
+  function whileCounterAccesses(loop: IrStmt & { kind: "while" }) {
+    return findWhileCounterAccesses(loop, {
+      counterLocal: (id) => {
+        const local = facts.locals.get(id);
+        return (
+          local?.type.kind === "f64" &&
+          !local.boxed &&
+          !local.tdz &&
+          !facts.captured.has(id) &&
+          facts.initialized.has(id)
+        );
+      },
+      boundKey: (e) => {
+        if (e.type.kind !== "f64") return null;
+        // The guard evaluates the bound before the loop: a typed array's
+        // length needs the binding itself, not a temp the body declares.
+        if (
+          e.kind === "bytesIntrinsic" &&
+          e.method === "length" &&
+          e.receiver.kind === "varRef" &&
+          e.receiver.type.kind === "bytes" &&
+          invariant(e.receiver.localId)
+        )
+          return `bytes-length:${e.receiver.localId}`;
+        if (
+          e.kind === "strIntrinsic" &&
+          e.method === "length" &&
+          e.args.length === 0 &&
+          e.receiver.kind === "varRef" &&
+          e.receiver.type.kind === "string" &&
+          invariant(e.receiver.localId)
+        )
+          return `string-length:${e.receiver.localId}`;
+        return pureKey(e);
+      },
+    });
+  }
   const receivers: { ref: VarRef; global: boolean }[] = [];
   const receiverIndex = new Map<string, number>();
   const fixed = new Map<number, { check: FixedCheck; nodes: Node[] }>();
@@ -356,6 +415,11 @@ export function matchBytesLoopVersion(
   // Worker programs: accesses through invariant receivers, proven or not,
   // become thread-owned when the guard finds no shared storage.
   const owned = new Map<number, Node[]>();
+  // While loops: accesses at a counter below a bound, by receiver and bound.
+  const boundedChecks = new Map<
+    string,
+    { check: BoundedCheck; nodes: { node: Node; access: WhileAccess }[] }
+  >();
   const access = (node: Node, array: IrExpr, index: IrExpr, width: number): void => {
     const ref = root(array);
     if (!ref) return;
@@ -390,11 +454,30 @@ export function matchBytesLoopVersion(
         });
       return;
     }
+    const bounded = whileAccesses?.accesses.get(node);
+    const loopCounter = bounded && whileAccesses!.counters.get(bounded.counter);
+    if (bounded && loopCounter && (!loopCounter.wraps || loopCounter.bounds !== null)) {
+      const receiver = receiverOf(ref);
+      const key = `${receiver}|${bounded.key}|${bounded.adj}`;
+      const slack = bounded.offset + width - 1 + bounded.adj;
+      const group = boundedChecks.get(key);
+      if (group) {
+        group.check.slack = Math.max(group.check.slack, slack);
+        group.nodes.push({ node, access: bounded });
+      } else
+        boundedChecks.set(key, {
+          check: { receiver, bound: bounded.bound, slack },
+          nodes: [{ node, access: bounded }],
+        });
+      return;
+    }
     const range = facts.ranges.get(index);
     if (!range) return;
     const end = range.max + width - 1;
-    // Only a length a buffer can have: an outer loop must not pay for a
-    // guard that bounds an inner loop's unbounded counter.
+    // Only an end below 2^31: an outer loop must not pay for a guard that
+    // bounds an inner loop's unbounded counter, and the counter of a loop to
+    // an array's length reaches 2^32 - 2, a length no typed array in
+    // practice has.
     if (range.min < 0 || end >= MAX_FIXED_END) return;
     const receiver = receiverOf(ref);
     const group = fixed.get(receiver);
@@ -416,6 +499,7 @@ export function matchBytesLoopVersion(
     return everyStmtChild(s, visitExpr, visitStmt);
   };
   loop.body.every(visitStmt);
+  if (loop.kind === "while") visitExpr(loop.cond);
   // Bound the counter so every counter operand stays a safe integer.
   let bound = SAFE;
   const affineGroups = [...affineChecks.values()].filter((group) => {
@@ -431,8 +515,18 @@ export function matchBytesLoopVersion(
   });
   const fixedKept = [...fixed.values()].slice(0, MAX_CHECKS);
   const affineKept = affineGroups.slice(0, MAX_CHECKS - fixedKept.length);
+  const boundedKept = [...boundedChecks.values()].slice(
+    0,
+    Math.max(0, MAX_CHECKS - fixedKept.length - affineKept.length),
+  );
   const ownedKept = [...owned.entries()].slice(0, MAX_CHECKS);
-  if (fixedKept.length === 0 && affineKept.length === 0 && ownedKept.length === 0) return null;
+  if (
+    fixedKept.length === 0 &&
+    affineKept.length === 0 &&
+    boundedKept.length === 0 &&
+    ownedKept.length === 0
+  )
+    return null;
   const bounds = new Set<Node>();
   const ownedNodes = new Set<Node>();
   const used = new Set<number>();
@@ -477,6 +571,39 @@ export function matchBytesLoopVersion(
   for (const group of affineKept)
     for (const term of group.check.terms) if (guardedRanges.has(term.expr)) guarded.add(term.key);
   for (const [node, key] of guardedNodes) if (guarded.has(key)) ranges.set(node, GUARDED_TERM);
+  // While loop counters of the kept checks: each entry value is an exact
+  // integer that no access offset takes below 0. In the fast copy each
+  // access's index is below its length, and its counter read at least the
+  // most negative offset.
+  const narrow = (e: IrExpr, range: IntegerRange): void => {
+    const known = ranges.get(e) ?? facts.ranges.get(e);
+    ranges.set(
+      e,
+      known ? { min: Math.max(known.min, range.min), max: Math.min(known.max, range.max) } : range,
+    );
+  };
+  const counterMins = new Map<string, number>();
+  for (const group of boundedKept) {
+    used.add(group.check.receiver);
+    for (const { node, access } of group.nodes) {
+      bounds.add(node);
+      const min = 0 - access.offset;
+      counterMins.set(access.counter, Math.max(counterMins.get(access.counter) ?? -SAFE, min));
+      narrow(access.ref, { min: -MAX_OFFSET, max: SAFE });
+      if (access.index !== access.ref) narrow(access.index, { min: 0, max: SAFE });
+    }
+  }
+  const counterChecks: CounterCheck[] = [];
+  for (const [id, min] of counterMins) {
+    const info = whileAccesses!.counters.get(id)!;
+    counterChecks.push({
+      ref: { kind: "varRef", localId: id, type: F64, loc: loop.loc },
+      min,
+      wrapBounds: info.wraps
+        ? info.bounds!.map((b) => ({ bound: b.bound, limit: 2 ** 31 - info.sum - b.adj }))
+        : [],
+    });
+  }
   // Renumber the receivers the kept checks use.
   const order = [...used].sort((a, b) => a - b);
   const renumber = new Map(order.map((old, index) => [old, index]));
@@ -500,6 +627,12 @@ export function matchBytesLoopVersion(
       minOffset: g.check.minOffset,
       maxEnd: g.check.maxEnd,
     })),
+    bounded: boundedKept.map((g) => ({
+      receiver: renumber.get(g.check.receiver)!,
+      bound: g.check.bound,
+      slack: g.check.slack,
+    })),
+    counters: counterChecks,
     guarded,
     counted: affineKept.length > 0 ? counted : null,
     counterBound: needsBound ? bound : null,
@@ -507,6 +640,307 @@ export function matchBytesLoopVersion(
   };
 }
 
+/** A typed-array access at a while loop counter: `counter + offset`, read
+ * or written while `counter < bound` (`<=` when `adj` is 1) holds. */
+interface WhileAccess {
+  counter: string;
+  /** The index, and the counter's read in it. */
+  index: IrExpr;
+  ref: IrExpr;
+  bound: IrExpr;
+  key: string;
+  adj: number;
+  offset: number;
+}
+
+interface WhileCounter {
+  /** Some increment is `c = (c + k) | 0`. */
+  wraps: boolean;
+  /** The sum of every increment literal in the loop. */
+  sum: number;
+  /** One bound from the condition of the loop and of every nested while
+   * loop that advances the counter; null when one of them has none. */
+  bounds: { bound: IrExpr; adj: number }[] | null;
+}
+
+/** `counter = counter + k`, `counter = (counter + k) | 0` or `counter++`. */
+function counterStep(s: IrStmt): { id: string; k: number; wraps: boolean } | null {
+  if (s.kind === "exprStmt" && s.expr.kind === "incDec" && s.expr.op === "+")
+    return { id: s.expr.localId, k: 1, wraps: false };
+  if (s.kind !== "assign") return null;
+  let value = s.value,
+    wraps = false;
+  if (
+    value.kind === "bin" &&
+    value.op === "|" &&
+    value.right.kind === "numLit" &&
+    value.right.value === 0 &&
+    value.left.kind === "bin"
+  ) {
+    value = value.left;
+    wraps = true;
+  }
+  if (
+    value.kind !== "bin" ||
+    value.op !== "+" ||
+    value.left.kind !== "varRef" ||
+    value.left.localId !== s.localId ||
+    value.right.kind !== "numLit" ||
+    !Number.isSafeInteger(value.right.value) ||
+    value.right.value <= 0 ||
+    value.right.value > MAX_OFFSET
+  )
+    return null;
+  return { id: s.localId, k: value.right.value, wraps };
+}
+
+/** Typed-array accesses of a while loop that happen at a counter below a
+ * loop-invariant bound. A counter is a number local that the loop changes
+ * only by positive integer increments, written as statements of the body,
+ * of its blocks and branches, or of nested while loops. The condition of
+ * the loop, an `if` or a `&&` chain that compares the counter below a bound
+ * (`while (pos < end && bytes[pos] === 32) pos++`) makes it below that
+ * bound until the next increment on the path. The counter starts at its
+ * entry value and only grows (a `| 0` increment never wraps when every
+ * bound is at most 2^31 minus the increments), so an access at `counter +
+ * d` is in bounds for the whole loop when the entry value is at least `-d`
+ * and `bound + d < length` (`+ 1` for `<=`). */
+function findWhileCounterAccesses(
+  loop: IrStmt & { kind: "while" },
+  host: {
+    counterLocal: (id: string) => boolean;
+    /** A key for a bound expression whose value the loop never changes. */
+    boundKey: (e: IrExpr) => string | null;
+  },
+): { accesses: Map<Node, WhileAccess>; counters: Map<string, WhileCounter> } {
+  const steps = new Map<string, { sum: number; wraps: boolean }>();
+  const refused = new Set<string>();
+  // Increments count only as statements of the loop's straight-line
+  // structure; any other write refuses the binding.
+  const writes = (list: readonly IrStmt[], direct: boolean): void => {
+    for (const s of list) {
+      const step = counterStep(s);
+      if (step && direct) {
+        const prior = steps.get(step.id);
+        steps.set(step.id, {
+          sum: (prior?.sum ?? 0) + step.k,
+          wraps: (prior?.wraps ?? false) || step.wraps,
+        });
+        continue;
+      }
+      if (step) refused.add(step.id);
+      const nested =
+        direct &&
+        (s.kind === "if" || s.kind === "while" || (s.kind === "block" && !s.labels?.length));
+      if (s.kind === "assign" || s.kind === "varDecl" || s.kind === "forOf") refused.add(s.localId);
+      if (s.kind === "tryCatch" && s.catchLocalId) refused.add(s.catchLocalId);
+      everyStmtChild(
+        s,
+        (e) => (exprWrites(e), true),
+        (c) => (writes([c], nested), true),
+      );
+    }
+  };
+  const exprWrites = (e: IrExpr): void => {
+    if (e.kind === "assignExpr" || e.kind === "incDec") refused.add(e.localId);
+    everyExprChild(
+      e,
+      (c) => (exprWrites(c), true),
+      (c) => (writes([c], false), true),
+    );
+  };
+  exprWrites(loop.cond);
+  writes(loop.body, true);
+  const counter = (id: string): boolean =>
+    steps.has(id) && !refused.has(id) && host.counterLocal(id);
+
+  type Fact = { bound: IrExpr; key: string; adj: number };
+  type Clean = ReadonlyMap<string, Fact>;
+  // The bounds a condition puts on counters while it holds.
+  const facts = (cond: IrExpr, into: Map<string, Fact>): void => {
+    if (cond.kind === "toBool") return facts(cond.operand, into);
+    if (cond.kind === "logical" && cond.op === "&&") {
+      facts(cond.left, into);
+      facts(cond.right, into);
+      return;
+    }
+    if (cond.kind !== "bin") return;
+    const below = cond.op === "<" || cond.op === "<=",
+      above = cond.op === ">" || cond.op === ">=";
+    const c = below ? cond.left : above ? cond.right : null,
+      bound = below ? cond.right : cond.left;
+    if (c?.kind !== "varRef" || !counter(c.localId) || into.has(c.localId)) return;
+    const key = host.boundKey(bound);
+    if (key === null) return;
+    into.set(c.localId, { bound, key, adj: cond.op === "<=" || cond.op === ">=" ? 1 : 0 });
+  };
+  // The facts that hold after either of two paths.
+  const meet = (a: Clean | null, b: Clean | null): Clean | null => {
+    if (a === null) return b;
+    if (b === null) return a;
+    const out = new Map<string, Fact>();
+    for (const [id, fact] of a) if (b.get(id) === fact) out.set(id, fact);
+    return out;
+  };
+  const withFacts = (clean: Clean, cond: IrExpr): Clean => {
+    const out = new Map(clean);
+    facts(cond, out);
+    return out;
+  };
+  const accesses = new Map<Node, WhileAccess>();
+  const record = (node: Node, index: IrExpr, clean: Clean): void => {
+    let ref: IrExpr = index,
+      offset = 0;
+    if (index.kind === "bin" && (index.op === "+" || index.op === "-")) {
+      const literal =
+        index.right.kind === "numLit" ? index.right : index.op === "+" ? index.left : null;
+      const other = literal === index.right ? index.left : index.right;
+      if (literal?.kind !== "numLit" || !Number.isSafeInteger(literal.value)) return;
+      if (Math.abs(literal.value) > MAX_OFFSET || Object.is(literal.value, -0)) return;
+      offset = index.op === "-" ? -literal.value : literal.value;
+      ref = other;
+    }
+    if (ref.kind !== "varRef") return;
+    const fact = clean.get(ref.localId);
+    if (fact) accesses.set(node, { counter: ref.localId, index, ref, ...fact, offset });
+  };
+  const expr = (e: IrExpr, clean: Clean): void => {
+    if (e.kind === "bytesIntrinsic") {
+      const numeric = byteNumberAccess(e);
+      const index = numeric
+        ? e.args[numeric.offsetArg]
+        : e.method === "get"
+          ? e.args[0]
+          : undefined;
+      if (index) record(e, index, clean);
+    }
+    if (e.kind === "logical" && e.op === "&&") {
+      expr(e.left, clean);
+      expr(e.right, withFacts(clean, e.left));
+      return;
+    }
+    if (e.kind === "ternary") {
+      expr(e.cond, clean);
+      expr(e.then, withFacts(clean, e.cond));
+      expr(e.else_, clean);
+      return;
+    }
+    everyExprChild(
+      e,
+      (c) => (expr(c, clean), true),
+      (c) => (stmts([c], clean), true),
+    );
+  };
+  // Walks statements in order; the result is the clean state where the
+  // list falls through, or null when it never does.
+  const stmts = (list: readonly IrStmt[], clean: Clean): Clean | null => {
+    let current: Clean | null = clean;
+    for (const s of list) {
+      if (current === null) return null;
+      const step = counterStep(s);
+      if (step) {
+        const next: Map<string, Fact> = new Map(current);
+        next.delete(step.id);
+        current = next;
+        continue;
+      }
+      switch (s.kind) {
+        case "if": {
+          expr(s.cond, current);
+          const then = stmts(s.then, withFacts(current, s.cond));
+          const otherwise = stmts(s.else_ ?? [], current);
+          current = meet(then, otherwise);
+          break;
+        }
+        case "while": {
+          // A pass of the nested loop may follow its own increments.
+          const advanced: Map<string, Fact> = new Map(current);
+          for (const id of writtenCounters(s.body)) advanced.delete(id);
+          expr(s.cond, advanced);
+          stmts(s.body, withFacts(advanced, s.cond));
+          current = advanced;
+          break;
+        }
+        case "block":
+          if (!s.labels?.length) {
+            current = stmts(s.body, current);
+            break;
+          }
+          everyStmtChild(
+            s,
+            (e) => (expr(e, current!), true),
+            (c) => (stmts([c], current!), true),
+          );
+          break;
+        case "bytesSet":
+          record(s, s.index, current);
+          everyStmtChild(
+            s,
+            (e) => (expr(e, current!), true),
+            (c) => (stmts([c], current!), true),
+          );
+          break;
+        case "break":
+        case "continue":
+        case "return":
+        case "throw":
+          everyStmtChild(
+            s,
+            (e) => (expr(e, current!), true),
+            (c) => (stmts([c], current!), true),
+          );
+          current = null;
+          break;
+        default:
+          everyStmtChild(
+            s,
+            (e) => (expr(e, current!), true),
+            (c) => (stmts([c], current!), true),
+          );
+          break;
+      }
+    }
+    return current;
+  };
+  const writtenCounters = (list: readonly IrStmt[]): Set<string> => {
+    const out = new Set<string>();
+    const visit = (s: IrStmt): boolean => {
+      const step = counterStep(s);
+      if (step) out.add(step.id);
+      return everyStmtChild(s, () => true, visit);
+    };
+    list.every(visit);
+    return out;
+  };
+  expr(loop.cond, new Map());
+  stmts(loop.body, withFacts(new Map(), loop.cond));
+
+  // Each counter's bounds: the loop's and those of nested while loops that
+  // advance it.
+  const counters = new Map<string, WhileCounter>();
+  const loops: (IrStmt & { kind: "while" })[] = [loop];
+  const collectLoops = (s: IrStmt): boolean => {
+    if (s.kind === "while") loops.push(s);
+    return everyStmtChild(s, () => true, collectLoops);
+  };
+  loop.body.every(collectLoops);
+  for (const [id, step] of steps) {
+    if (!counter(id)) continue;
+    let bounds: WhileCounter["bounds"] = [];
+    for (const l of loops) {
+      if (l !== loop && !writtenCounters(l.body).has(id)) continue;
+      const fact = withFacts(new Map(), l.cond).get(id);
+      if (!fact) {
+        bounds = null;
+        break;
+      }
+      bounds.push({ bound: fact.bound, adj: fact.adj });
+    }
+    counters.set(id, { ...step, bounds });
+  }
+  return { accesses, counters };
+}
+
 /** Branches to `fast` when every access the version proves is in bounds for
  * every iteration and every view it proves thread-owned has no shared
  * storage, and to `slow` otherwise (including a counted loop whose
@@ -553,6 +987,27 @@ export function emitBytesLoopVersionGuard(
     B.line(`${ok} = icmp ugt i64 ${lengths[check.receiver]}, ${check.end}`);
     conditions.push(ok);
   }
+  for (const counter of version.counters) {
+    const value = host.emitExpr(counter.ref).name;
+    conditions.push(emitExactIntegerTest(host, value, counter.min, SAFE).ok);
+    for (const wrap of counter.wrapBounds) {
+      const ok = B.tmp();
+      B.line(`${ok} = fcmp ole double ${host.emitExpr(wrap.bound).name}, ${f64Lit(wrap.limit)}`);
+      conditions.push(ok);
+    }
+  }
+  for (const check of version.bounded) {
+    // `bound <= length - slack` in doubles: both sides are exact (a length
+    // is below 2^53 and a slack below 2^32), and NaN fails.
+    const bound = host.emitExpr(check.bound).name;
+    const length = B.tmp(),
+      room = B.tmp(),
+      ok = B.tmp();
+    B.line(`${length} = uitofp i64 ${lengths[check.receiver]} to double`);
+    B.line(`${room} = fsub double ${length}, ${f64Lit(check.slack)}`);
+    B.line(`${ok} = fcmp ole double ${bound}, ${room}`);
+    conditions.push(ok);
+  }
   if (counted && version.affine.length > 0) {
     // The first and last counter values, in ascending order. The last value
     // of a non-unit step is exact: `start + ⌊(limit ∓ 1 - start) / step⌋ *
@@ -599,7 +1054,7 @@ export function emitBytesLoopVersionGuard(
       for (const term of check.terms) {
         const value = host.emitExpr(term.expr);
         if (version.guarded.has(term.key)) {
-          const tested = emitGuardedTerm(host, value.name);
+          const tested = emitExactIntegerTest(host, value.name, GUARDED_TERM.min, GUARDED_TERM.max);
           conditions.push(tested.ok);
           const sum = B.tmp();
           B.line(`${sum} = ${term.sign > 0 ? "add" : "sub"} i64 ${base}, ${tested.integer}`);
@@ -653,10 +1108,15 @@ export function emitBytesLoopVersionGuard(
   else B.condBr(valid, fast, slow);
 }
 
-/** Tests that a guarded term is an exact integer within GUARDED_TERM (not
- * -0; NaN fails the ordered comparisons) and gives its i64. The integer is
- * frozen, so it is a fixed value even when the test fails. */
-function emitGuardedTerm(host: LlvmEmitterContext, value: string): { ok: string; integer: string } {
+/** Tests that a double is an exact integer within [min, max] (not -0; NaN
+ * fails the ordered comparisons) and gives its i64. The integer is frozen,
+ * so it is a fixed value even when the test fails. */
+function emitExactIntegerTest(
+  host: LlvmEmitterContext,
+  value: string,
+  min: number,
+  max: number,
+): { ok: string; integer: string } {
   const B = host.B;
   const low = B.tmp(),
     high = B.tmp(),
@@ -668,8 +1128,8 @@ function emitGuardedTerm(host: LlvmEmitterContext, value: string): { ok: string;
     backBits = B.tmp(),
     exact = B.tmp(),
     ok = B.tmp();
-  B.line(`${low} = fcmp oge double ${value}, ${f64Lit(GUARDED_TERM.min)}`);
-  B.line(`${high} = fcmp ole double ${value}, ${f64Lit(GUARDED_TERM.max)}`);
+  B.line(`${low} = fcmp oge double ${value}, ${f64Lit(min)}`);
+  B.line(`${high} = fcmp ole double ${value}, ${f64Lit(max)}`);
   B.line(`${inRange} = and i1 ${low}, ${high}`);
   B.line(`${raw} = fptosi double ${value} to i64`);
   B.line(`${integer} = freeze i64 ${raw}`);

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

@@ -51,7 +51,7 @@ import {
   emitCountedLoopSafety,
   emitPairedLoopSafety,
 } from "./counted-loops.js";
-import { emitBytesLoopVersionGuard } from "./bytes-loop-versions.js";
+import { emitBytesLoopVersionGuard, type BytesLoopVersion } from "./bytes-loop-versions.js";
 import { mangleLocal } from "../mangle.js";
 import { LlvmUnsupportedError } from "./unsupported.js";
 import type { LlValue } from "./expr-context.js";
@@ -795,6 +795,16 @@ export function emitStmtBody(emitter: LlEmitter, s: IrStmt): void {
       break;
     }
     case "while": {
+      const version =
+        emitter.bytesLoopVersionsEnabled &&
+        emitter.bytesLoopVersioning.size < MAX_BYTES_LOOP_VERSION_DEPTH &&
+        !emitter.bytesLoopVersioning.has(s)
+          ? emitter.matchBytesLoopVersion(s, null)
+          : null;
+      if (version) {
+        emitVersionedLoop(emitter, s, version);
+        break;
+      }
       const workerBudget = emitter.workerLoopBudget(s);
       const le = B.newLabel("loop.e");
       emitter.emitLoopCopies(s, true, le, (exit) => {
@@ -981,35 +991,7 @@ export function emitStmtBody(emitter: LlEmitter, s: IrStmt): void {
           : countedLoop;
       const version = versionable ? emitter.matchBytesLoopVersion(s, versionCounted) : null;
       if (version) {
-        const fast = B.newLabel("bytes.version.fast"),
-          slow = B.newLabel("bytes.version.slow"),
-          done = B.newLabel("bytes.version.done");
-        emitBytesLoopVersionGuard(emitter, version, fast, slow);
-        const originalBounds = emitter.bytesBounds,
-          originalOwned = emitter.bytesOwned,
-          originalRanges = emitter.integerRanges;
-        emitter.bytesLoopVersioning.add(s);
-        try {
-          B.startBlock(slow);
-          emitter.emitStmt(s);
-          B.br(done);
-          B.startBlock(fast);
-          emitter.bytesBounds = new Set([...originalBounds, ...version.bounds]);
-          if (version.owned.size > 0)
-            emitter.bytesOwned = new Set([...originalOwned, ...version.owned]);
-          if (version.ranges.size > 0)
-            emitter.integerRanges = new Map([...originalRanges, ...version.ranges]);
-          if (version.counted?.guarded) emitter.provenCountedGuards.add(s);
-          emitter.emitStmt(s);
-          B.br(done);
-        } finally {
-          emitter.bytesBounds = originalBounds;
-          emitter.bytesOwned = originalOwned;
-          emitter.integerRanges = originalRanges;
-          emitter.provenCountedGuards.delete(s);
-          emitter.bytesLoopVersioning.delete(s);
-        }
-        B.startBlock(done);
+        emitVersionedLoop(emitter, s, version);
         break;
       }
       let countedJoin: string | null = null;
@@ -1896,6 +1878,45 @@ const SPECULATION_BAILS = 8;
  * when speculative-integers.ts plans a fast integer copy. Both copies
  * share the loop's setup and exit; the ordinary copy comes first so the
  * fast copy is only emitted when every resume point exists. */
+/** A typed-array loop version: the guard, then the checked copy, which
+ * runs whenever the guard fails, and the copy whose accesses the guard
+ * proves in bounds (or thread-owned). */
+function emitVersionedLoop(
+  emitter: LlEmitter,
+  s: IrStmt & { kind: "for" | "while" },
+  version: BytesLoopVersion,
+): void {
+  const B = emitter.B;
+  const fast = B.newLabel("bytes.version.fast"),
+    slow = B.newLabel("bytes.version.slow"),
+    done = B.newLabel("bytes.version.done");
+  emitBytesLoopVersionGuard(emitter, version, fast, slow);
+  const originalBounds = emitter.bytesBounds,
+    originalOwned = emitter.bytesOwned,
+    originalRanges = emitter.integerRanges;
+  emitter.bytesLoopVersioning.add(s);
+  try {
+    B.startBlock(slow);
+    emitter.emitStmt(s);
+    B.br(done);
+    B.startBlock(fast);
+    emitter.bytesBounds = new Set([...originalBounds, ...version.bounds]);
+    if (version.owned.size > 0) emitter.bytesOwned = new Set([...originalOwned, ...version.owned]);
+    if (version.ranges.size > 0)
+      emitter.integerRanges = new Map([...originalRanges, ...version.ranges]);
+    if (version.counted?.guarded) emitter.provenCountedGuards.add(s);
+    emitter.emitStmt(s);
+    B.br(done);
+  } finally {
+    emitter.bytesBounds = originalBounds;
+    emitter.bytesOwned = originalOwned;
+    emitter.integerRanges = originalRanges;
+    emitter.provenCountedGuards.delete(s);
+    emitter.bytesLoopVersioning.delete(s);
+  }
+  B.startBlock(done);
+}
+
 export function emitLoopCopies(
   emitter: LlEmitter,
   loop: Extract<IrStmt, { kind: "while" | "doWhile" | "for" }>,

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

@@ -3843,7 +3843,7 @@ export class LlEmitter {
   /** A binding's storage: a module global, a plain local slot, or a boxed
    * local (the slot holds the capture BOX; access goes through it). */
   matchBytesLoopVersion(
-    s: IrStmt & { kind: "for" },
+    s: IrStmt & { kind: "for" | "while" },
     counted: IntegerCountedForLoop | null,
   ): ReturnType<typeof matchBytesLoopVersion> {
     const fn = this.bytesLoopFunction;

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/6998-typed-array-while-loop-versions.ts.json

@@ -0,0 +1,6 @@
+{
+ "order": [
+  "<repo>/tests/corpus/6998-typed-array-while-loop-versions.ts"
+ ],
+ "diags": []
+}

tests/corpus/6998-typed-array-while-loop-versions.ts

@@ -0,0 +1,149 @@
+// While loops that scan typed arrays with a counter below a bound: scanners
+// with nested loops and `&&` conditions, `| 0` increments, look-behind and
+// look-ahead reads, writes, DataView reads, string-length bounds, and
+// bounds or starts past the array, fractional, negative, -0 or NaN.
+
+function words(bytes: Uint8Array, end: number): number {
+  let pos = 0;
+  let count = 0;
+  while (pos < end) {
+    while (pos < end && (bytes[pos] === 32 || bytes[pos] === 10)) pos++;
+    if (pos < end) count++;
+    while (pos < end && bytes[pos] !== 32 && bytes[pos] !== 10) pos++;
+  }
+  return count;
+}
+
+function identifiers(chars: Uint16Array, start: number, end: number): string {
+  let pos = start;
+  let count = 0;
+  let mix = 0;
+  while (pos < end) {
+    let hash = 0x811c9dc5 | 0;
+    while (pos < end) {
+      const c = chars[pos]!;
+      if (!((c >= 97 && c <= 122) || c === 95)) break;
+      hash = Math.imul(hash ^ c, 0x01000193);
+      pos = (pos + 1) | 0;
+    }
+    mix = (mix + hash) | 0;
+    count++;
+    pos = (pos + 1) | 0;
+  }
+  return `${count} ${mix >>> 0} ${pos}`;
+}
+
+function around(values: Int16Array, start: number, end: number): number[] {
+  const out: number[] = [];
+  let i = start;
+  while (i <= end) {
+    out.push(values[i - 1]! + values[i + 2]! * 2);
+    i += 1;
+  }
+  return out;
+}
+
+function fill(out: Float32Array, end: number, value: number): number {
+  let i = 0;
+  while (i < end) {
+    out[i] = value + i;
+    if (i < end - 1 && out[i] > 2) {
+      i += 2;
+      continue;
+    }
+    i++;
+  }
+  return i;
+}
+
+function afterInner(bytes: Uint8Array): number {
+  let pos = 0;
+  let sum = 0;
+  while (pos < bytes.length) {
+    while (pos < bytes.length && bytes[pos] !== 0) pos++;
+    // The inner loop may stop at the length: this read stays checked.
+    sum += bytes[pos] ?? -1;
+    pos++;
+  }
+  return sum;
+}
+
+function firstAbove(samples: Float64Array, limit: number): number {
+  let i = 0;
+  while (i < samples.length) {
+    if (samples[i]! > limit) return i;
+    i = i + 1;
+  }
+  return -1;
+}
+
+function halves(view: DataView, end: number): number {
+  let at = 0;
+  let total = 0;
+  while (at < end) {
+    total += view.getUint16(at, true);
+    at += 2;
+  }
+  return total;
+}
+
+function letters(text: string, codes: Uint8Array): number {
+  let i = 0;
+  let sum = 0;
+  while (i < text.length) {
+    sum += codes[i]! * (text.charCodeAt(i) & 7);
+    i++;
+  }
+  return sum;
+}
+
+function notCounters(bytes: Uint8Array): string {
+  let i = 1;
+  let doubled = 0;
+  while (i < bytes.length) {
+    doubled += bytes[i]!;
+    i = i * 2;
+  }
+  let j = bytes.length - 1;
+  let down = 0;
+  while (j < bytes.length && j >= 0) {
+    down += bytes[j]!;
+    j--;
+  }
+  return `${doubled} ${down}`;
+}
+
+const text = "let  alpha = beta_gamma\nreturn delta  ";
+const bytes = new Uint8Array(text.length);
+const chars = new Uint16Array(text.length);
+for (let i = 0; i < text.length; i++) {
+  bytes[i] = text.charCodeAt(i);
+  chars[i] = text.charCodeAt(i);
+}
+console.log("words", words(bytes, bytes.length), words(bytes, 10), words(bytes, bytes.length + 5), words(bytes, 7.5), words(bytes, NaN));
+console.log("identifiers", identifiers(chars, 0, chars.length));
+console.log("identifiers past end", identifiers(chars, 4, chars.length + 3));
+console.log("identifiers starts", identifiers(chars, -2, 9), identifiers(chars, 2.5, 12), identifiers(chars, -0, 6));
+
+const values = new Int16Array([5, -6, 7, -8, 9, -10, 11]);
+console.log("around", around(values, 1, 4).join(","), around(values, 0, 6).join(","), around(values, -1, 1).join(","));
+
+const floats = new Float32Array(6);
+console.log("fill", fill(floats, 6, 0.5), floats.join(","));
+const short = new Float32Array(3);
+console.log("fill past end", fill(short, 5, 1), short.join(","));
+
+console.log("after inner", afterInner(new Uint8Array([1, 2, 0, 3, 0, 4, 5])));
+console.log("first above", firstAbove(new Float64Array([0.5, 1.5, 2.5]), 1), firstAbove(new Float64Array([0.5]), 1));
+
+const buffer = new Uint8Array([1, 2, 3, 4, 5, 6, 7, 8]);
+const view = new DataView(buffer.buffer);
+console.log("halves", halves(view, 8));
+try {
+  console.log("halves past end", halves(view, 10));
+} catch (error) {
+  console.log((error as Error).name, (error as Error).message);
+}
+
+console.log("letters", letters("abcdef", new Uint8Array([1, 2, 3, 4, 5, 6])), letters("abcdefgh", new Uint8Array([1, 2, 3])));
+console.log("not counters", notCounters(new Uint8Array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])));

tests/harness/corpus-costs/6998-typed-array-while-loop-versions.ts.json

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