Run two-counter for loops on integer counters (#1768)

packages/compiler/src/backend/llvm/counted-loops.ts

@@ -1,4 +1,5 @@
-import type { IntegerCountedForLoop } from "../../ir/integer-loops.js";
+import type { IntegerCountedForLoop, IntegerPairedForLoop } from "../../ir/integer-loops.js";
+import type { IrExpr } from "../../ir/ir.js";
 import type { LlvmEmitterContext } from "./expr-context.js";
 import { f64Lit } from "./common.js";
 
@@ -59,37 +60,51 @@ export function emitCountedLoopSafety(
       `${safe} = fcmp ${loop.step > 0 ? "ole" : "oge"} double ${bound.name}, ${f64Lit(loop.guardLimit)}`,
     );
   }
-  if (loop.startGuarded) {
-    // An exact integer of at most 2^53 - 1 in magnitude, and not -0
-    // (the bit patterns of the value and its round trip agree).
-    // NaN and infinities fail the ordered range tests.
-    const start = host.emitExpr(loop.start);
-    const low = B.tmp(),
-      high = B.tmp(),
-      raw = B.tmp(),
-      wide = B.tmp(),
-      back = B.tmp(),
-      bits = B.tmp(),
-      backBits = B.tmp(),
-      exact = B.tmp(),
-      inRange = B.tmp(),
-      ok = B.tmp();
-    B.line(`${low} = fcmp oge double ${start.name}, ${f64Lit(-Number.MAX_SAFE_INTEGER)}`);
-    B.line(`${high} = fcmp ole double ${start.name}, ${f64Lit(Number.MAX_SAFE_INTEGER)}`);
-    B.line(`${inRange} = and i1 ${low}, ${high}`);
-    B.line(`${raw} = fptosi double ${start.name} to i64`);
-    B.line(`${wide} = freeze i64 ${raw}`);
-    B.line(`${back} = sitofp i64 ${wide} to double`);
-    B.line(`${bits} = bitcast double ${start.name} to i64`);
-    B.line(`${backBits} = bitcast double ${back} to i64`);
-    B.line(`${exact} = icmp eq i64 ${bits}, ${backBits}`);
-    B.line(`${ok} = and i1 ${inRange}, ${exact}`);
-    if (safe === "true") safe = ok;
-    else {
-      const both = B.tmp();
-      B.line(`${both} = and i1 ${safe}, ${ok}`);
-      safe = both;
-    }
-  }
+  if (loop.startGuarded) safe = andSafe(host, safe, emitExactStart(host, loop.start));
   return safe;
 }
+
+/** The run-time test of a guarded paired loop: each unproven start is an
+ * exact safe integer. */
+export function emitPairedLoopSafety(host: LlvmEmitterContext, loop: IntegerPairedForLoop): string {
+  let safe = "true";
+  for (const counter of [loop.low, loop.high])
+    if (counter.startGuarded) safe = andSafe(host, safe, emitExactStart(host, counter.start));
+  return safe;
+}
+
+function andSafe(host: LlvmEmitterContext, safe: string, ok: string): string {
+  if (safe === "true") return ok;
+  const both = host.B.tmp();
+  host.B.line(`${both} = and i1 ${safe}, ${ok}`);
+  return both;
+}
+
+/** An exact integer of at most 2^53 - 1 in magnitude, and not -0 (the bit
+ * patterns of the value and its round trip agree). NaN and infinities fail
+ * the ordered range tests. */
+function emitExactStart(host: LlvmEmitterContext, expr: IrExpr): string {
+  const B = host.B;
+  const start = host.emitExpr(expr);
+  const low = B.tmp(),
+    high = B.tmp(),
+    raw = B.tmp(),
+    wide = B.tmp(),
+    back = B.tmp(),
+    bits = B.tmp(),
+    backBits = B.tmp(),
+    exact = B.tmp(),
+    inRange = B.tmp(),
+    ok = B.tmp();
+  B.line(`${low} = fcmp oge double ${start.name}, ${f64Lit(-Number.MAX_SAFE_INTEGER)}`);
+  B.line(`${high} = fcmp ole double ${start.name}, ${f64Lit(Number.MAX_SAFE_INTEGER)}`);
+  B.line(`${inRange} = and i1 ${low}, ${high}`);
+  B.line(`${raw} = fptosi double ${start.name} to i64`);
+  B.line(`${wide} = freeze i64 ${raw}`);
+  B.line(`${back} = sitofp i64 ${wide} to double`);
+  B.line(`${bits} = bitcast double ${start.name} to i64`);
+  B.line(`${backBits} = bitcast double ${back} to i64`);
+  B.line(`${exact} = icmp eq i64 ${bits}, ${backBits}`);
+  B.line(`${ok} = and i1 ${inRange}, ${exact}`);
+  return ok;
+}

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

