Convert the base of an offset typed-array index once

A typed array read or write at `base + c`, `c + base` or `base - c`,
where `base` is a number local and `c` an integer literal of magnitude
below 2^31, now converts `base` to an integer and adds `c` as an
integer, instead of converting the double sum. Accesses at several
offsets from one base, such as `x[at]` through `x[at + 7]` with a number
parameter `at`, then share one conversion and integer test. Below a
length of at most 2^52 the sum is an integer index exactly when `base`
is an integer, so the accesses read and write the same elements as
before.

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

@@ -164,6 +164,43 @@ export function exactIndexLines(
   return { lines, ok, wide };
 }
 
+/** exactIndexLines for the index `base + offset`, where `offset` is an
+ * integer of magnitude below 2^31: the round trip converts `base` instead of
+ * the sum, so accesses at several offsets from one base share the
+ * conversion. A length is at most 2^52, and below it the sum is an integer
+ * exactly when `base` is one, and is then the integer sum: a fraction of
+ * `base` stays in a sum below 2^52, and a sum of 2^52 or more is out of
+ * range either way. */
+export function exactOffsetIndexLines(
+  base: string,
+  offset: number,
+  length: string,
+  sizeType: string,
+  tmp: () => string,
+): { lines: string[]; ok: string; wide: string } {
+  const raw = tmp(),
+    integer = tmp(),
+    back = tmp(),
+    integral = tmp(),
+    wide = tmp(),
+    below = tmp(),
+    ok = tmp();
+  const lines = [
+    `${raw} = fptosi double ${base} to i64`,
+    `${integer} = freeze i64 ${raw}`,
+    `${back} = sitofp i64 ${integer} to double`,
+    `${integral} = fcmp oeq double ${back}, ${base}`,
+    `${wide} = add i64 ${integer}, ${offset}`,
+  ];
+  let len64 = length;
+  if (sizeType !== "i64") {
+    len64 = tmp();
+    lines.push(`${len64} = zext ${sizeType} ${length} to i64`);
+  }
+  lines.push(`${below} = icmp ult i64 ${wide}, ${len64}`, `${ok} = and i1 ${integral}, ${below}`);
+  return { lines, ok, wide };
+}
+
 /** User-controlled text embedded after an LLVM `;` comment marker. Preserve
  * ordinary output unchanged, but encode control and line-separator code
  * units so a property name can never inject a line or invalid source byte. */

packages/compiler/src/backend/llvm/expr-bytes.ts

@@ -4,7 +4,7 @@ import { emitBorrowedInput } from "./borrowed-inputs.js";
 import type { BytesHeader } from "./bytes-headers.js";
 import { BYTES_ELEMENT_SIZE, F64, type IrBytesElem, type IrExpr } from "../../ir/ir.js";
 import type { LlvmEmitterContext, LlValue } from "./expr-context.js";
-import { F64_INF, exactIndexLines, f64Lit, sizeToDouble } from "./common.js";
+import { F64_INF, exactIndexLines, exactOffsetIndexLines, f64Lit, sizeToDouble } from "./common.js";
 import { alwaysExactInteger, exactInteger } from "./integer-values.js";
 import { byteNumberAccess } from "../../ir/byte-numbers.js";
 import { emitByteNumber } from "./byte-numbers.js";
@@ -77,6 +77,29 @@ export function emitIntegerLoopIndex(host: LlvmEmitterContext, expr: IrExpr): st
   return index;
 }
 
