Read the union payload inline when comparing it with a number or boolean

`u === 0`, `u !== true` and `Object.is(u, x)` against a scalar arm tested
the union's tag inline but read its payload through scr_union_get_f64 or
scr_union_get_bool. The call passed the union's stack box to the runtime,
so LLVM could not promote the box to registers, and an element read of a
`number[]` compared with a number kept its value, tag and box in memory.

The comparison now loads the payload slot directly, as other union
extractions already do, and the box no longer escapes.

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

@@ -1282,9 +1282,7 @@ function emitUnionEqWrappedScalar(
     B.startBlock(same);
     const eq = B.tmp();
     if (arm.kind === "f64") {
-      host.declare(`declare double @scr_union_get_f64(ptr)`);
-      const v = B.tmp();
-      B.line(`${v} = call double @scr_union_get_f64(ptr ${union.name})`);
+      const v = host.unionGetF64(union.name);
       if (e.sameValue) {
         host.declare(`declare zeroext i1 @scr_num_same_value(double, double)`);
         B.line(`${eq} = call zeroext i1 @scr_num_same_value(double ${v}, double ${plain.name})`);
@@ -1292,9 +1290,7 @@ function emitUnionEqWrappedScalar(
         B.line(`${eq} = fcmp oeq double ${v}, ${plain.name}`);
       }
     } else {
-      host.declare(`declare zeroext i1 @scr_union_get_bool(ptr)`);
-      const v = B.tmp();
-      B.line(`${v} = call zeroext i1 @scr_union_get_bool(ptr ${union.name})`);
+      const v = host.unionGetBool(union.name);
       B.line(`${eq} = icmp eq i1 ${v}, ${plain.name}`);
     }
     B.line(`store i1 ${eq}, ptr ${slot}`);

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/7322-union-scalar-equality.ts.json

@@ -0,0 +1,6 @@
+{
+ "order": [
+  "<repo>/tests/corpus/7322-union-scalar-equality.ts"
+ ],
+ "diags": []
+}

packages/compiler/test/union-scalar-equality-emission.test.ts

@@ -0,0 +1,45 @@
+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-union-scalar-eq-"));
+  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" });
+    expect(result.ok, JSON.stringify(result.diagnostics)).toBe(true);
+    return await readFile(outPath, "utf8");
+  } finally {
+    await rm(dir, { recursive: true, force: true });
+  }
+}
+
+const source = `
+function zeros(a: number[], at: number[]): number {
+  let n = 0;
+  for (const i of at) if (a[i] === 0) n++;
+  return n;
+}
+function flags(a: (boolean | undefined)[]): number {
+  let n = 0;
+  for (let i = 0; i < a.length; i++) if (a[i] === true) n++;
+  return n;
+}
+console.log(zeros([0, 1, 0], [0, 1, 2, 5]), flags([true, undefined, false]));
+`;
+
+test("comparing a union with a number or boolean reads its payload inline", async () => {
+  const llvm = await emitLlvm(source);
+  for (const name of ["zeros", "flags"]) {
+    const fn = new RegExp(`^define internal double @sc_bf_${name}\\([^]*?^}`, "m").exec(llvm)?.[0];
+    expect(fn, name).toBeDefined();
+    // The tag test and the payload load stay in the function, so the
+    // union's stack box has no escaping use.
+    expect(fn).not.toMatch(/@scr_union_get_(f64|bool)\(/);
+    expect(fn).toMatch(/getelementptr inbounds %ScrUnion, ptr %\S+, i64 0, i32 5/);
+  }
+});

tests/corpus/7322-union-scalar-equality.ts

@@ -0,0 +1,68 @@
+// Strict equality and inequality between a `number | undefined` or
+// `boolean | undefined` value and a number or boolean: element reads of
+// holes, past the end and of present undefined, NaN, -0 and +0, and
+// Object.is, in branches, loops and expressions.
+
+function zeros(a: number[], at: number[]): string {
+  const hits: string[] = [];
+  for (const i of at) {
+    const v = a[i];
+    hits.push(`${i}:${v === 0 ? "eq" : "ne"}${v !== 0 ? "!" : ""}${0 === v ? "r" : ""}`);
+  }
+  return hits.join(" ");
+}
+
+const holey: number[] = [0, -0, 1];
+holey[5] = 0;
+console.log(zeros(holey, [0, 1, 2, 3, 4, 5, 6, -1]));
+console.log(zeros([NaN, 0.5, -1, 0], [0, 1, 2, 3]));
+
+function nanAndZeros(a: (number | undefined)[]): string {
+  return a
+    .map((v) => `${v === -0}${v === 0}${Object.is(v, -0)}${Object.is(v, 0)}${Object.is(v, NaN)}`)
+    .join(" ");
+}
+console.log(nanAndZeros([0, -0, NaN, undefined, 1]));
+
+function flags(a: (boolean | undefined)[]): string {
+  let t = 0;
+  let f = 0;
+  let notTrue = 0;
+  for (let i = 0; i < a.length + 2; i++) {
+    if (a[i] === true) t++;
+    if (a[i] === false) f++;
+    if (a[i] !== true) notTrue++;
+  }
+  return `${t} ${f} ${notTrue}`;
+}
+const bools: (boolean | undefined)[] = [true, false, undefined, true];
+bools[7] = false;
+console.log(flags(bools));
+
+// A permutation walk that stops at a zero element, as the flip loops of a
+// counting benchmark do.
+function walk(p: number[], start: number): number {
+  let steps = 0;
+  let at = start;
+  while (p[at] !== 0 && steps < 20) {
+    at = p[at]!;
+    steps++;
+  }
+  return steps;
+}
+console.log(walk([0, 3, 1, 2], 1), walk([1, 2, 3, 0], 0), walk([2, 2, 2], 0));
+
+class Slot {
+  value: number | undefined = undefined;
+  on: boolean | undefined = undefined;
+}
+const slots = [new Slot(), new Slot(), new Slot()];
+slots[1]!.value = 0;
+slots[1]!.on = true;
+slots[2]!.value = -0;
+slots[2]!.on = false;
+console.log(
+  slots
+    .map((s) => `${s.value === 0}${s.value !== 0}${s.on === true}${s.on === false}${s.on !== false}`)
+    .join(" "),
+);

tests/harness/corpus-costs/7322-union-scalar-equality.ts.json

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