Find escaping functions with typed IR walks (#915)

packages/compiler/src/backend/llvm/constant-callbacks.ts

@@ -1,5 +1,6 @@
 import { isBorrowSafeMath } from "../../ir/analysis.js";
 import type { IrExpr, IrFunction, IrLibFn, IrModule, IrStmt } from "../../ir/ir.js";
+import { exprFunctionReference, moduleFunctionReferences } from "../../ir/function-references.js";
 import { everyStmtList } from "../../ir/traverse.js";
 import type { CallLifetimes } from "./call-lifetimes.js";
 
@@ -38,31 +39,14 @@ export function findConstantCallbacks(
   );
   const globals = new Map<string, IrExpr>();
   const writes = new Set<string>();
-  const escaped = new Set<string>([
-    mod.entry,
-    ...(mod.lib?.exports.map((entry) => entry.fnName) ?? []),
-  ]);
+  // A function named anywhere but as a direct call's callee can be entered
+  // by a generated adapter or the runtime, which a direct-call census cannot
+  // specialize: the entry, exports, class helpers and methods here, and
+  // closures and retag helpers in the bodies (collected by the walk below).
+  const escaped = new Set<string>(moduleFunctionReferences(mod));
   for (const cls of mod.classes ?? [])
     for (const method of ["constructor", ...(cls.methods ?? [])])
       escaped.add(`%${cls.name}.${method}`);
-  // Metadata and non-call expression references can expose an entry to a
-  // generated adapter or runtime callback. A direct-call census cannot
-  // specialize those entries. Conservatively reject every such reference,
-  // including future metadata that names a function.
-  function externalReferences(value: unknown): void {
-    if (typeof value === "string") {
-      if (functions.has(value)) escaped.add(value);
-    } else if (Array.isArray(value)) value.forEach(externalReferences);
-    else if (value && typeof value === "object") {
-      const object = value as Record<string, unknown>;
-      for (const key in object) {
-        if (object["kind"] === "call" && key === "callee") continue;
-        externalReferences(object[key]);
-      }
-    }
-  }
-  const { functions: _functions, ...metadata } = mod;
-  externalReferences(metadata);
   const incoming = new Map<string, { owner: string; args: IrExpr[] }[]>();
   // Where each lifted function's environment is created, for resolving its
   // immutable captures in the creating scope.
@@ -100,8 +84,9 @@ export function findConstantCallbacks(
         }
         if (expr.kind === "callValue" && expr.callee.kind === "varRef")
           callees.add(expr.callee.localId);
+        const referenced = exprFunctionReference(expr);
+        if (referenced !== undefined) escaped.add(referenced);
         if (expr.kind === "closure") {
-          escaped.add(expr.fnName);
           const sites = closureSites.get(expr.fnName) ?? [];
           sites.push({ owner: fn.name, captures: expr.captures });
           closureSites.set(expr.fnName, sites);
@@ -111,7 +96,6 @@ export function findConstantCallbacks(
         return true;
       },
     });
-    externalReferences(fn.body);
     definitions.set(fn.name, defs);
     constDefinitions.set(fn.name, consts);
     indirectCallees.set(fn.name, callees);

packages/compiler/src/ir/function-references.test.ts

