Speculate integer loops through pure helper calls, powers of two and byte reads (#2100)

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

@@ -2179,6 +2179,7 @@ export function emitLoopCopies(
             );
           },
           (id) => emitter.speculativeLocal(id) || emitter.speculativeInvariantGlobal(id),
+          emitter.speculationContext(),
         )
       : null;
   if (!plan) {
@@ -2237,6 +2238,20 @@ export function emitLoopCopies(
     B.line(`${both} = and i1 ${ok}, ${trying}`);
     ok = both;
   }
+  // Hoisted leaves live in locals of their own, which the resume copy of
+  // their write reads (see SpeculativeValue).
+  for (const { leaves } of plan.values.values())
+    for (const leaf of leaves) {
+      emitter.currentLocals.set(leaf.localId, {
+        id: leaf.localId,
+        name: "speculated leaf",
+        type: leaf.expr.type,
+        mutable: true,
+      });
+      B.entryAllocas.push(
+        `${emitter.binding(leaf.localId).slot} = alloca double ; speculated leaf`,
+      );
+    }
   const fast = B.newLabel("spec.fast"),
     slow = B.newLabel("spec.slow");
   B.condBr(ok, fast, slow);
@@ -2262,6 +2277,7 @@ export function emitLoopCopies(
     throw new InternalCompilerError("llvm emitter bug: speculative loop copy without a header");
   B.startBlock(B.newLabel("spec.tail"));
   emitter.speculationResume = resume;
+  emitter.speculationResumeValues = plan.values;
   emitter.speculationHeader = () => {
     B.br(header!);
     B.startBlock(B.newLabel("spec.unused"));
@@ -2270,6 +2286,7 @@ export function emitLoopCopies(
     emitCopy(exit);
   } finally {
     emitter.speculationResume = null;
+    emitter.speculationResumeValues = null;
     emitter.speculationHeader = null;
   }
   B.startBlock(fast);
@@ -2318,6 +2335,7 @@ export function emitLoopCopies(
     deferred,
     bails,
     ran,
+    values: plan.values,
   };
   try {
     const leave = B.newLabel("spec.leave");
@@ -2575,9 +2593,22 @@ export function emitSpeculativeWrite(
   // at once: a bailing path then owns nothing, and reads have no effect
   // whose reordering before the arithmetic could be observed.
   const reads = new Map<IrExpr, LlValue>();
-  if (target.value !== null) {
+  const rewritten = emitter.speculation!.values?.get(s);
+  const value = rewritten?.value ?? target.value;
+  if (rewritten !== undefined) {
+    // Hoisted leaves run first, in the write's order, and keep their
+    // doubles for the resume copy.
     emitter.frames.push([]);
-    for (const read of speculativeReads(target.value)) reads.set(read, emitter.emitExpr(read));
+    for (const leaf of rewritten.leaves) {
+      const evaluated = emitter.emitExpr(leaf.expr);
+      B.line(`store double ${evaluated.name}, ptr ${emitter.binding(leaf.localId).slot}`);
+      reads.set(leaf.ref, evaluated);
+    }
+    emitter.releaseTemps(emitter.frames.pop()!);
+  }
+  if (value !== null) {
+    emitter.frames.push([]);
+    for (const read of speculativeReads(value)) reads.set(read, emitter.emitExpr(read));
     emitter.releaseTemps(emitter.frames.pop()!);
   }
   emitter.frames.push([]);
@@ -2592,7 +2623,7 @@ export function emitSpeculativeWrite(
     B.line(`${biased} = add i64 ${result}, ${Number.MAX_SAFE_INTEGER}`);
     B.line(`${outside} = icmp ugt i64 ${biased}, ${2 * Number.MAX_SAFE_INTEGER}`);
     bail(outside);
-  } else result = emitSpeculativeInteger(emitter, target.value, bail, reads, resume === undefined);
+  } else result = emitSpeculativeInteger(emitter, value!, bail, reads, resume === undefined);
   // Speculable values read only numbers past this point.
   if (emitter.frames.pop()!.length > 0)
     throw new InternalCompilerError("llvm emitter bug: speculative write owns a temporary");

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

@@ -357,6 +357,13 @@ import {
 } from "./shapes.js";
 import type { ExprOf, LibCallExpr, LlStreamTypedRefAdapter, LlValue } from "./expr-context.js";
 import { INLINE_LITERAL_BYTES } from "./string-literal-equality.js";
+import {
+  inlinableNumberFunction,
+  withSpeculativeValue,
+  type InlinableFunction,
+  type SpeculationContext,
+  type SpeculativeValue,
+} from "./speculative-integers.js";
 import {
   emitFfiCallbackDefs,
   emitLibDefs,
@@ -1524,11 +1531,47 @@ export class LlEmitter {
      * integer loop has neither. */
     bails?: string;
     ran?: string;
+    /** Writes that speculate a rewritten value (rewriteSpeculable). */
+    values?: ReadonlyMap<IrStmt, SpeculativeValue>;
   } | null = null;
   /** Speculative loop sites numbered so far in the module. */
   speculativeLoopSites = 0;
   /** The ordinary copy being emitted: resume labels still to place. */
   speculationResume: Map<IrStmt, { label: string; placed: boolean }> | null = null;
+  /** The rewritten values of the writes `speculationResume` resumes at. */
+  speculationResumeValues: ReadonlyMap<IrStmt, SpeculativeValue> | null = null;
+  /** Hoisted speculative leaves numbered so far in the module. */
+  private speculativeLeaves = 0;
+  private readonly inlinableFunctions = new Map<string, InlinableFunction | undefined>();
+
+  /** What speculative loop planning may assume here (SpeculationContext). */
+  speculationContext(): SpeculationContext {
+    const inlinable = (callee: string): InlinableFunction | undefined => {
+      if (!this.inlinableFunctions.has(callee)) {
+        // Unset while its body is checked: a recursive helper is not inlinable.
+        this.inlinableFunctions.set(callee, undefined);
+        const fn = this.fnByName.get(callee);
+        this.inlinableFunctions.set(
+          callee,
+          fn ? inlinableNumberFunction(fn, inlinable) : undefined,
+        );
+      }
+      return this.inlinableFunctions.get(callee);
+    };
+    return {
+      inlinable,
+      unaliased: (id) => {
+        const local = this.currentLocals.get(id);
+        return (
+          local !== undefined &&
+          !local.boxed &&
+          !this.captureIds.has(id) &&
+          !this.lazyCaptureBoxes.has(id)
+        );
+      },
+      leafId: () => `%spec.leaf.${this.speculativeLeaves++}`,
+    };
+  }
   /** Run once at the top of the next loop copy's header
    * (`enterLoopHeader`): records the ordinary copy's header, or sends the
    * third copy's header there. */
@@ -5799,6 +5842,22 @@ export class LlEmitter {
     if (resume && !resume.placed && !this.B.isTerminated()) {
       // The fast copy's failed checks re-execute this statement here.
       resume.placed = true;
+      // A write with hoisted leaves resumed from its own failed check
+      // re-executes its arithmetic over the leaves' stored doubles: each
+      // leaf ran once, in the fast copy. Reached from an earlier statement
+      // of this copy (a failed check before it), its leaves never ran, so
+      // it runs as written.
+      const rewritten = this.speculationResumeValues?.get(s);
+      if (rewritten !== undefined && rewritten.leaves.length > 0) {
+        const join = this.B.newLabel("spec.resume.join");
+        this.emitStmt(s);
+        if (!this.B.isTerminated()) this.B.br(join);
+        this.B.startBlock(resume.label);
+        this.emitStmt(withSpeculativeValue(s, rewritten.value));
+        if (!this.B.isTerminated()) this.B.br(join);
+        this.B.startBlock(join);
+        return;
+      }
       this.B.br(resume.label);
       this.B.startBlock(resume.label);
     }

packages/compiler/src/backend/llvm/speculative-integers.test.ts

@@ -1,9 +1,11 @@
 import { expect, test } from "vitest";
 import { F64, type IrExpr, type IrStmt } from "../../ir/ir.js";
 import {
+  inlinableNumberFunction,
   isSpeculableValue,
   planProvenIntegerLoop,
   planSpeculativeLoop,
+  rewriteSpeculable,
   speculativeReads,
 } from "./speculative-integers.js";
 
@@ -266,3 +268,136 @@ test("bounded cursors with proven writes are shadowed in a proven integer loop",
     planProvenIntegerLoop(forLoop([...body, jump]) as never, cursors, all, () => false, all),
   ).toBeNull();
 });
+
+const libCall = (fn: string, args: IrExpr[]): IrExpr =>
+  ({ kind: "libCall", fn, args, type: F64, loc }) as never;
+const call = (callee: string, args: IrExpr[]): IrExpr => ({
+  kind: "call",
+  callee,
+  args,
+  type: F64,
+  loc,
+});
+const pow2 = (n: IrExpr): IrExpr => libCall("math.pow", [num(2), n]);
+/** `function trim(v, n) { return v - Math.floor(v / Math.pow(2, n)) * Math.pow(2, n); }` */
+const trim = {
+  name: "trim",
+  params: [
+    { localId: "v", name: "v", type: F64 },
+    { localId: "n", name: "n", type: F64 },
+  ],
+  returnType: F64,
+  locals: [],
+  body: [
+    {
+      kind: "return",
+      value: bin(
+        "-",
+        ref("v"),
+        bin("*", libCall("math.floor", [bin("/", ref("v"), pow2(ref("n")))]), pow2(ref("n"))),
+      ),
+      loc,
+    },
+  ],
+  loc,
+} as never;
+
+test("calls of small pure number functions expand into a speculated write", () => {
+  const inlinable = inlinableNumberFunction(trim)!;
+  expect(inlinable.params).toEqual(["v", "n"]);
+  const context = { inlinable: (callee: string) => (callee === "trim" ? inlinable : undefined) };
+  // key = trim(key, m) * 4 + c
+  const value = bin("+", bin("*", call("trim", [ref("key"), ref("m")]), num(4)), ref("c"));
+  expect(isSpeculableValue(value)).toBe(false);
+  const rewritten = rewriteSpeculable(value, context)!;
+  expect(rewritten.leaves).toEqual([]);
+  expect(JSON.stringify(rewritten.value)).not.toContain('"call"');
+  const write = assign("key", value);
+  const plan = planSpeculativeLoop(
+    loop([write, increment("i")]) as never,
+    all,
+    () => false,
+    () => false,
+    all,
+    context,
+  )!;
+  expect(plan.writes.has(write)).toBe(true);
+  expect(plan.values.get(write)).toEqual(rewritten);
+  // `m`, an input of the expanded body, is checked on entry with the others.
+  expect([...plan.locals.keys()].sort()).toEqual(["c", "i", "key", "m", "n"]);
+});
+
+test("Math.floor needs a power-of-two divisor, and expansion pure arguments", () => {
+  const context = { inlinable: () => inlinableNumberFunction(trim) };
+  expect(rewriteSpeculable(libCall("math.floor", [bin("/", ref("a"), num(8))]), {})).not.toBeNull();
+  expect(rewriteSpeculable(libCall("math.floor", [bin("/", ref("a"), num(6))]), {})).toBeNull();
+  expect(rewriteSpeculable(libCall("math.floor", [bin("/", ref("a"), ref("b"))]), {})).toBeNull();
+  // An argument that is itself a call is evaluated once: no expansion.
+  expect(rewriteSpeculable(call("trim", [call("other", []), ref("m")]), context)).toBeNull();
+  const effectful = {
+    ...(trim as object),
+    body: [{ kind: "exprStmt", expr: call("log", []), loc }],
+  };
+  expect(inlinableNumberFunction(effectful as never)).toBeUndefined();
+  // A divisor next to a power of two is not one, though Math.log2 rounds it to 52.
+  const near = (d: number): IrExpr => libCall("math.floor", [bin("/", ref("a"), num(d))]);
+  expect(rewriteSpeculable(near(2 ** 52), {})).not.toBeNull();
+  expect(rewriteSpeculable(near(2 ** 52 + 1), {})).toBeNull();
+  expect(rewriteSpeculable(near(2 ** 53), {})).toBeNull();
+  expect(rewriteSpeculable(near(0.5), {})).toBeNull();
+});
+
+test("a helper's own calls must be of helpers that expand too", () => {
+  // `function outer(v, n) { return trim(v, n) + other(v); }`
+  const outer = {
+    ...(trim as object),
+    name: "outer",
+    body: [
+      {
+        kind: "return",
+        value: bin("+", call("trim", [ref("v"), ref("n")]), call("other", [ref("v")])),
+        loc,
+      },
+    ],
+  };
+  const helpers = (callee: string) =>
+    callee === "trim" ? inlinableNumberFunction(trim) : undefined;
+  expect(inlinableNumberFunction(outer as never)).toBeUndefined();
+  expect(inlinableNumberFunction(outer as never, helpers)).toBeUndefined();
+  const pure = {
+    ...(outer as object),
+    body: [{ kind: "return", value: bin("*", call("trim", [ref("v"), ref("n")]), num(2)), loc }],
+  };
+  expect(inlinableNumberFunction(pure as never, helpers)).toBeDefined();
+});
+
+test("other leaves are hoisted only when no call can change the write's locals", () => {
+  let next = 0;
+  const leafId = (): string => `%leaf.${next++}`;
+  const value = bin(
+    "%",
+    bin("+", bin("*", ref("acc"), num(7)), call("noisy", [ref("i")])),
+    num(1000003),
+  );
+  const hoisted = rewriteSpeculable(value, { unaliased: () => true, leafId })!;
+  expect(hoisted.leaves.map((leaf) => leaf.expr)).toEqual([call("noisy", [ref("i")])]);
+  expect(rewriteSpeculable(value, { unaliased: (id) => id !== "acc", leafId })).toBeNull();
+  // A call whose argument writes a local stays where it is.
+  const bump: IrExpr = { kind: "incDec", localId: "i", op: "+", prefix: false, type: F64, loc };
+  expect(
+    rewriteSpeculable(bin("+", ref("i"), call("noisy", [bump])), { unaliased: () => true, leafId }),
+  ).toBeNull();
+  // Without leaf ids nothing is hoisted.
+  expect(rewriteSpeculable(value, { unaliased: () => true })).toBeNull();
+  const write = assign("acc", value);
+  const plan = planSpeculativeLoop(
+    loop([write, increment("i")]) as never,
+    (id) => !id.startsWith("%leaf"),
+    () => false,
+    () => false,
+    (id) => !id.startsWith("%leaf"),
+    { unaliased: () => true, leafId },
+  )!;
+  expect(plan.values.get(write)!.leaves).toHaveLength(1);
+  expect([...plan.locals.keys()].some((id) => id.startsWith("%leaf"))).toBe(false);
+});

packages/compiler/src/backend/llvm/speculative-integers.ts

@@ -1,5 +1,6 @@
-import type { IrExpr, IrStmt } from "../../ir/ir.js";
-import { everyExprChild, everyStmtChild } from "../../ir/traverse.js";
+import type { IrExpr, IrFunction, IrStmt } from "../../ir/ir.js";
+import { everyExpr, everyExprChild, everyStmtChild } from "../../ir/traverse.js";
+import { InternalCompilerError } from "../../errors.js";
 import { f64Lit } from "./common.js";
 import type { LlValue, LlvmEmitterContext } from "./expr-context.js";
 import { exactInteger, widenInteger } from "./integer-values.js";
@@ -30,6 +31,295 @@ export interface SpeculativeLoopPlan {
    * than through its condition, `break` or a failed check: no calls, no
    * throwing operations, no labeled jumps and no `return`. */
   observerFree: boolean;
+  /** The writes that speculate a rewritten value (rewriteSpeculable). */
+  values: Map<IrStmt, SpeculativeValue>;
+}
+
+/** A write's value as speculation evaluates it: calls of small pure number
+ * functions expanded in place, and other number leaves (`leaves`) replaced
+ * by references to fresh locals. The fast copy evaluates each leaf once,
+ * stores it to its local and checks it as an integer; a failed check
+ * resumes at a copy of the write that reads the locals, so a leaf is never
+ * evaluated twice. */
+export interface SpeculativeValue {
+  value: IrExpr;
+  leaves: { localId: string; expr: IrExpr; ref: IrExpr }[];
+}
+
+/** What the planner may assume about the code around a loop. */
+export interface SpeculationContext {
+  /** The parameters and the returned expression of `callee` when it is a
+   * small pure number function (inlinableNumberFunction). */
+  inlinable?: (callee: string) => InlinableFunction | undefined;
+  /** A local that no call can write: not captured, boxed or global. */
+  unaliased?: (localId: string) => boolean;
+  /** A fresh local id for a hoisted leaf. */
+  leafId?: () => string;
+}
+
+export interface InlinableFunction {
+  params: readonly string[];
+  body: IrExpr;
+}
+
+/** `s` (a speculated write) writing `value` instead. */
+export function withSpeculativeValue(s: IrStmt, value: IrExpr): IrStmt {
+  switch (s.kind) {
+    case "assign":
+      return { ...s, value };
+    case "varDecl":
+      return { ...s, init: value };
+    case "exprStmt":
+      if (s.expr.kind === "assignExpr") return { ...s, expr: { ...s.expr, value } };
+      break;
+    default:
+      break;
+  }
+  throw new InternalCompilerError(`llvm emitter bug: speculated ${s.kind} with a rewritten value`);
+}
+
+/** Calls nest at most this deep in one expansion. */
+const MAX_INLINE_DEPTH = 4;
+/** An expansion larger than this stays a call. */
+const MAX_INLINE_NODES = 64;
+
+/** A function that only returns arithmetic over its number parameters:
+ * `function trim(v, n) { return v - Math.floor(v / Math.pow(2, n)) * Math.pow(2, n); }`.
+ * Its call evaluates nothing but its arguments, so a speculated write may
+ * expand it in place. A call in its body must be of another such function
+ * (`resolve`): any other call could change what the expansion reads. */
+export function inlinableNumberFunction(
+  fn: IrFunction,
+  resolve?: (callee: string) => InlinableFunction | undefined,
+): InlinableFunction | undefined {
+  if (
+    fn.async ||
+    fn.generator ||
+    (fn.captures?.length ?? 0) > 0 ||
+    (fn.classCaptures?.length ?? 0) > 0 ||
+    fn.returnType.kind !== "f64" ||
+    fn.params.some((p) => p.type.kind !== "f64" && p.type.kind !== "union") ||
+    fn.body.length !== 1 ||
+    fn.body[0]!.kind !== "return" ||
+    fn.body[0]!.value === undefined
+  )
+    return undefined;
+  const params = fn.params.map((p) => p.localId);
+  const known = new Set(params);
+  let nodes = 0;
+  // A parameter of a union type (`number | undefined`) is read only
+  // through its tag test and narrowing, which a wrapped argument folds.
+  const union = (e: IrExpr): boolean => e.kind === "varRef" && known.has(e.localId);
+  const pure = (e: IrExpr): boolean => {
+    if (++nodes > MAX_INLINE_NODES) return false;
+    if (e.kind === "boolLit") return true;
+    if (e.kind === "unionIsTag") return union(e.value);
+    if (e.kind === "unionNarrow") return e.type.kind === "f64" && union(e.value);
+    if (e.type.kind !== "f64") return false;
+    switch (e.kind) {
+      case "numLit":
+        return true;
+      case "varRef":
+        return known.has(e.localId);
+      case "ternary":
+        return pure(e.cond) && pure(e.then) && pure(e.else_);
+      case "bin":
+        return (INTEGER_OPS.has(e.op) || e.op === "/") && pure(e.left) && pure(e.right);
+      case "unary":
+        return e.op === "-" && pure(e.operand);
+      case "libCall":
+        return (e.fn === "math.floor" || e.fn === "math.pow") && e.args.every(pure);
+      case "call":
+        return resolve?.(e.callee) !== undefined && e.args.every(pure);
+      default:
+        return false;
+    }
+  };
+  const body = fn.body[0]!.value!;
+  return pure(body) ? { params, body } : undefined;
+}
+
+/** `e` with each parameter reference replaced by its argument. A tag test
+ * or narrowing of a wrapped argument folds, and so does a ternary over a
+ * folded test. */
+function substitute(e: IrExpr, args: ReadonlyMap<string, IrExpr>): IrExpr {
+  switch (e.kind) {
+    case "varRef":
+      return args.get(e.localId) ?? e;
+    case "unionIsTag": {
+      const value = substitute(e.value, args);
+      if (value.kind !== "unionWrap" || value.unionId !== e.unionId) return { ...e, value };
+      return {
+        kind: "boolLit",
+        value: (value.tag === e.tag) !== (e.negated === true),
+        type: e.type,
+        loc: e.loc,
+      };
+    }
+    case "unionNarrow": {
+      const value = substitute(e.value, args);
+      return value.kind === "unionWrap" && value.unionId === e.unionId && value.tag === e.tag
+        ? value.value
+        : { ...e, value };
+    }
+    case "ternary": {
+      const cond = substitute(e.cond, args);
+      if (cond.kind === "boolLit") return substitute(cond.value ? e.then : e.else_, args);
+      return { ...e, cond, then: substitute(e.then, args), else_: substitute(e.else_, args) };
+    }
+    case "bin":
+      return { ...e, left: substitute(e.left, args), right: substitute(e.right, args) };
+    case "unary":
+      return { ...e, operand: substitute(e.operand, args) };
+    case "libCall":
+    case "call":
+      return { ...e, args: e.args.map((a) => substitute(a, args)) };
+    default:
+      return e;
+  }
+}
+
+/** k when `value` is exactly 2^k with k in [0, 52], else null.
+ * Math.log2 rounds: 2^52 + 1 gives 52, so the power is compared back. */
+function exactPowerOfTwo(value: number): number | null {
+  if (!(value >= 1)) return null;
+  const k = Math.round(Math.log2(value));
+  return k >= 0 && k <= 52 && 2 ** k === value ? k : null;
+}
+
+/** Whether `e` is a power of two speculation can shift by:
+ * `Math.pow(2, n)` or a literal 2^k with k in [0, 52]. */
+function powerOfTwo(e: IrExpr): boolean {
+  if (e.kind === "numLit") return exactPowerOfTwo(e.value) !== null;
+  return (
+    e.kind === "libCall" &&
+    e.fn === "math.pow" &&
+    e.args.length === 2 &&
+    e.args[0]!.kind === "numLit" &&
+    e.args[0]!.value === 2
+  );
+}
+
+/** Typed-array element kinds whose elements are integers. */
+const INTEGER_ELEMENTS = new Set(["u8", "u8c", "i8", "u16", "i16", "u32", "i32"]);
+
+/** A typed-array element read of an integer kind with no effect: an
+ * out-of-range read gives NaN, which fails its check. Float elements would
+ * only bail. */
+function isSpeculativeByteRead(e: IrExpr): boolean {
+  return (
+    e.kind === "bytesIntrinsic" &&
+    e.method === "get" &&
+    e.type.kind === "f64" &&
+    e.receiver.type.kind === "bytes" &&
+    INTEGER_ELEMENTS.has(e.receiver.type.elem) &&
+    e.args.length === 1 &&
+    isEffectFree(e.receiver) &&
+    isEffectFree(e.args[0]!)
+  );
+}
+
+/** Whether evaluating `e` assigns no binding: no `++`/`--`, no assignment
+ * expression, and no statement inside it. */
+function writesNoLocal(e: IrExpr): boolean {
+  return everyExpr(e, {
+    expr: (node) => node.kind !== "incDec" && node.kind !== "assignExpr",
+    stmt: () => false,
+  });
+}
+
+/** Rewrite a write's value that isSpeculableValue refuses into one it can
+ * speculate (SpeculativeValue), or null. Expansion and the new operations
+ * (`Math.pow(2, n)`, `Math.floor(a / 2^n)`, byte reads) evaluate nothing
+ * but numbers, so the fast copy may evaluate them in any order. Hoisted
+ * leaves run in the write's evaluation order before the arithmetic; that
+ * reorders nothing only when the write reads no heap and every local it
+ * reads is one no call can write. */
+export function rewriteSpeculable(
+  value: IrExpr,
+  context: SpeculationContext,
+): SpeculativeValue | null {
+  const leaves: SpeculativeValue["leaves"] = [];
+  let reads = false;
+  const locals: string[] = [];
+  const rewrite = (e: IrExpr, depth: number): IrExpr | null => {
+    if (e.type.kind !== "f64") return null;
+    switch (e.kind) {
+      case "numLit":
+        return Number.isSafeInteger(e.value) && !Object.is(e.value, -0) ? e : null;
+      case "varRef":
+        locals.push(e.localId);
+        return e;
+      case "bin": {
+        if (!ARITHMETIC_OPS.has(e.op)) break;
+        const left = rewrite(e.left, depth);
+        const right = left && rewrite(e.right, depth);
+        return left && right ? { ...e, left, right } : null;
+      }
+      case "unary": {
+        if (e.op !== "-") break;
+        const operand = rewrite(e.operand, depth);
+        return operand ? { ...e, operand } : null;
+      }
+      case "libCall": {
+        if (powerOfTwo(e) && e.kind === "libCall") {
+          const exponent = rewrite(e.args[1]!, depth);
+          return exponent ? { ...e, args: [e.args[0]!, exponent] } : null;
+        }
+        const quotient = e.args[0];
+        if (
+          e.fn === "math.floor" &&
+          e.args.length === 1 &&
+          quotient?.kind === "bin" &&
+          quotient.op === "/" &&
+          powerOfTwo(quotient.right)
+        ) {
+          const left = rewrite(quotient.left, depth);
+          const right = left && rewrite(quotient.right, depth);
+          return left && right ? { ...e, args: [{ ...quotient, left, right }] } : null;
+        }
+        break;
+      }
+      case "call": {
+        const fn = depth < MAX_INLINE_DEPTH ? context.inlinable?.(e.callee) : undefined;
+        // Arguments are evaluated before the body: an expansion may read
+        // them in another order, or twice, only when they are pure.
+        const plain = (a: IrExpr): boolean =>
+          isSpeculableValue(a, false) ||
+          (a.kind === "unionWrap" && isSpeculableValue(a.value, false));
+        if (fn !== undefined && fn.params.length === e.args.length && e.args.every(plain)) {
+          const args = new Map(fn.params.map((id, i) => [id, e.args[i]!]));
+          return rewrite(substitute(fn.body, args), depth + 1);
+        }
+        break;
+      }
+      default:
+        if (isSpeculativeRead(e) || isSpeculativeByteRead(e)) {
+          reads = true;
+          return e;
+        }
+        break;
+    }
+    // Only calls of program functions are hoisted: other leaves (a string's
+    // length, a Map read) are cheap where they stand, and would turn every
+    // loop that adds one into a second copy.
+    if (context.leafId === undefined || e.kind !== "call") return null;
+    // Hoisting runs the arguments before the rest of the write: one that
+    // assigns a local would move that write ahead of the write's earlier
+    // reads of it.
+    if (!e.args.every(writesNoLocal)) return null;
+    const ref: IrExpr = { kind: "varRef", localId: context.leafId(), type: e.type, loc: e.loc };
+    leaves.push({ localId: (ref as { localId: string }).localId, expr: e, ref });
+    return ref;
+  };
+  const rewritten = rewrite(value, 0);
+  if (rewritten === null) return null;
+  if (leaves.length > 0) {
+    const leafIds = new Set(leaves.map((leaf) => leaf.localId));
+    if (reads || !locals.every((id) => leafIds.has(id) || context.unaliased?.(id) === true))
+      return null;
+  }
+  return { value: rewritten, leaves };
 }
 
 type SpeculativeLoop = Extract<IrStmt, { kind: "for" | "while" | "doWhile" }>;
@@ -170,13 +460,16 @@ function isEffectFree(e: IrExpr): boolean {
   }
 }
 
-/** The speculative reads of a speculable value in evaluation order. */
+/** The speculative reads of a speculable value in evaluation order, those
+ * of a rewritten value (rewriteSpeculable) included. */
 export function speculativeReads(e: IrExpr, out: IrExpr[] = []): IrExpr[] {
-  if (e.kind === "bin" && ARITHMETIC_OPS.has(e.op)) {
+  if (e.kind === "bin" && (ARITHMETIC_OPS.has(e.op) || e.op === "/")) {
     speculativeReads(e.left, out);
     speculativeReads(e.right, out);
   } else if (e.kind === "unary" && e.op === "-") speculativeReads(e.operand, out);
-  else if (isSpeculativeRead(e)) out.push(e);
+  else if (e.kind === "libCall" && (e.fn === "math.floor" || e.fn === "math.pow"))
+    for (const arg of e.args) speculativeReads(arg, out);
+  else if (isSpeculativeRead(e) || isSpeculativeByteRead(e)) out.push(e);
   return out;
 }
 
@@ -195,6 +488,8 @@ function hasProduct(e: IrExpr): boolean {
   if (e.kind === "bin")
     return e.op === "*" || e.op === "%" || hasProduct(e.left) || hasProduct(e.right);
   if (e.kind === "unary") return hasProduct(e.operand);
+  // Math.floor(a / 2^n) and Math.pow(2, n) replace a division and a call.
+  if (e.kind === "libCall") return e.fn === "math.floor" || e.fn === "math.pow";
   return false;
 }
 
@@ -202,6 +497,7 @@ function hasArithmetic(e: IrExpr): boolean {
   if (e.kind === "bin")
     return ARITHMETIC_OPS.has(e.op) || hasArithmetic(e.left) || hasArithmetic(e.right);
   if (e.kind === "unary") return e.op === "-" || hasArithmetic(e.operand);
+  if (e.kind === "libCall") return e.fn === "math.floor" || e.fn === "math.pow";
   return false;
 }
 
@@ -236,8 +532,11 @@ export function planSpeculativeLoop(
   integral: (localId: string) => boolean = () => false,
   proven: (e: IrExpr) => boolean = () => false,
   invariantEligible: (localId: string) => boolean = eligible,
+  context: SpeculationContext = {},
 ): SpeculativeLoopPlan | null {
   const writes = new Map<string, IrStmt[]>();
+  const values = new Map<IrStmt, SpeculativeValue>();
+  const leafIds = new Set<string>();
   const unspeculable = new Set<string>();
   const reads = new Set<string>();
   // Locals with at least one read the static analysis does not prove.
@@ -263,11 +562,9 @@ export function planSpeculativeLoop(
   // make the body's work integer.
   let inCondition = false;
 
-  const expr = (e: IrExpr): boolean => {
-    if (rejected) return false;
-    nodes++;
-    if (!OBSERVER_FREE_EXPRESSIONS.has(e.kind) || (e.kind === "bin" && e.left.type.kind !== "f64"))
-      observerFree = false;
+  // What the integer writes read and consume. A rewritten value is noted
+  // here alone: its original, which `expr` walks, has its effects.
+  const note = (e: IrExpr): void => {
     if (e.kind === "varRef") {
       reads.add(e.localId);
       if (!proven(e)) unproven.add(e.localId);
@@ -281,6 +578,20 @@ export function planSpeculativeLoop(
     } else if (e.kind === "arrayGet" || e.kind === "arrayHas" || e.kind === "arrayState")
       consume(e.index);
     else if (e.kind === "bytesIntrinsic" || e.kind === "strIntrinsic") e.args.forEach(consume);
+  };
+  // An expanded call is emitted in the fast copy: it counts toward the
+  // duplicated code.
+  const noteAll = (e: IrExpr): boolean => {
+    nodes++;
+    note(e);
+    return everyExprChild(e, noteAll, () => true);
+  };
+  const expr = (e: IrExpr): boolean => {
+    if (rejected) return false;
+    nodes++;
+    if (!OBSERVER_FREE_EXPRESSIONS.has(e.kind) || (e.kind === "bin" && e.left.type.kind !== "f64"))
+      observerFree = false;
+    note(e);
     return everyExprChild(e, expr, nested);
   };
   // Statements inside expressions (sequence expressions) run mid-expression,
@@ -307,7 +618,11 @@ export function planSpeculativeLoop(
     if (s.kind === "arraySet" || s.kind === "bytesSet") consume(s.index);
     const target = speculativeWriteTarget(s);
     if (target) {
-      if (target.value === null || isSpeculableValue(target.value)) {
+      const rewritten =
+        target.value === null || isSpeculableValue(target.value)
+          ? null
+          : rewriteSpeculable(target.value, context);
+      if (target.value === null || isSpeculableValue(target.value) || rewritten !== null) {
         const list = writes.get(target.localId) ?? [];
         list.push(s);
         writes.set(target.localId, list);
@@ -315,6 +630,11 @@ export function planSpeculativeLoop(
           reads.add(target.localId);
           if (s.kind !== "exprStmt" || !proven(s.expr)) unproven.add(target.localId);
         } else expr(target.value);
+        if (rewritten !== null) {
+          values.set(s, rewritten);
+          for (const leaf of rewritten.leaves) leafIds.add(leaf.localId);
+          noteAll(rewritten.value);
+        }
         return !rejected;
       }
       unspeculable.add(target.localId);
@@ -341,25 +661,28 @@ export function planSpeculativeLoop(
   // write runs: speculated itself, already an integer, or loop-invariant
   // (checked on entry). A value computed from doubles the loop writes
   // would only fail its check on every entry.
+  const valueOf = (s: IrStmt): IrExpr | null =>
+    values.get(s)?.value ?? speculativeWriteTarget(s)!.value;
   const grounded = (e: IrExpr): boolean => {
     if (e.kind === "varRef") {
       const id = e.localId;
       return (
+        leafIds.has(id) ||
         proven(e) ||
         locals.has(id) ||
         integral(id) ||
         (!writes.has(id) && !unspeculable.has(id) && invariantEligible(id))
       );
     }
-    if (isSpeculativeRead(e)) return true;
+    if (isSpeculativeRead(e) || isSpeculativeByteRead(e)) return true;
     return everyExprChild(e, grounded, () => false);
   };
   for (let changed = true; changed;) {
     changed = false;
     for (const id of [...locals.keys()]) {
       for (const s of writes.get(id)!) {
-        const target = speculativeWriteTarget(s)!;
-        if (target.value !== null && !grounded(target.value)) {
+        const value = valueOf(s);
+        if (value !== null && !grounded(value)) {
           locals.delete(id);
           changed = true;
           break;
@@ -376,6 +699,7 @@ export function planSpeculativeLoop(
     if (e.kind === "varRef") {
       const id = e.localId;
       if (
+        !leafIds.has(id) &&
         !proven(e) &&
         !locals.has(id) &&
         !writes.has(id) &&
@@ -388,11 +712,11 @@ export function planSpeculativeLoop(
     return everyExprChild(e, invariant, () => true);
   };
   for (const s of speculated) {
-    const target = speculativeWriteTarget(s)!;
-    if (target.value === null) arithmetic = true;
+    const value = valueOf(s);
+    if (value === null) arithmetic = true;
     else {
-      if (hasArithmetic(target.value)) arithmetic = true;
-      invariant(target.value);
+      if (hasArithmetic(value)) arithmetic = true;
+      invariant(value);
     }
   }
   if (!arithmetic) return null;
@@ -401,13 +725,12 @@ export function planSpeculativeLoop(
     [...locals.keys()].some((id) => multiplied.has(id)) ||
     [...speculated].some((s) => {
       const target = speculativeWriteTarget(s)!;
-      return (
-        target.value !== null &&
-        (hasProduct(target.value) || isVariableStep(target.localId, target.value))
-      );
+      const value = valueOf(s);
+      return value !== null && (hasProduct(value) || isVariableStep(target.localId, value));
     });
   if (!heavy) return null;
-  return { locals, writes: speculated, observerFree };
+  for (const s of values.keys()) if (!speculated.has(s)) values.delete(s);
+  return { locals, writes: speculated, observerFree, values };
 }
 
 /** A proven integer loop. The bindings integer-ranges.ts bounds over a for
@@ -682,6 +1005,19 @@ export function emitSpeculativeInteger(
       bail(negativeZero);
       return result;
     }
+    case "libCall": {
+      // A rewritten value's Math.pow(2, n) and Math.floor(a / 2^n)
+      // (rewriteSpeculable): with n an integer in [0, 52], 2^n is exact and
+      // so is a / 2^n for a safe integer a, whose floor the arithmetic shift
+      // computes, also for negative a. No result is -0.
+      if (e.fn === "math.pow") return shiftOne(host, exponentOf(host, e, bail, reads, integral));
+      if (e.fn !== "math.floor" || e.args[0]?.kind !== "bin") break;
+      const a = emitSpeculativeInteger(host, e.args[0].left, bail, reads, integral);
+      const n = exponentOf(host, e.args[0].right, bail, reads, integral);
+      const result = B.tmp();
+      B.line(`${result} = ashr i64 ${a}, ${n}`);
+      return result;
+    }
     case "unary": {
       if (e.op !== "-") break;
       const a = emitSpeculativeInteger(host, e.operand, bail, reads, integral);
@@ -699,6 +1035,35 @@ export function emitSpeculativeInteger(
   return speculativeLeaf(host, reads.get(e) ?? host.emitExpr(e), e, bail);
 }
 
+/** The exponent of a power of two (Math.pow(2, n) or a literal 2^k) as an
+ * i64 in [0, 52]; other exponents bail. */
+function exponentOf(
+  host: LlvmEmitterContext,
+  e: IrExpr,
+  bail: (condition: string) => void,
+  reads: ReadonlyMap<IrExpr, LlValue>,
+  integral: boolean,
+): string {
+  if (e.kind === "numLit") {
+    const k = exactPowerOfTwo(e.value);
+    if (k === null) throw new InternalCompilerError("llvm emitter bug: speculative power of two");
+    return String(k);
+  }
+  if (e.kind !== "libCall")
+    throw new InternalCompilerError("llvm emitter bug: speculative power of two");
+  const n = emitSpeculativeInteger(host, e.args[1]!, bail, reads, integral);
+  const outside = host.B.tmp();
+  host.B.line(`${outside} = icmp ugt i64 ${n}, 52`);
+  bail(outside);
+  return n;
+}
+
+function shiftOne(host: LlvmEmitterContext, n: string): string {
+  const result = host.B.tmp();
+  host.B.line(`${result} = shl i64 1, ${n}`);
+  return result;
+}
+
 /** The safe i64 of an evaluated number: its proven integer, or a checked
  * conversion. */
 function speculativeLeaf(

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/7757-speculative-key-chains.ts.json

@@ -0,0 +1,6 @@
+{
+ "order": [
+  "<repo>/tests/corpus/7757-speculative-key-chains.ts"
+ ],
+ "diags": []
+}

tests/corpus/7757-speculative-key-chains.ts

@@ -0,0 +1,124 @@
+// Speculative integer loops over writes that call small pure number
+// functions, use Math.pow(2, n) and Math.floor(a / 2^n), read bytes, or
+// call other functions: every fallback gives Node's numbers.
+function trim(value: number, length: number): number {
+  return value - Math.floor(value / Math.pow(2, length)) * Math.pow(2, length);
+}
+function maybeTrim(value: number | undefined, length: number): number {
+  return (value ?? 0) - Math.floor((value ?? 0) / Math.pow(2, length)) * Math.pow(2, length);
+}
+
+function chain(codes: Uint8Array, length: number): number {
+  const mask = length * 2 - 2;
+  let key = 0;
+  let sum = 0;
+  for (let i = 0; i < codes.length; ++i) {
+    sum = (sum + key) % 1000000007;
+    key = trim(key, mask) * 4 + codes[i]!;
+  }
+  return sum * 31 + key;
+}
+
+function optionalChain(codes: Uint8Array, length: number): number {
+  let key = 0;
+  let sum = 0;
+  for (let i = 0; i < codes.length; ++i) {
+    sum = (sum * 3 + key) % 1000003;
+    key = maybeTrim(key, length * 2) * 4 + codes[i]!;
+  }
+  return sum + key;
+}
+
+// Exponents past 52, negative and fractional values, and NaN reads leave
+// the integer copy and finish in the double one.
+function shifts(start: number, exponent: number, steps: number): number[] {
+  const out: number[] = [];
+  let v = start;
+  for (let i = 0; i < steps; i++) {
+    v = Math.floor(v / Math.pow(2, exponent)) * 3 + i * 7;
+    out.push(v);
+  }
+  return out;
+}
+
+let calls = 0;
+function noisy(x: number): number {
+  calls++;
+  return x % 5 === 0 ? x / 2 : x * 2;
+}
+// A call that the write cannot expand runs once per iteration, also on
+// the iteration whose check fails.
+function hoisted(n: number): number {
+  let acc = 1;
+  for (let i = 0; i < n; i++) acc = (acc * 7 + noisy(i)) % 1000003;
+  return acc;
+}
+
+// A check that fails before a write with a hoisted call: that write still
+// makes its call on the iteration that bails.
+function tally(x: number): number {
+  calls++;
+  return x * 2 + 1;
+}
+function chained(n: number): number {
+  let a = 1;
+  let b = 1;
+  for (let i = 0; i < n; i++) {
+    a = a * 3 + i;
+    b = (b * 7 + tally(i) + (a % 11)) % 1000003;
+  }
+  return b;
+}
+
+// A divisor next to a power of two is not one: 2^52 + 1 divides as itself.
+function nearPowers(start: number, steps: number): number[] {
+  const out: number[] = [];
+  let v = start;
+  for (let i = 0; i < steps; i++) {
+    v = Math.floor(v / 4503599627370497) * 4 + Math.floor(v / 4503599627370496) + i;
+    out.push(v);
+  }
+  return out;
+}
+
+// A helper that calls a function with effects is not expanded: its
+// argument is read before the call runs.
+let level = 0;
+function raise(): number {
+  level += 100;
+  return 1;
+}
+function lifted(x: number): number {
+  return raise() + x;
+}
+function effectfulHelper(n: number): number {
+  let acc = 0;
+  for (let i = 0; i < n; i++) acc = (acc * 3 + lifted(level)) % 1000003;
+  return acc;
+}
+
+// A call whose argument assigns a local the write also reads stays where it
+// is: `i` is read before `i++` runs.
+function stepped(n: number): number {
+  let acc = 1;
+  let i = 0;
+  while (i < n) acc = (acc * 3 + i + tally(i++)) % 1000003;
+  return acc;
+}
+
+const codes = new Uint8Array(5000);
+let seed = 7;
+for (let i = 0; i < codes.length; i++) {
+  seed = (seed * 1103515245 + 12345) % 2147483648;
+  codes[i] = seed % 4;
+}
+for (const length of [1, 2, 3, 6, 12, 18, 27]) console.log(length, chain(codes, length), optionalChain(codes, length));
+console.log(shifts(1000, 3, 6), shifts(-1000, 2, 6), shifts(5, 53, 3), shifts(7.5, 1, 3), shifts(2 ** 52, 1, 3));
+console.log(hoisted(20), calls);
+const short = codes.subarray(0, 3);
+console.log(chain(short, 2), chain(new Uint8Array(0), 2));
+calls = 0;
+console.log(chained(34), chained(60), calls);
+console.log(nearPowers(4503599627370496, 3), nearPowers(9007199254740991, 2));
+console.log(effectfulHelper(10), level);
+console.log(stepped(40));

tests/harness/corpus-costs/7757-speculative-key-chains.ts.json

@@ -0,0 +1 @@
+{ "plain": 2200, "san": 2900 }