Read the discriminant again after an arm's check fails

An arm's field check can run a getter that changes the discriminant. The
checked cast now reads the discriminant again after each failed arm;
when its case changed, the tests continue arm by arm after the failed
arm, so the cast picks the arm the arm-by-arm tests pick.

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

@@ -440,8 +440,12 @@ export class LlDyn {
    * discriminant is read once (scr_dyn_discriminant_case, against the
    * constant `table`), and only the arms whose literal it holds and the arms
    * without a literal are tested, in `order`: every other arm's literal test
-   * fails. A value the runtime cannot classify that way (a typed capsule, an
-   * accessor, an inherited discriminant) takes the tests arm by arm. */
+   * fails. An arm's failed check can run a getter that changes the
+   * discriminant, so it is read again after each one; when its case
+   * changed, the tests continue arm by arm after the failed arm, as they
+   * would have run from the first arm. A value the runtime cannot classify
+   * that way (a typed capsule, an accessor, an inherited discriminant) takes
+   * the tests arm by arm. */
   private unionArmsByDiscriminant(
     B: BlockBuilder,
     def: IrUnionDef,
@@ -455,47 +459,66 @@ export class LlDyn {
     const discriminant = def.discriminant!;
     const literalsOf = (tag: number): readonly (string | number | boolean)[] | null =>
       discriminant.cases.find((candidate) => candidate.tag === tag)?.values ?? null;
-    const tryArms = (tags: readonly number[], literal: boolean): void => {
-      if (tags.length === 0) {
-        B.br(none);
-        return;
-      }
-      tags.forEach((tag, index) => {
-        const next = index === tags.length - 1 ? none : B.newLabel("du.c.n");
-        if (literal) this.unionArmMatch(B, def, tag, hits[tag]!, next, preserveRefs);
-        else this.unionArmStructure(B, def, tag, hits[tag]!, next, preserveRefs);
-        if (next !== none) B.startBlock(next);
-      });
-    };
-    const armsFor = (value: string | null): number[] =>
-      order.filter((tag) => {
-        const literals = literalsOf(tag);
-        return literals === null || (value !== null && literals.includes(value));
-      });
     const host = this.host;
     this.defs.push(
       `${table} = private unnamed_addr constant [${values.length} x ptr] [${values.map((value) => `ptr ${host.internLiteral(value)}`).join(", ")}]`,
       ``,
     );
     host.declare(`declare i32 @scr_dyn_discriminant_case(ptr, ptr, ptr, ${host.sizeType})`);
-    const index = B.tmp();
-    B.line(
-      `${index} = call i32 @scr_dyn_discriminant_case(ptr %d, ptr ${host.internLiteral(discriminant.field)}, ptr ${table}, ${host.sizeType} ${values.length})`,
-    );
+    const classify = (): string => {
+      const index = B.tmp();
+      B.line(
+        `${index} = call i32 @scr_dyn_discriminant_case(ptr %d, ptr ${host.internLiteral(discriminant.field)}, ptr ${table}, ${host.sizeType} ${values.length})`,
+      );
+      return index;
+    };
+    // The arm-by-arm tests, entered at any arm (`from[i]` tests order[i]
+    // and the arms after it).
     const tests = B.newLabel("du.tests");
+    const from = order.map((_, i) => (i === 0 ? tests : B.newLabel("du.t")));
+    const after = (position: number): string =>
+      position + 1 < order.length ? from[position + 1]! : none;
+    // Tests the arms a value in case `k` (-1: no literal) can match.
+    const tryCase = (k: number): void => {
+      const value = k < 0 ? null : values[k]!;
+      const positions = order
+        .map((_, position) => position)
+        .filter((position) => {
+          const literals = literalsOf(order[position]!);
+          return literals === null || (value !== null && literals.includes(value));
+        });
+      if (positions.length === 0) {
+        B.br(none);
+        return;
+      }
+      positions.forEach((position, index) => {
+        const tag = order[position]!;
+        const failed = B.newLabel("du.c.f");
+        this.unionArmStructure(B, def, tag, hits[tag]!, failed, preserveRefs);
+        B.startBlock(failed);
+        const next = index === positions.length - 1 ? none : B.newLabel("du.c.n");
+        const same = B.tmp();
+        B.line(`${same} = icmp eq i32 ${classify()}, ${k}`);
+        B.condBr(same, next, after(position));
+        if (next !== none) B.startBlock(next);
+      });
+    };
+    const index = classify();
     const unmatched = B.newLabel("du.none");
     const cases = values.map(() => B.newLabel("du.case"));
     B.terminate(
       `switch i32 ${index}, label %${tests} [ i32 -1, label %${unmatched} ${cases.map((label, k) => `i32 ${k}, label %${label}`).join(" ")} ]`,
     );
-    values.forEach((value, k) => {
+    values.forEach((_, k) => {
       B.startBlock(cases[k]!);
-      tryArms(armsFor(value), false);
+      tryCase(k);
     });
     B.startBlock(unmatched);
-    tryArms(armsFor(null), false);
-    B.startBlock(tests);
-    tryArms(order, true);
+    tryCase(-1);
+    order.forEach((tag, position) => {
+      B.startBlock(from[position]!);
+      this.unionArmMatch(B, def, tag, hits[tag]!, after(position), preserveRefs);
+    });
   }
 
   /** Branch to a variant only after its literal discriminator and complete

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/6306-discriminant-getter-dispatch.ts.json

@@ -0,0 +1,6 @@
+{
+ "order": [
+  "<repo>/tests/corpus/6306-discriminant-getter-dispatch.ts"
+ ],
+ "diags": []
+}

tests/corpus/6306-discriminant-getter-dispatch.ts

@@ -0,0 +1,55 @@
+// A checked cast into a union with many string discriminant literals reads
+// the discriminant once, but a getter that runs while an arm's fields are
+// checked can change it. The cast then continues with the arms after the
+// one that failed, testing each arm's literal again, and matches the arm
+// the new discriminant names. Each getter changes the discriminant on its
+// first call only; the program reads the getter once itself, so Node (where
+// the cast runs no code) reaches the same state. Valid casts only (see
+// 1002-json-parse-cast.ts).
+
+type Shape =
+  | { kind: "circle"; radius: number }
+  | { kind: "square"; side: number }
+  | { kind: "rect"; height: number; width: number }
+  | { kind: "tri"; a: number; b: number; c: number };
+
+function area(shape: Shape): number {
+  switch (shape.kind) {
+    case "circle":
+      return Math.round(Math.PI * shape.radius * shape.radius);
+    case "square":
+      return shape.side * shape.side;
+    case "rect":
+      return shape.width * shape.height;
+    case "tri": {
+      const s = (shape.a + shape.b + shape.c) / 2;
+      return Math.sqrt(s * (s - shape.a) * (s - shape.b) * (s - shape.c));
+    }
+  }
+}
+
+// The circle arm reads `radius`, whose getter turns the value into a square.
+let squareReads = 0;
+const square: any = {
+  kind: "circle",
+  get radius() {
+    squareReads++;
+    if (squareReads === 1) square.kind = "square";
+    return "none";
+  },
+  side: 3,
+};
+const asSquare = square as Shape;
+void square.radius;
+console.log(square.kind, asSquare.kind, area(asSquare));
+
+// Values without getters take the arm their discriminant names.
+const plain: any[] = [
+  { kind: "rect", width: 2, height: 5 },
+  { kind: "tri", a: 3, b: 4, c: 5 },
+  { kind: "circle", radius: 2 },
+];
+for (const value of plain) {
+  const shape = value as Shape;
+  console.log(shape.kind, area(shape));
+}

tests/harness/corpus-costs/6306-discriminant-getter-dispatch.ts.json

@@ -0,0 +1 @@
+{ "plain": 3800, "san": 4900 }

tests/harness/dyncheck.test.ts

@@ -444,6 +444,47 @@ for (const input of inputs) {
     );
   });
 
+  test("a getter that changes the discriminant during an arm's check continues with the later arms", async () => {
+    // Each getter sets `kind` and fails its arm's check: to a later arm's
+    // literal (that arm matches), to an earlier arm's literal (arm-by-arm
+    // tests do not go back), and to a literal no arm holds.
+    const r = await compileAndRun(
+      "union-discriminant-getter",
+      `type Shape =
+  | { kind: "circle"; radius: number }
+  | { kind: "square"; side: number }
+  | { kind: "rect"; height: number; width: number }
+  | { kind: "tri"; a: number; b: number; c: number };
+const flip = (value: any, to: string): string => {
+  value.kind = to;
+  return "none";
+};
+const later: any = { kind: "circle", get radius() { return flip(later, "square"); }, side: 3 };
+const earlier: any = { kind: "square", radius: 2, get side() { return flip(earlier, "circle"); } };
+const unknown: any = { kind: "square", get side() { return flip(unknown, "hexagon"); } };
+const inputs = [later, earlier, unknown];
+for (const input of inputs) {
+  try {
+    const shape = input as Shape;
+    console.log("matched", shape.kind);
+  } catch (error) {
+    console.log((error as Error).message);
+  }
+}
+`,
+    );
+    const any = "object | object | object | object";
+    expect(r.exitCode).toBe(0);
+    expect(r.stdout).toBe(
+      [
+        "matched square",
+        `expected ${any} at $, got object`,
+        `expected ${any} at $, got object`,
+        "",
+      ].join("\n"),
+    );
+  });
+
   test("nested paths point at the offending element", async () => {
     const r = await compileAndRun(
       "nested-path",