@@ -0,0 +1,340 @@
+import { expect, test } from "vitest";
+import { exprFunctionReference, moduleFunctionReferences } from "./function-references.js";
+import { functionCallFacts } from "./int32-slots.js";
+import {
+  F64,
+  VOID,
+  type IrExpr,
+  type IrFunction,
+  type IrModule,
+  type IrStmt,
+  type IrType,
+  type SrcLoc,
+} from "./ir.js";
+import { everyStmtList } from "./traverse.js";
+
+/* Every plain `string` field of the IR, classified. The analyses that ask
+ * which functions escape (constant-callbacks.ts, int32-slots.ts) read only
+ * the fields classified "reference" here, through function-references.ts.
+ * A new string field anywhere in the IR fails to compile until it is
+ * classified, and a new "reference" fails the tests below until the shared
+ * helpers report it. */
+type Role =
+  /** Names a module function as a value or a runtime entry. */
+  | "reference"
+  /** A direct call's callee. */
+  | "callee"
+  /** A class method name: the function is `%<class>.<method>`. */
+  | "method"
+  /** A verified contract fact about a function, not a way to reach it. */
+  | "fact"
+  /** Never a module function's name. */
+  | "other";
+
+type Leaf = IrExpr | IrStmt | IrType | SrcLoc;
+type PathsOf<V, P extends string> = V extends string
+  ? string extends V
+    ? P
+    : never
+  : V extends readonly (infer E)[]
+    ? PathsOf<E, `${P}[]`>
+    : V extends Leaf
+      ? never
+      : V extends object
+        ? StringFields<V, `${P}.`>
+        : never;
+type StringFields<T, P extends string = ""> = {
+  [K in keyof T & string]-?: PathsOf<NonNullable<T[K]>, `${P}${K}`>;
+}[keyof T & string];
+type Classified<U extends { kind: string }> = {
+  [K in U["kind"] as [StringFields<Extract<U, { kind: K }>>] extends [never] ? never : K]: Record<
+    StringFields<Extract<U, { kind: K }>>,
+    Role
+  >;
+};
+
+const EXPR_FIELDS: Classified<IrExpr> = {
+  numLit: { spelling: "other" },
+  strLit: { value: "other" },
+  moduleNsRef: { moduleId: "other" },
+  varRef: { localId: "other" },
+  incDec: { localId: "other" },
+  fieldIncDec: { className: "other", field: "other" },
+  assignExpr: { localId: "other" },
+  orDefault: { retag: "reference" },
+  optChain: { id: "other" },
+  chainRecv: { id: "other" },
+  regexLit: { pattern: "other", flags: "other" },
+  templateStrings: { key: "other", "cooked[]": "other" },
+  call: { callee: "callee" },
+  ffiCall: { import: "other" },
+  closure: { fnName: "reference", "captures[]": "other" },
+  new: { className: "other" },
+  classRef: { className: "other", "captures[]": "other" },
+  instanceOf: { className: "other" },
+  virtualCall: { className: "other", method: "method" },
+  fieldGet: { className: "other", field: "other" },
+  recordLit: { "fields[].name": "other" },
+  recordClone: { "overrides[].name": "other" },
+  recordGet: { shapeId: "other", field: "other" },
+  recordKeyGet: { shapeId: "other" },
+  dynFrom: { fnName: "other" },
+  dynCall: { calleeName: "other", "spreads[].what": "other" },
+  dynInvoke: { method: "other", calleeName: "other" },
+  dynHasKey: { key: "other" },
+  dynDestrCheck: { spelling: "other", firstProp: "other" },
+  dynIterN: { notIterableMessage: "other" },
+  recordOvfKeys: { shapeId: "other" },
+  recordOvfHas: { shapeId: "other" },
+  recordHas: { shapeId: "other", field: "other" },
+  fieldAbsent: { unionId: "other" },
+  unionWrap: { unionId: "other" },
+  unionRetype: { unionId: "other", fromId: "other" },
+  unionFuncEq: { unionId: "other" },
+  caughtTest: { className: "other" },
+  caughtCheck: { className: "other" },
+  unionNarrow: { unionId: "other" },
+  unionDisc: { unionId: "other", field: "other" },
+  unionKeyGet: { unionId: "other" },
+  unionIsTag: { unionId: "other" },
+  unionEq: { unionId: "other" },
+  jsonStringify: { indent: "other" },
+};
+
+const STMT_FIELDS: Classified<IrStmt> = {
+  varDecl: { localId: "other" },
+  assign: { localId: "other" },
+  while: { "labels[]": "other" },
+  doWhile: { "labels[]": "other" },
+  switch: { "labels[]": "other" },
+  for: { "labels[]": "other" },
+  forOf: { localId: "other", "labels[]": "other" },
+  fieldSet: { className: "other", field: "other" },
+  recordSet: { shapeId: "other", field: "other" },
+  recordKeySet: { shapeId: "other" },
+  recordKeyDelete: { shapeId: "other" },
+  break: { label: "other" },
+  continue: { label: "other" },
+  block: { "labels[]": "other" },
+  runtimeFence: { code: "other", message: "other" },
+  rethrow: { localId: "other" },
+  tryCatch: { catchLocalId: "other" },
+};
+
+const TYPE_FIELDS: Classified<IrType> = {
+  object: { className: "other" },
+  record: { shapeId: "other" },
+  classval: { className: "other" },
+  moduleNs: { moduleId: "other" },
+  union: { unionId: "other" },
+};
+
+const FUNCTION_FIELDS: Record<StringFields<Omit<IrFunction, "body">>, Role> = {
+  name: "other",
+  sourceName: "other",
+  asyncCacheGlobal: "other",
+  asyncCycleCacheGlobal: "other",
+  "params[].name": "other",
+  "params[].localId": "other",
+  "locals[].name": "other",
+  "locals[].id": "other",
+  "captures[].name": "other",
+  "captures[].localId": "other",
+  "classCaptures[].name": "other",
+  "classCaptures[].localId": "other",
+};
+
+const MODULE_FIELDS: Record<StringFields<Omit<IrModule, "functions">>, Role> = {
+  sourceFile: "other",
+  entry: "reference",
+  "classes[].name": "other",
+  "classes[].base": "other",
+  "classes[].jsName": "other",
+  "classes[].genericOf": "other",
+  "classes[].baseValueGlobal": "other",
+  "classes[].prototypeDataHelper": "reference",
+  "classes[].instancePrototypeHelper": "reference",
+  "classes[].localCaptures[].name": "other",
+  "classes[].localCaptures[].localId": "other",
+  "classes[].symbolFields[].field": "other",
+  "classes[].symbolFields[].globalId": "other",
+  "classes[].fields[].name": "other",
+  "classes[].methods[]": "method",
+  "classes[].abstractMethods[]": "method",
+  "globals[].id": "other",
+  "globals[].name": "other",
+  "globals[].initFlag": "other",
+  "records[].id": "other",
+  "records[].fields[].name": "other",
+  "records[].declaredOrder[]": "other",
+  "unions[].id": "other",
+  "unions[].discriminant.field": "other",
+  "unions[].discriminant.cases[].values[]": "other",
+  "ffiImports[].name": "other",
+  "ffiImports[].symbol": "other",
+  "ffiImports[].library": "other",
+  "ffiImports[].callbackTarget": "other",
+  "ffiImports[].params[].context": "other",
+  "ffiImports[].params[].callback.id": "other",
+  "ffiImports[].params[].callback.release": "other",
+  "ffiImports[].params[].callback.params[].context": "other",
+  "lib.profileName": "other",
+  "lib.prefix": "other",
+  "lib.initSymbol": "other",
+  "lib.sinkRegisterSymbol": "other",
+  "lib.collectSymbol": "other",
+  "lib.resultResetSymbol": "other",
+  "lib.callbackRegisterSymbol": "other",
+  "lib.callbacks[].name": "other",
+  "lib.callbacks[].unregisteredTrap": "other",
+  "lib.exports[].symbol": "other",
+  "lib.exports[].fnName": "reference",
+  "lib.exports[].inboundBytesTrap": "other",
+  "lib.exports[].inboundIntTrap": "other",
+  "lib.trapOverlays[].code": "other",
+  "lib.trapOverlays[].teaching": "other",
+  "lib.trapOverlays[].remediation": "other",
+  "lib.identity.buildIdSymbol": "other",
+  "lib.identity.abiVersionSymbol": "other",
+  "lib.identity.buildId": "other",
+  "contracts.int32Fields[]": "other",
+  "contracts.int32Params[]": "other",
+  "contracts.int32Returns[]": "fact",
+  "contracts.int32Locals[]": "other",
+  "contracts.noThrow[]": "fact",
+  "contracts.acyclic[]": "other",
+  "contracts.immortalGlobals[]": "other",
+  "contracts.programLifetime[]": "other",
+};
+
+const loc = { file: "references.ts", start: 0, end: 0 };
+const TARGET = "%m0.target";
+const target: IrFunction = {
+  name: TARGET,
+  params: [{ localId: "n", name: "n", type: F64 }],
+  locals: [{ id: "n", name: "n", type: F64, mutable: true }],
+  returnType: VOID,
+  body: [],
+  loc,
+};
+const entry = (body: IrStmt[] = []): IrFunction => ({
+  name: "%main",
+  params: [],
+  locals: [],
+  returnType: VOID,
+  body,
+  loc,
+});
+const module = (body: IrStmt[] = []): IrModule => ({
+  irVersion: 16,
+  sourceFile: loc.file,
+  entry: "%main",
+  functions: [target, entry(body)],
+});
+const exprStmt = (expr: IrExpr): IrStmt => ({ kind: "exprStmt", expr, loc });
+
+/** A module whose only mention of TARGET, other than its definition, is
+ * through the given reference field. */
+const REFERENCE_FIXTURES: Record<string, () => IrModule> = {
+  "closure.fnName": () =>
+    module([
+      exprStmt({
+        kind: "closure",
+        fnName: TARGET,
+        captures: [],
+        type: { kind: "func", params: [F64], ret: VOID },
+        loc,
+      }),
+    ]),
+  "orDefault.retag": () =>
+    module([
+      exprStmt({
+        kind: "orDefault",
+        left: { kind: "numLit", value: 0, type: F64, loc },
+        right: { kind: "numLit", value: 1, type: F64, loc },
+        retag: TARGET,
+        type: F64,
+        loc,
+      }),
+    ]),
+  entry: () => ({ ...module(), entry: TARGET }),
+  "classes[].prototypeDataHelper": () => ({
+    ...module(),
+    classes: [{ name: "%m0.Cell", fields: [], prototypeDataHelper: TARGET, loc }],
+  }),
+  "classes[].instancePrototypeHelper": () => ({
+    ...module(),
+    classes: [{ name: "%m0.Cell", fields: [], instancePrototypeHelper: TARGET, loc }],
+  }),
+  "lib.exports[].fnName": () => ({
+    ...module(),
+    lib: {
+      profileName: "references",
+      prefix: "ref_",
+      initSymbol: "ref_init",
+      sinkRegisterSymbol: "ref_sink",
+      collectSymbol: null,
+      resultResetSymbol: null,
+      threadInstances: false,
+      exports: [{ symbol: "ref_target", fnName: TARGET, params: ["f64"], returns: "void" }],
+      trapOverlays: [],
+    },
+  }),
+};
+
+function referencesOf(mod: IrModule): Set<string> {
+  const names = new Set(moduleFunctionReferences(mod));
+  for (const fn of mod.functions)
+    everyStmtList(fn.body, {
+      expr: (expr) => {
+        const name = exprFunctionReference(expr);
+        if (name !== undefined) names.add(name);
+        return true;
+      },
+      stmt: () => true,
+    });
+  return names;
+}
+
+const referencePaths = (table: object, prefix: string): string[] =>
+  Object.entries(table).flatMap(([key, value]) =>
+    typeof value === "string"
+      ? value === "reference"
+        ? [prefix + key]
+        : []
+      : referencePaths(value as object, `${key}.`),
+  );
+
+test("every IR field that can name a function has a reference fixture", () => {
+  const references = [
+    ...referencePaths(EXPR_FIELDS, ""),
+    ...referencePaths(STMT_FIELDS, ""),
+    ...referencePaths(TYPE_FIELDS, ""),
+    ...referencePaths(FUNCTION_FIELDS, "functions[]."),
+    ...referencePaths(MODULE_FIELDS, ""),
+  ];
+  expect(references.sort()).toEqual(Object.keys(REFERENCE_FIXTURES).sort());
+});
+
+test.each(Object.keys(REFERENCE_FIXTURES))("%s escapes its function", (path) => {
+  const mod = REFERENCE_FIXTURES[path]!();
+  expect(referencesOf(mod).has(TARGET)).toBe(true);
+  const facts = functionCallFacts(mod, new Map(), (name) => [name]);
+  expect(facts.valueEscaped.has(TARGET)).toBe(true);
+});
+
+test("direct calls and look-alike strings do not escape a function", () => {
+  const mod = module([
+    exprStmt({ kind: "call", callee: TARGET, args: [], type: VOID, loc }),
+    exprStmt({ kind: "strLit", value: TARGET, type: { kind: "string" }, loc }),
+    exprStmt({
+      kind: "dynFrom",
+      value: { kind: "numLit", value: 0, type: F64, loc },
+      fnName: TARGET,
+      type: { kind: "dyn" },
+      loc,
+    }),
+  ]);
+  expect(referencesOf(mod).has(TARGET)).toBe(false);
+  expect(functionCallFacts(mod, new Map(), (name) => [name]).valueEscaped.has(TARGET)).toBe(false);
+});

