Keep every field of a record converted into a union arm (#2304)

packages/compiler/src/frontend/lowering/coercions/structural-plans.ts

@@ -1,6 +1,6 @@
 import { planRecordUnionLift, planRecordUnionWrap } from "../union-retag.js";
 import type { WidthLift } from "../width-lift.js";
-import type { IrType } from "../../../ir/ir.js";
+import type { IrRecordShape, IrType } from "../../../ir/ir.js";
 import { canDynCheckTo, isUnitType, typeEquals, UNDEFINED_T } from "../../../ir/ir.js";
 import { type ClassInfo, findMethodOn, findGenericMethodOn } from "../lower-classes.js";
 import type { Lowerer } from "../lowerer.js";
@@ -104,7 +104,15 @@ export function widthLiftPlan(lowerer: Lowerer, src: IrType, dst: IrType): Width
           })
         );
       });
-      if (absentOnly.length === 1) selected = absentOnly;
+      // That layout does not win when it holds the record only by dropping
+      // fields and another arm keeps every field (`{ kind; type }` beside
+      // `{ kind; name; adapts?: true; type }`): the value is not its variant.
+      const dropsFields = (arm: IrType): boolean => !keepsEveryField(lowerer, source, arm);
+      if (
+        absentOnly.length === 1 &&
+        !(dropsFields(absentOnly[0]!.arm) && selected.some(({ arm }) => !dropsFields(arm)))
+      )
+        selected = absentOnly;
     }
     // A record that fits several arms of a discriminated union selects one
     // by its discriminant at run time. Without this, an arm that a nested
@@ -131,21 +139,19 @@ export function widthLiftPlan(lowerer: Lowerer, src: IrType, dst: IrType): Width
     // A record whose fields some arm keeps EVERY one of, each lifting
     // (`{ id; limit; opts?: A | B }` into `{ id; limit; opts?: A | B | C } |
     // { id; opts?: C }`): that arm is the value's own layout, widened; the
-    // others hold it only by dropping fields. Prefer the unique lossless arm.
+    // others hold it only by dropping fields. Prefer the unique lossless arm
+    // with the record's own fields, then the unique one that also declares
+    // optional fields the record lacks, which complete as absent.
     if (selected.length > 1 && src.kind === "record") {
       const source = lowerer.shapes.get(src.shapeId);
-      const lossless = selected.filter(({ arm }) => {
-        const target = arm.kind === "record" ? lowerer.shapes.get(arm.shapeId) : undefined;
-        return (
-          source !== undefined &&
-          target !== undefined &&
-          !source.tuple &&
-          !target.tuple &&
-          target.fields.length === source.fields.length &&
-          source.fields.every((field) => target.fields.some((kept) => kept.name === field.name))
-        );
-      });
-      if (lossless.length === 1) selected = lossless;
+      const lossless = selected.filter(({ arm }) => keepsEveryField(lowerer, source, arm));
+      const exact = lossless.filter(
+        ({ arm }) =>
+          arm.kind === "record" &&
+          lowerer.shapes.get(arm.shapeId)?.fields.length === source?.fields.length,
+      );
+      if (exact.length === 1) selected = exact;
+      else if (exact.length === 0 && lossless.length === 1) selected = lossless;
     }
     if (selected.length !== 1) return null;
     return { how: "liftWrap", tag: selected[0]!.tag, arm: selected[0]!.arm };
@@ -241,6 +247,23 @@ export function widthLiftPlan(lowerer: Lowerer, src: IrType, dst: IrType): Width
   return null;
 }
 
