Read and write record fields through keys whose literal type names fields of one type

packages/compiler/src/frontend/lowering/expressions/element-access.ts

@@ -84,7 +84,13 @@ import { isSafeToDiscard, isSafeToRepeat } from "./evaluation-safety.js";
 import { withInheritedMethods } from "../record-inherited.js";
 import { lowerClassPrototypeComputedAssignment } from "../class-prototypes.js";
 import { isClassCallback } from "../class-callbacks.js";
-import { isKeyedFieldAccess } from "./keyed-field-targets.js";
+import {
+  isKeyedFieldAccess,
+  keyedFieldRead,
+  keyedFieldWrite,
+  sameTypedFieldKeys,
+  type KeyedFieldTarget,
+} from "./keyed-field-targets.js";
 import { lowerFieldCompoundValue } from "./updates.js";
 import { readsHeldEvolvingArray } from "../lower-stmts.js";
 import {
@@ -1548,6 +1554,36 @@ export function lowerRecordKeyRead(
     recordKeyResultOk(lowerer, shape, DYN)
   )
     declared = DYN;
+  // A key whose literal type names only declared fields of one type, on a
+  // record whose other fields have other types, reads the field it names.
+  // A called access (`r[key]()`) keeps the refusal.
+  const named =
+    !declared && litKey === null && !isCalledAccess(expr)
+      ? sameTypedFieldKeys(lowerer, shape, keyNode)
+      : null;
+  if (named) {
+    const stmts: IrStmt[] = [];
+    const settle = (value: IrExpr, name: string, force = false): IrExpr => {
+      if (!force && isSafeToRepeat(value)) return value;
+      const temp = lowerer.declareHiddenLocal(name, value.type);
+      stmts.push({ kind: "varDecl", localId: temp.id, init: value, loc });
+      return varRef(temp.id, temp.type, loc);
+    };
+    const receiver = settle(obj, "%keyedReceiver", !isSafeToDiscard(key));
+    const target: KeyedFieldTarget = {
+      container: "recordKeyed",
+      obj: receiver,
+      shapeId,
+      field: named.keys[0]!,
+      keys: named.keys,
+      key: settle(key, "%keyedKey"),
+      fieldType: named.fieldType,
+    };
+    const get = keyedFieldRead(lowerer, target, loc, true);
+    const read: IrExpr =
+      stmts.length === 0 ? get : { kind: "seqExpr", stmts, result: get, type: get.type, loc };
+    return includeUndefined ? read : lowerer.maybeNarrow(read, expr);
+  }
   if (!declared) {
     lowerer.unsupported(
       "SC1090",
@@ -2571,6 +2607,48 @@ export function lowerElementWrite(lowerer: Lowerer, expr: ts.BinaryExpression):
           };
         }
       }
+      // A key whose literal type names only declared fields of one type
+      // writes the field it names. JavaScript evaluates the receiver, then
+      // the key, then the value; the dispatch reads the saved receiver and
+      // key once per named field.
+      const named =
+        litKey === null ? sameTypedFieldKeys(lowerer, shape, target.argumentExpression) : null;
+      if (named && !shape.fields.every((f) => typeEquals(f.type, named.fieldType))) {
+        const loc = locOf(expr);
+        const body: IrStmt[] = [];
+        const save = (value: IrExpr, name: string): IrExpr => {
+          if (value.kind === "varRef") return value;
+          const temp = lowerer.declareHiddenLocal(name, value.type);
+          body.push({ kind: "varDecl", localId: temp.id, init: value, loc });
+          return varRef(temp.id, temp.type, loc);
+        };
+        const obj = save(lowerer.lowerExpr(target.expression), "%keyedReceiver");
+        const rawKey = lowerRecordPropertyKey(
+          lowerer,
+          lowerer.lowerExpr(target.argumentExpression),
+          target.argumentExpression,
+        );
+        if (rawKey.type.kind !== "string") {
+          lowerer.unsupported(
+            "SC1090",
+            target.argumentExpression,
+            "indexing records with non-string or non-number keys",
+          );
+        }
+        const key = save(rawKey, "%keyedKey");
+        const value = lowerer.lowerExprExpecting(expr.right, named.fieldType);
+        const keyed: KeyedFieldTarget = {
+          container: "recordKeyed",
+          obj,
+          shapeId: receiverIr.shapeId,
+          field: named.keys[0]!,
+          keys: named.keys,
+          key,
+          fieldType: named.fieldType,
+        };
+        body.push(keyedFieldWrite(lowerer, keyed, value, loc));
+        return { kind: "block", body, loc };
+      }
       if (!shape.indexValue) {
         // A SIGNATURE-FREE shape whose declared fields share ONE type
         // writes through the same per-shape dispatch (the mockable-clock

packages/compiler/src/frontend/lowering/expressions/keyed-field-targets.ts

@@ -5,9 +5,12 @@ import {
   DYN,
   F64,
   STRING,
+  shapeHasAccessorSlots,
   typeEquals,
   type IrExpr,
+  type IrRecordShape,
   type IrStmt,
+  type IrType,
   type SrcLoc,
 } from "../../../ir/ir.js";
 import { isJsSourceFile } from "../../program.js";
@@ -141,6 +144,32 @@ export function keyedFieldTarget(
   };
 }
 
