Revisit only the changed sites of the optional-read fixed point (#1703)

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

@@ -121,6 +121,7 @@ import { ArrayElementStates } from "./runtime-optional-elements.js";
 import { indexReadInBounds } from "./runtime-optional-bounds.js";
 import { ArrayOwnership } from "./runtime-optional-ownership.js";
 import { RecordFlowClasses } from "./runtime-optional-records.js";
+import { SiteDependencies } from "./site-dependencies.js";
 import { StringIndexBounds } from "./string-index-bounds.js";
 import { NanCodedReads } from "./nan-coded-reads.js";
 import { RuntimeOptionalLocals } from "./runtime-optional-locals.js";
@@ -4688,10 +4689,15 @@ export class Lowerer {
     // resolveValueSymbol also flushes deferred diagnostics and applies
     // merged-namespace fences, whose authority remains ordinary lowering.
     this.checker.prefetchSymbolRoots(sourceFiles, true);
+    // The facts each fixed-point site reads, by the symbols, nodes and
+    // element keys they belong to (created with the site list below).
+    let siteDependencies: SiteDependencies | null = null;
     const symbolOf = (node: ts.Node): ts.Symbol | null => {
       const s = this.checker.getSymbolAtLocation(node);
       if (!s) return null;
-      return s.flags & ts.SymbolFlags.Alias ? this.checker.getAliasedSymbol(s) : s;
+      const symbol = s.flags & ts.SymbolFlags.Alias ? this.checker.getAliasedSymbol(s) : s;
+      siteDependencies?.read(symbol);
+      return symbol;
     };
     const signatureOf = (node: ts.Identifier): FnSig | null => {
       const symbol = symbolOf(node);
@@ -4876,7 +4882,10 @@ export class Lowerer {
     const callableSymbols = new Map<ts.Expression, ts.Symbol | null>();
     const callableSymbolOf = (node: ts.Expression): ts.Symbol | null => {
       const cached = callableSymbols.get(node);
-      if (cached !== undefined) return cached;
+      if (cached !== undefined) {
+        if (cached !== null) siteDependencies?.read(cached);
+        return cached;
+      }
       let symbol = symbolOf(ts.isPropertyAccessExpression(node) ? node.name : node);
       if (symbol === null || signatureBySymbol.has(symbol)) {
         callableSymbols.set(node, symbol);
@@ -5165,7 +5174,10 @@ export class Lowerer {
       if (ts.isPropertyAccessExpression(e) && isSettledIteratorValueRead(this, e)) return true;
       if (ts.isPropertyAccessExpression(e) && isOsConstantsSymbol(symbolOf(e.expression), this))
         return true;
-      if (ts.isCallExpression(e) && this.runtimeOptionalReduceTypes.has(e)) return true;
+      if (ts.isCallExpression(e)) {
+        siteDependencies?.read(e);
+        if (this.runtimeOptionalReduceTypes.has(e)) return true;
+      }
       if (
         ts.isCallExpression(e) &&
         e.arguments.nodes.length === 0 &&
@@ -5663,6 +5675,7 @@ export class Lowerer {
         if (fnType?.kind === "func" && bodyReturnsOptional(fn)) {
           const widenedReturn = this.runtimeOptionalType(fnType.ret);
           provenance?.note("callback-return", fn, provenanceSite);
+          siteDependencies?.read(fn);
           const previousReturn = this.runtimeOptionalFunctionReturnType(fn, fnType.ret);
           this.runtimeOptionalFunctionReturns.set(fn, widenedReturn);
           if (!typeEquals(previousReturn, widenedReturn)) changed = true;
@@ -5898,8 +5911,10 @@ export class Lowerer {
           arm.elem.kind === "union"
             ? this.armTag(arm.elem.unionId, UNDEFINED_T) >= 0
             : isUnitType(arm.elem) || arm.elem.kind === "void"
-        )
+        ) {
+          if (!optionalElementKeys.has(key)) siteDependencies?.changed(key);
           optionalElementKeys.add(key);
+        }
         keys.push(key);
       }
       return keys;
@@ -6011,7 +6026,10 @@ export class Lowerer {
       return keys;
     };
     const elementStates = new ArrayElementStates({
-      elementAdmitsUndefined: (key) => optionalElementKeys.has(key),
+      elementAdmitsUndefined: (key) => {
+        siteDependencies?.read(key);
+        return optionalElementKeys.has(key);
+      },
       ownerOf: (node) => ownership.ownerOf(node),
       allocationOwner: (node) => ownership.allocationOwner(node),
       arrayKeys: arrayKeysOf,
@@ -6084,6 +6102,8 @@ export class Lowerer {
       note: provenance
         ? (fact, key, site) => provenance.note(`element-${fact}`, site, site, key)
         : null,
+      read: (fact) => siteDependencies?.read(fact),
+      changed: (fact) => siteDependencies?.changed(fact),
     });
     // Each pass revisits every array callback site; the last visit holds
     // the decision at the fixed point.
@@ -6385,6 +6405,7 @@ export class Lowerer {
               ) {
                 if (promoteHofCallback(callback, [0, 1])) changed = true;
                 const result = this.runtimeOptionalType(this.irTypeOf(node));
+                siteDependencies?.read(node);
                 const previous = this.runtimeOptionalReduceTypes.get(node);
                 if (!previous || !typeEquals(previous, result)) {
                   this.runtimeOptionalReduceTypes.set(node, result);
@@ -6452,6 +6473,7 @@ export class Lowerer {
               )
                 continue;
               for (const target of familyBySymbol.get(symbol) ?? [symbol]) {
+                siteDependencies?.read(target);
                 const shape = signatureBySymbol.get(target)?.params[i];
                 if (!shape || shape.mode !== "required") continue;
                 const widened = addUndefined(shape.type);
@@ -6496,6 +6518,7 @@ export class Lowerer {
             if (!ts.isSpreadElement(arg) && shape && i !== restAt && fieldsOfValue(arg)) {
               const fields = fieldsOfValue(arg)!;
               for (const target of familyBySymbol.get(symbol) ?? [symbol]) {
+                siteDependencies?.read(target);
                 const parameterShape = signatureBySymbol.get(target)?.params[i];
                 const parameter = functionDeclBySymbol.get(target)?.parameters?.nodes[i];
                 if (!parameterShape || parameterShape.type.kind !== "record") continue;
@@ -6541,6 +6564,7 @@ export class Lowerer {
               // argument can make its elements undefined, not the pack.
               // Keep every virtual implementation on the same array ABI.
               for (const target of familyBySymbol.get(symbol) ?? [symbol]) {
+                siteDependencies?.read(target);
                 const rest = signatureBySymbol.get(target)?.params[i];
                 if (rest?.mode !== "rest" || rest.type.kind !== "array") continue;
                 const widened = arrayOf(addUndefined(rest.type.elem));
@@ -6562,6 +6586,7 @@ export class Lowerer {
             // downstream calls and returns see the same optional value.
             const family = familyBySymbol.get(symbol);
             for (const target of family ?? [symbol]) {
+              siteDependencies?.read(target);
               const parameterShape = signatureBySymbol.get(target)?.params[i];
               if (!parameterShape) continue;
               parameterShape.type = addUndefined(parameterShape.type);
@@ -6592,8 +6617,21 @@ export class Lowerer {
           });
         }
       };
-      for (const node of sitesOf(sf)) scan(node);
-      return changed;
+      // A site whose facts are unchanged since its last visit would compute
+      // that visit's result again, except in a full pass.
+      const base = siteBase.get(sf)!;
+      const dependencies = siteDependencies!;
+      const sites = sitesOf(sf);
+      let fileChanged = false;
+      for (let i = 0; i < sites.length; i++) {
+        if (!fullPass && !dependencies.isStale(base + i)) continue;
+        changed = false;
+        dependencies.begin(base + i);
+        scan(sites[i]!);
+        dependencies.end(changed);
+        if (changed) fileChanged = true;
+      }
+      return fileChanged;
     };
     const scanReturns = (symbol: ts.Symbol, decl: ts.FunctionLikeDeclaration): boolean => {
       let changed = false;
@@ -6643,23 +6681,41 @@ export class Lowerer {
       optionalSymbols.add(symbol);
       this.earlyReadVars.add(symbol);
     }
-    let changed = true;
-    while (changed) {
+    // Number the sites of every file. The first and the last pass visit
+    // every site; the passes between visit only the sites whose facts
+    // changed (SiteDependencies). When such a pass changes nothing, a full
+    // pass confirms the fixed point, and the iteration resumes if it
+    // changes a fact after all.
+    const siteBase = new Map<ts.SourceFile, number>();
+    let siteCount = 0;
+    for (const sf of sourceFiles) {
+      siteBase.set(sf, siteCount);
+      siteCount += sitesOf(sf).length;
+    }
+    siteDependencies = new SiteDependencies(siteCount);
+    let fullPass = true;
+    for (let changed = true, first = true; changed || !fullPass; first = false) {
+      fullPass = first || !changed;
       changed = false;
       for (const [symbol, field] of staticFieldsBySymbol) {
         if (!optionalSymbols.has(symbol) && mayBeOptional(field.initializer)) {
           optionalSymbols.add(symbol);
+          siteDependencies.changed(symbol);
           provenance?.note("static-field", field.initializer.parent, field.initializer);
           changed = true;
         }
       }
       for (const sf of sourceFiles) if (scanFile(sf)) changed = true;
-      for (const [symbol, decl] of functionDeclBySymbol)
-        if (scanReturns(symbol, decl)) changed = true;
+      for (const [symbol, decl] of functionDeclBySymbol) {
+        if (!scanReturns(symbol, decl)) continue;
+        siteDependencies.changed(symbol);
+        changed = true;
+      }
       const noteGenericReturn = (symbol: ts.Symbol, decl: ts.FunctionLikeDeclaration): void => {
         if (optionalReturns.has(symbol) || isJsSourceFile(ts.getSourceFileOfNode(decl))) return;
         if (!returnsOf(decl).some(genericMayBeOptional)) return;
         optionalReturns.add(symbol);
+        siteDependencies!.changed(symbol);
         this.runtimeOptionalGenericReturns.add(decl);
         provenance?.note("return", decl.name ?? decl, decl);
         changed = true;
@@ -6696,6 +6752,7 @@ export class Lowerer {
               : previous;
         if (!typeEquals(previous, next)) {
           this.runtimeOptionalFunctionReturns.set(fn, next);
+          siteDependencies.changed(fn);
           provenance?.note("lambda-return", fn, fn.body);
           changed = true;
         }
@@ -6713,10 +6770,12 @@ export class Lowerer {
           if (symbol && next.kind === "union" && this.armTag(next.unionId, UNDEFINED_T) >= 0) {
             if (!optionalReturns.has(symbol)) {
               optionalReturns.add(symbol);
+              siteDependencies.changed(symbol);
               changed = true;
             }
           } else if (symbol && !arithmeticReturns.has(symbol)) {
             arithmeticReturns.set(symbol, next);
+            siteDependencies.changed(symbol);
             changed = true;
           }
         }
@@ -6733,6 +6792,7 @@ export class Lowerer {
             : this.runtimeOptionalType(previous);
         if (!typeEquals(previous, next)) {
           this.runtimeOptionalFunctionReturns.set(fn, next);
+          siteDependencies.changed(fn);
           changed = true;
         }
       }
@@ -6746,6 +6806,7 @@ export class Lowerer {
         for (const symbol of optional ? family : []) {
           if (!optionalReturns.has(symbol)) {
             optionalReturns.add(symbol);
+            siteDependencies.changed(symbol);
             provenance?.note(
               "return",
               functionDeclBySymbol.get(symbol),
@@ -6768,10 +6829,14 @@ export class Lowerer {
           const previous = arithmeticReturns.get(symbol);
           if (previous && typeEquals(previous, joined)) continue;
           arithmeticReturns.set(symbol, joined);
+          siteDependencies.changed(symbol);
           changed = true;
         }
       }
     }
+    // The site facts serve only the fixed point; release them for the rest
+    // of lowering.
+    siteDependencies = null;
     // NaN-coded slots convert optional values on entry instead of holding
     // the union, so promotion never applies to them.
     for (const symbol of [...optionalSymbols])

packages/compiler/src/frontend/lowering/runtime-optional-elements.ts

@@ -39,8 +39,16 @@ export interface ElementStateHost {
   allocationOwner(node: ts.Expression): ts.Symbol | null;
   /** Debug provenance for a new fact. */
   note: ((fact: "undefined" | "holes", key: string, site: ts.Node | undefined) => void) | null;
+  /** The scan read, or added to, the facts of an element key, of an owning
+   * slot, or of every array (ANY_ARRAY), for dependency tracking
+   * (SiteDependencies). */
+  read?(fact: string | ts.Symbol): void;
+  changed?(fact: string | ts.Symbol): void;
 }
 
+/** The fact name of the facts shared by every array. */
+const ANY_ARRAY = "*";
+
 /** What an array value can contain besides present values. */
 interface Contents {
   undefined: boolean;
@@ -179,8 +187,10 @@ export class ArrayElementStates {
       if (built !== undefined) return built;
     }
     const owner = this.host.ownerOf(e);
-    if (owner !== null)
+    if (owner !== null) {
+      this.host.read?.(owner);
       return { undefined: this.undefinedOwners.has(owner), holes: this.holeyOwners.has(owner) };
+    }
     return this.stored(this.host.arrayKeys(e));
   }
 
@@ -317,7 +327,9 @@ export class ArrayElementStates {
   /** The facts recorded for arrays with these element ABIs. */
   private stored(keys: readonly string[] | null): Contents | null {
     if (this.ownFactsOnly) return NOTHING;
+    this.host.read?.(ANY_ARRAY);
     if (keys === null) return null;
+    for (const key of keys) this.host.read?.(key);
     return {
       undefined:
         this.anyUndefined ||
@@ -333,6 +345,7 @@ export class ArrayElementStates {
     const keys = this.host.arrayKeys(node);
     if (keys !== null && keys.length === 0) return false;
     const contents = this.contents(node);
+    this.host.read?.(ANY_ARRAY);
     return contents === null
       ? this.anyUndefined || this.anyHoley
       : contents.undefined || contents.holes;
@@ -368,6 +381,7 @@ export class ArrayElementStates {
         this.holeyOwners.add(owner);
         changed = true;
       }
+      if (changed) this.host.changed?.(owner);
       return changed;
     }
     let changed = false;
@@ -406,6 +420,7 @@ export class ArrayElementStates {
     if (owner === null) return this.addUndefined(this.host.arrayKeys(receiver));
     if (this.undefinedOwners.has(owner)) return false;
     this.undefinedOwners.add(owner);
+    this.host.changed?.(owner);
     return true;
   }
 
@@ -415,6 +430,7 @@ export class ArrayElementStates {
     if (owner === null) return this.addHoles(this.host.arrayKeys(receiver));
     if (this.holeyOwners.has(owner)) return false;
     this.holeyOwners.add(owner);
+    this.host.changed?.(owner);
     return true;
   }
 
@@ -422,6 +438,7 @@ export class ArrayElementStates {
     if (keys === null) {
       if (this.anyUndefined) return false;
       this.anyUndefined = true;
+      this.host.changed?.(ANY_ARRAY);
       this.host.note?.("undefined", "*", this.site);
       return true;
     }
@@ -429,6 +446,7 @@ export class ArrayElementStates {
     for (const key of keys) {
       if (this.undefinedKeys.has(key)) continue;
       this.undefinedKeys.add(key);
+      this.host.changed?.(key);
       this.host.note?.("undefined", key, this.site);
       changed = true;
     }
@@ -439,6 +457,7 @@ export class ArrayElementStates {
     if (keys === null) {
       if (this.anyHoley) return false;
       this.anyHoley = true;
+      this.host.changed?.(ANY_ARRAY);
       this.host.note?.("holes", "*", this.site);
       return true;
     }
@@ -446,6 +465,7 @@ export class ArrayElementStates {
     for (const key of keys) {
       if (this.holeyKeys.has(key)) continue;
       this.holeyKeys.add(key);
+      this.host.changed?.(key);
       this.host.note?.("holes", key, this.site);
       changed = true;
     }

packages/compiler/src/frontend/lowering/site-dependencies.test.ts

@@ -0,0 +1,56 @@
+import { expect, test } from "vitest";
+import { SiteDependencies } from "./site-dependencies.js";
+
+test("a changed fact makes the sites that read it stale", () => {
+  const dependencies = new SiteDependencies(3);
+  const a = {};
+  const b = "key";
+  dependencies.begin(0);
+  dependencies.read(a);
+  dependencies.end(false);
+  dependencies.begin(1);
+  dependencies.read(b);
+  dependencies.end(false);
+  expect([0, 1, 2].map((site) => dependencies.isStale(site))).toEqual([false, false, false]);
+  dependencies.changed(b);
+  expect([0, 1, 2].map((site) => dependencies.isStale(site))).toEqual([false, true, false]);
+  dependencies.begin(1);
+  expect(dependencies.isStale(1)).toBe(false);
+});
+
+test("a visit that changes a fact makes every reader of its facts stale", () => {
+  const dependencies = new SiteDependencies(3);
+  const shared = {};
+  const own = {};
+  dependencies.begin(0);
+  dependencies.read(shared);
+  dependencies.end(false);
+  dependencies.begin(2);
+  dependencies.read(own);
+  dependencies.end(false);
+  dependencies.begin(1);
+  dependencies.read(shared);
+  dependencies.read(shared);
+  dependencies.end(true);
+  expect([0, 1, 2].map((site) => dependencies.isStale(site))).toEqual([true, true, false]);
+});
+
+test("reads outside a visit are not recorded", () => {
+  const dependencies = new SiteDependencies(1);
+  const fact = {};
+  dependencies.read(fact);
+  dependencies.changed(fact);
+  expect(dependencies.isStale(0)).toBe(false);
+});
+
+test("a site visited again is listed once per fact", () => {
+  const dependencies = new SiteDependencies(3);
+  const fact = {};
+  for (const site of [0, 2, 1, 2, 0]) {
+    dependencies.begin(site);
+    dependencies.read(fact);
+    dependencies.end(false);
+  }
+  dependencies.changed(fact);
+  expect([0, 1, 2].map((site) => dependencies.isStale(site))).toEqual([true, true, true]);
+});

packages/compiler/src/frontend/lowering/site-dependencies.ts

@@ -0,0 +1,79 @@
+/** The facts each site of a fixed-point scan read on its last visit, so that
+ * a pass visits again only the sites whose facts changed since.
+ *
+ * Sites are numbered from 0. A fact is any value usable as a Map key: a
+ * symbol, a syntax node, or a string that names a derived fact. A visit
+ * reports what it reads (`read`). When a visit changes a fact, every site
+ * that read one of the facts this visit read is visited again: the facts a
+ * visit changes belong to what it read. Facts changed outside a visit, or
+ * not read before they change, are reported with `changed`.
+ *
+ * A site that is not visited again would compute what its last visit
+ * computed, as long as every fact it reads is reported. The scan confirms
+ * that with a pass over every site once a pass changes nothing, and resumes
+ * when that pass changes a fact: a fact this tracking misses costs passes,
+ * not a different result. */
+export class SiteDependencies {
+  /** Each fact's entry: the visit that last read it, so that a visit lists a
+   * fact once, followed by the sites that read it, ascending. */
+  private readonly readers = new Map<unknown, number[]>();
+  private readonly stale: Uint8Array;
+  private site = -1;
+  private visit = 0;
+  private facts: unknown[] = [];
+
+  constructor(siteCount: number) {
+    this.stale = new Uint8Array(siteCount);
+  }
+
+  /** A fact that site `index` read changed since its last visit. */
+  isStale(index: number): boolean {
+    return this.stale[index] === 1;
+  }
+
+  /** Starts a visit of site `index`. */
+  begin(index: number): void {
+    this.site = index;
+    this.visit++;
+    this.stale[index] = 0;
+    this.facts = [];
+  }
+
+  /** Ends the visit; `changed` when it changed a fact. */
+  end(changed: boolean): void {
+    this.site = -1;
+    if (changed) for (const fact of this.facts) this.changed(fact);
+  }
+
+  /** The site being visited read `fact`. */
+  read(fact: unknown): void {
+    if (this.site < 0) return;
+    const entry = this.readers.get(fact);
+    if (entry === undefined) {
+      this.readers.set(fact, [this.visit, this.site]);
+      this.facts.push(fact);
+      return;
+    }
+    if (entry[0] === this.visit) return;
+    entry[0] = this.visit;
+    this.facts.push(fact);
+    // A pass visits sites in ascending order, so a site is new at the end;
+    // a site visited again is usually listed already.
+    const site = this.site;
+    let low = 1;
+    let high = entry.length;
+    if (entry[high - 1]! < site) low = high;
+    while (low < high) {
+      const middle = (low + high) >>> 1;
+      if (entry[middle]! < site) low = middle + 1;
+      else high = middle;
+    }
+    if (entry[low] !== site) entry.splice(low, 0, site);
+  }
+
+  /** `fact` changed: every site that read it is visited again. */
+  changed(fact: unknown): void {
+    const entry = this.readers.get(fact);
+    if (entry !== undefined) for (let i = 1; i < entry.length; i++) this.stale[entry[i]!] = 1;
+  }
+}