packages/compiler/src/ir/function-references.ts

@@ -0,0 +1,26 @@
+import type { IrExpr, IrModule } from "./ir.js";
+
+/** The function an expression names as a value rather than as a direct
+ * call's callee: a closure's code, or the union→union retag helper an
+ * `orDefault` hands its left box to. These are the only expression fields
+ * that hold a module function's name besides `call.callee`
+ * (function-references.test.ts classifies every string field of the IR). */
+export function exprFunctionReference(expr: IrExpr): string | undefined {
+  if (expr.kind === "closure") return expr.fnName;
+  if (expr.kind === "orDefault") return expr.retag;
+  return undefined;
+}
+
+/** Functions the module's metadata names: the entry, library exports and
+ * the class prototype helpers the runtime calls. Class methods and
+ * constructors follow the `%<class>.<method>` convention instead of being
+ * named here; callers that need them derive them from `mod.classes`. */
+export function moduleFunctionReferences(mod: IrModule): string[] {
+  const names = [mod.entry];
+  for (const entry of mod.lib?.exports ?? []) names.push(entry.fnName);
+  for (const cls of mod.classes ?? []) {
+    if (cls.prototypeDataHelper !== undefined) names.push(cls.prototypeDataHelper);
+    if (cls.instancePrototypeHelper !== undefined) names.push(cls.instancePrototypeHelper);
+  }
+  return names;
+}