+/** The declared fields a computed key names when the key's literal type
+ * allows only names of declared fields stored with one type
+ * (`stats[counter]` with `counter: "hits" | "misses"` on a record whose
+ * other fields have other types). Null for any other key or shape:
+ * accessor-carrying shapes, index-signature shapes and JavaScript sources
+ * keep their own keyed access. */
+export function sameTypedFieldKeys(
+  lowerer: Lowerer,
+  shape: IrRecordShape,
+  keyNode: ts.Expression,
+): { keys: string[]; fieldType: IrType } | null {
+  if (
+    shape.tuple ||
+    shape.indexValue !== undefined ||
+    shapeHasAccessorSlots(shape) ||
+    isJsSourceFile(ts.getSourceFileOfNode(keyNode))
+  )
+    return null;
+  const keys = literalFieldKeys(lowerer, keyNode);
+  if (!keys) return null;
+  const fields = keys.map((k) => shape.fields.find((f) => f.name === k));
+  const first = fields[0];
+  if (!first || fields.some((f) => !f || !typeEquals(f.type, first.type))) return null;
+  return { keys, fieldType: first.type };
+}
+
 /** Does `key` name `field`? */
 function keyIs(key: IrExpr, field: string, loc: SrcLoc): IrExpr {
   return {

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/record-literal-key-unions.ts.json

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

tests/corpus/record-literal-key-unions.ts

@@ -0,0 +1,78 @@
+// Keyed reads and writes whose key is a union of literal names of declared
+// fields stored with one type, on records whose other fields have other
+// types: plain reads, plain writes, compound updates, keys from parameters
+// and loops, and JavaScript evaluation order with effectful receivers and
+// keys.
+type Stats = { hits: number; misses: number; bytes: number; name: string; ok: boolean };
+type Counter = "hits" | "misses" | "bytes";
+type Flags = { name: string; verbose: boolean; strict: boolean; level: number };
+type Labels = { id: number; title: string; summary: string; tags: string[] };
+
+function read(s: Stats, key: Counter): number {
+  return s[key];
+}
+
+function bump(s: Stats, key: Counter, by: number): void {
+  s[key] = s[key] + by;
+}
+
+function toggle(flags: Flags, key: "verbose" | "strict"): boolean {
+  flags[key] = !flags[key];
+  return flags[key];
+}
+
+function relabel(labels: Labels, key: "title" | "summary", text: string): string {
+  const before = labels[key];
+  labels[key] = text.toUpperCase();
+  return `${before} -> ${labels[key]}`;
+}
+
+const stats: Stats[] = [];
+for (let i = 0; i < 4; i++) stats.push({ hits: 0, misses: 0, bytes: 0, name: "s" + i, ok: true });
+
+let seed = 7;
+for (let i = 0; i < 1000; i++) {
+  seed = (Math.imul(seed, 1103515245) + 12345) >>> 0;
+  const s = stats[seed & 3]!;
+  bump(s, seed % 5 === 0 ? "misses" : "hits", 1);
+  bump(s, "bytes", (seed >>> 8) & 255);
+}
+for (const s of stats) console.log(s.name, read(s, "hits"), read(s, "misses"), read(s, "bytes"), s.ok);
+
+// Keys from a loop over a constant list of names.
+const counters = ["hits", "misses", "bytes"] as const;
+const totals: Stats = { hits: 0, misses: 0, bytes: 0, name: "total", ok: false };
+for (const s of stats) for (const key of counters) totals[key] += s[key];
+console.log(totals.name, totals.hits, totals.misses, totals.bytes);
+
+// Booleans and strings.
+const flags: Flags = { name: "build", verbose: false, strict: true, level: 2 };
+console.log(toggle(flags, "verbose"), toggle(flags, "strict"), toggle(flags, "verbose"), flags.level);
+const labels: Labels = { id: 1, title: "draft", summary: "none", tags: ["a"] };
+console.log(relabel(labels, "title", "final"), relabel(labels, "summary", "short"), labels.tags.length);
+
+// The receiver, then the key, then the value, each once.
+const order: string[] = [];
+function pick(i: number): Stats {
+  order.push("receiver " + i);
+  return stats[i]!;
+}
+function name(key: Counter): Counter {
+  order.push("key " + key);
+  return key;
+}
+function amount(n: number): number {
+  order.push("value " + n);
+  return n;
+}
+pick(1)[name("misses")] = amount(40);
+console.log(order.join(", "), stats[1]!.misses);
+order.length = 0;
+console.log(pick(2)[name("bytes")], order.join(", "));
+
+// A key that names one field.
+function only(s: Stats, key: "bytes"): number {
+  s[key] = s[key] * 2;
+  return s[key];
+}
+console.log(only(stats[3]!, "bytes"));

tests/harness/corpus-costs/record-literal-key-unions.ts.json

@@ -0,0 +1 @@
+{ "plain": 5800, "san": 62700 }