Keep object literals' key order through IR rewrites and on shapes typed JSON.parse also builds (#1607)

packages/compiler/src/frontend/lowering/key-order-witness.test.ts

@@ -1,5 +1,6 @@
 import { expect, test } from "vitest";
 import { F64, type IrExpr, type IrFunction, type IrStmt, type IrType } from "../../ir/ir.js";
+import { observeRecordLiteralRebuilds, transformExpr } from "../../ir/traverse.js";
 import { ShapeRegistry } from "../type-mapper.js";
 import {
   applyKeyOrderWitnesses,
@@ -156,7 +157,24 @@ test("a synthesized construction pins today's table", () => {
   expect(f.order(ab)).toEqual(["a", "b"]);
 });
 
-test("a checked conversion pins today's table", () => {
+test("a checked conversion alone pins today's table", () => {
+  const f = fixture();
+  const ab = f.record(
+    [
+      ["a", OPT],
+      ["b", F64],
+    ],
+    ["a", "b"],
+  );
+  const parsed: IrExpr = { kind: "dynCheck", value: ref("j", { kind: "dyn" }), type: ab, loc };
+  // A late write to `a` would move it after `b`; values only a checked
+  // conversion builds keep today's table, so parsed values in that order
+  // need no source view.
+  f.run([decl("p", parsed), set(ref("p", ab), "a"), observe(ref("p", ab))]);
+  expect(f.order(ab)).toEqual(["a", "b"]);
+});
+
+test("a shape literals and a checked conversion both build takes the literals' order", () => {
   const f = fixture();
   const ab = f.record(
     [
@@ -166,8 +184,33 @@ test("a checked conversion pins today's table", () => {
     ["a", "b"],
   );
   const parsed: IrExpr = { kind: "dynCheck", value: ref("j", { kind: "dyn" }), type: ab, loc };
+  // The literal's values can enumerate in no other order; parsed values in
+  // another order keep a record-source view.
   f.run([observe(f.literal(ab, ["b", "a"])), observe(parsed)]);
-  expect(f.order(ab)).toEqual(["a", "b"]);
+  expect(f.order(ab)).toEqual(["b", "a"]);
+});
+
+test("a literal an IR rewrite rebuilds keeps its witness", () => {
+  const f = fixture();
+  const ab = f.record(
+    [
+      ["a", F64],
+      ["b", F64],
+    ],
+    ["a", "b"],
+  );
+  const identity = { expr: (e: IrExpr): IrExpr => e, stmt: (st: IrStmt): IrStmt => st };
+  const original = f.literal(ab, ["b", "a"]);
+  observeRecordLiteralRebuilds((from, to) => f.notes.follow(from, to));
+  let rebuilt: IrExpr;
+  try {
+    rebuilt = transformExpr(original, identity);
+  } finally {
+    observeRecordLiteralRebuilds(null);
+  }
+  expect(rebuilt).not.toBe(original);
+  f.run([observe(rebuilt)]);
+  expect(f.order(ab)).toEqual(["b", "a"]);
 });
 
 test("a spread contributes its source's order, and an override keeps its position", () => {

packages/compiler/src/frontend/lowering/key-order-witness.ts

@@ -23,10 +23,12 @@
  *  - any other construction (synthesized builtin results, opaque runtime
  *    producers) pins today's table;
  *  - checked conversions from dynamic values (typed JSON.parse, values
- *    from JavaScript) pin today's table: the record-source table keeps
- *    their own key order relative to it, and another table would make
- *    values parsed in today's order keep a source view. Their order is
- *    unknown, so a copy out of such a shape cannot be ordered statically;
+ *    from JavaScript) contribute nothing: the record-source table keeps
+ *    their own key order relative to the chosen table. A shape only they
+ *    build keeps today's table, since another table would make values
+ *    parsed in today's order keep a source view; a shape object literals
+ *    also build takes the literals' order. Their order is unknown, so a
+ *    copy out of such a shape cannot be ordered statically;
  *  - a write to an optional key that may be absent appends it in Node. A
  *    key no construction creates ("late-only") goes after every key a
  *    construction can create; any other moving write goes after every key.
@@ -56,8 +58,10 @@ export type KeyOrderSegment =
   | { kind: "opaque" };
 
 /** Witness notes the lowering attaches to the IR nodes it builds. Keyed by
- * node identity, so nodes a speculative lowering discards are never walked,
- * and a node some later pass rebuilds falls back to pinning. */
+ * node identity, so nodes a speculative lowering discards are never walked.
+ * Literals the IR rewriters rebuild (traverse.ts mapExprChildren) carry
+ * their notes over (follow); a node rebuilt any other way falls back to
+ * pinning. */
 export class KeyOrderNotes {
   private readonly literals = new Map<IrExpr, KeyOrderSegment[]>();
   private readonly copies = new Map<IrExpr, string>();
@@ -77,6 +81,15 @@ export class KeyOrderNotes {
     return node;
   }
 
+  /** A pass rebuilt `from` as `to` (traverse.ts mapExprChildren): the
+   * rebuilt literal is the same construction, so it keeps the witness. */
+  follow(from: IrExpr, to: IrExpr): void {
+    const segments = this.literals.get(from);
+    if (segments !== undefined) this.literals.set(to, segments);
+    const source = this.copies.get(from);
+    if (source !== undefined) this.copies.set(to, source);
+  }
+
   observe(type: IrType): void {
     if (type.kind === "record") this.observed.add(type.shapeId);
   }
@@ -388,8 +401,13 @@ export function applyKeyOrderWitnesses(lowerer: Lowerer, fns: readonly IrFunctio
   // A checked conversion builds values in their source's own order, which
   // the record-source table reproduces relative to the shape's table. A new
   // table would make every value parsed in today's order keep a source
-  // view, so such shapes keep today's table.
-  for (const f of facts.values()) if (f.runtime) f.pinned = true;
+  // view, so a shape only checked conversions build keeps today's table.
+  // A shape the program also builds with object literals takes their order:
+  // the literal's values have no other way to enumerate in it, and values
+  // parsed from JSON those literals wrote (a message sent to a child or
+  // another service of the program) arrive in it too. Parsed values in any
+  // other order keep a source view, exactly as before.
+  for (const f of facts.values()) if (f.runtime && f.literals.length === 0) f.pinned = true;
 
   // Nothing to choose: every witness agrees with today's table and no
   // optional key is written late.

packages/compiler/src/frontend/lowering/lowerer.ts

@@ -35,6 +35,7 @@ import {
 import { identityPreservingWidening } from "./coercions/identity.js";
 import { scanDivergences } from "./divergences.js";
 import { applyKeyOrderWitnesses, KeyOrderNotes } from "./key-order-witness.js";
+import { observeRecordLiteralRebuilds } from "../../ir/traverse.js";
 import { checkedClassAssertion } from "./class-assertions.js";
 import { narrowStoredClassValue, narrowGenericClassValue } from "./class-unions.js";
 import {
@@ -7441,6 +7442,10 @@ export class Lowerer {
    * body; both fire edges through the same hooks and are not units
    * themselves. */
   emitReachable(extraRoots?: readonly string[]): { reachable: Set<string>; result: LowerResult } {
+    // Frontend passes below rewrite function bodies through the IR mapper
+    // (caught-value and class dispatch); key-order notes follow the record
+    // literals they rebuild until the witness pass has read them.
+    observeRecordLiteralRebuilds((from, to) => this.keyOrder.follow(from, to));
     this.deferredModuleInitializers = new DeferredModuleInitializers();
     const parts = this.splitFiles();
     this.narrowing.setFiles(parts.map((fp) => fp.sf));
@@ -7871,6 +7876,7 @@ export class Lowerer {
       ...this.liftedFns,
       ...this.implicitFns,
     ]);
+    observeRecordLiteralRebuilds(null);
     for (const finalize of this.shapeOrderHelperFinalizers) finalize();
     const functions = [
       ...orderedUnits,

packages/compiler/src/ir/traverse.ts

@@ -466,6 +466,18 @@ export function everyStmtChild(
   throw new Error("unhandled IR node");
 }
 
+/** Called with each record literal mapExprChildren rebuilds and its
+ * replacement, so a side table keyed by node identity (the lowering's
+ * key-order witness notes) follows the literal into the rewritten tree. */
+let recordLiteralRebuilt: ((from: IrExpr, to: IrExpr) => void) | null = null;
+
+/** Installs (or with null removes) the record-literal rebuild observer. */
+export function observeRecordLiteralRebuilds(
+  observer: ((from: IrExpr, to: IrExpr) => void) | null,
+): void {
+  recordLiteralRebuilt = observer;
+}
+
 /** Copy a node with transformed immediate children; preserve all metadata. */
 export function mapExprChildren(
   node: IrExpr,
@@ -653,11 +665,14 @@ export function mapExprChildren(
       return { ...node, value: expr(node.value) };
     case "virtualCall":
       return { ...node, args: mapOf(node.args, expr) };
-    case "recordLit":
-      return {
+    case "recordLit": {
+      const rebuilt: IrExpr = {
         ...node,
         fields: node.fields.map((child) => ({ ...child!, value: expr(child.value) })),
       };
+      if (recordLiteralRebuilt !== null) recordLiteralRebuilt(node, rebuilt);
+      return rebuilt;
+    }
     case "recordClone":
       return {
         ...node,

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/6495-key-order-rewritten-literals.ts.json

@@ -0,0 +1,6 @@
+{
+ "order": [
+  "<repo>/tests/corpus/6495-key-order-rewritten-literals.ts"
+ ],
+ "diags": []
+}

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/6496-key-order-built-and-parsed.ts.json

@@ -0,0 +1,6 @@
+{
+ "order": [
+  "<repo>/tests/corpus/6496-key-order-built-and-parsed.ts"
+ ],
+ "diags": []
+}

tests/corpus/6495-key-order-rewritten-literals.ts

@@ -0,0 +1,54 @@
+// Object literals keep their own key order through JSON.stringify when the
+// function that builds them is rewritten by a later lowering pass (here a
+// catch block that hands the caught value to an `unknown` parameter). The
+// declared type lists the keys in another order; Node writes the literal's.
+interface StatusBody {
+  pool?: string;
+  cpuModel?: string;
+  cellId?: string;
+  hiveId?: string;
+  cluster?: string;
+  commit?: string;
+  idleSince?: number;
+}
+interface PatchBody {
+  readyState: string;
+  status: StatusBody;
+}
+
+function describe(value: unknown): string {
+  return value instanceof Error ? `Error(${value.message})` : String(value);
+}
+
+function send(path: string, body: string): string {
+  return `${path} ${body}`;
+}
+
+function write(status: StatusBody): string {
+  return JSON.stringify(status);
+}
+
+function report(commit: string, idleSince: number | undefined, fail: boolean): string[] {
+  const lines: string[] = [];
+  try {
+    if (fail) throw new Error("probe failed");
+  } catch (e) {
+    lines.push(`caught ${describe(e)}`);
+  }
+  lines.push(
+    send(
+      "/patch",
+      JSON.stringify({
+        readyState: "DOWNLOADING",
+        status: { cellId: "c1", cluster: "k1", commit, cpuModel: "m1", hiveId: "h1", idleSince },
+      } satisfies PatchBody),
+    ),
+  );
+  lines.push(write({ cellId: "c2", cluster: "k2", commit, cpuModel: "m2", hiveId: "h2" }));
+  const body = { hiveId: "h3", commit, cellId: "c3" } satisfies StatusBody;
+  lines.push(send("/status", JSON.stringify(body)), Object.keys(body).join(","));
+  return lines;
+}
+
+for (const line of report("abc", 5, true)) console.log(line);
+for (const line of report("def", undefined, false)) console.log(line);

tests/corpus/6496-key-order-built-and-parsed.ts

@@ -0,0 +1,35 @@
+// A record type the program both builds with an object literal and reads
+// back with a typed JSON.parse: the literal enumerates in its own key order
+// (not the declaration's), values parsed from the JSON it wrote keep that
+// order, and values parsed from JSON in another order keep theirs.
+interface Order {
+  id: string;
+  region: string;
+}
+interface Message {
+  order: Order;
+  outputs: string[];
+  traceCarrier?: Record<string, string>;
+  pollMs: number;
+}
+
+function encode(order: Order, outputs: string[], pollMs: number): string {
+  const message: Message = { outputs, order, pollMs };
+  return JSON.stringify(message);
+}
+
+function decode(text: string): Message {
+  return JSON.parse(text) as Message;
+}
+
+const wire = encode({ region: "iad1", id: "o1" }, ["a", "b"], 250);
+console.log(wire);
+const back = decode(wire);
+console.log(JSON.stringify(back), Object.keys(back).join(","), Object.keys(back.order).join(","));
+
+const external = decode('{"pollMs":5,"order":{"id":"o2","region":"sfo1"},"outputs":[],"traceCarrier":{"t":"1"}}');
+console.log(JSON.stringify(external), Object.keys(external).join(","));
+
+const declared = decode('{"order":{"id":"o3","region":"fra1"},"outputs":["x"],"pollMs":1}');
+console.log(JSON.stringify(declared), Object.keys(declared).join(","));
+for (const [key, value] of Object.entries(back)) console.log(key, value !== undefined ? JSON.stringify(value) : "-");