+/** The base, read again, and the offset of an index `base + c`, `c + base`
+ * or `base - c`, where `base` is a number local and `c` an integer literal
+ * of magnitude below 2^31. Null for other indexes. */
+function offsetIndex(
+  host: LlvmEmitterContext,
+  expr: IrExpr | undefined,
+): { base: string; offset: number } | null {
+  if (expr?.kind !== "bin" || expr.type.kind !== "f64" || (expr.op !== "+" && expr.op !== "-"))
+    return null;
+  const literal = (e: IrExpr): number | null =>
+    e.kind === "numLit" && Number.isInteger(e.value) && Math.abs(e.value) < 2 ** 31
+      ? e.value
+      : null;
+  const local = (e: IrExpr): boolean => e.kind === "varRef" && e.type.kind === "f64";
+  const right = literal(expr.right);
+  const left = literal(expr.left);
+  if (right !== null && local(expr.left))
+    return { base: host.emitExpr(expr.left).name, offset: expr.op === "+" ? right : -right };
+  if (expr.op === "+" && left !== null && local(expr.right))
+    return { base: host.emitExpr(expr.right).name, offset: left };
+  return null;
+}
+
 export function emitBytesIndex(
   host: LlvmEmitterContext,
   receiver: string,
@@ -134,7 +157,10 @@ export function emitBytesIndex(
     B.line(`${ok} = and i1 ${nanIndex.ordered}, ${below}`);
     wide = nanIndex.wide;
   } else {
-    const exact = exactIndexLines(index, len, host.sizeType, () => B.tmp());
+    const offset = offsetIndex(host, expr);
+    const exact = offset
+      ? exactOffsetIndexLines(offset.base, offset.offset, len, host.sizeType, () => B.tmp())
+      : exactIndexLines(index, len, host.sizeType, () => B.tmp());
     for (const line of exact.lines) B.line(line);
     ok = exact.ok;
     wide = exact.wide;

packages/compiler/test/offset-index-emission.test.ts

@@ -0,0 +1,39 @@
+import { mkdtemp, readFile, rm, writeFile } from "node:fs/promises";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import { expect, test } from "vitest";
+import { compile } from "../src/index.js";
+
+async function emitLlvm(source: string): Promise<string> {
+  const dir = await mkdtemp(join(tmpdir(), "scriptc-offset-index-"));
+  try {
+    const entry = join(dir, "main.ts");
+    const outPath = join(dir, "main.ll");
+    await writeFile(entry, source);
+    const result = await compile(entry, { outDir: dir, outPath, outputKind: "llvm" });
+    if (!result.ok)
+      throw new Error(result.diagnostics.map((d) => `${d.code}: ${d.message}`).join("\n"));
+    return await readFile(outPath, "utf8");
+  } finally {
+    await rm(dir, { recursive: true, force: true });
+  }
+}
+
+test("typed array accesses at a number base plus an integer convert the base, not the sum", async () => {
+  const llvm = await emitLlvm(`
+function window(values: Float64Array, at: number): number {
+  return values[at + 0]! + values[at + 1]! + values[2 + at]! + values[at - 1]!;
+}
+const values = new Float64Array([1, 2, 3, 4, 5, 6]);
+console.log(window(values, Number(process.argv[2] ?? "1")));
+`);
+  const start = llvm.indexOf("@sc_bf_window(");
+  const body = llvm.slice(start, llvm.indexOf("\n}\n", start));
+  const sums = new Set([...body.matchAll(/(%t\d+) = fadd double/g)].map((m) => m[1]));
+  const converted = [...body.matchAll(/fptosi double (%[\w.]+) to i64/g)].map((m) => m[1]);
+  expect(converted.length).toBeGreaterThanOrEqual(4);
+  // Each access converts the base and adds its offset as an integer.
+  expect(converted.filter((name) => sums.has(name))).toEqual([]);
+  for (const offset of [1, 2, -1])
+    expect(body).toMatch(new RegExp(`= add i64 %t\\d+, ${offset}\\n`));
+});

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/typed-array-offset-indexes.ts.json

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

tests/corpus/typed-array-offset-indexes.ts

@@ -0,0 +1,33 @@
+// Typed array reads and writes at `base + c`, `c + base` and `base - c`,
+// where `base` is a number that may be an integer, a fraction, negative,
+// -0, NaN, infinite or past the length.
+
+function readGroup(values: Float64Array, at: number): number[] {
+  return [values[at + 0]!, values[at + 1]!, values[2 + at]!, values[at - 1]!, values[at + 7]!];
+}
+
+function writeGroup(values: Int32Array, at: number, value: number): void {
+  values[at + 0] = value;
+  values[1 + at] = value + 1;
+  values[at - 2] = value - 2;
+}
+
+function sumWindow(values: Uint8Array, at: number): number {
+  return (values[at] ?? 0) + (values[at + 1] ?? 0) + (values[at + 2] ?? 0) + (values[at + 3] ?? 0);
+}
+
+const floats = new Float64Array(16);
+for (let i = 0; i < floats.length; i++) floats[i] = i * 1.5;
+const bases = [0, 1, 4, 8.5, -0, -1, -0.5, 15, 16, 2 ** 31, 2 ** 52, 2 ** 53, -(2 ** 53), NaN, Infinity, -Infinity];
+for (const base of bases) console.log(String(base), JSON.stringify(readGroup(floats, base)));
+
+const ints = new Int32Array(8);
+for (const base of [2, 3.25, -0, 7, 9, NaN, 2 ** 52 + 1]) {
+  writeGroup(ints, base, 10);
+  console.log(String(base), ints.join(","));
+}
+
+const bytes = new Uint8Array([5, 4, 3, 2, 1, 250, 251, 252]);
+let total = 0;
+for (let i = 0; i < 8; i++) total += sumWindow(bytes, i);
+console.log(total, sumWindow(bytes, 6), sumWindow(bytes, 4.5));

tests/harness/corpus-costs/typed-array-offset-indexes.ts.json

@@ -0,0 +1 @@
+{ "plain": 2100, "san": 2500 }