Run counted loops from an unproven start as integer loops

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

@@ -6051,14 +6051,49 @@ export class LlEmitter {
         }
         let countedJoin: string | null = null;
         if (countedLoop?.guarded) {
-          const bound = this.emitExpr(countedLoop.limit);
-          const safe = B.tmp();
+          let safe = "true";
+          if (countedLoop.limitGuarded) {
+            const bound = this.emitExpr(countedLoop.limit);
+            safe = B.tmp();
+            B.line(
+              `${safe} = fcmp ${countedLoop.step > 0 ? "ole" : "oge"} double ${bound.name}, ${f64Lit(countedLoop.guardLimit)}`,
+            );
+          }
+          if (countedLoop.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 = this.emitExpr(countedLoop.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;
+            }
+          }
           const fast = B.newLabel("counted.fast");
           const slow = B.newLabel("counted.slow");
           countedJoin = B.newLabel("counted.done");
-          B.line(
-            `${safe} = fcmp ${countedLoop.step > 0 ? "ole" : "oge"} double ${bound.name}, ${f64Lit(countedLoop.guardLimit)}`,
-          );
           B.condBr(safe, fast, slow);
           B.startBlock(slow);
           this.countedLoopsEnabled = false;

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

@@ -164,3 +164,21 @@ test("versioning is bounded and requires an integer consumer", () => {
   expect(matchIntegerCountedForLoop(nested, withLimit)).toBeNull();
   expect(matchIntegerCountedForLoop(counted(number(4)), locals)?.guarded).toBe(false);
 });
+
+test("an unproven start versions the loop on an exact integer start", () => {
+  const withStart = new Map(locals);
+  withStart.set("start", { id: "start", name: "start", type: F64, mutable: false });
+  const s = counted(number(8));
+  s.init = { kind: "varDecl", localId: "i", init: ref("start"), loc };
+  expect(matchIntegerCountedForLoop(s, withStart)).toMatchObject({
+    guarded: true,
+    startGuarded: true,
+    limitGuarded: false,
+    range: { min: -Number.MAX_SAFE_INTEGER, max: 7 },
+  });
+  // The versioned body is duplicated, so the usual size bound applies.
+  const large = counted(number(8));
+  large.init = { kind: "varDecl", localId: "i", init: ref("start"), loc };
+  large.body = Array.from({ length: 81 }, () => large.body[0]!);
+  expect(matchIntegerCountedForLoop(large, withStart)).toBeNull();
+});

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

@@ -20,6 +20,11 @@ export interface IntegerCountedForLoop {
   step: number;
   inclusive: boolean;
   guarded: boolean;
+  /** The limit needs the run-time `guardLimit` test (part of `guarded`). */
+  limitGuarded: boolean;
+  /** The start has no proven integer range: the versioned loop also tests
+   * that it is an exact integer (not -0) of at most 2^53 in magnitude. */
+  startGuarded: boolean;
   guardLimit: number;
   range: IntegerRange;
 }
@@ -104,10 +109,13 @@ export function matchIntegerCountedForLoop(
   const local = locals.get(init.localId);
   if (local?.type.kind !== "f64" || !local.mutable || local.boxed || local.tdz) return null;
   const start = init.init;
-  const startRange =
+  let startRange =
     start.kind === "numLit" && Number.isSafeInteger(start.value) && !Object.is(start.value, -0)
       ? { min: start.value, max: start.value }
       : ranges.get(start);
+  // An unproven start is tested at run time, on a versioned loop.
+  const startGuarded = !startRange && start.kind !== "numLit" && start.type.kind === "f64";
+  if (startGuarded) startRange = { min: -Number.MAX_SAFE_INTEGER, max: Number.MAX_SAFE_INTEGER };
   if (
     !startRange ||
     startRange.min < -Number.MAX_SAFE_INTEGER ||
@@ -132,9 +140,10 @@ export function matchIntegerCountedForLoop(
   const guardLimit = step > 0 ? 2 ** 53 - (inclusive ? 1 : 0) : -(2 ** 53) + (inclusive ? 1 : 0);
   const boundRange =
     limit.kind === "numLit" ? { min: limit.value, max: limit.value } : ranges.get(limit);
-  const guarded =
+  const limitGuarded =
     !boundRange || (step > 0 ? !(boundRange.max <= guardLimit) : !(boundRange.min >= guardLimit));
-  if (guarded && limit.kind === "numLit") return null;
+  if (limitGuarded && limit.kind === "numLit") return null;
+  const guarded = limitGuarded || startGuarded;
   let useful = false;
   let cost = 0;
   const counter = (e: IrExpr): boolean =>
@@ -204,7 +213,18 @@ export function matchIntegerCountedForLoop(
           ),
           max: startRange.max,
         };
-  return { localId: local.id, start, limit, step, inclusive, guarded, guardLimit, range };
+  return {
+    localId: local.id,
+    start,
+    limit,
+    step,
+    inclusive,
+    guarded,
+    limitGuarded,
+    startGuarded,
+    guardLimit,
+    range,
+  };
 }
 
 /** True when a lowered subtree writes `localId`. Local ids are unique per

tests/corpus/exact-index-round-trip.ts

@@ -96,3 +96,17 @@ console.log(
   rangeHash(cursor.units, 0.5, 4),
   rangeHash(cursor.units, opaque(odd, 1), opaque(odd, 3)),
 );
+
+// Counted loops from an unproven start run as integer loops only when the
+// start is an exact integer; -0, fractions, NaN and huge starts keep the
+// double loop, and the counter's own value stays observable.
+function observe(chars: Uint16Array, start: number, end: number): string {
+  const seen: string[] = [];
+  for (let i = start; i < end; i++) {
+    seen.push(`${Object.is(i, -0) ? "-0" : i}:${chars[i] ^ 1}:${(i * 2) | 0}`);
+    if (seen.length > 6) break;
+  }
+  return seen.join(" ");
+}
+for (const start of [0, -0, 1, 2.5, -2, NaN, 2 ** 53 - 2, -(2 ** 53), 1e300, -Infinity])
+  console.log(String(start), observe(cursor.units, opaque([start], 0), opaque([start + 4, 9], start === -Infinity ? 1 : 0)));