+/** True when the record arm `arm` has a field of every name the record
+ * `source` has. Tuples keep their own positional layouts. */
+function keepsEveryField(
+  lowerer: Lowerer,
+  source: IrRecordShape | undefined,
+  arm: IrType,
+): boolean {
+  const target = arm.kind === "record" ? lowerer.shapes.get(arm.shapeId) : undefined;
+  return (
+    source !== undefined &&
+    target !== undefined &&
+    !source.tuple &&
+    !target.tuple &&
+    source.fields.every((field) => target.fields.some((kept) => kept.name === field.name))
+  );
+}
+
 /** Two record arms a source record fits, where one arm declares a field
  * the other lacks and the source types that field `T | undefined`
  * (`{ count, errors: errs.length ? errs : undefined }` into `{ count;

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/union-record-kept-fields-arm.ts.json

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

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/union-record-kept-fields-no-discriminant.ts.json

@@ -0,0 +1,6 @@
+{
+ "order": [
+  "<repo>/tests/corpus/union-record-kept-fields-no-discriminant.ts"
+ ],
+ "diags": []
+}

tests/corpus/union-record-kept-fields-arm.ts

@@ -0,0 +1,56 @@
+// A record flowing into a discriminated union where one record arm holds it
+// by dropping fields and another keeps every field but declares an optional
+// field the record lacks. The record's discriminant selects the arm that keeps
+// its fields.
+
+type Loc = { file: string; start: number };
+
+type Expr =
+  | { kind: "ref"; type: number; loc: Loc }
+  | {
+      kind: "closure";
+      fnName: string;
+      captures: string[];
+      adapts?: true;
+      type: number;
+      loc: Loc;
+    }
+  | { kind: "call"; callee: string; args: Expr[]; type: number; loc: Loc };
+
+function closureOf(name: string, loc: Loc): Expr {
+  const closure = {
+    kind: "closure" as const,
+    fnName: name,
+    captures: [],
+    type: 1,
+    loc,
+  };
+  return closure;
+}
+
+function describe(expr: Expr): string {
+  if (expr.kind === "closure")
+    return `closure ${expr.fnName}/${expr.captures.length} ${expr.adapts ?? "-"}`;
+  if (expr.kind === "call")
+    return `call ${expr.callee}(${expr.args.map(describe).join(", ")})`;
+  return `ref ${expr.type}`;
+}
+
+const loc: Loc = { file: "a.ts", start: 3 };
+const first = closureOf("pick", loc);
+console.log(first.kind, describe(first), JSON.stringify(first));
+
+const exprs: Expr[] = [first];
+const captured = {
+  kind: "closure" as const,
+  fnName: "inner",
+  captures: ["x"],
+  type: 2,
+  loc,
+};
+exprs.push(captured);
+exprs.push({ kind: "ref", type: 7, loc });
+console.log(
+  describe({ kind: "call", callee: "run", args: exprs, type: 0, loc }),
+);
+console.log(exprs.map((e) => e.kind).join(","));

tests/corpus/union-record-kept-fields-no-discriminant.ts

@@ -0,0 +1,39 @@
+// A record flowing into a union without a discriminant where one record arm
+// holds it by dropping fields and another keeps every field but declares an
+// optional field the record lacks. The record converts into the arm that
+// keeps its fields, and the optional field stays absent.
+
+type Loc = { file: string; start: number };
+
+type Slot =
+  | { type: number; loc: Loc }
+  | {
+      fnName: string;
+      captures: string[];
+      adapts?: true;
+      type: number;
+      loc: Loc;
+    };
+
+const loc: Loc = { file: "a.ts", start: 3 };
+
+function slotOf(name: string): Slot {
+  const slot = { fnName: name, captures: [], type: 1, loc };
+  return slot;
+}
+
+const slot = slotOf("stored");
+console.log("fnName" in slot, "adapts" in slot, JSON.stringify(slot));
+if ("fnName" in slot)
+  console.log(slot.fnName, slot.adapts, slot.captures.length);
+
+function describe(value: Slot): string {
+  return "fnName" in value
+    ? `${value.fnName}/${value.captures.length}`
+    : `type ${value.type}`;
+}
+
+const inner = { fnName: "inner", captures: ["x"], type: 2, loc };
+const slots: Slot[] = [slot, { type: 7, loc }];
+slots.push(inner);
+console.log(slots.map(describe).join(", "), JSON.stringify(slots[2]));

tests/harness/corpus-costs/union-record-kept-fields-arm.ts.json

@@ -0,0 +1 @@
+{ "plain": 4793, "san": 7135 }

tests/harness/corpus-costs/union-record-kept-fields-no-discriminant.ts.json

@@ -0,0 +1 @@
+{ "plain": 4466, "san": 5253 }