Key Sets and Maps of identity-arm unions by their pointer representation (#2487)

benchmarks/runtime/README.md

@@ -27,6 +27,7 @@ Application-shaped workloads for measuring the speed of compiled executables. Ea
 | `log-summary`, `inventory-report` | cpu | the multi-module applications from `benchmarks/builds`, on 300k-line inputs |
 | `number-maps` | cpu | number-keyed Map/Set: a dense-id table, sparse ids in many small maps, get/has/set mixes |
 | `numeric-kernels` | cpu | sieve, hashing, `number[][]` matmul, `Float64Array` |
+| `optional-child-sets` | cpu | syntax-tree walks through class nodes whose getters return `TreeNode \| undefined` children, in a program that also keeps a `Set<TreeNode \| undefined>` |
 | `optional-fields` | cpu | records and class instances with `number \| undefined` and other scalar-union fields |
 | `published-graphs` | cpu | syntax trees with parent links and a symbol table, published with `@scriptc/threads` and then read |
 | `record-untyped-writes` | cpu | typed records written through untyped views: a counter helper taking `unknown` bumps fields of a stats record, and `Object.assign` merges patches into a settings record, with typed reads in between |

benchmarks/runtime/workloads/optional-child-sets.json

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

benchmarks/runtime/workloads/optional-child-sets.ts

@@ -0,0 +1,48 @@
+// Syntax-tree walks through class nodes whose getters return optional
+// children, in a program that also keeps a `Set` of the optional node type.
+class TreeNode {
+  readonly kind: number;
+  private readonly kids: (TreeNode | undefined)[];
+  parentNode: TreeNode | undefined;
+  constructor(kind: number, kids: (TreeNode | undefined)[]) {
+    this.kind = kind;
+    this.kids = kids;
+    this.parentNode = undefined;
+    for (const kid of kids) if (kid !== undefined) kid.parentNode = this;
+  }
+  get left(): TreeNode | undefined {
+    return this.kids[0];
+  }
+  get right(): TreeNode | undefined {
+    return this.kids[1];
+  }
+  get parent(): TreeNode | undefined {
+    return this.parentNode;
+  }
+}
+
+function build(depth: number, seed: number): TreeNode {
+  if (depth === 0) return new TreeNode(seed % 7, [undefined, undefined]);
+  const right = seed % 11 === 0 ? undefined : build(depth - 1, seed * 7 + 3);
+  return new TreeNode(seed % 7, [build(depth - 1, seed * 5 + 1), right]);
+}
+
+function sumKinds(node: TreeNode | undefined): number {
+  let total = 0;
+  while (node !== undefined) {
+    total += node.kind;
+    const up = node.parent;
+    if (up !== undefined && up.kind === node.kind) total++;
+    total += sumKinds(node.right);
+    node = node.left;
+  }
+  return total;
+}
+
+const root = build(20, 1);
+// The one Set of the optional node type.
+const marked = new Set<TreeNode | undefined>([root, root.left, undefined]);
+let total = 0;
+const rounds = 30 * Number(process.argv[2] ?? "1");
+for (let i = 0; i < rounds; i++) total += sumKinds(root);
+console.log(total, marked.size, marked.has(root.right));

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

@@ -561,6 +561,10 @@ export interface LlvmTargetOptions {
   debugSources?: ReadonlyMap<string, string>;
   /** Pointer width of the target C ABI. Native targets are 64-bit today. */
   pointerBits?: 32 | 64;
+  /** Unions emitted as tagged boxes whatever their arms, for tests of the
+   * box paths: in a program, a union that has a pointer representation
+   * (nullable-unions.ts) never holds a box. */
+  boxedUnions?: ReadonlySet<string>;
   /** Select the WASI libc entry-point convention. */
   wasi?: boolean;
   /** Library archive assembly may move the volatile identity getters into a
@@ -2089,6 +2093,7 @@ export class LlEmitter {
       options.debugSources === undefined,
       this.nodeIndex,
       options.pointerBits !== 32,
+      options.boxedUnions,
     );
     this.heapCensus = heapCensusOf(mod, this.nullableUnions);
     for (const decl of this.heapCensus?.declarations() ?? []) this.declare(decl);

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

@@ -1327,7 +1327,7 @@ function collectionOps(host: LlvmEmitterContext, key: IrType, value: IrType): st
       : { retain: "null", release: "null", trace: "null" };
   const arms = key.kind === "union" ? host.unionsById.get(key.unionId)?.arms : undefined;
   let unionKeys = "null";
-  if (mapKeyKindNum(key, arms) === 7) {
+  if (mapKeyKindNum(key, arms) === 7 && host.nullableUnions.of(key) === null) {
     if (!arms) throw new InternalCompilerError("llvm emitter bug: collection key union missing");
     unionKeys = host.cstr(arms.map((arm) => String.fromCharCode(mapKeyKindNum(arm))).join(""));
   }
@@ -1346,9 +1346,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;
-  // 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);
+  // SCR_MAP_KEY_REF (2): a union key with a pointer representation is that
+  // pointer (boxedUnionIds).
+  const kind = host.nullableUnions.of(key) !== null ? 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/local-array-reads.test.ts

@@ -153,7 +153,7 @@ function candidates(mod: IrModule) {
 }
 function workBody(mod: IrModule, pointerBits: 32 | 64 = 64) {
   return /^define internal [^\n]*@sc_(?:b)?f_work(?:\.i32)?\([^]*?^}/m.exec(
-    emitLlvmModule(keepUnionsBoxed(mod), { pointerBits }),
+    emitLlvmModule(mod, { pointerBits, ...keepUnionsBoxed(mod) }),
   )![0];
 }
 

packages/compiler/src/backend/llvm/local-union-storage.test.ts

@@ -104,7 +104,7 @@ function facts(mod: IrModule) {
 function body(mod: IrModule, pointerBits: 32 | 64): string {
   expect(validateModule(mod)).toEqual([]);
   return /^define internal [^\n]*@sc_(?:b)?f_work(?:\.i32)?\([^]*?^}/m.exec(
-    emitLlvmModule(keepUnionsBoxed(mod), { pointerBits }),
+    emitLlvmModule(mod, { pointerBits, ...keepUnionsBoxed(mod) }),
   )![0];
 }
 

packages/compiler/src/backend/llvm/map-read-lifetimes.test.ts

@@ -93,7 +93,10 @@ function body(mod: IrModule, bits: 32 | 64 = 64): string {
   expect(validateModule(mod)).toEqual([]);
   // Record values are nullable pointers; the other fixtures pin stack boxes.
   const nullable = mod.unions?.[0]?.arms[0]?.kind === "record";
-  const llvm = emitLlvmModule(nullable ? mod : keepUnionsBoxed(mod), { pointerBits: bits });
+  const llvm = emitLlvmModule(mod, {
+    pointerBits: bits,
+    ...(nullable ? {} : keepUnionsBoxed(mod)),
+  });
   return /^define internal [^\n]*@sc_(?:b)?f_work(?:\.i32)?\([^]*?^}/m.exec(llvm)![0];
 }
 
@@ -255,11 +258,11 @@ test("reference-preserving keys borrow receivers across unrelated later assignme
 });
 
 test("string-key lookups test the absent-key filter inline on 64-bit targets only", () => {
-  const mod = keepUnionsBoxed(fixture(F64, STRING));
+  const mod = fixture(F64, STRING);
   expect(validateModule(mod)).toEqual([]);
   const helper = (bits: 32 | 64) =>
     /^define internal ptr @sc_map_entry_str(?:\.i32)?\([^]*?^}/m.exec(
-      emitLlvmModule(mod, { pointerBits: bits }),
+      emitLlvmModule(mod, { pointerBits: bits, ...keepUnionsBoxed(mod) }),
     )![0];
   const wide = helper(64);
   // The cached hash is the capacity word's high half; the filter bits are its

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

@@ -78,6 +78,36 @@ test("records of runtime-inspected unions keep every union over them boxed", ()
   expect(unions.shapeNumber("c")).toBe(0);
 });
 
+test("Set elements and Map keys over identity arms keep their pointer representation", () => {
+  const union = (unionId: string): IrType => ({ kind: "union", unionId });
+  const globals = [
+    { id: "%g.s", name: "s", type: { kind: "set", elem: union("au") } as IrType, mutable: false },
+    {
+      id: "%g.m",
+      name: "m",
+      type: { kind: "map", key: union("abnu"), value: F64 } as IrType,
+      mutable: false,
+    },
+    { id: "%g.t", name: "t", type: { kind: "set", elem: union("su") } as IrType, mutable: false },
+  ];
+  const unions: IrUnionDef[] = [
+    { id: "au", arms: [a, UNDEFINED_T] },
+    { id: "abnu", arms: [a, b, NULL_T, UNDEFINED_T] },
+    { id: "su", arms: [{ kind: "string" }, UNDEFINED_T] },
+  ];
+  const program = new NullableUnions(module(unions, { globals }));
+  expect(program.single("au")).not.toBeNull();
+  expect(program.shapes("abnu")).toMatchObject({ unitTag: 3, nullTag: 2 });
+  // Strings compare by value: such keys stay boxed.
+  expect(program.get("su")).toBeNull();
+  // Library builds keep them boxed.
+  const library = new NullableUnions(
+    module(unions, { globals, lib: { exports: [] } as unknown as NonNullable<IrModule["lib"]> }),
+  );
+  expect(library.get("au")).toBeNull();
+  expect(library.get("abnu")).toBeNull();
+});
+
 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 };

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

@@ -1,4 +1,5 @@
 import {
+  isIdentityCollectionKey,
   isUnitType,
   type IrClassDef,
   type IrModule,
@@ -190,13 +191,21 @@ export class NullableUnions {
   /** `shapes`: allow shape and class unions (debug builds keep boxes,
    * whose debug-info layouts describe every arm). `classIds`: the target
    * keeps class ids in instance headers (64-bit targets), which class
-   * unions read their arm from. */
-  constructor(mod: IrModule, shapes = true, nodes?: IrNodeIndex, classIds = true) {
+   * unions read their arm from. `kept`: unions that stay boxes whatever
+   * their arms (LlvmTargetOptions.boxedUnions). */
+  constructor(
+    mod: IrModule,
+    shapes = true,
+    nodes?: IrNodeIndex,
+    classIds = true,
+    kept: ReadonlySet<string> = new Set(),
+  ) {
     const emitted = new Set<string>();
     for (const cls of mod.classes ?? []) if (!cls.runtime) emitted.add(cls.name);
     const records = new Set<string>();
     for (const shape of mod.records ?? []) records.add(shape.id);
     const boxed = boxedUnionIds(mod, nodes);
+    for (const id of kept) boxed.add(id);
     // Shape unions stay out of library builds (their host ABI marshals
     // boxes) and of programs with lifetime regions (whose walks report a
     // reference slot by its type's release entry point).
@@ -500,19 +509,21 @@ function classUnion(def: IrUnionDef, classes: ClassForest): ClassUnion | null {
 
 /** 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. 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). */
+ * tag and payload inside the C map. A union whose arms all compare by
+ * identity (records, class instances, arrays, functions), with at most an
+ * undefined and a null arm, is the exception: a pointer representation
+ * carries that identity, with NULL and the null sentinel as two more
+ * distinct keys, so the map keys it by reference. Library builds keep
+ * such unions boxed, except unions of records alone. */
 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>();
+  const identityKeys = 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);
+    if (mod.lib === undefined ? identityKeyUnion(def) : recordsOnly(def)) identityKeys.add(def.id);
   const mark = (t: IrType): void => {
-    if (t.kind === "union" && !recordsOnly.has(t.unionId)) boxed.add(t.unionId);
+    if (t.kind === "union" && !identityKeys.has(t.unionId)) boxed.add(t.unionId);
   };
   if (nodes !== undefined) {
     for (const t of nodes.types) {
@@ -533,6 +544,25 @@ function boxedUnionIds(mod: IrModule, nodes?: IrNodeIndex): Set<string> {
   return boxed;
 }
 
+function recordsOnly(def: IrUnionDef): boolean {
+  return def.arms.length >= 2 && def.arms.every((arm) => arm.kind === "record");
+}
+
+/** Every arm compares by identity as a key, beside at most one undefined
+ * and one null arm. */
+function identityKeyUnion(def: IrUnionDef): boolean {
+  let refs = 0,
+    undefinedArms = 0,
+    nullArms = 0;
+  for (const arm of def.arms) {
+    if (arm.kind === "undefinedT") undefinedArms++;
+    else if (arm.kind === "nullT") nullArms++;
+    else if (isIdentityCollectionKey(arm)) refs++;
+    else return false;
+  }
+  return refs > 0 && undefinedArms <= 1 && nullArms <= 1;
+}
+
 /** Emitted classes keyed by name, for callers that only hold the class
  * list (nullable class fields). */
 export function emittedClassNames(classes: readonly IrClassDef[]): Set<string> {

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

@@ -58,6 +58,37 @@ console.log(seen.size, depth.get(nodes[1]!), seen.has({ kind: "num", value: 1 })
   expect(ir).not.toMatch(/call ptr @scr_union_new_ref/);
 });
 
+test("Sets and Maps of a class, null or undefined key the instance pointer", async () => {
+  const ir = await llvmFor(`
+class Leaf {
+  name: string;
+  constructor(name: string) {
+    this.name = name;
+  }
+}
+const a = new Leaf("a");
+const leaves = new Set<Leaf | undefined>([a, undefined, new Leaf("b")]);
+const counts = new Map<Leaf | null | undefined, number>();
+for (const leaf of [a, null, undefined, a]) counts.set(leaf, (counts.get(leaf) ?? 0) + 1);
+const child = (i: number): Leaf | undefined => (i < 2 ? a : undefined);
+console.log(leaves.has(undefined), counts.get(a), counts.get(null), child(1)?.name);
+`);
+  expect(ir).toMatch(/call ptr @scr_map_new\(i32 2,/);
+  expect(ir).not.toMatch(/call ptr @scr_map_new\(i32 [37],/);
+  expect(ir).not.toMatch(/call ptr @scr_union_new_ref/);
+});
+
+test("Maps keyed by a string or undefined keep union keys, which compare by value", async () => {
+  const ir = await llvmFor(`
+const names = new Map<string | undefined, number>([["a", 1], [undefined, 2]]);
+names.set(["a"].join(""), 3);
+console.log(names.size, names.get("a"), names.get(undefined));
+`);
+  // SCR_MAP_KEY_UNION_VALUE (7): the strings compare by their text.
+  expect(ir).toMatch(/call ptr @scr_map_new\(i32 7,/);
+  expect(ir).not.toMatch(/call ptr @scr_map_new\(i32 2,/);
+});
+
 const recordType = (shapeId: string): IrType => ({ kind: "record", shapeId });
 const wide: IrType = { kind: "union", unionId: "wide" };
 

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

@@ -90,7 +90,7 @@ function body(ir: string, name: string): string {
 
 function emit(mod: IrModule, pointerBits: 32 | 64 = 64): string {
   expect(validateModule(mod)).toEqual([]);
-  return emitLlvmModule(keepUnionsBoxed(mod), { pointerBits });
+  return emitLlvmModule(mod, { pointerBits, ...keepUnionsBoxed(mod) });
 }
 
 function declareLocal(mod: IrModule, init: IrExpr): void {

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

@@ -1,29 +1,26 @@
-import { F64, type IrModule, type IrType, type IrUnionDef } from "../../ir/ir.js";
+import type { IrModule } from "../../ir/ir.js";
+import { everyModuleNode } from "../../ir/traverse.js";
+import type { LlvmTargetOptions } from "./emitter.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. 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 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,
-    };
+/** Test helper: emission options that keep every union of `mod` a tagged
+ * `%ScrUnion` box, so stack-box tests exercise the paths that
+ * nullable-pointer unions (nullable-unions.ts) bypass. */
+export function keepUnionsBoxed(mod: IrModule): LlvmTargetOptions {
+  return { boxedUnions: new Set((mod.unions ?? []).map((union) => union.id)) };
+}
+
+/** Test helper: emission options that keep each union the Sets of `mod`
+ * hold a tagged box. Source tests of the stack-box paths declare a
+ * `Set<T | undefined>` of the union they exercise. */
+export function keepSetElementsBoxed(mod: IrModule): LlvmTargetOptions {
+  const boxedUnions = new Set<string>();
+  everyModuleNode(mod, {
+    expr: () => true,
+    stmt: () => true,
+    type: (type) => {
+      if (type.kind === "set" && type.elem.kind === "union") boxedUnions.add(type.elem.unionId);
+      return true;
+    },
   });
-  return {
-    ...mod,
-    unions: [...(mod.unions ?? []), ...companions],
-    globals: [...(mod.globals ?? []), ...keep],
-  };
+  return { boxedUnions };
 }

packages/compiler/src/backend/llvm/union-retype.test.ts

@@ -103,7 +103,7 @@ function fixture(): IrModule {
 
 function emit(mod: IrModule): string {
   expect(validateModule(mod)).toEqual([]);
-  return emitLlvmModule(keepUnionsBoxed(mod), { pointerBits: 64 });
+  return emitLlvmModule(mod, { pointerBits: 64, ...keepUnionsBoxed(mod) });
 }
 
 test("an identity retype retains the input box instead of allocating", () => {
@@ -196,7 +196,7 @@ async function llvmFor(source: string, boxed = true): Promise<{ module: IrModule
       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 { module, ir: emitLlvmModule(boxed ? keepUnionsBoxed(module) : module) };
+    return { module, ir: emitLlvmModule(module, boxed ? keepUnionsBoxed(module) : {}) };
   } finally {
     await rm(dir, { recursive: true, force: true });
   }

packages/compiler/test/call-lifetimes.test.ts

@@ -5,6 +5,7 @@ import { expect, test } from "vitest";
 import { compile, deserializeModule, validateModule } from "../src/index.js";
 import { analyzeCallLifetimes } from "../src/backend/llvm/call-lifetimes.js";
 import { emitLlvmModule } from "../src/backend/llvm/emitter.js";
+import { keepSetElementsBoxed } from "../src/backend/llvm/test-boxed-unions.js";
 import { type IrModule, type IrFunction } from "../src/ir/ir.js";
 
 async function lower(source: string): Promise<IrModule> {
@@ -170,7 +171,7 @@ function work(items: Item[]): number {
 }
 console.log(work([new Item(7)]));
 
-// Keeps the optional element union a tagged box (Set elements are boxes),
+// The Set's element union stays a tagged box under keepSetElementsBoxed,
 // so the stack-box paths stay exercised; nullable pointers bypass them.
 const keepBoxed = new Set<Item | undefined>();
 void keepBoxed;
@@ -181,7 +182,7 @@ void keepBoxed;
   const facts = analyzeCallLifetimes(new Map(mod.functions.map((item) => [item.name, item])));
   expect(facts.parameters.get("read")).toEqual(new Set([0]));
   expect(facts.parameters.get("relay")).toEqual(new Set([0]));
-  const llvm = emitLlvmModule(mod);
+  const llvm = emitLlvmModule(mod, keepSetElementsBoxed(mod));
   const work = body(llvm, "sc_bf_work");
   expect(work).toContain("alloca %ScrUnion");
   expect(work).toContain("@sc_bf_relay");

packages/compiler/test/string-and-array-lifetimes.test.ts

@@ -5,6 +5,7 @@ import { expect, test } from "vitest";
 import { compile, deserializeModule, serializeModule, validateModule } from "../src/index.js";
 import { analyzeCallLifetimes } from "../src/backend/llvm/call-lifetimes.js";
 import { emitLlvmModule } from "../src/backend/llvm/emitter.js";
+import { keepSetElementsBoxed } from "../src/backend/llvm/test-boxed-unions.js";
 import { type IrModule } from "../src/ir/ir.js";
 
 /** A pending-exception check: the inline active-cell test or a runtime call. */
@@ -114,14 +115,14 @@ function value(item: Item): number { return item.value; }
 function work(items: Item[], index: number): number { return value(items[index]); }
 console.log(work([new Item(7)], 0));
 
-// Keeps the optional element union a tagged box (Set elements are boxes),
+// The Set's element union stays a tagged box under keepSetElementsBoxed,
 // so the stack-box paths stay exercised; nullable pointers bypass them.
 const keepBoxed = new Set<Item | undefined>();
 void keepBoxed;
 `);
   const value = module.functions.find((fn) => fn.name === "value")!;
   expect(value.params[0]!.type.kind).toBe("union");
-  const llvm = emitLlvmModule(module);
+  const llvm = emitLlvmModule(module, keepSetElementsBoxed(module));
   const work = body(llvm, "sc_bf_work");
   expect(work).toContain("alloca %ScrUnion");
   expect(work).toContain("@sc_bf_value");
@@ -139,12 +140,12 @@ function clear(items: Item[]): number { items.length = 0; return 1; }
 function work(items: Item[]): number { return value(items[0], clear(items)); }
 console.log(work([new Item(7)]));
 
-// Keeps the optional element union a tagged box (Set elements are boxes),
+// The Set's element union stays a tagged box under keepSetElementsBoxed,
 // so the stack-box paths stay exercised; nullable pointers bypass them.
 const keepBoxed = new Set<Item | undefined>();
 void keepBoxed;
 `);
-  const llvm = emitLlvmModule(module);
+  const llvm = emitLlvmModule(module, keepSetElementsBoxed(module));
   const work = body(llvm, "sc_bf_work");
   expect(work).toContain("alloca %ScrUnion");
   expect(work).toContain("@sc_retain_Item");
@@ -172,12 +173,12 @@ function label(item: Item): string { return item.label; }
 function work(items: Item[]): string { return label(items[0]); }
 console.log(work([new Item("kept")]));
 
-// Keeps the optional element union a tagged box (Set elements are boxes),
+// The Set's element union stays a tagged box under keepSetElementsBoxed,
 // so the stack-box paths stay exercised; nullable pointers bypass them.
 const keepBoxed = new Set<Item | undefined>();
 void keepBoxed;
 `);
-  const llvm = emitLlvmModule(module);
+  const llvm = emitLlvmModule(module, keepSetElementsBoxed(module));
   expect(body(llvm, "sc_bf_work")).toContain("alloca %ScrUnion");
   expect(body(llvm, "sc_bf_label")).toContain("@scr_str_retain_v");
   expect(body(llvm, "sc_bf_label")).not.toContain("@scr_union_retain_v");
@@ -202,12 +203,12 @@ function pair(left: Item, right: Item): number { return left.value + right.value
 function work(items: Item[]): number { return pair(items[0], items[1]) + pair(items[1], items[0]); }
 console.log(work([new Item(), new Item()]));
 
-// Keeps the optional element union a tagged box (Set elements are boxes),
+// The Set's element union stays a tagged box under keepSetElementsBoxed,
 // so the stack-box paths stay exercised; nullable pointers bypass them.
 const keepBoxed = new Set<Item | undefined>();
 void keepBoxed;
 `);
-  const work = body(emitLlvmModule(module), "sc_bf_work");
+  const work = body(emitLlvmModule(module, keepSetElementsBoxed(module)), "sc_bf_work");
   expect(work.match(/alloca %ScrUnion/g)).toHaveLength(4);
   expect(work.match(/@sc_bf_pair/g)).toHaveLength(2);
   expect(work).not.toContain("@scr_union_new_ref");
@@ -347,14 +348,13 @@ function count(parts: Part[], probe: Part): number {
 const parts = [new Part()];
 console.log(same(parts, 0, parts[0]), missing(parts, 3), weight(parts, 2), count(parts, parts[0]));
 `;
-  // Keeps the optional element union a tagged box (Set elements are boxes),
+  // The Set's element union stays a tagged box under keepSetElementsBoxed,
   // so the stack-tag paths stay exercised; nullable pointers bypass them.
-  const boxed = emitLlvmModule(
-    await lower(`${source}
+  const kept = await lower(`${source}
 const keepBoxed = new Set<Part | undefined>();
 void keepBoxed;
-`),
-  );
+`);
+  const boxed = emitLlvmModule(kept, keepSetElementsBoxed(kept));
   const nullable = emitLlvmModule(await lower(source));
   for (const name of ["same", "missing", "weight", "count"]) {
     const work = body(boxed, `sc_bf_${name}`);

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/6313-identity-union-collection-keys.ts.json

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

tests/corpus/6313-identity-union-collection-keys.ts

@@ -0,0 +1,86 @@
+// Sets and Maps keyed by unions whose arms compare by identity, beside
+// undefined and null keys.
+
+class Leaf {
+  readonly name: string;
+  constructor(name: string) {
+    this.name = name;
+  }
+}
+
+class Branch {
+  readonly name: string;
+  readonly children: Leaf[];
+  constructor(name: string, children: Leaf[]) {
+    this.name = name;
+    this.children = children;
+  }
+}
+
+const a = new Leaf("a");
+const b = new Leaf("b");
+const twin = new Leaf("a");
+
+// A class or undefined.
+const leaves = new Set<Leaf | undefined>();
+for (const item of [a, undefined, b, a, undefined, twin]) leaves.add(item);
+console.log(leaves.size, leaves.has(a), leaves.has(undefined), leaves.has(new Leaf("b")));
+console.log([...leaves].map((leaf) => (leaf === undefined ? "undefined" : leaf.name)).join(","));
+leaves.delete(undefined);
+console.log(leaves.size, leaves.has(undefined));
+console.log(leaves);
+
+// A class, null or undefined, as Map keys.
+const counts = new Map<Leaf | null | undefined, number>();
+for (const item of [a, null, undefined, b, null, a])
+  counts.set(item, (counts.get(item) ?? 0) + 1);
+console.log(counts.get(a), counts.get(null), counts.get(undefined), counts.get(twin));
+for (const [key, count] of counts)
+  console.log(key === null ? "null" : key === undefined ? "undefined" : key.name, count);
+console.log(Array.from(counts.keys()).length, counts.has(null), counts.delete(null), counts.has(null));
+console.log(counts);
+
+// Two classes or undefined.
+const nodes: (Leaf | Branch | undefined)[] = [a, new Branch("root", [a, b]), undefined, b];
+const seen = new Set(nodes);
+seen.add(nodes[1]);
+seen.add(undefined);
+console.log(seen.size);
+seen.forEach((node) => {
+  if (node === undefined) console.log("undefined");
+  else if (node instanceof Branch) console.log("branch", node.name, node.children.length);
+  else console.log("leaf", node.name);
+});
+
+// Records, arrays and functions, or undefined.
+type Point = { x: number; y: number };
+const origin: Point = { x: 0, y: 0 };
+const points = new Map<Point | undefined, string>([
+  [origin, "origin"],
+  [undefined, "none"],
+  [{ x: 0, y: 0 }, "copy"],
+]);
+console.log(points.size, points.get(origin), points.get(undefined), points.get({ x: 0, y: 0 }));
+
+const row = [1, 2, 3];
+const rows = new Set<number[] | undefined>([row, undefined, [1, 2, 3], row]);
+console.log(rows.size, rows.has(row), rows.has([1, 2, 3]));
+
+const greet = (): string => "hi";
+const handlers = new Map<(() => string) | undefined, number>();
+handlers.set(greet, 1);
+handlers.set(undefined, 2);
+handlers.set(greet, 3);
+console.log(handlers.size, handlers.get(greet), handlers.get(undefined));
+
+// Strings compare by value, beside undefined.
+const names = new Map<string | undefined, number>([
+  ["a", 1],
+  [undefined, 2],
+]);
+names.set(["a"].join(""), 3);
+console.log(names.size, names.get("a"), names.get(undefined));
+
+// The values keep their identity through the collection.
+const copied = new Set<Leaf | undefined>(leaves);
+console.log([...copied][0] === a, copied.has(b), copied.size);