Keep module-level lets unboxed in the dynamic tier and return receivers from dispatch helpers (#2471)

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

@@ -225,6 +225,7 @@ import { recordTextCodecClass } from "../../ir/ir.js";
 import { assignLegacyTiers, type DynamicFunction } from "../../ir/tiers.js";
 import { dynamicSlotOperationCounts, dynamicTierWarning } from "./dynamic-tier.js";
 import { promoteJsvals } from "../../ir/jsval.js";
+import { returnDispatchReceivers } from "./receiver-returns.js";
 import {
   analyzeTiers,
   dynamicStatementCount,
@@ -9464,6 +9465,7 @@ export class Lowerer {
     applyHandleProperties(this, functions);
     installBufferValue(functions);
     installFetchBodyWebStreams(this, functions);
+    returnDispatchReceivers(this, functions);
 
     // Globals typed by a class that never REGISTERED (a JS class whose
     // collection fenced — Symbol-keyed fields, an unsupported base): the

packages/compiler/src/frontend/lowering/receiver-returns.ts

@@ -0,0 +1,131 @@
+/* Methods that return their receiver, called through an untyped value.
+ *
+ * A JavaScript method that ends every path with `return this` (a fluent
+ * setter such as `step(dx, dy) { this.x += dx; return this; }`) boxes its
+ * receiver on every call: a lookup of the instance's live capsule, or a new
+ * capsule when it has none. A class dispatch helper
+ * (`%dyn.class.call.N`) that calls such a method already holds a capsule
+ * of the same instance at its exact class: its receiver after
+ * normalization. The helper now calls a copy of the method that returns
+ * nothing and answers its own receiver, so the call boxes nothing. Other
+ * callers (vtables, typed code, method values) keep the method itself. */
+import type { IrExpr, IrFunction, IrStmt } from "../../ir/ir.js";
+import { VOID } from "../../ir/ir.js";
+import { varRef } from "../../ir/build.js";
+import { everyStmtList, transformStmtList } from "../../ir/traverse.js";
+import type { Lowerer } from "./lowerer.js";
+
+const DISPATCH_PREFIX = "%dyn.class.call.";
+const BOX_EXACT_PREFIX = "%dyn.class.boxExact:";
+
+/** The name of a method's copy that returns nothing. */
+function receiverCopyName(method: string): string {
+  return `%dyn.class.receiver:${method.slice(1)}`;
+}
+
+export function returnDispatchReceivers(lowerer: Lowerer, functions: IrFunction[]): void {
+  const byName = new Map(functions.map((fn) => [fn.name, fn]));
+  const copies = new Map<string, IrFunction | null>();
+  const copyOf = (name: string): IrFunction | null => {
+    const known = copies.get(name);
+    if (known !== undefined) return known;
+    const method = byName.get(name);
+    const copy = method && returnsReceiver(lowerer, method) ? receiverCopy(method) : null;
+    copies.set(name, copy);
+    if (copy) functions.push(copy);
+    return copy;
+  };
+  for (const helper of functions) {
+    if (!helper.name.startsWith(DISPATCH_PREFIX)) continue;
+    const receiver = helper.params[0]?.localId;
+    if (receiver === undefined || !normalizedOnly(helper, receiver)) continue;
+    for (const stmt of helper.body) {
+      if (stmt.kind !== "if" || stmt.then.length !== 1) continue;
+      const only = stmt.then[0]!;
+      if (only.kind !== "return" || only.value === null || only.value.kind !== "call") continue;
+      const call = only.value;
+      if (!receiverArgument(call.args[0], receiver)) continue;
+      const copy = copyOf(call.callee);
+      if (!copy) continue;
+      stmt.then = [
+        {
+          kind: "exprStmt",
+          expr: { ...call, callee: copy.name, type: VOID },
+          loc: only.loc,
+        },
+        { kind: "return", value: varRef(receiver, only.value.type, only.loc), loc: only.loc },
+      ];
+    }
+  }
+}
+
+/** The helper's receiver is assigned only by its normalization, so it is
+ * the capsule of the instance's own class wherever a branch reads it. */
+function normalizedOnly(helper: IrFunction, receiver: string): boolean {
+  return everyStmtList(helper.body, {
+    stmt: (s) =>
+      !(
+        (s.kind === "assign" || s.kind === "forOf") &&
+        s.localId === receiver &&
+        !(
+          s.kind === "assign" &&
+          s.value.kind === "call" &&
+          s.value.callee === "%dyn.class.normalize"
+        )
+      ),
+    expr: (e) => !((e.kind === "assignExpr" || e.kind === "incDec") && e.localId === receiver),
+  });
+}
+
+/** `dynCheck(receiver)`, possibly upcast to the method's class. */
+function receiverArgument(arg: IrExpr | undefined, receiver: string): boolean {
+  if (arg?.kind === "upcast") arg = arg.value;
+  return arg?.kind === "dynCheck" && arg.value.kind === "varRef" && arg.value.localId === receiver;
+}
+
+/** Whether every path of a synchronous method returns `this` boxed at its
+ * exact class. */
+function returnsReceiver(lowerer: Lowerer, fn: IrFunction): boolean {
+  if (fn.async !== undefined || fn.generator !== undefined) return false;
+  if (fn.captures?.length || fn.classCaptures?.length) return false;
+  if (fn.returnType.kind !== "dyn") return false;
+  const self = fn.params[0];
+  if (self?.type.kind !== "object" || self.name !== "this") return false;
+  const last = fn.body[fn.body.length - 1];
+  if (last?.kind !== "return") return false;
+  const isSelf = (value: IrExpr | null): boolean => {
+    if (value === null) return false;
+    if (value.kind === "call" && value.callee.startsWith(BOX_EXACT_PREFIX))
+      return thisRef(value.args[0], self.localId);
+    // Without exact boxing the capsule names the static class, which is
+    // the instance's own class only when nothing extends it.
+    if (value.kind === "dynFrom" && value.liveRef === undefined && value.fnName === undefined) {
+      const className = self.type.kind === "object" ? self.type.className : "";
+      return (
+        thisRef(value.value, self.localId) &&
+        (lowerer.classes.get(className)?.subclasses.length ?? 1) === 0
+      );
+    }
+    return false;
+  };
+  return everyStmtList(fn.body, {
+    stmt: (s) => s.kind !== "return" || isSelf(s.value),
+    expr: () => true,
+  });
+}
+
+function thisRef(value: IrExpr | undefined, thisId: string): boolean {
+  return value?.kind === "varRef" && value.localId === thisId;
+}
+
+/** A copy of `method` whose returns return nothing. */
+function receiverCopy(method: IrFunction): IrFunction {
+  const copy = structuredClone(method);
+  copy.name = receiverCopyName(method.name);
+  copy.returnType = VOID;
+  copy.body = transformStmtList(copy.body, {
+    stmt: (s): IrStmt => (s.kind === "return" ? { ...s, value: null } : s),
+    expr: (e) => e,
+  });
+  return copy;
+}

packages/compiler/src/ir/jsval.test.ts

@@ -344,3 +344,66 @@ test("a parameter read by dyn consumers or passed cells keeps the dyn ABI", () =
   expect(twice.returnType).toEqual(DYN);
   expect(validateModule(mod)).toEqual([]);
 });
+
+/** `for (...) %g.sum = %g.sum + o.step` in a dynamic function. */
+function globalLoop(name: string, global: string): IrFunction {
+  return fn(name, [
+    {
+      kind: "while",
+      cond: { kind: "boolLit", value: false, type: BOOL, loc },
+      body: [
+        {
+          kind: "assign",
+          localId: global,
+          value: op("dyn.add", ref(global), keyGet(ref("o"), "step")),
+          loc,
+        },
+      ],
+      loc,
+    },
+    { kind: "return", value: ref("o"), loc },
+  ]);
+}
+
+test("a dyn module global written by operators in a loop becomes a jsval global", () => {
+  const loop = globalLoop("loop", "%g.e.sum");
+  // %g.sum.x, once: a dyn consumer.
+  const report = fn("report", [
+    { kind: "exprStmt", expr: keyGet(ref("%g.e.sum"), "x"), loc },
+    { kind: "return", value: ref("o"), loc },
+  ]);
+  const mod = moduleOf(loop, report);
+  mod.globals = [{ id: "%g.e.sum", name: "sum", type: DYN, mutable: true }];
+  promoteJsvals(mod);
+  expect(mod.globals[0]!.type).toEqual(JSVAL);
+  // The loop's own result (o, a jsval result) is the second js.fromDyn.
+  expect(jsCalls(loop.body)).toEqual(["js.add", "js.fromDyn", "js.fromDyn"]);
+  // One box for the read; the second call is the jsval result o.
+  expect(jsCalls(report.body)).toEqual(["js.toDyn", "js.fromDyn"]);
+  expect(validateModule(mod)).toEqual([]);
+});
+
+test("a module global a static function reads stays dyn", () => {
+  const loop = globalLoop("loop", "%g.e.sum");
+  const reader: IrFunction = {
+    name: "reader",
+    params: [],
+    returnType: DYN,
+    locals: [],
+    body: [{ kind: "return", value: ref("%g.e.sum"), loc }],
+    loc,
+  };
+  const mod = moduleOf(loop, reader);
+  expect(reader.tier).toBeUndefined();
+  mod.globals = [{ id: "%g.e.sum", name: "sum", type: DYN, mutable: true }];
+  promoteJsvals(mod);
+  expect(mod.globals[0]!.type).toEqual(DYN);
+  expect(jsCalls(loop.body)).toEqual([
+    "js.toDyn",
+    "js.add",
+    "js.fromDyn",
+    "js.fromDyn",
+    "js.fromDyn",
+  ]);
+  expect(validateModule(mod)).toEqual([]);
+});

packages/compiler/src/ir/jsval.ts

@@ -32,7 +32,9 @@
  *   of the parameter by dyn consumers box more often than its callers do;
  * - a dyn local becomes a jsval local when the boxing it saves on writes
  *   of operator results outweighs the boxing its reads by dyn consumers add
- *   (loop depth weights both).
+ *   (loop depth weights both);
+ * - a dyn module global becomes a jsval global by the same measure, summed
+ *   over every function, when only dynamic-tier functions reference it.
  *
  * The rewrite changes representation only: every js builtin implements its
  * dyn twin's JS semantics, including ToPrimitive order and error messages
@@ -116,10 +118,12 @@ export function promoteJsvals(mod: IrModule): number {
   const storesJs = (e: IrExpr): boolean =>
     e.kind === "call" && sites.site(e.callee)?.kind === "set";
   const params = jsvalParams(mod, direct, results, storesJs);
+  const globals = promotableGlobals(mod, results, params, storesJs);
   let count = 0;
   for (const fn of mod.functions)
-    if (jsvalEligible(fn) && promoteFunction(fn, results, params, storesJs)) count++;
+    if (jsvalEligible(fn) && promoteFunction(fn, results, params, storesJs, globals)) count++;
   for (const fn of mod.functions) if (results.has(fn.name)) fn.returnType = JSVAL;
+  for (const global of mod.globals ?? []) if (globals.has(global.id)) global.type = JSVAL;
   return count;
 }
 
@@ -273,6 +277,22 @@ function promotableLocals(
   credits: ReadonlyMap<string, number> = new Map(),
   jsParams: ReadonlyMap<string, readonly boolean[]> = new Map(),
 ): Set<string> {
+  const candidates = localCandidates(fn, credits);
+  if (candidates.size === 0) return new Set();
+  const blocked = new Set<string>();
+  scanBindingCosts(fn, candidates, blocked, producer, jsResult, storesJs, jsParams);
+  const promoted = new Set<string>();
+  for (const [id, cost] of candidates)
+    if (!blocked.has(id) && cost.saved > cost.added) promoted.add(id);
+  return promoted;
+}
+
+/** The dyn locals of a function that may become jsval locals, parameters
+ * with their credit. */
+function localCandidates(
+  fn: IrFunction,
+  credits: ReadonlyMap<string, number>,
+): Map<string, LocalCost> {
   const candidates = new Map<string, LocalCost>();
   for (const local of fn.locals) {
     if (local.type.kind === "dyn" && local.boxed === undefined && local.tdz === undefined)
@@ -286,8 +306,69 @@ function promotableLocals(
   }
   for (const c of fn.captures ?? []) candidates.delete(c.localId);
   for (const c of fn.classCaptures ?? []) candidates.delete(c.localId);
+  return candidates;
+}
+
+/** The dyn module globals that become jsval globals: those that only
+ * dynamic-tier functions reference, by reads and plain writes, and whose
+ * saved boxing outweighs the added boxing summed over every function, the
+ * same measure as for locals. Libraries keep every global (the host reads
+ * exported bindings), and so do bindings with a temporal dead zone. */
+function promotableGlobals(
+  mod: IrModule,
+  results: ReadonlySet<string>,
+  jsParams: ReadonlyMap<string, readonly boolean[]>,
+  storesJs: (e: IrExpr) => boolean,
+): Set<string> {
+  const candidates = new Map<string, LocalCost>();
+  if (mod.lib !== undefined) return new Set();
+  for (const global of mod.globals ?? [])
+    if (global.type.kind === "dyn" && global.tdz === undefined && global.initFlag === undefined)
+      candidates.set(global.id, { saved: 0, added: 0 });
   if (candidates.size === 0) return new Set();
   const blocked = new Set<string>();
+  const producer = (e: IrExpr): boolean =>
+    isProducerCall(e) || (e.kind === "call" && results.has(e.callee));
+  for (const fn of mod.functions) {
+    if (!jsvalEligible(fn)) {
+      // A static function reads and writes the binding as a cell (I1).
+      everyStmtList(fn.body, {
+        expr: (e) => {
+          if (
+            (e.kind === "varRef" || e.kind === "assignExpr" || e.kind === "incDec") &&
+            candidates.has(e.localId)
+          )
+            blocked.add(e.localId);
+          return true;
+        },
+        stmt: (st) => {
+          if ((st.kind === "assign" || st.kind === "forOf") && candidates.has(st.localId))
+            blocked.add(st.localId);
+          return true;
+        },
+      });
+      continue;
+    }
+    scanBindingCosts(fn, candidates, blocked, producer, results.has(fn.name), storesJs, jsParams);
+  }
+  const promoted = new Set<string>();
+  for (const [id, cost] of candidates)
+    if (!blocked.has(id) && cost.saved > cost.added) promoted.add(id);
+  return promoted;
+}
+
+/** Adds to `candidates` the boxing a jsval binding saves and adds in `fn`,
+ * and to `blocked` the candidates `fn` references other than by reads and
+ * plain writes. */
+function scanBindingCosts(
+  fn: IrFunction,
+  candidates: ReadonlyMap<string, LocalCost>,
+  blocked: Set<string>,
+  producer: (e: IrExpr) => boolean,
+  jsResult: boolean,
+  storesJs: (e: IrExpr) => boolean,
+  jsParams: ReadonlyMap<string, readonly boolean[]>,
+): void {
   const write = (id: string, value: IrExpr | null, weight: number): void => {
     const cost = candidates.get(id);
     if (cost === undefined) return;
@@ -431,10 +512,6 @@ function promotableLocals(
     }
   };
   for (const s of fn.body) stmt(s, 1);
-  const promoted = new Set<string>();
-  for (const [id, cost] of candidates)
-    if (!blocked.has(id) && cost.saved > cost.added) promoted.add(id);
-  return promoted;
 }
 
 /** A statement-level sequence whose result nothing reads. */
@@ -449,6 +526,7 @@ function promoteFunction(
   results: ReadonlySet<string>,
   params: ReadonlyMap<string, readonly boolean[]>,
   storesJs: (e: IrExpr) => boolean,
+  globals: ReadonlySet<string>,
 ): boolean {
   const resultCall = (e: IrExpr): boolean => e.kind === "call" && results.has(e.callee);
   const operatorCall = (e: IrExpr): boolean => isProducerCall(e);
@@ -460,6 +538,7 @@ function promoteFunction(
     new Map(),
     params,
   );
+  for (const id of globals) promoted.add(id);
   const ownParams = params.get(fn.name);
   fn.params.forEach((p, i) => {
     if (ownParams?.[i] === true) promoted.add(p.localId);

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/dynamic-method-returns-receiver.mjs.json

@@ -0,0 +1,6 @@
+{
+ "order": [
+  "<repo>/tests/corpus/dynamic-method-returns-receiver.mjs"
+ ],
+ "diags": []
+}

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/dynamic-module-let-imports/main.ts.json

@@ -0,0 +1,7 @@
+{
+ "order": [
+  "<repo>/tests/corpus/dynamic-module-let-imports/counter.mjs",
+  "<repo>/tests/corpus/dynamic-module-let-imports/main.ts"
+ ],
+ "diags": []
+}

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/dynamic-module-let-numbers.mjs.json

@@ -0,0 +1,6 @@
+{
+ "order": [
+  "<repo>/tests/corpus/dynamic-module-let-numbers.mjs"
+ ],
+ "diags": []
+}

tests/corpus/dynamic-method-returns-receiver.mjs

@@ -0,0 +1,72 @@
+// Untyped methods that return their receiver, called through untyped
+// values: chains, identity, subclasses, overrides and partial returns.
+class Point {
+  constructor(x, y) {
+    this.x = x;
+    this.y = y;
+  }
+  move(dx, dy) {
+    this.x += dx;
+    this.y += dy;
+    return this;
+  }
+  scale(k) {
+    if (k === 0) {
+      this.x = 0;
+      this.y = 0;
+      return this;
+    }
+    this.x *= k;
+    this.y *= k;
+    return this;
+  }
+  maybe(flag) {
+    if (flag) return this;
+    return { x: -1 };
+  }
+  check(v) {
+    if (v < 0) throw new RangeError("negative " + v);
+    this.x += v;
+    return this;
+  }
+}
+class Tagged extends Point {
+  constructor(x, y, tag) {
+    super(x, y);
+    this.tag = tag;
+  }
+  move(dx, dy) {
+    super.move(dx, dy);
+    return this.tag;
+  }
+}
+class Plain extends Point {}
+const points = [new Point(1, 2), new Tagged(3, 4, "t"), new Plain(5, 6), new Point(7, 8)];
+let sum = 0;
+for (let r = 0; r < 200; r++) {
+  const p = points[r % points.length];
+  const moved = p.move(1, -1);
+  if (typeof moved === "object") {
+    sum += moved.x + moved.y;
+    if (moved !== p) throw new Error("identity lost");
+  } else sum += moved.length;
+  const scaled = p.scale(r % 3);
+  sum += scaled === p ? 1 : 0;
+  sum += p.maybe(r % 2 === 0).x;
+}
+console.log(sum);
+const chain = points[0].move(1, 1).scale(2).move(-1, 0);
+console.log(chain === points[0], chain.x, chain.y, chain instanceof Point);
+const plain = points[2];
+const r1 = plain.move(2, 2);
+r1.extra = "seen";
+console.log(plain.extra, r1 === plain, Object.keys(r1).join(","), r1.constructor.name);
+try {
+  points[3].check(1).check(-2);
+} catch (e) {
+  console.log(e instanceof RangeError, e.message, points[3].x);
+}
+const results = [];
+for (const p of points) results.push(p.scale(1));
+console.log(results.map((p, i) => p === points[i]).join(","));
+console.log(JSON.stringify(points));

tests/corpus/dynamic-module-let-imports/counter.mjs

@@ -0,0 +1,13 @@
+// Module-level bindings of an untyped module that a typed module imports.
+export let total = 0;
+let hidden = 0;
+export function run(rounds) {
+  for (let i = 0; i < rounds; i++) {
+    total += i * 3;
+    hidden += i;
+  }
+  return hidden;
+}
+export function reset(value) {
+  total = value;
+}

tests/corpus/dynamic-module-let-imports/main.ts

@@ -0,0 +1,6 @@
+// A typed module reads a live binding that untyped code updates in a loop.
+import { reset, run, total } from "./counter.mjs";
+
+console.log(run(100), total);
+reset(2.5);
+console.log(run(10), total);

tests/corpus/dynamic-module-let-numbers.mjs

@@ -0,0 +1,53 @@
+// Module-level `let` bindings that untyped code updates in loops: numbers,
+// then values of other kinds, read by functions declared before and after.
+function readTotal() {
+  return total;
+}
+let total = 0;
+let label = "n";
+let mixed = 1;
+let big = 10n;
+let counter = 0;
+var hoisted;
+function bump(n) {
+  return n * 2 + 1;
+}
+for (let i = 0; i < 1000; i++) {
+  total += bump(i);
+  label += i % 10 === 0 ? i : "";
+  mixed = mixed * 1.5 + i;
+  big += BigInt(i);
+  counter++;
+}
+console.log(total, readTotal(), label.length, mixed, big, counter);
+total = total / 0;
+console.log(total, readTotal());
+total = -0;
+console.log(Object.is(total, -0), 1 / total);
+total = NaN;
+console.log(total === total, String(total));
+mixed = "text";
+mixed += 5;
+console.log(mixed, typeof mixed);
+mixed = { value: 3 };
+mixed.value += 4;
+console.log(mixed.value, JSON.stringify(mixed));
+mixed = undefined;
+console.log(mixed, typeof mixed, mixed ?? "fallback");
+mixed = null;
+console.log(mixed === null, `${mixed}`);
+hoisted = 0;
+for (let i = 0; i < 10; i++) hoisted += i * 0.1;
+console.log(hoisted);
+const add = (n) => {
+  total = (Number.isNaN(total) ? 0 : total) + n;
+  return total;
+};
+for (let i = 0; i < 5; i++) add(i);
+console.log(total, readTotal());
+let product = 1;
+for (let i = 1; i <= 25; i++) product *= i;
+console.log(product, product > Number.MAX_SAFE_INTEGER);
+let text = 0;
+for (let i = 0; i < 3; i++) text = text + "" + i;
+console.log(text, typeof text);