Represent record-only unions used as Map keys and Set elements as record pointers (#2160)

benchmarks/runtime/README.md

@@ -36,6 +36,7 @@ Application-shaped workloads for measuring the speed of compiled executables. Ea
 | `typed-array-kernels` | cpu | element-wise loops that run to one typed array's length and index others at the same position: `Float64Array` axpy, `Float32Array` blend, `Uint8Array` mask, wrapping `Int32Array` add, `Uint8ClampedArray` brighten, RGBA to gray, a row-major `Float32Array` matrix-vector product, and `reduce`/`some`/`map` callbacks |
 | `typed-array-scan`, `typed-array-scan-workers` | cpu | element-by-element scan loops over `Uint8Array`, `Uint16Array`, `Int32Array` and `Float64Array`; the second entry is a worker-enabled executable (it can start a worker), so each access also checks for shared storage |
 | `unicode-lexer` | cpu | a hand-written lexer over a 2.7M-unit bundled module whose comments and strings are in several languages: `charCodeAt` scans, line starts |
+| `union-node-sets` | cpu | syntax-tree passes over a discriminated union of object literals: a `Set` of visited nodes, a `Map` from node to value, conversions to a sub-union of leaves |
 | `validate-errors` | cpu | throw/catch on a fraction of inputs, small call chains |
 | `value-objects` | cpu | small immutable records and class instances: vector helpers, method chains, particle updates |
 | `value-records` | cpu | records that hold strings, spread state updates and destructured tuples: token helpers, an immutable reducer, tuple returns |

benchmarks/runtime/workloads/union-node-sets.json

@@ -0,0 +1,6 @@
+{
+  "entry": "workloads/union-node-sets.ts",
+  "args": [
+    "1"
+  ]
+}

benchmarks/runtime/workloads/union-node-sets.ts

@@ -0,0 +1,141 @@
+// Syntax-tree passes over a discriminated union of object literals, the way
+// compilers walk their IR: a DAG of shared subexpressions, a `Set` of
+// visited nodes, a `Map` from node to its value, and conversions between
+// the whole union and the sub-unions passes narrow it to.
+type Num = { kind: "num"; value: number };
+type Var = { kind: "var"; slot: number };
+type Neg = { kind: "neg"; operand: Node };
+type Add = { kind: "add"; left: Node; right: Node };
+type Mul = { kind: "mul"; left: Node; right: Node };
+type Pick = { kind: "pick"; cond: Node; then: Node; else_: Node };
+type Leaf = Num | Var;
+type Node = Leaf | Neg | Add | Mul | Pick;
+
+function children(node: Node, visit: (child: Node) => void): void {
+  switch (node.kind) {
+    case "num":
+    case "var":
+      return;
+    case "neg":
+      visit(node.operand);
+      return;
+    case "add":
+    case "mul":
+      visit(node.left);
+      visit(node.right);
+      return;
+    case "pick":
+      visit(node.cond);
+      visit(node.then);
+      visit(node.else_);
+      return;
+  }
+}
+
+function leafValue(leaf: Leaf, env: number[]): number {
+  return leaf.kind === "num" ? leaf.value : env[leaf.slot]!;
+}
+
+function evaluate(node: Node, env: number[], memo: Map<Node, number>): number {
+  const cached = memo.get(node);
+  if (cached !== undefined) return cached;
+  let result: number;
+  switch (node.kind) {
+    case "num":
+    case "var":
+      result = leafValue(node, env);
+      break;
+    case "neg":
+      result = -evaluate(node.operand, env, memo);
+      break;
+    case "add":
+      result = (evaluate(node.left, env, memo) + evaluate(node.right, env, memo)) % 1000003;
+      break;
+    case "mul":
+      result = (evaluate(node.left, env, memo) * evaluate(node.right, env, memo)) % 1000003;
+      break;
+    case "pick":
+      result =
+        evaluate(node.cond, env, memo) % 2 === 0
+          ? evaluate(node.then, env, memo)
+          : evaluate(node.else_, env, memo);
+      break;
+  }
+  memo.set(node, result);
+  return result;
+}
+
+function reachable(root: Node): number {
+  const seen = new Set<Node>();
+  const stack: Node[] = [root];
+  while (stack.length > 0) {
+    const node = stack.pop()!;
+    if (seen.has(node)) continue;
+    seen.add(node);
+    children(node, (child) => {
+      if (!seen.has(child)) stack.push(child);
+    });
+  }
+  return seen.size;
+}
+
+function leaves(root: Node): Leaf[] {
+  const out: Leaf[] = [];
+  const seen = new Set<Node>();
+  const walk = (node: Node): void => {
+    if (seen.has(node)) return;
+    seen.add(node);
+    if (node.kind === "num" || node.kind === "var") out.push(node);
+    else children(node, walk);
+  };
+  walk(root);
+  return out;
+}
+
+let seed = 2463534242;
+function next(): number {
+  seed ^= seed << 13;
+  seed ^= seed >>> 17;
+  seed ^= seed << 5;
+  return seed >>> 0;
+}
+
+function build(size: number): Node {
+  const pool: Node[] = [];
+  for (let i = 0; i < 64; i++)
+    pool.push(i % 2 === 0 ? { kind: "num", value: i * 7 + 1 } : { kind: "var", slot: i % 16 });
+  // Operands come from the most recent nodes, so most of the DAG is
+  // reachable from the root and many nodes are shared.
+  const recent = (): Node => pool[pool.length - 1 - (next() % 48)]!;
+  for (let i = 0; i < size; i++) {
+    const a = recent();
+    const b = recent();
+    const r = next() % 10;
+    let node: Node;
+    if (r < 1) node = { kind: "neg", operand: a };
+    else if (r < 5) node = { kind: "add", left: a, right: b };
+    else if (r < 8) node = { kind: "mul", left: a, right: b };
+    else node = { kind: "pick", cond: a, then: b, else_: recent() };
+    pool.push(node);
+  }
+  return pool[pool.length - 1]!;
+}
+
+const scale = Number(process.argv[2] ?? "1");
+const rounds = Math.max(1, Math.round(40 * scale));
+let checksum = 0;
+let nodes = 0;
+let leafCount = 0;
+for (let round = 0; round < rounds; round++) {
+  const root = build(20_000);
+  nodes += reachable(root);
+  const found = leaves(root);
+  leafCount += found.length;
+  for (let pass = 0; pass < 4; pass++) {
+    const env: number[] = [];
+    for (let i = 0; i < 16; i++) env.push((round * 31 + pass * 7 + i) % 97);
+    checksum = (checksum + evaluate(root, env, new Map())) % 1000003;
+    for (const leaf of found) checksum = (checksum + leafValue(leaf, env)) % 1000003;
+  }
+}
+console.log("rounds", rounds, "nodes", nodes, "leaves", leafCount, "checksum", checksum);

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

@@ -1014,7 +1014,7 @@ export function emitModuleParts(emitter: LlEmitter): string[] {
     // LIBRARY mode: no @main — the profile-declared external
     // symbols specified by the library IR instead.
     for (const line of emitter.emitLibDefs(globals, globalReleaseLines, stamps)) out.push(line);
-    out.push(...emitter.inlineRcTail(flushedDecls));
+    for (const line of emitter.inlineRcTail(flushedDecls)) out.push(line);
     out.push(...functionAttributes(emitter, hasNoInlineRecordClone));
     out.push(``);
     out.push(...emitter.fieldAliasMetadata());
@@ -1160,7 +1160,9 @@ export function emitModuleParts(emitter: LlEmitter): string[] {
       : [`  %native_exit = call i32 @scr_exit_code_hint_get()`, `  ret i32 %native_exit`]),
     `}`,
     ``,
-    ...emitter.inlineRcTail(flushedDecls),
+  );
+  for (const line of emitter.inlineRcTail(flushedDecls)) out.push(line);
+  out.push(
     // sanitize_address is inert under the plain pipeline; the sanitized
     // lane's -fsanitize=address link activates instrumentation over the
     // emitted functions too (the runtime TUs get theirs from clang).

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

@@ -1219,7 +1219,9 @@ function collectionOps(host: LlvmEmitterContext, key: IrType, value: IrType): st
 function newCollection(host: LlvmEmitterContext, key: IrType, value: IrType): string {
   const B = host.B;
   const arms = key.kind === "union" ? host.unionsById.get(key.unionId)?.arms : undefined;
-  const kind = mapKeyKindNum(key, arms);
+  // SCR_MAP_KEY_REF (2): shape-union keys are record pointers.
+  const kind =
+    key.kind === "union" && host.nullableUnions.shapes(key.unionId) ? 2 : mapKeyKindNum(key, arms);
   const m = B.tmp();
   host.declare(`declare ptr @scr_map_new(i32, i32, ptr)`);
   // SCR_MAP_VAL_CREF (3): string keys and compact values (compactMapValues).

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

@@ -7,6 +7,8 @@ import { DYN_KIND } from "./dyn.js";
 import { LlvmUnsupportedError } from "./unsupported.js";
 import type { LlvmEmitterContext, ExprOf, LlValue } from "./expr-context.js";
 import { f64Lit } from "./common.js";
+import { NULLABLE_NULL, SHAPE_WORD_INDEX, shapeWordHeader } from "./nullable-unions.js";
+import { emitUnionTag } from "./union-repr.js";
 
 export function emitControlExpr(
   host: LlvmEmitterContext,
@@ -52,6 +54,8 @@ export function emitControlExpr(
     case "toBool":
       return { name: host.truthy(host.emitUnionProjection(e.operand)), type: e.type };
     case "logical": {
+      const tagSet = e.op === "||" ? unionTagSet(e) : null;
+      if (tagSet) return { name: emitUnionTagSet(host, tagSet), type: e.type };
       // JS value semantics: the result is the deciding operand itself.
       // Left evaluates once, ownership moves into the result slot; when
       // the branch takes the right operand the stale left releases first
@@ -614,3 +618,85 @@ function literalChoice(e: IrExpr): boolean {
   if (e.kind === "strLit") return true;
   return e.kind === "ternary" && literalChoice(e.then) && literalChoice(e.else_);
 }
+
+type TagSet = { unionId: string; value: IrExpr; tags: Set<number> };
+
+/** An `||` chain of two or more tag tests of one local (`isTag(u, 1) ||
+ * isTag(u, 4) || ...`, the shape of a union conversion's identity test). */
+function unionTagSet(e: IrExpr & { kind: "logical" }): TagSet | null {
+  const leaves: IrExpr[] = [];
+  const collect = (x: IrExpr): void => {
+    if (x.kind === "logical" && x.op === "||") {
+      collect(x.left);
+      collect(x.right);
+    } else leaves.push(x);
+  };
+  collect(e);
+  const first = leaves[0]!;
+  if (first.kind !== "unionIsTag" || first.value.kind !== "varRef") return null;
+  const localId = first.value.localId;
+  const tags = new Set<number>();
+  for (const leaf of leaves) {
+    if (
+      leaf.kind !== "unionIsTag" ||
+      leaf.negated ||
+      leaf.unionId !== first.unionId ||
+      leaf.value.kind !== "varRef" ||
+      leaf.value.localId !== localId
+    )
+      return null;
+    tags.add(leaf.tag);
+  }
+  return { unionId: first.unionId, value: first.value, tags };
+}
+
+/** i1: the union value holds one of the tags. One read of the tag and one
+ * switch over the set; a shape union switches on the record's shape word
+ * after its unit encodings. */
+function emitUnionTagSet(host: LlvmEmitterContext, set: TagSet): string {
+  const B = host.B;
+  const u = host.emitUnionProjection(set.value).name;
+  const nullable = host.nullableUnions.get(set.unionId);
+  const yes = B.newLabel("tags.y");
+  const no = B.newLabel("tags.n");
+  const join = B.newLabel("tags.j");
+  const target = (tag: number): string => (set.tags.has(tag) ? yes : no);
+  if (nullable?.repr === "shapes") {
+    if (nullable.unitTag >= 0) {
+      const isNull = B.tmp();
+      const next = B.newLabel("tags.r");
+      B.line(`${isNull} = icmp eq ptr ${u}, null`);
+      B.condBr(isNull, target(nullable.unitTag), next);
+      B.startBlock(next);
+    }
+    if (nullable.nullTag >= 0) {
+      const isNullArm = B.tmp();
+      const next = B.newLabel("tags.r");
+      B.line(`${isNullArm} = icmp eq ptr ${u}, ${NULLABLE_NULL}`);
+      B.condBr(isNullArm, target(nullable.nullTag), next);
+      B.startBlock(next);
+    }
+    const wp = B.tmp();
+    const w = B.tmp();
+    B.line(
+      `${wp} = getelementptr inbounds ${shapeWordHeader(host.sizeType)}, ptr ${u}, i32 0, i32 ${SHAPE_WORD_INDEX}`,
+    );
+    B.line(`${w} = load i32, ptr ${wp}`);
+    const cases = nullable.arms
+      .filter((arm) => set.tags.has(arm.tag))
+      .map((arm) => `i32 ${host.nullableUnions.shapeNumber(arm.shapeId)}, label %${yes}`);
+    B.terminate(`switch i32 ${w}, label %${no} [ ${cases.join(" ")} ]`);
+  } else {
+    const tag = nullable ? emitUnionTag(B, nullable, u) : host.unionTag(u, set.unionId);
+    const cases = [...set.tags].map((t) => `i32 ${t}, label %${yes}`);
+    B.terminate(`switch i32 ${tag}, label %${no} [ ${cases.join(" ")} ]`);
+  }
+  B.startBlock(yes);
+  B.br(join);
+  B.startBlock(no);
+  B.br(join);
+  B.startBlock(join);
+  const t = B.tmp();
+  B.line(`${t} = phi i1 [ true, %${yes} ], [ false, %${no} ]`);
+  return t;
+}

packages/compiler/src/backend/llvm/nullable-unions.test.ts

@@ -56,10 +56,11 @@ test("unions of two or more records become shape unions", () => {
 });
 
 test("records of runtime-inspected unions keep every union over them boxed", () => {
-  const set: IrType = { kind: "set", elem: { kind: "union", unionId: "ab" } };
+  const set: IrType = { kind: "set", elem: { kind: "union", unionId: "abf" } };
   const unions = new NullableUnions(
     module(
       [
+        { id: "abf", arms: [a, b, F64] },
         { id: "ab", arms: [a, b] },
         { id: "abu", arms: [a, b, UNDEFINED_T] },
         { id: "bc", arms: [b, c] },
@@ -68,6 +69,7 @@ test("records of runtime-inspected unions keep every union over them boxed", ()
       { globals: [{ id: "%g.s", name: "s", type: set, mutable: false }] },
     ),
   );
+  expect(unions.get("abf")).toBeNull();
   expect(unions.get("ab")).toBeNull();
   expect(unions.get("abu")).toBeNull();
   expect(unions.get("bc")).toBeNull();
@@ -75,6 +77,34 @@ test("records of runtime-inspected unions keep every union over them boxed", ()
   expect(unions.shapeNumber("c")).toBe(0);
 });
 
+test("unions of records alone stay shape unions as Map keys and Set elements", () => {
+  const set: IrType = { kind: "set", elem: { kind: "union", unionId: "ab" } };
+  const map: IrType = { kind: "map", key: { kind: "union", unionId: "abc" }, value: F64 };
+  const unions = new NullableUnions(
+    module(
+      [
+        { id: "ab", arms: [a, b] },
+        { id: "abc", arms: [a, b, c] },
+        { id: "ba", arms: [b, a] },
+        { id: "abu", arms: [a, b, UNDEFINED_T] },
+      ],
+      {
+        globals: [
+          { id: "%g.s", name: "s", type: set, mutable: false },
+          { id: "%g.m", name: "m", type: map, mutable: false },
+        ],
+      },
+    ),
+  );
+  expect(unions.shapes("ab")).not.toBeNull();
+  expect(unions.shapes("abc")).not.toBeNull();
+  expect(unions.shapes("abu")).not.toBeNull();
+  // Unions over the same records share one release and trace helper.
+  expect(unions.helperOwner("ba")).toBe("ab");
+  expect(unions.helperOwner("abu")).toBe("ab");
+  expect(unions.helperOwner("abc")).toBe("abc");
+});
+
 test("library, lifetime-region and debug builds keep record unions boxed", () => {
   const unions = [{ id: "ab", arms: [a, b] }];
   expect(new NullableUnions(module(unions), false).get("ab")).toBeNull();

packages/compiler/src/backend/llvm/nullable-unions.ts

@@ -118,6 +118,9 @@ export class NullableUnions {
   readonly retainWrappers = new Map<string, string>();
   /** Some nullable union encodes `null` with NULLABLE_NULL. */
   readonly usesNullSentinel: boolean = false;
+  /** The shape union whose release and trace helpers each shape union
+   * uses: the first one over the same record arms. */
+  private readonly helperOwners = new Map<string, string>();
 
   /** `shapes`: allow shape unions (debug builds keep boxes, whose
    * debug-info layouts describe every arm). */
@@ -151,6 +154,24 @@ export class NullableUnions {
           if (!this.shapeNumbers.has(arm.shapeId))
             this.shapeNumbers.set(arm.shapeId, this.shapeNumbers.size + 2);
     }
+    // Release and trace dispatch on the shape word, which names a record
+    // shape module-wide, so unions over the same records share them.
+    const owners = new Map<string, string>();
+    for (const nullable of this.byId.values()) {
+      if (nullable.repr !== "shapes") continue;
+      const key = nullable.arms
+        .map((arm) => arm.shapeId)
+        .sort()
+        .join(" ");
+      const owner = owners.get(key) ?? nullable.unionId;
+      owners.set(key, owner);
+      this.helperOwners.set(nullable.unionId, owner);
+    }
+  }
+
+  /** The shape union whose release and trace helpers this one uses. */
+  helperOwner(unionId: string): string {
+    return this.helperOwners.get(unionId) ?? unionId;
   }
 
   get(unionId: string): PointerUnion | null {
@@ -277,13 +298,19 @@ function shapeUnion(
 
 /** Unions the runtime itself inspects as tagged boxes keep that
  * representation: Map keys and Set elements hash and compare union keys by
- * tag and payload inside the C map. */
+ * tag and payload inside the C map. A union of records alone is the
+ * exception: its keys compare by record identity, which a shape union's
+ * record pointer carries as well (the map keys it by reference). */
 function boxedUnionIds(mod: IrModule, nodes?: IrNodeIndex): Set<string> {
   const boxed = new Set<string>();
   // Only the module's union definitions consult the result.
   if (!mod.unions?.length) return boxed;
+  const recordsOnly = new Set<string>();
+  for (const def of mod.unions)
+    if (def.arms.length >= 2 && def.arms.every((arm) => arm.kind === "record"))
+      recordsOnly.add(def.id);
   const mark = (t: IrType): void => {
-    if (t.kind === "union") boxed.add(t.unionId);
+    if (t.kind === "union" && !recordsOnly.has(t.unionId)) boxed.add(t.unionId);
   };
   if (nodes !== undefined) {
     for (const t of nodes.types) {

packages/compiler/src/backend/llvm/record-union-keys.test.ts

@@ -0,0 +1,169 @@
+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, deserializeModule } from "../../index.js";
+import { BOOL, F64, VOID, type IrExpr, type IrModule, type IrType } from "../../ir/ir.js";
+import { validateModule } from "../../ir/validate.js";
+import { emitLlvmModule } from "./emitter.js";
+
+const loc = { file: "record-union-keys.ts", start: 0, end: 0 };
+
+async function llvmFor(source: string): Promise<string> {
+  const dir = await mkdtemp(join(tmpdir(), "scriptc-record-union-keys-"));
+  try {
+    const entry = join(dir, "main.ts"),
+      output = join(dir, "main.ir.json");
+    await writeFile(entry, source);
+    const result = await compile(entry, {
+      outDir: dir,
+      outPath: output,
+      outputKind: "ir",
+    });
+    if (!result.ok)
+      throw new Error(result.diagnostics.map((item) => `${item.code}: ${item.message}`).join("\n"));
+    const module = deserializeModule(await readFile(output, "utf8"));
+    expect(validateModule(module)).toEqual([]);
+    return emitLlvmModule(module);
+  } finally {
+    await rm(dir, { recursive: true, force: true });
+  }
+}
+
+function body(ir: string, name: string): string {
+  const found = new RegExp(`^define internal [^\\n]*@${name}(?:\\.i32)?\\([^]*?^}`, "m").exec(ir);
+  expect(found, name).not.toBeNull();
+  return found![0];
+}
+
+test("Sets and Maps of record unions key the record pointer", async () => {
+  const ir = await llvmFor(`
+type Num = { kind: "num"; value: number };
+type Neg = { kind: "neg"; operand: Node };
+type Add = { kind: "add"; left: Node; right: Node };
+type Node = Num | Neg | Add;
+const one: Node = { kind: "num", value: 1 };
+const nodes: Node[] = [one, { kind: "neg", operand: one }, { kind: "add", left: one, right: one }];
+const seen = new Set<Node>();
+const depth = new Map<Node, number>();
+for (const node of nodes) {
+  seen.add(node);
+  depth.set(node, node.kind === "num" ? 0 : 1);
+}
+console.log(seen.size, depth.get(nodes[1]!), seen.has({ kind: "num", value: 1 }));
+`);
+  // SCR_MAP_KEY_REF (2), never SCR_MAP_KEY_UNION_REF (3), and no union box.
+  expect(ir).toMatch(/call ptr @scr_map_new\(i32 2,/);
+  expect(ir).not.toMatch(/call ptr @scr_map_new\(i32 3,/);
+  expect(ir).not.toMatch(/call ptr @scr_union_new_ref/);
+});
+
+const recordType = (shapeId: string): IrType => ({ kind: "record", shapeId });
+const wide: IrType = { kind: "union", unionId: "wide" };
+
+/** `pick(u)` answers whether `u` holds arm 0, 2 or 3 of a five-record union
+ * through an `||` chain of tag tests; `same` is a second union over the same
+ * records in another order. */
+function tagChainModule(): IrModule {
+  const u: IrExpr = { kind: "varRef", localId: "u", type: wide, loc };
+  const isTag = (tag: number): IrExpr => ({
+    kind: "unionIsTag",
+    unionId: "wide",
+    tag,
+    negated: false,
+    value: u,
+    type: BOOL,
+    loc,
+  });
+  const or = (left: IrExpr, right: IrExpr): IrExpr => ({
+    kind: "logical",
+    op: "||",
+    left,
+    right,
+    type: BOOL,
+    loc,
+  });
+  const shapes = ["a", "b", "c", "d", "e"];
+  return {
+    irVersion: 17,
+    sourceFile: loc.file,
+    entry: "main",
+    records: shapes.map((id) => ({ id, fields: [{ name: id, type: F64 }] })),
+    unions: [
+      { id: "wide", arms: shapes.map(recordType) },
+      { id: "same", arms: [...shapes].reverse().map(recordType) },
+    ],
+    functions: [
+      {
+        name: "main",
+        params: [],
+        returnType: VOID,
+        locals: [],
+        body: [],
+        loc,
+      },
+      {
+        name: "pick",
+        params: [{ localId: "u", name: "u", type: wide }],
+        returnType: BOOL,
+        locals: [{ id: "u", name: "u", type: wide, mutable: false }],
+        body: [{ kind: "return", value: or(or(isTag(0), isTag(2)), isTag(3)), loc }],
+        loc,
+      },
+      {
+        name: "keep",
+        params: [{ localId: "s", name: "s", type: { kind: "union", unionId: "same" } }],
+        returnType: VOID,
+        locals: [
+          {
+            id: "s",
+            name: "s",
+            type: { kind: "union", unionId: "same" },
+            mutable: false,
+          },
+        ],
+        body: [],
+        loc,
+      },
+    ],
+  };
+}
+
+test("an || chain of tag tests is one switch on the shape word", () => {
+  const mod = tagChainModule();
+  expect(validateModule(mod)).toEqual([]);
+  const ir = emitLlvmModule(mod, { pointerBits: 64 });
+  const pick = body(ir, "sc_bf_pick");
+  const switches = pick.match(/switch i32 [^\n]*/g) ?? [];
+  expect(switches).toHaveLength(1);
+  expect(switches[0]!.match(/label %tags\.y/g)).toHaveLength(3);
+  expect(pick).not.toMatch(/@sc_stag\./);
+});
+
+test("unions over the same records share one release", () => {
+  const ir = emitLlvmModule(tagChainModule(), { pointerBits: 64 });
+  expect(ir.match(/^define internal void @sc_srelease\.[^\n]*noinline/gm)).toHaveLength(1);
+  expect(ir).toMatch(/^define internal void @sc_srelease\.wide\(/m);
+  expect(ir).not.toMatch(/^define internal void @sc_srelease\.same\(/m);
+});
+
+test("a union with a freezable arm compares its count without the frozen bit", async () => {
+  const ir = await llvmFor(`
+type Name = { kind: "name"; text: string };
+type Count = { kind: "count"; value: number };
+type Key = Name | Count;
+function name(text: string): Key {
+  const built: Name = { kind: "name", text };
+  return Object.freeze(built);
+}
+const seen = new Set<Key>();
+seen.add(name("a"));
+seen.add({ kind: "count", value: 1 });
+console.log(seen.size);
+`);
+  const release = /^define internal void @sc_srelease\.[^\n]*\n[^]*?^}/m.exec(ir);
+  expect(release).not.toBeNull();
+  // SCR_RC_FROZEN is bit 62: a frozen record's last reference has count 1.
+  expect(release![0]).toContain("and i64 %rc, -4611686018427387905");
+  expect(release![0]).toContain("icmp ult i64 %rcm2, -4611686018427387907");
+});

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

@@ -233,7 +233,7 @@ export { computeTraced } from "../cycle-analysis.js";
  * `_v`-shaped. */
 export function vAdapters(host: ShapeHost, t: IrType): { retain: string; release: string } {
   const nullable = host.nullableUnions.of(t);
-  if (nullable?.repr === "shapes") return shapeUnionRc(nullable);
+  if (nullable?.repr === "shapes") return shapeUnionRc(host, nullable);
   if (nullable) {
     const arm = vAdapters(host, nullable.arm);
     return { retain: nullableRetainSym(host, nullable.arm), release: arm.release };
@@ -290,7 +290,7 @@ export function vAdapters(host: ShapeHost, t: IrType): { retain: string; release
 export function retainSym(host: ShapeHost, t: IrType): string {
   if (t.kind === "jsval") return jsvalRcSym(host, "retain");
   const nullable = host.nullableUnions.of(t);
-  if (nullable?.repr === "shapes") return shapeUnionRc(nullable).retain;
+  if (nullable?.repr === "shapes") return shapeUnionRc(host, nullable).retain;
   if (nullable) return nullableRetainSym(host, nullable.arm);
   const stem = runtimeRcStem(t);
   if (host.rcHelpers !== null && stem !== null && inlineRcFamily(host, stem, "retain") !== null) {
@@ -305,10 +305,10 @@ export function retainSym(host: ShapeHost, t: IrType): string {
 /** A shape union's RC entry points: every record retain is the same
  * NULL- and immortal-skipping increment (the null sentinel is immortal),
  * so the first arm's serves; release dispatches on the arm. */
-function shapeUnionRc(union: ShapeUnion): { retain: string; release: string } {
+function shapeUnionRc(host: ShapeHost, union: ShapeUnion): { retain: string; release: string } {
   return {
     retain: `@${mangleRecordRetain(union.arms[0]!.shapeId)}`,
-    release: shapeReleaseFn(union.unionId),
+    release: shapeReleaseFn(host.nullableUnions.helperOwner(union.unionId)),
   };
 }
 
@@ -352,38 +352,81 @@ export function emitShapeUnionHelpers(host: ShapeHost): string[] {
       `}`,
       ``,
     );
-    // release: unit encodings are NULL or immortal; each arm releases
-    // through its own (NULL- and immortal-tolerant) helper.
-    out.push(
-      `define internal void ${shapeReleaseFn(union.unionId)}(ptr %u) ${FN_ATTRS} { ; shape union release`,
-      `entry:`,
-      `  %t = call i32 ${tagFn}(ptr %u)`,
-      `  switch i32 %t, label %done [ ${union.arms.map((arm) => `i32 ${arm.tag}, label %a${arm.tag}`).join(" ")} ]`,
-      ...union.arms.flatMap((arm) => [
-        `a${arm.tag}:`,
-        `  call void @${mangleRecordRelease(arm.shapeId)}(ptr %u)`,
-        `  br label %done`,
-      ]),
+    // Release and trace: one definition for the unions over the same
+    // record arms (helperOwner). Unit encodings are NULL or the immortal
+    // null sentinel; the shape word picks the arm.
+    if (host.nullableUnions.helperOwner(union.unionId) !== union.unionId) continue;
+    const units = [
+      `  %isnull = icmp eq ptr %o, null`,
+      `  br i1 %isnull, label %done, label %${host.nullableUnions.usesNullSentinel ? "notnull" : "rec"}`,
+      ...(host.nullableUnions.usesNullSentinel
+        ? [
+            `notnull:`,
+            `  %isnullarm = icmp eq ptr %o, ${NULLABLE_NULL}`,
+            `  br i1 %isnullarm, label %done, label %rec`,
+          ]
+        : []),
+      `rec:`,
+    ];
+    const armSwitch = (call: (shapeId: string) => string): string[] => [
+      `  %wp = getelementptr inbounds ${header}, ptr %o, i32 0, i32 ${SHAPE_WORD_INDEX}`,
+      `  %w = load i32, ptr %wp`,
+      `  switch i32 %w, label %done [ ${union.arms.map((arm) => `i32 ${host.nullableUnions.shapeNumber(arm.shapeId)}, label %a${arm.tag}`).join(" ")} ]`,
+      ...union.arms.flatMap((arm) => [`a${arm.tag}:`, call(arm.shapeId), `  br label %done`]),
       `done:`,
       `  ret void`,
       `}`,
       ``,
-    );
+    ];
+    // Every record arm shares the release fast path (the count at the same
+    // offset, one collector header choice for the whole union), so only a
+    // release that may free dispatches on the arm. The release stays a call:
+    // inlined at every release site, it grows code more than it saves.
+    const traced = union.arms.map((arm) => host.tracedShapes.has(`record:${arm.shapeId}`));
+    // A freezable arm's count is the word without its frozen bit, as in
+    // releaseFastPath: a frozen record's last reference takes the slow path.
+    const frozen = union.arms.some((arm) => host.recordsById.get(arm.shapeId)?.freezable === true);
+    const notFrozen = ~recordFrozenBit(S);
+    const release = shapeReleaseFn(union.unionId);
+    if (traced.every((t) => t === traced[0]))
+      out.push(
+        `define internal void ${release}(ptr %o) noinline ${FN_ATTRS} { ; shape union release`,
+        `entry:`,
+        ...units,
+        `  %rc = load ${S}, ptr %o`,
+        ...(frozen
+          ? [
+              `  %refs = and ${S} %rc, ${sizeLiteral(S, notFrozen)}`,
+              `  %rcm2 = sub ${S} %refs, 2`,
+              `  %shared = icmp ult ${S} %rcm2, ${sizeLiteral(S, notFrozen - 2n)}`,
+            ]
+          : [`  %rcm2 = sub ${S} %rc, 2`, `  %shared = icmp ult ${S} %rcm2, -3`]),
+        `  br i1 %shared, label %dec, label %slow`,
+        `dec:`,
+        `  %n = sub ${S} %rc, 1`,
+        `  store ${S} %n, ptr %o`,
+        ...(traced[0] ? cycleRootLines(host, "done") : [`  br label %done`]),
+        `slow:`,
+        `  %imm = icmp eq ${S} %rc, -1`,
+        `  br i1 %imm, label %done, label %free`,
+        `free:`,
+        ...armSwitch((shapeId) => `  call void @${mangleRecordRelease(shapeId)}_slow(ptr %o)`),
+      );
+    else
+      out.push(
+        `define internal void ${release}(ptr %o) ${FN_ATTRS} { ; shape union release`,
+        `entry:`,
+        ...units,
+        ...armSwitch((shapeId) => `  call void @${mangleRecordRelease(shapeId)}(ptr %o)`),
+      );
     if (shapeUnionTraced(host, union))
       out.push(
-        `define internal void ${shapeTraceFn(union.unionId)}(ptr %u, ptr %visit, ptr %ctx) ${FN_ATTRS} { ; shape union trace`,
+        `define internal void ${shapeTraceFn(union.unionId)}(ptr %o, ptr %visit, ptr %ctx) ${FN_ATTRS} { ; shape union trace`,
         `entry:`,
-        `  %t = call i32 ${tagFn}(ptr %u)`,
-        `  switch i32 %t, label %done [ ${union.arms.map((arm) => `i32 ${arm.tag}, label %a${arm.tag}`).join(" ")} ]`,
-        ...union.arms.flatMap((arm) => [
-          `a${arm.tag}:`,
-          `  call void @${mangleRecordTrace(arm.shapeId)}(ptr %u, ptr %visit, ptr %ctx)`,
-          `  br label %done`,
-        ]),
-        `done:`,
-        `  ret void`,
-        `}`,
-        ``,
+        ...units,
+        ...armSwitch(
+          (shapeId) => `  call void @${mangleRecordTrace(shapeId)}(ptr %o, ptr %visit, ptr %ctx)`,
+        ),
       );
   }
   return out;
@@ -435,7 +478,7 @@ export function emitNullableRetainWrappers(host: ShapeHost): string[] {
 export function releaseSym(host: ShapeHost, t: IrType): string {
   if (t.kind === "jsval") return jsvalRcSym(host, "release");
   const nullable = host.nullableUnions.of(t);
-  if (nullable?.repr === "shapes") return shapeUnionRc(nullable).release;
+  if (nullable?.repr === "shapes") return shapeUnionRc(host, nullable).release;
   if (nullable) return releaseSym(host, nullable.arm);
   const stem = runtimeRcStem(t);
   if (host.rcHelpers !== null && stem !== null && inlineRcFamily(host, stem, "release") !== null) {
@@ -985,7 +1028,9 @@ export function traceAdapter(host: ShapeHost, t: IrType): string | null {
     case "union": {
       const nullable = host.nullableUnions.get(t.unionId);
       if (nullable?.repr === "shapes")
-        return shapeUnionTraced(host, nullable) ? shapeTraceFn(nullable.unionId) : null;
+        return shapeUnionTraced(host, nullable)
+          ? shapeTraceFn(host.nullableUnions.helperOwner(nullable.unionId))
+          : null;
       if (nullable) return traceAdapter(host, nullable.arm);
       if (!host.tracedUnions.has(t.unionId)) return null;
       host.declare(`declare void @scr_union_trace_v(ptr, ptr, ptr)`);

packages/compiler/src/backend/llvm/test-boxed-unions.ts

@@ -1,14 +1,29 @@
-import type { IrModule, IrType } from "../../ir/ir.js";
+import { F64, type IrModule, type IrType, type IrUnionDef } from "../../ir/ir.js";
 
 /** Test helper: a Set of each union keeps it a tagged `%ScrUnion` box (Set
  * elements are runtime-inspected boxes), so stack-box tests exercise the
- * paths that nullable-pointer unions (nullable-unions.ts) bypass. */
+ * paths that nullable-pointer unions (nullable-unions.ts) bypass. A union of
+ * records alone is a Set element by record identity, so for those the Set
+ * holds a companion union with a number arm, which keeps the records out of
+ * every shape union. */
 export function keepUnionsBoxed(mod: IrModule): IrModule {
-  const keep = (mod.unions ?? []).map((union) => ({
-    id: `%g.keep.${union.id}`,
-    name: `keep_${union.id}`,
-    type: { kind: "set", elem: { kind: "union", unionId: union.id } } as IrType,
-    mutable: false,
-  }));
-  return { ...mod, globals: [...(mod.globals ?? []), ...keep] };
+  const companions: IrUnionDef[] = [];
+  const keep = (mod.unions ?? []).map((union) => {
+    let elem = union.id;
+    if (union.arms.every((arm) => arm.kind === "record")) {
+      elem = `%keep.${union.id}`;
+      companions.push({ id: elem, arms: [...union.arms, F64] });
+    }
+    return {
+      id: `%g.keep.${union.id}`,
+      name: `keep_${union.id}`,
+      type: { kind: "set", elem: { kind: "union", unionId: elem } } as IrType,
+      mutable: false,
+    };
+  });
+  return {
+    ...mod,
+    unions: [...(mod.unions ?? []), ...companions],
+    globals: [...(mod.globals ?? []), ...keep],
+  };
 }

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/8291-frozen-record-union-keys.ts.json

@@ -0,0 +1,6 @@
+{
+ "order": [
+  "<repo>/tests/corpus/8291-frozen-record-union-keys.ts"
+ ],
+ "diags": []
+}

tests/corpus/8291-frozen-record-union-keys.ts

@@ -0,0 +1,65 @@
+// Frozen records in record-only unions used as Set elements and Map keys:
+// the union's release frees a frozen record when it drops the last
+// reference, for unions the cycle collector traces and ones it does not.
+type Leaf = { kind: "leaf"; label: string };
+type Pair = { kind: "pair"; left: Node; right: Node };
+type Node = Leaf | Pair;
+
+type Name = { kind: "name"; text: string };
+type Count = { kind: "count"; value: number; unit: string };
+type Key = Name | Count;
+
+function leaf(label: string): Node {
+  const built: Leaf = { kind: "leaf", label };
+  return Object.freeze(built);
+}
+
+function pair(left: Node, right: Node): Node {
+  const built: Pair = { kind: "pair", left, right };
+  return Object.freeze(built);
+}
+
+function key(i: number): Key {
+  if (i % 2 === 0) {
+    const name: Name = { kind: "name", text: "n" + i };
+    return Object.freeze(name);
+  }
+  const count: Count = { kind: "count", value: i, unit: "u" + i };
+  return Object.freeze(count);
+}
+
+function nodes(n: number): number {
+  const seen = new Set<Node>();
+  const weight = new Map<Node, number>();
+  let last = leaf("start");
+  for (let i = 0; i < n; i++) {
+    const next = i % 3 === 0 ? leaf("l" + i) : pair(last, leaf("r" + i));
+    seen.add(next);
+    weight.set(next, i);
+    last = next;
+  }
+  let total = 0;
+  for (const node of seen) total += (weight.get(node) ?? 0) + (node.kind === "pair" ? 1 : 0);
+  seen.clear();
+  weight.clear();
+  return total + (last.kind === "pair" ? 1 : 0);
+}
+
+function keys(n: number): number {
+  const seen = new Set<Key>();
+  const size = new Map<Key, number>();
+  for (let i = 0; i < n; i++) {
+    const k = key(i);
+    seen.add(k);
+    size.set(k, k.kind === "name" ? k.text.length : k.value);
+  }
+  let total = 0;
+  for (const k of seen) total += size.get(k) ?? 0;
+  seen.clear();
+  size.clear();
+  return total;
+}
+
+let sum = 0;
+for (let r = 0; r < 50; r++) sum += nodes(40) + keys(40);
+console.log(sum);

tests/harness/corpus-costs/8291-frozen-record-union-keys.ts.json

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