@@ -42,10 +42,15 @@ import {
   matchIntegerArrayForLoop,
   matchIntegerBytesForLoop,
   matchIntegerCountedForLoop,
+  matchIntegerPairedForLoop,
 } from "../../ir/integer-loops.js";
 import { emitByteWindowGuard, matchByteWindow } from "./byte-windows.js";
 import { countedDenseShape, emitDenseLoopGuard, matchDenseArrayLoop } from "./dense-array-loops.js";
-import { emitCountedLoopLimit, emitCountedLoopSafety } from "./counted-loops.js";
+import {
+  emitCountedLoopLimit,
+  emitCountedLoopSafety,
+  emitPairedLoopSafety,
+} from "./counted-loops.js";
 import { emitBytesLoopVersionGuard } from "./bytes-loop-versions.js";
 import { mangleLocal } from "../mangle.js";
 import { LlvmUnsupportedError } from "./unsupported.js";
@@ -866,6 +871,12 @@ export function emitStmtBody(emitter: LlEmitter, s: IrStmt): void {
               emitter.unaliasedArrays,
             )
           : null;
+      // Two counters closing in on each other (`lo++, hi--`) take the same
+      // integer storage as a counted loop's one.
+      const pairedLoop =
+        !integerLoop && !countedLoop && emitter.countedLoopsEnabled
+          ? matchIntegerPairedForLoop(s, emitter.numericLocals, emitter.integerRanges)
+          : null;
       const window =
         countedLoop && emitter.byteWindowsEnabled
           ? matchByteWindow(
@@ -998,8 +1009,10 @@ export function emitStmtBody(emitter: LlEmitter, s: IrStmt): void {
         break;
       }
       let countedJoin: string | null = null;
-      if (countedLoop?.guarded && !emitter.provenCountedGuards.has(s)) {
-        const safe = emitCountedLoopSafety(emitter, countedLoop);
+      if ((countedLoop?.guarded && !emitter.provenCountedGuards.has(s)) || pairedLoop?.guarded) {
+        const safe = countedLoop
+          ? emitCountedLoopSafety(emitter, countedLoop)
+          : emitPairedLoopSafety(emitter, pairedLoop!);
         const fast = B.newLabel("counted.fast");
         const slow = B.newLabel("counted.slow");
         countedJoin = B.newLabel("counted.done");
@@ -1014,6 +1027,7 @@ export function emitStmtBody(emitter: LlEmitter, s: IrStmt): void {
       emitter.scopes.push([]);
       let integerSlot: string | null = null;
       let countedLimit: string | null = null;
+      const pairedSlots: string[] = [];
       if (integerLoop) {
         integerSlot = B.slot();
         B.entryAllocas.push(
@@ -1058,6 +1072,24 @@ export function emitStmtBody(emitter: LlEmitter, s: IrStmt): void {
               : countedLoop.range,
         });
         countedLimit = emitCountedLoopLimit(emitter, countedLoop);
+      } else if (pairedLoop) {
+        for (const counter of [pairedLoop.low, pairedLoop.high]) {
+          const slot = B.slot();
+          B.entryAllocas.push(
+            `${slot} = alloca i64 ; integer induction ${emitter.currentLocals.get(counter.localId)!.name}`,
+          );
+          const value = emitter.emitExpr(counter.start);
+          const start = B.tmp();
+          B.line(`${start} = fptosi double ${value.name} to i64`);
+          B.line(`store i64 ${start}, ptr ${slot}`);
+          pairedSlots.push(slot);
+          emitter.integerLoopBindings.set(counter.localId, {
+            slot,
+            type: "i64",
+            signed: true,
+            range: pairedLoop.range,
+          });
+        }
       } else if (s.init) {
         // A multi-declarator head shares the loop's scope. Emitting its
         // IR block as an ordinary block would release captured/ref locals
@@ -1139,6 +1171,14 @@ export function emitStmtBody(emitter: LlEmitter, s: IrStmt): void {
             `${inBounds} = icmp ${countedLoop.step > 0 ? "slt" : "sgt"} i64 ${index}, ${countedLimit}`,
           );
           B.condBr(inBounds, lb, exit);
+        } else if (pairedLoop) {
+          const low = B.tmp(),
+            high = B.tmp(),
+            open = B.tmp();
+          B.line(`${low} = load i64, ptr ${pairedSlots[0]}`);
+          B.line(`${high} = load i64, ptr ${pairedSlots[1]}`);
+          B.line(`${open} = icmp ${pairedLoop.inclusive ? "sle" : "slt"} i64 ${low}, ${high}`);
+          B.condBr(open, lb, exit);
         } else if (s.cond) B.condBr(emitter.emitCondition(s.cond), lb, exit);
         else B.br(lb);
         B.startBlock(lb);
@@ -1170,6 +1210,14 @@ export function emitStmtBody(emitter: LlEmitter, s: IrStmt): void {
             B.line(`${old} = load i64, ptr ${integerSlot}`);
             B.line(`${next} = add nsw i64 ${old}, ${countedLoop.step}`);
             B.line(`store i64 ${next}, ptr ${integerSlot}`);
+          } else if (pairedLoop) {
+            [pairedLoop.low, pairedLoop.high].forEach((counter, i) => {
+              const old = B.tmp();
+              const next = B.tmp();
+              B.line(`${old} = load i64, ptr ${pairedSlots[i]}`);
+              B.line(`${next} = add nsw i64 ${old}, ${counter.step}`);
+              B.line(`store i64 ${next}, ptr ${pairedSlots[i]}`);
+            });
           } else if (s.update) emitter.emitStmt(s.update);
           B.br(lc);
         }
@@ -1178,6 +1226,10 @@ export function emitStmtBody(emitter: LlEmitter, s: IrStmt): void {
       emitter.releaseScope(emitter.scopes.pop()!);
       if (integerLoop) emitter.integerLoopBindings.delete(integerLoop.localId);
       if (countedLoop) emitter.integerLoopBindings.delete(countedLoop.localId);
+      if (pairedLoop) {
+        emitter.integerLoopBindings.delete(pairedLoop.low.localId);
+        emitter.integerLoopBindings.delete(pairedLoop.high.localId);
+      }
       if (integerArrayLoop) emitter.integerArrayBindings.delete(integerArrayLoop.localId);
       if (countedJoin) {
         B.br(countedJoin);

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

@@ -1,6 +1,10 @@
 import { expect, test } from "vitest";
 import { BOOL, F64, STRING, arrayOf, type IrExpr, type IrLocal, type IrStmt } from "./ir.js";
-import { matchIntegerArrayForLoop, matchIntegerCountedForLoop } from "./integer-loops.js";
+import {
+  matchIntegerArrayForLoop,
+  matchIntegerCountedForLoop,
+  matchIntegerPairedForLoop,
+} from "./integer-loops.js";
 
 const loc = { file: "integer-array-loop.ts", start: 0, end: 0 };
 const number = (value: number): IrExpr => ({ kind: "numLit", value, type: F64, loc });
@@ -355,3 +359,70 @@ test("a descending loop from an array's last index takes an integer counter for
     range: { min: 0, max: 4294967294 },
   });
 });
+
+test("paired counters closing in on each other run as integers", () => {
+  const vars = new Map(locals);
+  for (const id of ["lo", "hi", "head"]) vars.set(id, { id, name: id, type: F64, mutable: true });
+  const step = (id: string, op: "+" | "-", value = 1): IrStmt => ({
+    kind: "assign",
+    localId: id,
+    value: { kind: "bin", op, left: ref(id), right: number(value), type: F64, loc },
+    loc,
+  });
+  const element = (index: IrExpr): IrExpr => ({
+    kind: "arrayGet",
+    arr: { kind: "varRef", localId: "a", type: arrayOf(F64), loc },
+    index,
+    type: F64,
+    loc,
+  });
+  const paired = (
+    lowStart: IrExpr,
+    highStart: IrExpr,
+    update = [step("lo", "+"), step("hi", "-")],
+  ) =>
+    ({
+      kind: "for",
+      loc,
+      init: {
+        kind: "block",
+        body: [
+          { kind: "varDecl", localId: "lo", init: lowStart, loc },
+          { kind: "varDecl", localId: "hi", init: highStart, loc },
+        ],
+        loc,
+      },
+      cond: { kind: "bin", op: "<", left: ref("lo"), right: ref("hi"), type: BOOL, loc },
+      update: { kind: "block", body: update, loc },
+      body: [{ kind: "exprStmt", expr: element(ref("hi")), loc }],
+    }) as IrStmt & { kind: "for" };
+  expect(matchIntegerPairedForLoop(paired(number(0), number(9)), vars)).toMatchObject({
+    low: { localId: "lo", step: 1, startGuarded: false },
+    high: { localId: "hi", step: -1, startGuarded: false },
+    guarded: false,
+    range: { min: 0, max: 9 },
+  });
+  // A start without a known range is tested at run time.
+  expect(matchIntegerPairedForLoop(paired(number(0), ref("head")), vars)).toMatchObject({
+    high: { startGuarded: true },
+    guarded: true,
+  });
+  // Counters moving the same way, written by the body, captured, started
+  // from each other, or never used as an index are refused.
+  expect(
+    matchIntegerPairedForLoop(
+      paired(number(0), number(9), [step("lo", "+"), step("hi", "+")]),
+      vars,
+    ),
+  ).toBeNull();
+  const writing = paired(number(0), number(9));
+  writing.body.push(step("hi", "-", 2));
+  expect(matchIntegerPairedForLoop(writing, vars)).toBeNull();
+  const boxed = new Map(vars);
+  boxed.set("hi", { ...vars.get("hi")!, boxed: true });
+  expect(matchIntegerPairedForLoop(paired(number(0), number(9)), boxed)).toBeNull();
+  expect(matchIntegerPairedForLoop(paired(number(0), ref("lo")), vars)).toBeNull();
+  const unused = paired(number(0), number(9));
+  unused.body = [];
+  expect(matchIntegerPairedForLoop(unused, vars)).toBeNull();
+});

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

@@ -281,6 +281,120 @@ export function matchIntegerCountedForLoop(
   };
 }
 
+/** A for loop whose two counters close in on each other:
+ *
+ *   for (let lo = a, hi = b; lo < hi; lo++, hi--) { ... }
+ *
+ * The first counter only grows and the second only shrinks, by positive
+ * literal steps, the body writes neither, and the loop runs while the first
+ * is below (or at most) the second. Every body entry therefore holds both
+ * between the first start and the second start: exact integers when the
+ * starts are. A start without a proven integer range is tested at run time
+ * on a versioned loop, as a counted loop's is. */
+export interface IntegerPairedCounter {
+  localId: string;
+  start: IrExpr;
+  /** The signed step of each update. */
+  step: number;
+  startGuarded: boolean;
+}
+export interface IntegerPairedForLoop {
+  low: IntegerPairedCounter;
+  high: IntegerPairedCounter;
+  inclusive: boolean;
+  guarded: boolean;
+  /** Both counters' values on every body entry. */
+  range: IntegerRange;
+}
+
+export function matchIntegerPairedForLoop(
+  stmt: IrStmt & { kind: "for" },
+  locals: ReadonlyMap<string, IrLocal>,
+  ranges: IntegerRanges = new Map(),
+): IntegerPairedForLoop | null {
+  const init = stmt.init,
+    update = stmt.update,
+    cond = stmt.cond;
+  if (init?.kind !== "block" || init.body.length !== 2) return null;
+  if (update?.kind !== "block" || update.body.length !== 2) return null;
+  if (cond?.kind !== "bin" || (cond.op !== "<" && cond.op !== "<=")) return null;
+  if (cond.left.kind !== "varRef" || cond.right.kind !== "varRef") return null;
+  const lowId = cond.left.localId,
+    highId = cond.right.localId;
+  if (lowId === highId) return null;
+  const counter = (id: string): IntegerPairedCounter | null => {
+    const decl = init.body.find((d) => d.kind === "varDecl" && d.localId === id);
+    if (decl?.kind !== "varDecl" || !decl.init) return null;
+    const local = locals.get(id);
+    if (local?.type.kind !== "f64" || !local.mutable || local.boxed || local.tdz) return null;
+    const start = decl.init;
+    // Each start is evaluated once in the head and again on a fallback path.
+    if (
+      !stableLoopValue(start, locals, stmt.body, false, lowId) ||
+      !stableLoopValue(start, locals, stmt.body, false, highId)
+    )
+      return null;
+    const steps = update.body.map((u) => constantStep(u, id)).filter((v) => v !== null);
+    if (steps.length !== 1 || writesLocal(stmt.body, id)) return null;
+    const known =
+      start.kind === "numLit" && Number.isSafeInteger(start.value) && !Object.is(start.value, -0)
+        ? { min: start.value, max: start.value }
+        : ranges.get(start);
+    const startGuarded = !known && start.kind !== "numLit" && start.type.kind === "f64";
+    if (!known && !startGuarded) return null;
+    if (known && (known.min < -Number.MAX_SAFE_INTEGER || known.max > Number.MAX_SAFE_INTEGER))
+      return null;
+    return { localId: id, start, step: steps[0]!, startGuarded };
+  };
+  const low = counter(lowId),
+    high = counter(highId);
+  if (!low || !high || low.step <= 0 || high.step >= 0) return null;
+  const startRange = (c: IntegerPairedCounter): IntegerRange =>
+    c.startGuarded
+      ? { min: -Number.MAX_SAFE_INTEGER, max: Number.MAX_SAFE_INTEGER }
+      : c.start.kind === "numLit"
+        ? { min: c.start.value, max: c.start.value }
+        : ranges.get(c.start)!;
+  const guarded = low.startGuarded || high.startGuarded;
+  // Only a loop whose counters index something gains from integer storage.
+  let useful = false,
+    cost = 0;
+  const isCounter = (e: IrExpr): boolean =>
+    (e.kind === "varRef" && (e.localId === lowId || e.localId === highId)) ||
+    (e.kind === "bin" &&
+      (e.op === "+" || e.op === "-") &&
+      (isCounter(e.left) || isCounter(e.right)));
+  const expr = (e: IrExpr): boolean => {
+    if (e.kind === "arrayGet" && isCounter(e.index)) useful = true;
+    if (e.kind === "strIntrinsic" && e.method === "charCodeAt" && e.args[0] && isCounter(e.args[0]))
+      useful = true;
+    if (e.kind === "bytesIntrinsic" && e.method === "get" && e.args[0] && isCounter(e.args[0]))
+      useful = true;
+    return (!guarded || ++cost <= 80) && everyExprChild(e, expr, bodyStmt);
+  };
+  const bodyStmt = (s: IrStmt): boolean => {
+    if ((s.kind === "arraySet" || s.kind === "bytesSet") && isCounter(s.index)) useful = true;
+    if (
+      guarded &&
+      (s.kind === "for" || s.kind === "forOf" || s.kind === "while" || s.kind === "doWhile")
+    )
+      return false;
+    return (!guarded || ++cost <= 80) && everyStmtChild(s, expr, bodyStmt);
+  };
+  if (!stmt.body.every(bodyStmt) || !useful) return null;
+  return {
+    low,
+    high,
+    inclusive: cond.op === "<=",
+    guarded,
+    // A loop whose starts already fail the condition never enters the body.
+    range: {
+      min: startRange(low).min,
+      max: Math.max(startRange(low).min, startRange(high).max),
+    },
+  };
+}
+
 /** True when a lowered subtree writes `localId`. Local ids are unique per
  * function, so typed traversal includes writes
  * nested in expressions, branches, nested loops, and try/finally bodies. */

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/7140-paired-loop-counters.ts.json

@@ -0,0 +1,6 @@
+{
+ "order": [
+  "<repo>/tests/corpus/7140-paired-loop-counters.ts"
+ ],
+ "diags": []
+}

tests/corpus/7140-paired-loop-counters.ts

@@ -0,0 +1,81 @@
+// For loops with two counters closing in on each other (`lo++, hi--`):
+// in-place reversals, inclusive bounds, larger steps, early exits, starts
+// that are fractions, negative zero, NaN, huge or already crossed, nested
+// reversals, and counters a closure captures or the body writes.
+function reverse(a: number[], from: number, to: number): number[] {
+  for (let lo = from, hi = to; lo < hi; lo++, hi--) {
+    const t = a[lo]!;
+    a[lo] = a[hi]!;
+    a[hi] = t;
+  }
+  return a;
+}
+const base = (): number[] => [0, 1, 2, 3, 4, 5, 6, 7, 8, 9];
+console.log(reverse(base(), 0, 9).join(","));
+console.log(reverse(base(), 2, 6).join(","));
+console.log(reverse(base(), 6, 2).join(","));
+console.log(reverse(base(), 3, 3).join(","));
+// Unproven starts: fractions, negative zero, NaN and huge values keep JS
+// semantics (charCodeAt truncates its index; out-of-range reads are NaN).
+function codes(s: string, from: number, to: number): string {
+  const out: number[] = [];
+  for (let lo = from, hi = to; lo < hi; lo++, hi--) out.push(s.charCodeAt(lo), s.charCodeAt(hi), lo, hi);
+  return out.join(",");
+}
+const word = "abcdefgh";
+console.log(codes(word, 0.5, 4.5));
+console.log(codes(word, -0, 3), Object.is(-0 + 0, 0));
+console.log(codes(word, NaN, 4), "|", codes(word, 0, NaN));
+console.log(codes(word, 2 ** 53 - 3, 2 ** 53));
+console.log(codes(word, -3, 2));
+
+function pairs(n: number): string {
+  const out: string[] = [];
+  for (let i = 0, j = n; i <= j; i += 2, j -= 3) out.push(`${i}:${j}`);
+  return out.join(" ");
+}
+console.log(pairs(20), "|", pairs(0), "|", pairs(-1));
+
+function firstGap(a: number[]): number {
+  let found = -1;
+  for (let lo = 0, hi = a.length - 1; lo < hi; lo++, hi--) {
+    if (a[lo] === a[hi]) continue;
+    found = lo;
+    break;
+  }
+  return found;
+}
+console.log(firstGap([1, 2, 3, 2, 1]), firstGap([1, 2, 3, 4, 2, 1]), firstGap([]));
+
+// The flip loop of a permutation puzzle: the next index comes from the array.
+function flips(perm: number[]): number {
+  const work = perm.slice();
+  let count = 0;
+  let head = work[0]!;
+  while (head !== 0) {
+    for (let lo = 0, hi = head; lo < hi; lo++, hi--) {
+      const t = work[lo]!;
+      work[lo] = work[hi]!;
+      work[hi] = t;
+    }
+    count++;
+    head = work[0]!;
+  }
+  return count;
+}
+console.log(flips([3, 1, 0, 2]), flips([4, 3, 2, 1, 0]), flips([0, 1, 2]));
+
+// Counters a closure reads, or the body writes, stay ordinary bindings.
+const readers: (() => number)[] = [];
+for (let lo = 0, hi = 4; lo < hi; lo++, hi--) readers.push(() => lo * 10 + hi);
+console.log(readers.map((f) => f()).join(","));
+const seen: number[] = [];
+for (let lo = 0, hi = 10; lo < hi; lo++, hi--) {
+  seen.push(lo, hi);
+  if (lo === 2) hi -= 3;
+}
+console.log(seen.join(","));
+const text = "abcdefgh";
+let mirrored = "";
+for (let lo = 0, hi = text.length - 1; lo <= hi; lo++, hi--) mirrored += text.charAt(hi) + text.charCodeAt(lo);
+console.log(mirrored);

tests/harness/corpus-costs/7140-paired-loop-counters.ts.json

@@ -0,0 +1 @@
+{ "plain": 3880, "san": 6000 }