packages/compiler/src/ir/int32-slots.ts

@@ -7,8 +7,9 @@ import {
   type IntegerRanges,
   type IntegerSlotFacts,
 } from "./integer-ranges.js";
+import { exprFunctionReference, moduleFunctionReferences } from "./function-references.js";
 import { findInitializerBindings, withInitializerBindings } from "./initializer-bindings.js";
-import { everyExprChild, everyStmtChild, everyTypeChild } from "./traverse.js";
+import { everyExprChild, everyStmtChild, everyStmtList, everyTypeChild } from "./traverse.js";
 
 /** Whole-program int32 specialization of numeric storage.
  *
@@ -265,18 +266,6 @@ const DYN_READ_ONLY = new Set<string>([
   "text.decodeStream",
   "worker.new",
 ]);
-/** IR metadata keys whose strings never name a function. */
-const NON_REFERENCE_KEYS = new Set([
-  "loc",
-  "source",
-  "name",
-  "sourceName",
-  "jsName",
-  "localId",
-  "type",
-  "returnType",
-]);
-
 function walkFunction(
   fn: IrFunction,
   visit: { expr?: (e: IrExpr) => void; stmt?: (s: IrStmt) => void },
@@ -480,36 +469,8 @@ export function functionCallFacts(
   virtualTargets: Set<string>;
   runtimeVirtual: (name: string) => boolean;
 } {
-  const functions = new Map(mod.functions.map((fn) => [fn.name, fn]));
-  const escaped = new Set<string>();
+  const escaped = new Set<string>(moduleFunctionReferences(mod));
   const virtualTargets = new Set<string>();
-  const seenObjects = new Set<object>();
-  const scan = (
-    value: unknown,
-    key: string | null,
-    parent: Record<string, unknown> | null,
-  ): void => {
-    if (typeof value === "string") {
-      if (functions.has(value)) {
-        const direct =
-          key === "callee" && parent?.kind === "call"
-            ? true
-            : key === "value" && parent?.kind === "strLit";
-        if (!direct) escaped.add(value);
-      }
-      return;
-    }
-    if (value === null || typeof value !== "object" || seenObjects.has(value)) return;
-    seenObjects.add(value);
-    if (Array.isArray(value)) {
-      for (const item of value) scan(item, null, null);
-      return;
-    }
-    const record = value as Record<string, unknown>;
-    for (const [k, v] of Object.entries(record)) if (!NON_REFERENCE_KEYS.has(k)) scan(v, k, record);
-  };
-  scan(mod, null, null);
-  if (typeof mod.entry === "string") escaped.add(mod.entry);
   // Virtual dispatch: a `virtualCall` can reach the method's implementation
   // in any class of the receiver's subtree or the one it inherits. Classes
   // under a runtime-provided root keep every method escaped: the runtime
@@ -525,6 +486,8 @@ export function functionCallFacts(
   for (const fn of mod.functions)
     walkFunction(fn, {
       expr: (e) => {
+        const referenced = exprFunctionReference(e);
+        if (referenced !== undefined) escaped.add(referenced);
         if (e.kind !== "virtualCall") return;
         for (const name of subtree(e.className)) virtualTargets.add(`%${name}.${e.method}`);
         for (let c = classes.get(e.className); c; c = c.base ? classes.get(c.base) : undefined)
@@ -825,13 +788,47 @@ export function analyzeInt32Slots(mod: IrModule): Int32Slots {
           queue.push(fnName);
         }
   };
+  const candidateArgs = (callee: IrFunction | undefined, args: number, offset: number): boolean => {
+    if (!callee) return false;
+    for (let i = 0; i < args; i++) {
+      const param = callee.params[i + offset];
+      if (param && slots.hasParam(callee.name, param.localId)) return true;
+    }
+    return false;
+  };
+  const hasCandidateWrite = (fn: IrFunction): boolean => {
+    const box = (localId: string): boolean => slots.hasBox(slots.boxGroup(fn.name, localId));
+    const returns = slots.hasReturn(fn.name);
+    return !everyStmtList(fn.body, {
+      stmt: (s) => {
+        if (s.kind === "assign") return !box(s.localId);
+        if (s.kind === "varDecl") return s.init === null || !box(s.localId);
+        if (s.kind === "fieldSet") {
+          const family = slots.family(s.className, s.field);
+          return !family || !slots.hasFamily(family);
+        }
+        return !(s.kind === "return" && returns);
+      },
+      expr: (e) => {
+        if (e.kind === "call") return !candidateArgs(functions.get(e.callee), e.args.length, 0);
+        if (e.kind === "new")
+          return !candidateArgs(functions.get(`%${e.className}.constructor`), e.args.length, 1);
+        if (e.kind === "assignExpr") return !box(e.localId);
+        return true;
+      },
+    });
+  };
   const provenByFunction = new Map<string, IrExpr[]>();
   while (queue.length) {
     const fn = functions.get(queue.pop()!)!;
     queued.delete(fn.name);
     const params = slots.facts(fn.name).params;
     if (params) for (const id of params.keys()) reads(`param:${fn.name}\0${id}`, fn.name);
-    const ranges: IntegerRanges = analyzeIntegerRanges(numeric(fn), slots.facts(fn.name));
+    // Range facts are consulted only at writes into candidate slots; slots
+    // are only ever dropped, so a body without such a write needs none.
+    const ranges: IntegerRanges = hasCandidateWrite(fn)
+      ? analyzeIntegerRanges(numeric(fn), slots.facts(fn.name))
+      : new Map();
     const proven: IrExpr[] = [];
     const int32 = (e: IrExpr): boolean => withinInt32(ranges.get(e));
     const args = (callee: IrFunction | undefined, list: IrExpr[], offset: number): void => {