Keep union locals copied between locals in private storage (#2457)

packages/compiler/src/backend/llvm/call-lifetimes.test.ts

@@ -594,6 +594,35 @@ test("closure bodies get local facts but never parameter facts", () => {
   expect(analyze(fn).locals.get("closure")).toEqual(new Set(["item"]));
 });
 
+test("union locals copied to other locals are reported with their copies", () => {
+  const fn = helper("copies");
+  fn.locals.push(local("source"), { ...local("first"), mutable: true }, local("second"));
+  fn.body = [
+    { kind: "varDecl", localId: "source", init: ref("value"), loc },
+    { kind: "varDecl", localId: "first", init: ref("source"), loc },
+    { kind: "assign", localId: "first", value: ref("source"), loc },
+    { kind: "varDecl", localId: "second", init: ref("first"), loc },
+    { kind: "exprStmt", expr: tag("second"), loc },
+  ];
+  const result = analyze(fn);
+  expect(result.projectedLocals.get("copies")).toEqual(new Set(["second"]));
+  expect(result.copiedLocals.get("copies")).toEqual(
+    new Map([
+      ["source", ["first", "first"]],
+      ["first", ["second"]],
+    ]),
+  );
+  // A copied parameter stays rejected: callers may lend it a stack box.
+  expect(result.parameters.has("copies")).toBe(false);
+  // A copy inside an expression is an ordinary whole-value use.
+  fn.body.push({
+    kind: "exprStmt",
+    expr: { kind: "assignExpr", localId: "first", value: ref("source"), type: optional, loc },
+    loc,
+  });
+  expect(analyze(fn).copiedLocals.get("copies")?.has("source") ?? false).toBe(false);
+});
+
 test("stores into nullable-pointer class fields project the stored union", () => {
   const self: IrType = { kind: "object", className: "Node" };
   const store = (className: string, field: string, value: IrExpr): IrStmt => ({

packages/compiler/src/backend/llvm/call-lifetimes.ts

@@ -34,6 +34,10 @@ interface Uses {
   /** Narrowed payloads of union bindings passed directly as call arguments:
    * the payload stays in place if the callee's parameter is proved. */
   payloadForwards: Map<string, ForwardedUse[]>;
+  /** Union bindings whose whole value initializes or is assigned to
+   * another union local (`const x = y`, `x = y` as a statement), with the
+   * locals it is copied to. */
+  copies: Map<string, string[]>;
 }
 
 interface Parameter {
@@ -59,6 +63,11 @@ export interface CallLifetimes {
   /** Projection-only local uses, allowing statement assignments. Each
    * initializer and assignment still needs a separate storage proof. */
   projectedLocals: Map<string, Set<string>>;
+  /** Locals that projectedLocals leaves out only because their whole value
+   * is copied to other union locals, with those locals. Such a local is
+   * only projected when every local it is copied to snapshots the copy
+   * (local-union-storage.ts) instead of keeping the box. */
+  copiedLocals: Map<string, Map<string, readonly string[]>>;
   /** Generator bodies own their parameters for the whole body, including
    * every suspension: the frame releases them only on its final exit. A
    * reference parameter nothing rebinds, declares or captures is therefore
@@ -132,6 +141,7 @@ function collectUses(
     forwards: new Map(),
     payloadKept: new Set(),
     payloadForwards: new Map(),
+    copies: new Map(),
   };
   const forward = (
     table: Map<string, ForwardedUse[]>,
@@ -292,6 +302,18 @@ function collectUses(
     }
     return everyExprChild(node, expr, (node) => stmt(node, true));
   }
+  /** `const x = y` or a statement `x = y` of union locals: the copy is
+   * recorded instead of rejecting `y`. A parameter's copy still rejects it,
+   * since callers may lend projection-only parameters a stack box. */
+  const params = new Set(fn.params.map((param) => param.localId));
+  const unionCopy = (target: string, value: IrExpr | null | undefined): boolean => {
+    if (value?.kind !== "varRef" || value.type.kind !== "union" || params.has(value.localId))
+      return false;
+    let targets = uses.copies.get(value.localId);
+    if (!targets) uses.copies.set(value.localId, (targets = []));
+    targets.push(target);
+    return true;
+  };
   function stmt(node: IrStmt, inExpression: boolean): boolean {
     switch (node.kind) {
       case "return":
@@ -300,6 +322,7 @@ function collectUses(
         break;
       case "varDecl":
         uses.declarations.set(node.localId, (uses.declarations.get(node.localId) ?? 0) + 1);
+        if (unionCopy(node.localId, node.init)) return true;
         // A receiver temporary (`%compoundReceiver = this`) projects its
         // source as long as the temporary itself is only projected.
         if (
@@ -342,6 +365,7 @@ function collectUses(
         // Sequence expressions can rebind a previous call argument before
         // its consumer runs. Those locals need an owned heap snapshot.
         if (inExpression) uses.invalid.add(node.localId);
+        else if (unionCopy(node.localId, node.value)) return true;
         break;
       case "forOf":
         // The binding is initialized afresh for each iteration. Its uses
@@ -429,6 +453,7 @@ export function analyzeCallLifetimes(
     borrowed: new Map(),
     bindings: new Map(),
     projectedLocals: new Map(),
+    copiedLocals: new Map(),
     generatorParameters: new Map(),
   };
   for (const fn of functions.values()) {
@@ -602,6 +627,7 @@ export function analyzeCallLifetimes(
     const params = new Set(fn.params.map((param) => param.localId));
     const safe = new Set<string>();
     const projected = new Set<string>();
+    const copied = new Map<string, readonly string[]>();
     for (const local of fn.locals) {
       if (
         params.has(local.id) ||
@@ -611,13 +637,18 @@ export function analyzeCallLifetimes(
       )
         continue;
       const forwards = uses.forwards.get(local.id) ?? [];
-      if (forwards.every((use) => target(use)?.safe === true)) {
-        projected.add(local.id);
-        if (eligible(local) && !uses.written.has(local.id)) safe.add(local.id);
+      if (!forwards.every((use) => target(use)?.safe === true)) continue;
+      const copies = uses.copies.get(local.id);
+      if (copies) {
+        copied.set(local.id, copies);
+        continue;
       }
+      projected.add(local.id);
+      if (eligible(local) && !uses.written.has(local.id)) safe.add(local.id);
     }
     result.locals.set(fn.name, safe);
     result.projectedLocals.set(fn.name, projected);
+    if (copied.size > 0) result.copiedLocals.set(fn.name, copied);
   }
   // Consumed parameters take the owned convention, unless only projected.
   for (const { fn, index } of consumed) {

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

@@ -130,9 +130,10 @@ export function findLocalUnionStorage(
   const types = new Map(fn.locals.map((local) => [local.id, local.type]));
   const result = new Map<string, LocalUnionStorageProof>();
   const safe = lifetimes.projectedLocals.get(fn.name);
-  if (!safe?.size) return result;
+  const copied = lifetimes.copiedLocals.get(fn.name);
+  if (!safe?.size && !copied?.size) return result;
   for (const local of fn.locals) {
-    if (local.type.kind !== "union" || !safe.has(local.id)) continue;
+    if (local.type.kind !== "union" || !(safe?.has(local.id) || copied?.has(local.id))) continue;
     const arms = unions.get(local.type.unionId)?.arms;
     if (!arms) continue;
     const refs = arms.filter(isRefCounted);
@@ -149,6 +150,18 @@ export function findLocalUnionStorage(
       return true;
     },
   });
+  // A copied local's box stays private only while every local it is copied
+  // to snapshots the copy's tag and bits: one that keeps the box pointer
+  // makes the source a heap box again, which can cascade along the copies.
+  for (let changed = copied !== undefined; changed;) {
+    changed = false;
+    for (const [id, targets] of copied!) {
+      if (result.has(id) && !targets.every((target) => result.has(target))) {
+        result.delete(id);
+        changed = true;
+      }
+    }
+  }
   return result;
 }
 

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

@@ -172,8 +172,11 @@ test.each(["return", "alias", "capture", "store"])(
       work.body.push(ret(ref("item")));
     }
     if (reason === "alias") {
+      // The alias escapes, so it holds a heap box and so does its source.
+      work.returnType = optional;
       work.locals.push({ id: "alias", name: "alias", type: optional, mutable: false });
       work.body.push({ kind: "varDecl", localId: "alias", init: ref("item"), loc });
+      work.body.push(ret(ref("alias")));
     }
     if (reason === "capture") {
       work.locals[0]!.boxed = true;
@@ -204,6 +207,27 @@ test.each(["return", "alias", "capture", "store"])(
   },
 );
 
+test("copies to a projection-only alias keep both locals in private storage", () => {
+  const mod = fixture();
+  const work: IrFunction = {
+    name: "work",
+    params: [],
+    locals: [],
+    body: [],
+    returnType: VOID,
+    loc,
+  };
+  mod.functions.unshift(work);
+  declareLocal(mod, wrap(fresh()));
+  work.locals.push({ id: "alias", name: "alias", type: optional, mutable: true });
+  work.body.push({ kind: "varDecl", localId: "alias", init: ref("item"), loc });
+  work.body.push({ kind: "assign", localId: "alias", value: ref("item"), loc });
+  work.body.push(effect(call("read", [ref("alias")], record)));
+  const text = body(emit(mod), "sc_f_work");
+  expect(text).toContain("alloca %ScrUnion");
+  expect(text).not.toMatch(/@scr_union_(?:new|retain|release)/);
+});
+
 test("forwarded parameters never retain or release a borrowed stack box", () => {
   const mod = fixture();
   const outer: IrFunction = {

packages/compiler/src/frontend/lowering/array-values.ts

@@ -260,6 +260,9 @@ const PURE_KINDS = new Set<IrExpr["kind"]>([
   "unionIsTag",
   "unionNarrow",
   "unionWrap",
+  // Its statements are checked on their own (declarations and assignments
+  // are recorded writes).
+  "seqExpr",
 ]);
 
 /** The locals evaluating `value` may assign, and whether it may run other
@@ -273,7 +276,11 @@ function valueEffects(lowerer: Lowerer, value: IrExpr): { writes: Set<string>; o
   everyExpr(value, {
     expr: (e) => {
       if (e.kind === "assignExpr" || e.kind === "incDec") writes.add(e.localId);
-      else if (e.kind === "call" ? !readers.has(e.callee) : !PURE_KINDS.has(e.kind)) opaque = true;
+      else if (e.kind === "libCall") {
+        // The "undefined" key of an undefined index is a constant.
+        if (e.fn !== "array.undefinedKey") opaque = true;
+      } else if (e.kind === "call" ? !readers.has(e.callee) : !PURE_KINDS.has(e.kind))
+        opaque = true;
       return true;
     },
     stmt: (s) => {
@@ -314,6 +321,16 @@ function readsMutableLocal(lowerer: Lowerer, e: IrExpr): boolean {
   });
 }
 
+/** Whether `receiver` still reads the same value after `index` runs: the
+ * index assigns none of the bindings it reads, and runs no code that could
+ * assign them through a closure. */
+export function unaffectedByIndex(lowerer: Lowerer, receiver: IrExpr, index: IrExpr): boolean {
+  const effects = valueEffects(lowerer, index);
+  return (
+    unaffectedBy(receiver, effects) && !(effects.opaque && readsMutableLocal(lowerer, receiver))
+  );
+}
+
 /** A number, string or boolean element type. Arrays of it hold a present
  * undefined in the element's state, so `(T | undefined)[]` stores T
  * elements (the type mapper) and a T[] seen as one stays the same array. */

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

@@ -66,6 +66,7 @@ import {
   lowerSafeIndexRead,
   primitiveElement,
   unionArrayValueRead,
+  unaffectedByIndex,
 } from "../array-values.js";
 import { lowerOptionalStringIndex, lowerUnionStringIndex } from "../string-index.js";
 import { InternalCompilerError } from "../../../errors.js";
@@ -2475,14 +2476,35 @@ export function lowerElementWrite(lowerer: Lowerer, expr: ts.BinaryExpression):
     (value.type.kind === "union" && lowerer.armTag(value.type.unionId, UNDEFINED_T) >= 0)
       ? arrayValueType(lowerer, arr.type.elem)
       : arr.type.elem;
-  return arrayValueStore(
-    lowerer,
-    arr,
-    index,
-    lowerer.coerceInto(expr.right, value, expected),
-    arr.type.elem,
-    locOf(expr),
-  );
+  const stored = lowerer.coerceInto(expr.right, value, expected);
+  // A binding's `T | undefined` value is stored by two writes that both
+  // read the index (arrayValueStore). A computed index, such as the key
+  // conversion of a union index that declares a temporary, is held once
+  // ahead of them. The stores then read the receiver after the index, so
+  // an index that can rebind the receiver (`a[(a = b, 0)] = u`) keeps the
+  // stores as they are, reading the receiver first.
+  if (
+    stored.kind === "varRef" &&
+    stored.type.kind === "union" &&
+    lowerer.armTag(stored.type.unionId, UNDEFINED_T) >= 0 &&
+    !typeEquals(stored.type, arr.type.elem) &&
+    arr.kind === "varRef" &&
+    index.kind !== "varRef" &&
+    index.kind !== "numLit" &&
+    unaffectedByIndex(lowerer, arr, index)
+  ) {
+    const loc = locOf(expr);
+    const held = lowerer.declareHiddenLocal("%arrayIndex", index.type);
+    return {
+      kind: "block",
+      body: [
+        { kind: "varDecl", localId: held.id, init: index, loc },
+        arrayValueStore(lowerer, arr, varRef(held.id, held.type, loc), stored, arr.type.elem, loc),
+      ],
+      loc,
+    };
+  }
+  return arrayValueStore(lowerer, arr, index, stored, arr.type.elem, locOf(expr));
 }
 
 /** Environment property names use ToPrimitive with the string hint, so

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/union-local-copies.ts.json

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

tests/corpus/union-local-copies.ts

@@ -0,0 +1,195 @@
+// Array reads held in locals and copied between them: `number | undefined`,
+// `boolean | undefined` and `string | undefined` values copied by
+// declarations and assignments, read back, and stored into arrays, with
+// holes and present undefined values along the way.
+
+function flips(perm: number[]): number {
+  const pp = perm.slice();
+  let count = 0;
+  let p0 = pp[0];
+  while (p0 !== 0 && p0 !== undefined) {
+    count++;
+    for (let i = 1, j = p0 - 1; i < j; i++, j--) {
+      const t = pp[i];
+      pp[i] = pp[j];
+      pp[j] = t;
+    }
+    const t = pp[p0];
+    pp[p0] = p0;
+    p0 = t;
+  }
+  return count;
+}
+
+console.log(flips([3, 1, 2, 0]), flips([2, 0, 1, 3]), flips([0, 1, 2]));
+
+// A hole read into a local, copied, and written back as present undefined.
+const sparse: number[] = [];
+sparse[0] = 1;
+sparse[2] = 3;
+const out: number[] = [];
+for (let i = 0; i < 4; i++) {
+  const read = sparse[i];
+  let copy = read;
+  if (i === 2) copy = sparse[0];
+  out[i] = copy;
+}
+console.log(out.length, 1 in out, out[1], out[2], 3 in out, out[3]);
+
+// The copy feeds an index and is written back at it.
+const keys: number[] = [2];
+keys[2] = 0;
+const values: number[] = [10, 20, 30];
+const picked: (number | undefined)[] = [];
+for (let i = 0; i < keys.length; i++) {
+  let k = keys[i];
+  const again = k;
+  k = again;
+  if (k === undefined) {
+    picked.push(k);
+    continue;
+  }
+  picked.push(values[k]);
+  values[k] = k;
+}
+console.log(picked, values);
+
+// Boolean and string elements.
+const flags: boolean[] = [true, false];
+let flag = flags[1];
+const flagCopy = flag;
+flag = flags[5];
+console.log(flagCopy, flag, typeof flag);
+
+const names: string[] = ["ada"];
+names[2] = "lin";
+const renamed: string[] = [];
+for (let i = 0; i < 3; i++) {
+  const name = names[i];
+  let held = name;
+  if (held === undefined) held = names[0];
+  renamed[2 - i] = held;
+}
+console.log(renamed.join("/"), renamed.length);
+
+// A copy that escapes: returned, captured, or stored in an object.
+function firstOr(list: number[], fallback: number): number | undefined {
+  const head = list[0];
+  const result = head;
+  return result ?? fallback;
+}
+console.log(firstOr([], 7), firstOr([4], 7));
+
+function captured(list: number[]): () => number | undefined {
+  const head = list[1];
+  const kept = head;
+  return () => kept;
+}
+console.log(captured([1, 2])(), captured([1])());
+
+const holder: { value: number | undefined } = { value: 0 };
+const element = sparse[1];
+const moved = element;
+holder.value = moved;
+console.log(holder.value, "value" in holder);
+
+// Reference elements, exception regions, generators, async functions,
+// closures over a copy, and nested arrays.
+class Box { v: number; constructor(v: number) { this.v = v; } }
+function chain(xs: Box[], n: number): string {
+  let out = "";
+  let cur = xs[0];
+  for (let i = 0; i < n; i++) {
+    const next = xs[(i * 7) % (xs.length + 2)];
+    let held = next;
+    if (held === undefined) held = cur;
+    cur = held;
+    out += cur === undefined ? "u" : String(cur.v);
+    if (i % 5 === 4) xs.pop();
+  }
+  return out;
+}
+console.log(chain([new Box(1), new Box(2), new Box(3), new Box(4)], 12));
+
+function strings(words: string[]): string[] {
+  const res: string[] = [];
+  let last = words[0];
+  for (let i = 0; i < words.length + 2; i++) {
+    const w = words[i];
+    const copy = w;
+    last = copy ?? last;
+    res[i] = copy;
+    try {
+      if (i === 3) throw new Error("x" + last);
+    } catch (e) {
+      res.push((e as Error).message);
+    } finally {
+      last = words[i + 1];
+    }
+  }
+  return res;
+}
+console.log(strings(["a", "bb", "ccc", "dddd"]));
+
+function* gen(xs: number[]): Generator<number | undefined> {
+  const a = xs[5];
+  const b = a;
+  yield b;
+  const c = xs[0];
+  let d = c;
+  d = xs[9];
+  yield d;
+}
+console.log([...gen([1, 2])]);
+
+async function asy(xs: number[]): Promise<string> {
+  const a = xs[1];
+  const b = a;
+  await null;
+  const c = b;
+  return String(c);
+}
+asy([4, 5]).then((v) => console.log(v));
+
+function closures(xs: number[]): number[] {
+  const fns: (() => number | undefined)[] = [];
+  for (let i = 0; i < 3; i++) {
+    const v = xs[i];
+    let copy = v;
+    fns.push(() => copy);
+    copy = xs[i + 1];
+  }
+  return fns.map((f) => f() ?? -1);
+}
+console.log(closures([10, 20]));
+
+function nested(grid: number[][]): number {
+  let s = 0;
+  for (let r = 0; r < 3; r++) {
+    const row = grid[r];
+    const rowCopy = row;
+    if (rowCopy === undefined) continue;
+    for (let c = 0; c < 3; c++) {
+      const cell = rowCopy[c];
+      let x = cell;
+      x = x === undefined ? 100 : x;
+      s += x;
+    }
+  }
+  return s;
+}
+console.log(nested([[1, 2], [3]]));
+
+// The index rebinds the receiver: the store goes to the array the receiver
+// held before the index ran, as for any element assignment.
+let target: number[] = [1, 2, 3];
+const original = target;
+const other: number[] = [7, 8, 9];
+const keyed: number[] = [1];
+const missing = sparse[1];
+const kept = sparse[0];
+target[(target = other, keyed[0])] = missing;
+console.log(original, target === other, 1 in original, other);
+target = original;
+target[(target = other, keyed[0])] = kept;
+console.log(original, other);

tests/harness/corpus-costs/union-local-copies.ts.json

@@ -0,0 +1 @@
+{ "plain": 9441, "san": 15140 }