Keep integer index arithmetic in integers across nested counted loops (#2113)

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

@@ -50,8 +50,8 @@ export function emitCountedLoopLimit(
 }
 
 /** The run-time test of a guarded counted loop: an unproven limit is within
- * the integer form's finite range and an unproven start is an exact safe
- * integer. Returns `true` when nothing needs testing. */
+ * the integer form's finite range and an unproven start is an exact integer
+ * within the loop's start bound. Returns `true` when nothing needs testing. */
 export function emitCountedLoopSafety(
   host: LlvmEmitterContext,
   loop: IntegerCountedForLoop,
@@ -65,7 +65,11 @@ export function emitCountedLoopSafety(
       `${safe} = fcmp ${loop.step > 0 ? "ole" : "oge"} double ${bound.name}, ${f64Lit(loop.guardLimit)}`,
     );
   }
-  if (loop.startGuarded) safe = andSafe(host, safe, emitExactStart(host, loop.start));
+  if (loop.startGuarded) {
+    const start = host.emitExpr(loop.start);
+    const bound = loop.startBound;
+    safe = andSafe(host, safe, emitExactIntegerTest(host, start.name, -bound, bound).ok);
+  }
   for (const follower of loop.followers) {
     if (follower.startGuarded) safe = andSafe(host, safe, emitExactStart(host, follower.start));
     if (follower.tripGuarded) safe = andSafe(host, safe, emitFollowerReach(host, loop, follower));

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

@@ -122,6 +122,7 @@ const counted = (
   guarded: false,
   limitGuarded: false,
   startGuarded: false,
+  startBound: Number.MAX_SAFE_INTEGER,
   guardLimit: 0,
   followers: [],
   ...fields,

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

@@ -80,6 +80,7 @@ import {
 } from "./emit-journal.js";
 import { specializeNumericCalls } from "../../ir/numeric-call-specialization.js";
 import { versionLoopBoundParameters } from "../../ir/parameter-versions.js";
+import { inlineExpressionFunctions } from "../../ir/expression-inlining.js";
 import {
   allocatedLiteral,
   fillAllocationLength,
@@ -564,13 +565,15 @@ function moduleEmitter(input: IrModule, options: LlvmTargetOptions): LlEmitter {
               scalarizeNumericRecords(
                 specializeCallbackCalls(
                   specializeNumericCalls(
-                    foldConstantGlobals(
-                      copyStableCaptures(
-                        resolveSelfCaptures(
-                          hoistFinalFieldReads(
-                            freezeInitOnlyGlobals(
-                              holdCharLocalsAsUnits(
-                                dropUnwrittenPropertyBags(runPairedWhileLoopsAsFor(mod)),
+                    inlineExpressionFunctions(
+                      foldConstantGlobals(
+                        copyStableCaptures(
+                          resolveSelfCaptures(
+                            hoistFinalFieldReads(
+                              freezeInitOnlyGlobals(
+                                holdCharLocalsAsUnits(
+                                  dropUnwrittenPropertyBags(runPairedWhileLoopsAsFor(mod)),
+                                ),
                               ),
                             ),
                           ),

packages/compiler/src/backend/llvm/wrapping-sums.ts

@@ -4,13 +4,15 @@ import { integerArithmeticRange, type IntegerRange } from "../../ir/integer-rang
 import type { LlValue, LlvmEmitterContext } from "./expr-context.js";
 import type { LlInteger } from "./integer-values.js";
 
-/** Marks the `+` and `-` operands whose number only reaches ToInt32/ToUint32
- * consumers (`| 0`, the other bitwise operators, `Math.imul`, `Math.clz32`),
- * directly or through further sums, in the emitter's `wrappingSums`. Called
- * by the consumer before its operands emit. */
+/** Marks the `+`, `-` and `*` operands whose number only reaches
+ * ToInt32/ToUint32 consumers (`| 0`, the other bitwise operators,
+ * `Math.imul`, `Math.clz32`), directly or through further sums and
+ * products, in the emitter's `wrappingSums`. Called by the consumer before
+ * its operands emit. */
 export function markWrappingOperands(host: LlvmEmitterContext, ...operands: IrExpr[]): void {
   for (const operand of operands) {
-    if (operand.kind !== "bin" || (operand.op !== "+" && operand.op !== "-")) continue;
+    if (operand.kind !== "bin" || (operand.op !== "+" && operand.op !== "-" && operand.op !== "*"))
+      continue;
     if (operand.type.kind !== "f64" || host.wrappingSums.has(operand)) continue;
     host.wrappingSums.add(operand);
     markWrappingOperands(host, operand.left, operand.right);
@@ -128,12 +130,15 @@ function wrappingOperand(
   return { low, valid };
 }
 
-/** ToUint32 of `left op right` (`+` or `-`) in 32-bit arithmetic when the
- * expression only reaches ToInt32/ToUint32 consumers and both operands are
- * exact integers or NaN with an exact result. Both conversions are
- * arithmetic mod 2^32, so the wrapping i32 sum of the operands' low halves
- * is the converted exact sum; a NaN operand makes the sum NaN, which
- * converts to 0. Returns undefined, emitting nothing, for any other sum. */
+/** ToUint32 of `left op right` (`+`, `-` or `*`) in 32-bit arithmetic when
+ * the expression only reaches ToInt32/ToUint32 consumers and both operands
+ * are exact integers or NaN with an exact result. Both conversions are
+ * arithmetic mod 2^32, so the wrapping i32 sum or product of the operands'
+ * low halves is the converted exact result; a NaN operand makes the result
+ * NaN, which converts to 0. A product is exact when the operands' magnitude
+ * bounds multiply to at most 2^53; it may be -0 (`0 * -1`), which converts
+ * to 0 like the integer 0. Returns undefined, emitting nothing, for any
+ * other expression. */
 export function emitWrappingSum(
   host: LlvmEmitterContext,
   e: IrExpr & { kind: "bin" },
@@ -142,15 +147,23 @@ export function emitWrappingSum(
   li: LlInteger | null,
   ri: LlInteger | null,
 ): string | undefined {
-  if ((e.op !== "+" && e.op !== "-") || !host.wrappingSums.has(e)) return undefined;
+  if ((e.op !== "+" && e.op !== "-" && e.op !== "*") || !host.wrappingSums.has(e)) return undefined;
   const leftRange = li ? li.range : integerOrNaNRange(host, e.left);
   const rightRange = leftRange && (ri ? ri.range : integerOrNaNRange(host, e.right));
-  if (!rightRange || !integerArithmeticRange(e.op, leftRange!, rightRange)) return undefined;
+  if (!rightRange) return undefined;
+  if (
+    e.op === "*"
+      ? !exactProduct(leftRange!, rightRange)
+      : !integerArithmeticRange(e.op, leftRange!, rightRange)
+  )
+    return undefined;
   const B = host.B;
   const left = wrappingOperand(host, l, li, leftRange!);
   const right = wrappingOperand(host, r, ri, rightRange);
   const sum = B.tmp();
-  B.line(`${sum} = ${e.op === "+" ? "add" : "sub"} i32 ${left.low}, ${right.low}`);
+  B.line(
+    `${sum} = ${e.op === "+" ? "add" : e.op === "-" ? "sub" : "mul"} i32 ${left.low}, ${right.low}`,
+  );
   let valid = left.valid ?? right.valid;
   if (valid === null) return sum;
   if (left.valid && right.valid) {
@@ -162,6 +175,17 @@ export function emitWrappingSum(
   return out;
 }
 
+/** Whether every product of a value in `left` and one in `right` is an
+ * integer of magnitude below 2^53, which the double product holds exactly.
+ * The bound's own product is computed in a double: a bound of at least
+ * 2^53 never rounds below 2^53, while 2^53 + 1 (3 * 3002399751580331)
+ * rounds to 2^53, so the test is strict. */
+function exactProduct(left: IntegerRange, right: IntegerRange): boolean {
+  const magnitude = (range: IntegerRange): number =>
+    Math.max(Math.abs(range.min), Math.abs(range.max));
+  return magnitude(left) * magnitude(right) < 2 ** 53;
+}
+
 /** The number locals of `fn` that only hold exact integers (never -0) or NaN, with the range of
  * those integers: locals declared once with an initializer, neither parameters nor captured or
  * boxed, whose initializer and every later write `integerOrNaNRange` covers. A local joins once

packages/compiler/src/ir/expression-inlining.test.ts

@@ -0,0 +1,145 @@
+import { expect, test } from "vitest";
+import { inlineExpressionFunctions } from "./expression-inlining.js";
+import { BOOL, F64, VOID, type IrExpr, type IrFunction, type IrModule, type IrStmt } from "./ir.js";
+
+const loc = { file: "expression-inlining.ts", start: 0, end: 0 };
+const num = (value: number): IrExpr => ({ loc, kind: "numLit", value, type: F64 });
+const ref = (localId: string): IrExpr => ({ loc, kind: "varRef", localId, type: F64 });
+const bin = (op: "+" | "*" | ">>>", left: IrExpr, right: IrExpr): IrExpr => ({
+  loc,
+  kind: "bin",
+  op,
+  left,
+  right,
+  type: F64,
+});
+const call = (callee: string, args: IrExpr[]): IrExpr => ({
+  loc,
+  kind: "call",
+  callee,
+  args,
+  type: F64,
+});
+const local = (id: string, extra: object = {}) => ({
+  id,
+  name: id,
+  type: F64,
+  mutable: true,
+  ...extra,
+});
+
+// a(i, j) = ((i + j) * (i + j + 1)) >>> 1
+const a: IrFunction = {
+  loc,
+  name: "a",
+  params: [
+    { localId: "i", name: "i", type: F64 },
+    { localId: "j", name: "j", type: F64 },
+  ],
+  locals: [local("i"), local("j")],
+  returnType: F64,
+  body: [
+    {
+      loc,
+      kind: "return",
+      value: bin(
+        ">>>",
+        bin("*", bin("+", ref("i"), ref("j")), bin("+", bin("+", ref("i"), ref("j")), num(1))),
+        num(1),
+      ),
+    },
+  ],
+};
+
+/** `caller` running `body` in a loop. */
+function caller(body: IrStmt[], extra: object = {}): IrFunction {
+  return {
+    loc,
+    name: "caller",
+    params: [],
+    locals: [local("x", extra), local("y"), local("out")],
+    returnType: VOID,
+    body: [{ loc, kind: "while", cond: { loc, kind: "boolLit", value: true, type: BOOL }, body }],
+  };
+}
+
+const loopBody = (fn: IrFunction): IrStmt[] => (fn.body[0] as IrStmt & { kind: "while" }).body;
+
+const mod = (fn: IrFunction): IrModule => ({
+  irVersion: 17,
+  sourceFile: "expression-inlining.ts",
+  entry: "caller",
+  functions: [a, fn],
+});
+
+test("calls passing literals and locals take the callee's expression over them", () => {
+  const result = inlineExpressionFunctions(
+    mod(caller([{ loc, kind: "assign", localId: "out", value: call("a", [ref("x"), num(3)]) }])),
+  );
+  const assign = loopBody(result.functions[1]!)[0] as IrStmt & { kind: "assign" };
+  expect(assign.value).toMatchObject({
+    kind: "bin",
+    op: ">>>",
+    left: {
+      kind: "bin",
+      op: "*",
+      left: {
+        kind: "bin",
+        op: "+",
+        left: { kind: "varRef", localId: "x" },
+        right: { kind: "numLit", value: 3 },
+      },
+    },
+  });
+  // Every substituted read is a fresh node.
+  const left = (assign.value as IrExpr & { kind: "bin" }).left as IrExpr & { kind: "bin" };
+  const factors = [left.left, (left.right as IrExpr & { kind: "bin" }).left] as (IrExpr & {
+    kind: "bin";
+  })[];
+  expect(factors[0]!.left).not.toBe(factors[1]!.left);
+  expect(result.functions[0]).toBe(a);
+});
+
+test("calls with other arguments, boxed locals or callees with other bodies keep their call", () => {
+  const nested = call("a", [call("a", [ref("x"), ref("y")]), num(1)]);
+  const kept = inlineExpressionFunctions(
+    mod(caller([{ loc, kind: "assign", localId: "out", value: nested }])),
+  );
+  const outer = (loopBody(kept.functions[1]!)[0] as IrStmt & { kind: "assign" }).value;
+  // The inner call inlines; the outer one passes a computed argument.
+  expect(outer).toMatchObject({
+    kind: "call",
+    callee: "a",
+    args: [{ kind: "bin", op: ">>>" }, { kind: "numLit" }],
+  });
+  const boxed = caller(
+    [{ loc, kind: "assign", localId: "out", value: call("a", [ref("x"), ref("y")]) }],
+    {
+      boxed: true,
+    },
+  );
+  expect(inlineExpressionFunctions(mod(boxed)).functions[1]).toBe(boxed);
+  const twoStatements: IrFunction = {
+    ...a,
+    body: [{ loc, kind: "exprStmt", expr: ref("i") }, ...a.body],
+  };
+  const unchanged: IrModule = {
+    ...mod(
+      caller([{ loc, kind: "assign", localId: "out", value: call("a", [ref("x"), ref("y")]) }]),
+    ),
+    functions: [
+      twoStatements,
+      caller([{ loc, kind: "assign", localId: "out", value: call("a", [ref("x"), ref("y")]) }]),
+    ],
+  };
+  expect(inlineExpressionFunctions(unchanged)).toBe(unchanged);
+});
+
+test("calls outside loops keep their call", () => {
+  const outside: IrFunction = {
+    ...caller([]),
+    body: [{ loc, kind: "assign", localId: "out", value: call("a", [ref("x"), num(3)]) }],
+  };
+  const m = mod(outside);
+  expect(inlineExpressionFunctions(m)).toBe(m);
+});

packages/compiler/src/ir/expression-inlining.ts

@@ -0,0 +1,157 @@
+import type { IrExpr, IrFunction, IrModule, IrStmt } from "./ir.js";
+import { everyExprChild, everyStmtChild, mapExprChildren, transformStmtList } from "./traverse.js";
+
+/** Callee bodies above this many nodes keep their call. */
+const MAX_EXPRESSION_NODES = 32;
+
+/** Number operators that run no other code and cannot throw. */
+const PURE_BINARY = new Set<string>([
+  "+",
+  "-",
+  "*",
+  "/",
+  "%",
+  "**",
+  "&",
+  "|",
+  "^",
+  "<<",
+  ">>",
+  ">>>",
+]);
+
+/** The returned expression of a function whose body is `return <expr>` over
+ * its number parameters: literals, parameter reads, and the number
+ * operators of PURE_BINARY, unary `-` and `~`. Null for any other function. */
+function inlinableExpression(fn: IrFunction): IrExpr | null {
+  if (fn.async || fn.generator || fn.captures !== undefined || fn.classCaptures?.length)
+    return null;
+  if (fn.tier === "dynamic" || fn.returnType.kind !== "f64") return null;
+  if (fn.body.length !== 1) return null;
+  const only = fn.body[0]!;
+  if (only.kind !== "return" || only.value === null) return null;
+  const params = new Set<string>();
+  for (const param of fn.params) {
+    if (param.type.kind !== "f64") return null;
+    params.add(param.localId);
+  }
+  const locals = new Map(fn.locals.map((local) => [local.id, local]));
+  for (const id of params) {
+    const local = locals.get(id);
+    if (local === undefined || local.boxed || local.tdz) return null;
+  }
+  let nodes = 0;
+  const pure = (e: IrExpr): boolean => {
+    if (++nodes > MAX_EXPRESSION_NODES || e.type.kind !== "f64") return false;
+    switch (e.kind) {
+      case "numLit":
+        return true;
+      case "varRef":
+        return params.has(e.localId);
+      case "bin":
+        return PURE_BINARY.has(e.op) && pure(e.left) && pure(e.right);
+      case "unary":
+        return (e.op === "-" || e.op === "~") && pure(e.operand);
+      default:
+        return false;
+    }
+  };
+  return pure(only.value) ? only.value : null;
+}
+
+/** A call argument that reads the same value wherever the inlined body
+ * reads it: a number literal or a read of one of the caller's own unboxed
+ * locals outside the temporal dead zone. The body writes nothing, so every
+ * read sees the value the call would have passed. */
+function duplicableArgument(e: IrExpr, caller: ReadonlyMap<string, IrFunction["locals"][number]>) {
+  if (e.type.kind !== "f64") return false;
+  // A literal whose source spelling does not round-trip keeps its call.
+  if (e.kind === "numLit") return e.spelling === undefined;
+  if (e.kind !== "varRef") return false;
+  const local = caller.get(e.localId);
+  return local !== undefined && !local.boxed && !local.tdz;
+}
+
+/** A fresh copy of `e`: facts downstream are keyed by node identity. */
+function copy(e: IrExpr, substitute: (ref: IrExpr & { kind: "varRef" }) => IrExpr): IrExpr {
+  if (e.kind === "varRef") return substitute(e);
+  return mapExprChildren(
+    { ...e },
+    (child) => copy(child, substitute),
+    (stmt) => stmt,
+  );
+}
+
+/** The direct calls of `body` that run inside a loop body or condition. */
+function callsInLoops(body: IrStmt[]): Set<IrExpr> {
+  const found = new Set<IrExpr>();
+  let depth = 0;
+  const expr = (e: IrExpr): boolean => {
+    if (depth > 0 && e.kind === "call") found.add(e);
+    return everyExprChild(e, expr, stmt);
+  };
+  const stmt = (s: IrStmt): boolean => {
+    const loop =
+      s.kind === "for" || s.kind === "while" || s.kind === "doWhile" || s.kind === "forOf";
+    if (loop) depth++;
+    everyStmtChild(s, expr, stmt);
+    if (loop) depth--;
+    return true;
+  };
+  body.forEach(stmt);
+  return found;
+}
+
+/** Inlines functions whose body is one number expression over their
+ * parameters (`function a(i, j) { return (((i + j) * (i + j + 1)) >>> 1) +
+ * i + 1; }`) at direct call sites that pass literals or reads of the
+ * caller's locals.
+ *
+ * The callee runs no other code, allocates nothing and cannot throw, so
+ * the call's only effect is its result. Its arguments are reads that
+ * nothing between them and the callee's reads can change, so substituting
+ * them gives the same value. The caller's loop and integer analyses then
+ * see the arithmetic on its own counters (`a(i, j)` inside a counted loop
+ * over `i` and `j`). Only calls inside loops are inlined: elsewhere the
+ * integer analyses gain nothing. The function itself stays for its other
+ * uses. */
+export function inlineExpressionFunctions(mod: IrModule): IrModule {
+  const members = new Set((mod.classes ?? []).flatMap((cls) => cls.methods ?? []));
+  const expressions = new Map<string, { fn: IrFunction; value: IrExpr }>();
+  for (const fn of mod.functions) {
+    if (fn.name === mod.entry || members.has(fn.name)) continue;
+    const value = inlinableExpression(fn);
+    if (value !== null) expressions.set(fn.name, { fn, value });
+  }
+  if (expressions.size === 0) return mod;
+  let changed = false;
+  const functions = mod.functions.map((fn) => {
+    const locals = new Map(fn.locals.map((local) => [local.id, local]));
+    const inLoops = callsInLoops(fn.body);
+    if (inLoops.size === 0) return fn;
+    let inlined = false;
+    const body = transformStmtList(fn.body, {
+      stmt: (s) => s,
+      expr: (e) => {
+        if (e.kind !== "call" || e.callee === fn.name || !inLoops.has(e)) return e;
+        const target = expressions.get(e.callee);
+        if (target === undefined || target.fn.params.length !== e.args.length) return e;
+        if (!e.args.every((arg) => duplicableArgument(arg, locals))) return e;
+        const args = new Map(target.fn.params.map((param, i) => [param.localId, e.args[i]!]));
+        inlined = true;
+        return copy(target.value, (ref) => {
+          const arg = args.get(ref.localId);
+          if (arg?.kind === "numLit")
+            return { kind: "numLit", value: arg.value, type: arg.type, loc: ref.loc };
+          if (arg?.kind === "varRef")
+            return { kind: "varRef", localId: arg.localId, type: arg.type, loc: ref.loc };
+          return ref;
+        });
+      },
+    });
+    if (!inlined) return fn;
+    changed = true;
+    return { ...fn, body };
+  });
+  return changed ? { ...mod, functions } : mod;
+}

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

@@ -552,3 +552,71 @@ test("counters stepped beside the tested one run as integers", () => {
     { startGuarded: true, tripGuarded: true },
   ]);
 });
+
+/** `counted(limit)` with the body `x = ((i + k) * (i + k + 1)) >>> 1`
+ * inside a nested `for (let j = 0; j < 4; j++)` loop. */
+function productLoop(limit: IrExpr, start: IrExpr = number(0)): IrStmt & { kind: "for" } {
+  const s = counted(limit);
+  if (s.init?.kind === "varDecl") s.init.init = start;
+  const sum = (extra: number): IrExpr => {
+    const base: IrExpr = { kind: "bin", op: "+", left: ref("i"), right: ref("k"), type: F64, loc };
+    return extra === 0
+      ? base
+      : { kind: "bin", op: "+", left: base, right: number(extra), type: F64, loc };
+  };
+  const product: IrExpr = { kind: "bin", op: "*", left: sum(0), right: sum(1), type: F64, loc };
+  const inner = counted(number(4));
+  if (inner.init?.kind === "varDecl") inner.init.localId = "j";
+  if (inner.cond?.kind === "bin") inner.cond.left = ref("j");
+  if (inner.update?.kind === "assign") {
+    inner.update.localId = "j";
+    if (inner.update.value.kind === "bin") inner.update.value.left = ref("j");
+  }
+  inner.body = [
+    {
+      kind: "assign",
+      localId: "x",
+      value: { kind: "bin", op: ">>>", left: product, right: number(1), type: F64, loc },
+      loc,
+    },
+  ];
+  s.body = [inner];
+  return s;
+}
+
+test("guarded counted loops whose products the counter bound proves exact may hold nested loops", () => {
+  const scope = new Map(locals);
+  for (const id of ["limit", "start", "k", "x", "j"])
+    scope.set(id, { id, name: id, type: F64, mutable: true });
+  const tight = matchIntegerCountedForLoop(productLoop(ref("limit"), ref("start")), scope);
+  expect(tight).toMatchObject({
+    guarded: true,
+    limitGuarded: true,
+    startGuarded: true,
+    startBound: 2 ** 25,
+    guardLimit: 2 ** 25,
+    range: { min: -(2 ** 25), max: 2 ** 25 - 1 },
+  });
+  // A literal trip count past the bound keeps the safe-integer bound, and
+  // a guarded loop under that bound holds no nested loop.
+  expect(matchIntegerCountedForLoop(productLoop(number(1e9), ref("start")), scope)).toBeNull();
+  const plain = counted(ref("limit"));
+  plain.body = [productLoop(number(4)).body[0]!];
+  const inner = plain.body[0] as IrStmt & { kind: "for" };
+  inner.body = [
+    {
+      kind: "assign",
+      localId: "x",
+      value: { kind: "bin", op: "|", left: ref("i"), right: number(0), type: F64, loc },
+      loc,
+    },
+  ];
+  expect(matchIntegerCountedForLoop(plain, scope)).toBeNull();
+  // Without a nested loop the product bound still applies to the limit.
+  const flat = productLoop(ref("limit"), ref("start"));
+  flat.body = (flat.body[0] as IrStmt & { kind: "for" }).body;
+  expect(matchIntegerCountedForLoop(flat, scope)).toMatchObject({
+    startBound: 2 ** 25,
+    guardLimit: 2 ** 25,
+  });
+});

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

@@ -24,8 +24,12 @@ export interface IntegerCountedForLoop {
   /** 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. */
+   * that it is an exact integer (not -0) of at most `startBound` in
+   * magnitude. */
   startGuarded: boolean;
+  /** The magnitude the guarded start is tested against: 2^53 - 1, or
+   * PRODUCT_COUNTER_BOUND when the body multiplies sums of the counter. */
+  startBound: number;
   guardLimit: number;
   range: IntegerRange;
   /** Counters declared and stepped beside this one that the condition does
@@ -50,6 +54,64 @@ export interface IntegerLoopFollower {
   range: IntegerRange;
 }
 
+/** The counter magnitude below which a product of two sums of counters and
+ * small literals stays exact in a double: `(i + j) * (i + j + 1)` with both
+ * counters within 2^25 is below 2^52. A guarded loop whose body multiplies
+ * such sums for a bitwise operator tests its counter against this bound, so
+ * the product's integer form is proven exact. */
+export const PRODUCT_COUNTER_BOUND = 2 ** 25;
+
+/** Whether `e` is a sum or difference of locals and safe integer literals
+ * of magnitude at most 2^20, and whether it reads `counter`. */
+function localSum(e: IrExpr, counter: string): { sum: boolean; reads: boolean } {
+  let reads = false;
+  const sum = (x: IrExpr): boolean => {
+    if (x.kind === "varRef") {
+      if (x.localId === counter) reads = true;
+      return true;
+    }
+    if (x.kind === "numLit") return Number.isSafeInteger(x.value) && Math.abs(x.value) <= 2 ** 20;
+    return x.kind === "bin" && (x.op === "+" || x.op === "-") && sum(x.left) && sum(x.right);
+  };
+  return { sum: sum(e), reads };
+}
+
+/** Whether a bitwise operand `e` multiplies two sums of locals and
+ * literals, one of which reads `counter` (`((i + j) * (i + j + 1)) >>> 1`). */
+function multipliesCounterSums(e: IrExpr, counter: string): boolean {
+  if (e.kind !== "bin") return false;
+  if (e.op === "*") {
+    const left = localSum(e.left, counter),
+      right = localSum(e.right, counter);
+    return left.sum && right.sum && (left.reads || right.reads);
+  }
+  return (
+    (e.op === "+" || e.op === "-") &&
+    (multipliesCounterSums(e.left, counter) || multipliesCounterSums(e.right, counter))
+  );
+}
+
+const BITWISE = new Set(["&", "|", "^", "<<", ">>", ">>>"]);
+
+/** Whether `body`, nested loops included, takes a bitwise operator of a
+ * product that multipliesCounterSums matches. */
+function bodyMultipliesCounter(body: IrStmt[], counter: string): boolean {
+  let found = false;
+  everyStmtList(body, {
+    stmt: () => !found,
+    expr: (e) => {
+      if (
+        e.kind === "bin" &&
+        BITWISE.has(e.op) &&
+        (multipliesCounterSums(e.left, counter) || multipliesCounterSums(e.right, counter))
+      )
+        found = true;
+      return !found;
+    },
+  });
+  return found;
+}
+
 /** The bound of the run-time follower test. */
 export const FOLLOWER_GUARD_LIMIT = 2 ** 52;
 /** The magnitude a follower that passed the test can reach. The test's four
@@ -171,6 +233,23 @@ export function matchIntegerCountedForLoop(
   // The condition tests the driving counter; every other declaration
   // follows it.
   const driver = declared(cond.left) ?? declared(cond.right);
+  // A body that multiplies sums of the driver tests it against the bound
+  // under which those products are exact. A limit proven past that bound
+  // (a literal trip count) keeps the safe-integer bound instead.
+  const limitSide = driver !== null && declared(cond.left) === driver ? cond.right : cond.left;
+  const knownLimit =
+    limitSide.kind === "numLit"
+      ? { min: limitSide.value, max: limitSide.value }
+      : ranges.get(limitSide);
+  const driverBound =
+    driver !== null &&
+    !(
+      knownLimit &&
+      Math.max(Math.abs(knownLimit.min), Math.abs(knownLimit.max)) > PRODUCT_COUNTER_BOUND
+    ) &&
+    bodyMultipliesCounter(stmt.body, driver)
+      ? PRODUCT_COUNTER_BOUND
+      : Number.MAX_SAFE_INTEGER;
   const counterIds = decls.map((d) => (d.kind === "varDecl" ? d.localId : ""));
   const stepOf = (id: string): number | null => {
     const steps = updates.map((u) => constantStep(u, id)).filter((v) => v !== null);
@@ -191,7 +270,8 @@ export function matchIntegerCountedForLoop(
         ? { min: start.value, max: start.value }
         : ranges.get(start);
     const guarded = !known && start.kind !== "numLit" && start.type.kind === "f64";
-    const range = guarded ? { min: -Number.MAX_SAFE_INTEGER, max: Number.MAX_SAFE_INTEGER } : known;
+    const bound = id === driver ? driverBound : Number.MAX_SAFE_INTEGER;
+    const range = guarded ? { min: -bound, max: bound } : known;
     const step = stepOf(id);
     if (
       !range ||
@@ -248,7 +328,9 @@ export function matchIntegerCountedForLoop(
     );
   if (!stableLoopValue(limit, locals, stmt.body, true, local.id, stableLength)) return null;
   const inclusive = op === "<=" || op === ">=";
-  const guardLimit = step > 0 ? 2 ** 53 - (inclusive ? 1 : 0) : -(2 ** 53) + (inclusive ? 1 : 0);
+  const limitBound = driverBound === Number.MAX_SAFE_INTEGER ? 2 ** 53 : driverBound;
+  const guardLimit =
+    step > 0 ? limitBound - (inclusive ? 1 : 0) : -limitBound + (inclusive ? 1 : 0);
   const boundRange =
     limit.kind === "numLit" ? { min: limit.value, max: limit.value } : ranges.get(limit);
   const limitGuarded =
@@ -264,7 +346,9 @@ export function matchIntegerCountedForLoop(
               Number.MAX_SAFE_INTEGER,
               boundRange
                 ? Math.ceil(boundRange.max) - (inclusive ? 0 : 1)
-                : Number.MAX_SAFE_INTEGER,
+                : limitGuarded
+                  ? Math.ceil(guardLimit) - (inclusive ? 0 : 1)
+                  : Number.MAX_SAFE_INTEGER,
             ),
           ),
         }
@@ -275,7 +359,9 @@ export function matchIntegerCountedForLoop(
               -Number.MAX_SAFE_INTEGER,
               boundRange
                 ? Math.floor(boundRange.min) + (inclusive ? 0 : 1)
-                : -Number.MAX_SAFE_INTEGER,
+                : limitGuarded
+                  ? Math.floor(guardLimit) + (inclusive ? 0 : 1)
+                  : -Number.MAX_SAFE_INTEGER,
             ),
           ),
           max: startRange.max,
@@ -366,9 +452,15 @@ export function matchIntegerCountedForLoop(
   const bodyStmt = (s: IrStmt): boolean => {
     if (s.kind === "bytesSet" && counter(s.index)) useful = true;
     if (!guarded && s.kind === "arraySet" && counter(s.index)) useful = true;
+    // A guarded loop's copy duplicates nested loops too, so only a body
+    // whose counter products the bound proves exact may hold a nested
+    // `for` loop; the node budget counts its body.
     if (
       guarded &&
-      (s.kind === "for" || s.kind === "forOf" || s.kind === "while" || s.kind === "doWhile")
+      (s.kind === "forOf" ||
+        s.kind === "while" ||
+        s.kind === "doWhile" ||
+        (s.kind === "for" && driverBound === Number.MAX_SAFE_INTEGER))
     )
       return false;
     return (!guarded || ++cost <= 80) && everyStmtChild(s, expr, bodyStmt);
@@ -383,6 +475,7 @@ export function matchIntegerCountedForLoop(
     guarded,
     limitGuarded,
     startGuarded,
+    startBound: driverBound,
     guardLimit,
     range,
     followers,

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/8193-counted-loop-products.ts.json

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

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/8195-wrapping-products-near-2-53.ts.json

@@ -0,0 +1,6 @@
+{
+ "order": [
+  "<repo>/tests/corpus/8195-wrapping-products-near-2-53.ts"
+ ],
+ "diags": []
+}

tests/corpus/8193-counted-loop-products.ts

@@ -0,0 +1,59 @@
+// Nested counted loops whose bounds arrive at run time, with the counters'
+// products reaching bitwise operators. Starts and limits take the integer
+// copy or the double copy: negative, fractional, NaN, -0, and starts or
+// limits past 2^25, past which the guarded copy no longer proves products
+// exact.
+function a(i: number, j: number): number {
+  return (((i + j) * (i + j + 1)) >>> 1) + i + 1;
+}
+
+function rows(start: number, end: number, n: number): number[] {
+  const out: number[] = [];
+  for (let i = start; i < end; i++) {
+    let t = 0;
+    for (let j = 0; j < n; j++) t += 1 / a(i, j);
+    out.push(t);
+  }
+  return out;
+}
+
+function mixed(start: number, end: number, n: number): number {
+  let h = 0;
+  for (let i = start; i < end; i++) {
+    for (let j = start; j < n; j++) h = (h + ((i - j) * (j - i + 3) | 0)) | 0;
+  }
+  return h;
+}
+
+function wide(start: number, count: number): number[] {
+  const out: number[] = [];
+  for (let i = start; i < start + count; i++) out.push(((i * (i + 1)) >>> 0) ^ (i * i) | 0);
+  return out;
+}
+
+const args = process.argv.slice(2).map(Number);
+const big = 2 ** 25;
+for (const [start, end, n] of [
+  [0, 4, 6],
+  [-3, 2, 5],
+  [1.5, 4, 3],
+  [-0, 2, 2],
+  [NaN, 2, 2],
+  [big - 2, big + 2, 3],
+  [0, 3, args.length > 0 ? big + 1 : 4],
+]) {
+  console.log(start, end, n, rows(start!, end!, Math.min(n!, 8)).map((x) => x.toFixed(12)).join(" "));
+  console.log("mixed", mixed(start!, end!, Math.min(n!, 8)));
+}
+for (const start of [0, -5, 2 ** 26, 2 ** 31 - 2, 2 ** 52, -(2 ** 40), 0.5]) console.log(start, wide(start, 4).join(" "));
+console.log(Object.is(((-0 * 5) >>> 0), 0), a(-1, 0), a(0, -1), a(2 ** 30, 2 ** 30));
+
+// Single-expression functions inlined at calls that pass literals and
+// locals, also read as values and called with other arguments.
+function mix(x: number, y: number): number {
+  return ((x * 31 + y) ^ (y << 3)) % 1000;
+}
+const table = [mix, (x: number, y: number) => x - y];
+let acc = 0;
+for (let k = 0; k < 5; k++) acc += mix(k, 7) + mix(3, k) + table[k % 2]!(k, 2) + mix(k * 2, args.length);
+console.log(acc, mix(2 ** 40, -0), mix(NaN, 1), mix(Infinity, 2), [1, 2, 3].map((v) => mix(v, v)).join());

tests/corpus/8195-wrapping-products-near-2-53.ts

@@ -0,0 +1,25 @@
+// Products that reach bitwise operators with operands whose ranges are
+// known, around 2^53: below it the double product is exact, and from it on
+// the double product rounds, so the 32-bit result comes from the rounded
+// value (3 * 3002399751580331 is 2^53 + 1, which rounds to 2^53).
+const pairs: [number, number][] = [
+  [3, 3002399751580331],
+  [-3, 3002399751580331],
+  [2 ** 26, 2 ** 27],
+  [94906265, 94906267],
+  [94906267, 94906267],
+  [7, 1286742750677284],
+  [65535, 137443147777],
+];
+for (const [a, b] of pairs) console.log(a, b, (a * b) | 0, (a * b) >>> 0, (a * b) ^ 1, (a * b) << 3);
+
+// The same products with literal operands, which the compiler sees whole.
+console.log((3 * 3002399751580331) | 0, (3 * 3002399751580331) >>> 0);
+console.log((-3 * 3002399751580331) | 0, (2 ** 26 * 2 ** 27) | 0);
+console.log((94906267 * 94906267) | 0, (94906265 * 94906267) >>> 0);
+
+// A counted loop whose counter multiplies a constant past 2^53.
+const step = 1286742750677284;
+let h = 0;
+for (let k = 5; k < 9; k++) h = (h + ((k * step) | 0)) | 0;
+console.log(h);