astgen: cancel a type write TypeScript cannot bring to an end

astgen.js

@@ -43,8 +43,11 @@ import { lowerEmberFile } from "./emberAst.js";
 // parsed; the AST JSON carries the rewritten text as `source`. 4.5.1: files
 // with source phase or deferred imports (`import source`, `import defer`) are
 // parsed instead of dropped, and type maps include `node_modules/@types`
-// packages again when the project does not set `types`.
-const ASTGEN_VERSION = "4.5.1";
+// packages again when the project does not set `types`. 4.5.2: a project
+// whose types recursed through conditional types (Prisma 4's client) ran
+// astgen out of memory, leaving no ASTs at all; the empty parse must not be
+// reused.
+const ASTGEN_VERSION = "4.5.2";
 
 const HELP_TEXT = `Options:
   -i, --src      Source directory                                 [default: "."]
@@ -1203,6 +1206,35 @@ const TSC_FLAGS_BY_NAME = TSC_FLAGS & ~tsc.TypeFormatFlags.InTypeAlias;
 // type exhausts the heap after a few files. Real code stays far below this.
 const MAX_TYPE_STRING_LENGTH = 4096;
 
+// Most type nodes TypeScript may build while writing one type. Its node builder
+// stops honouring the length limit in places once it is truncating: it keeps
+// expanding substitution, conditional and template literal types, and a type
+// that recurses through them runs on until the heap is gone. One
+// `prisma.article.findMany` (Prisma 4.16) reached 17 million characters across
+// 300,000 builder calls and exhausted 4 GB, taking every file's types and ASTs
+// with it. The builder asks the checker's cancellation token once per node, so
+// a token that cancels past this many ends such a write, which is then a type
+// too long to write, as when the limit holds. A type astgen writes needs far
+// fewer: across 37 projects (juice-shop, taxonomy, the Angular, React, Vue,
+// Svelte and Solid realworld apps among them) none needed more than 305, and
+// every type map came out byte for byte as before. ASTGEN_TYPE_TIMING reports
+// the most a file's written types needed, and how many writes were cancelled.
+const MAX_TYPE_BUILDER_CALLS = 16384;
+
+/** A cancellation token that ends a type write past MAX_TYPE_BUILDER_CALLS. */
+const typeWriteBudget = () => {
+  const budget = {
+    calls: 0,
+    isCancellationRequested: () => budget.calls > MAX_TYPE_BUILDER_CALLS,
+    throwIfCancellationRequested: () => {
+      if (++budget.calls > MAX_TYPE_BUILDER_CALLS) {
+        throw new tsc.OperationCanceledException();
+      }
+    }
+  };
+  return budget;
+};
+
 /**
  * Convert a single vue file to AST.
  * When `tsInstance` is present also collect type inference from TSC and return both AST & types.
@@ -1643,8 +1675,23 @@ const instrumentTypeChecker = (typeChecker) => {
     timing.kinds = new Map();
     timing.repeats = 0;
     timing.slow = [];
+    timing.maxWrittenBuilderCalls = 0;
+    timing.cancelledWrites = 0;
   };
   reset();
+  // The builder calls a written type needed, and the writes cancelled for
+  // running past MAX_TYPE_BUILDER_CALLS.
+  timing.recordWrite = (builderCalls, written) => {
+    if (written) {
+      timing.maxWrittenBuilderCalls = Math.max(
+        timing.maxWrittenBuilderCalls,
+        builderCalls
+      );
+    }
+    if (builderCalls > MAX_TYPE_BUILDER_CALLS) {
+      timing.cancelledWrites++;
+    }
+  };
   const queried = new WeakSet();
   const wrap = (name, onCall) => {
     if (typeof typeChecker[name] !== "function") {
@@ -1690,7 +1737,7 @@ const instrumentTypeChecker = (typeChecker) => {
       .map(([kind, n]) => `${kind}:${n}`)
       .join(",");
     console.error(
-      `[astgen-type-timing] file=${fileName} ms=${ms.toFixed(0)} nodes=${timing.nodes} getTypeAtLocation=${count("getTypeAtLocation")} (ms=${own("getTypeAtLocation")}) repeatQueries=${timing.repeats} writeType=${count("writeType")} (ms=${own("writeType")}) ${TYPE_TIMING_COUNTED.map((name) => `${name}=${count(name)}`).join(" ")} topKinds=${topKinds}`
+      `[astgen-type-timing] file=${fileName} ms=${ms.toFixed(0)} nodes=${timing.nodes} getTypeAtLocation=${count("getTypeAtLocation")} (ms=${own("getTypeAtLocation")}) repeatQueries=${timing.repeats} writeType=${count("writeType")} (ms=${own("writeType")}) maxWrittenBuilderCalls=${timing.maxWrittenBuilderCalls} cancelledWrites=${timing.cancelledWrites} ${TYPE_TIMING_COUNTED.map((name) => `${name}=${count(name)}`).join(" ")} topKinds=${topKinds}`
     );
     for (const q of timing.slow.sort((a, b) => b.ms - a.ms).slice(0, 20)) {
       console.error(
@@ -1748,7 +1795,8 @@ function createTsc(srcFiles, src) {
     // between workers.
     const boundedTypeToString = (type, context, flags) => {
       const out = { canIncreaseExpansionDepth: false, truncated: false };
-      const text =
+      const budget = typeWriteBudget();
+      const write = () =>
         typeof typeChecker.writeType === "function"
           ? typeChecker.writeType(
               type,
@@ -1760,9 +1808,21 @@ function createTsc(srcFiles, src) {
               out
             )
           : typeChecker.typeToString(type, context, flags);
-      return out.truncated || text.length >= MAX_TYPE_STRING_LENGTH
-        ? undefined
-        : text;
+      let text;
+      try {
+        text = typeChecker.runWithCancellationToken
+          ? typeChecker.runWithCancellationToken(budget, write)
+          : write();
+      } catch (err) {
+        if (err instanceof tsc.OperationCanceledException) {
+          typeTiming?.recordWrite(budget.calls, false);
+          return undefined;
+        }
+        throw err;
+      }
+      const written = !out.truncated && text.length < MAX_TYPE_STRING_LENGTH;
+      typeTiming?.recordWrite(budget.calls, written);
+      return written ? text : undefined;
     };
     let oversizedTypes = new WeakMap();
     const typeToBoundedString = (type, context) => {

package.json

@@ -6,7 +6,7 @@
   "type": "module",
   "scripts": {
     "pretty": "prettier --write *.js lib/**/*.{js,mjs} --trailing-comma=none",
-    "test": "node test-fixtures/supervise-regression.js && node test-fixtures/astgen-type-regression.js && node test-fixtures/astgen-json-regression.js && node test-fixtures/astgen-vue-regression.js && node test-fixtures/astgen-svelte-regression.js && node test-fixtures/astgen-astro-regression.js && node test-fixtures/astgen-html-regression.js && node test-fixtures/astgen-ember-regression.js && node test-fixtures/astgen-node-types-regression.js && node test-fixtures/astgen-import-phases-regression.js && node test-fixtures/astgen-oversized-types-regression.js && node test-fixtures/astgen-ignored-dirs-regression.js && node test-fixtures/astgen-shape-snapshot.js && node test-fixtures/evaluate-astgen.js",
+    "test": "node test-fixtures/supervise-regression.js && node test-fixtures/astgen-type-regression.js && node test-fixtures/astgen-json-regression.js && node test-fixtures/astgen-vue-regression.js && node test-fixtures/astgen-svelte-regression.js && node test-fixtures/astgen-astro-regression.js && node test-fixtures/astgen-html-regression.js && node test-fixtures/astgen-ember-regression.js && node test-fixtures/astgen-node-types-regression.js && node test-fixtures/astgen-import-phases-regression.js && node test-fixtures/astgen-oversized-types-regression.js && node test-fixtures/astgen-recursive-types-regression.js && node test-fixtures/astgen-ignored-dirs-regression.js && node test-fixtures/astgen-shape-snapshot.js && node test-fixtures/evaluate-astgen.js",
     "test:evaluate": "node test-fixtures/evaluate-astgen.js",
     "test:ruby": "node test-fixtures/rbastgen-regression.js",
     "test:php": "node test-fixtures/phpastgen-cli.js && node test-fixtures/phpastgen-parser-info.js && node test-fixtures/phpastgen-discovery.js && node test-fixtures/phpastgen-provenance.js && node test-fixtures/phpastgen-framework-facts.js && node test-fixtures/phpastgen-legacy-regression.js && node test-fixtures/phpastgen-concurrency.js && node test-fixtures/phpastgen-chunk.js && node test-fixtures/phpastgen-regression.js && node test-fixtures/phpastgen-contract-snapshot.js && npm run test:php:pbt",

test-fixtures/astgen-recursive-types-regression.js

@@ -0,0 +1,74 @@
+// A type that recurses through conditional types must not take the heap with
+// it. Writing one, TypeScript's node builder keeps expanding past the length
+// limit once it is truncating; an ORM client's result types (Prisma 4's, in
+// the shape the fixture copies) run to tens of thousands of builder calls a
+// write,
+// and in a real project to millions of characters, until the process runs out
+// of memory and no file gets an AST. astgen cancels such a write and treats the
+// type as one too long to write, which it was going to be anyway.
+
+import assert from "node:assert/strict";
+import { spawnSync } from "node:child_process";
+import { readFileSync, rmSync } from "node:fs";
+import { dirname, join } from "node:path";
+import { fileURLToPath } from "node:url";
+import {
+  createOutputRoot,
+  expectType,
+  readJson
+} from "./ast-regression-helpers.js";
+
+const __dirname = dirname(fileURLToPath(import.meta.url));
+const repoRoot = dirname(__dirname);
+const fixtureRoot = join(__dirname, "projects", "recursive-payload");
+const outputRoot = createOutputRoot("atom-parsetools-recursive-types-");
+
+const load = (relativeName) => ({
+  code: readFileSync(join(fixtureRoot, relativeName), "utf8"),
+  typemap: readJson(join(outputRoot, `${relativeName}.typemap`)),
+  relativeName
+});
+
+try {
+  // Without the cancellation the fixture needs about 2 GB; 512 MB leaves the
+  // bounded writes ample room and fails the unbounded ones.
+  const run = spawnSync(
+    process.execPath,
+    [
+      "--max-old-space-size=512",
+      join(repoRoot, "astgen.js"),
+      "-i",
+      fixtureRoot,
+      "-o",
+      outputRoot,
+      "-t",
+      "ts"
+    ],
+    {
+      cwd: repoRoot,
+      encoding: "utf8",
+      env: { ...process.env, ASTGEN_TYPE_TIMING: "1" },
+      stdio: ["ignore", "pipe", "pipe"]
+    }
+  );
+  assert.equal(run.status, 0, `astgen failed:\n${run.stderr.slice(-2000)}`);
+  const cancelled = [...run.stderr.matchAll(/cancelledWrites=(\d+)/g)].reduce(
+    (sum, m) => sum + Number(m[1]),
+    0
+  );
+  assert.ok(cancelled > 0, "no type write was cancelled");
+
+  const service = load("service.ts");
+  load("client.ts");
+
+  // Everything else is written as before
+  expectType(service, "title =", "(name: string) => string");
+  expectType(service, "author: string", "string", { offset: 8 });
+  expectType(service, "include:", "{ author: boolean; tags: boolean; }", {
+    offset: 9
+  });
+
+  console.log("astgen recursive type regression tests passed");
+} finally {
+  rmSync(outputRoot, { recursive: true, force: true });
+}

test-fixtures/projects/recursive-payload/client.ts

@@ -0,0 +1,158 @@
+// The shape of an ORM client's result types (Prisma 4): each model's payload
+// lists its scalars, its relations and its embedded composites, generic over
+// the client's extensions; relations point back at each other, and a query's
+// result is a conditional mapped type over its arguments that recurses into
+// every relation they select. Nothing here is Prisma's code, only its shape.
+type Row = { [field in string]: unknown };
+
+export type Payload = {
+  scalars: Row;
+  objects: Row;
+  composites: Row;
+};
+
+export type Extensions = {
+  result: { [model in string]: { [field in string]: unknown } };
+};
+export type NoExtensions = { result: {} & Record<string, {}> };
+
+// A model's scalars with the fields an extension computes
+type Computed<Base extends Record<any, any>, R, KR extends keyof R = string extends keyof R ? never : keyof R> = {
+  [K in KR | keyof Base]: K extends KR
+    ? R[K] extends () => { compute: (...args: any) => infer C }
+      ? C
+      : never
+    : Base[K];
+};
+
+type Unwrap<P> = {
+  [K in keyof P]: P[K] extends Payload[]
+    ? Unwrap<P[K]>
+    : P[K] extends Payload
+      ? P[K]["scalars"] & Unwrap<P[K]["composites"]>
+      : P[K] extends (infer O extends Payload) | null
+        ? (O["scalars"] & Unwrap<O["composites"]>) | null
+        : P[K];
+} & unknown;
+
+type Defaults<P> = P extends Payload ? P["scalars"] & Unwrap<P["composites"]> : P;
+
+type Field<K extends PropertyKey, O> =
+  | { objects: { [k in K]: O } }
+  | { composites: { [k in K]: O } };
+
+type Nullable<P extends Field<any, any>, K extends PropertyKey> = P extends {
+  objects: Record<K, any>;
+}
+  ? P["objects"][K]
+  : P extends { composites: Record<K, any> }
+    ? P["composites"][K]
+    : never;
+
+type Count<O> = { [K in keyof O]: number } & {};
+
+export type Found<P extends Payload, A> = {} extends A
+  ? Defaults<P>
+  : A extends
+        | ({ select: infer S } & Record<string, unknown>)
+        | ({ include: infer S } & Record<string, unknown>)
+    ? S extends undefined
+      ? Defaults<P>
+      : {
+          [K in keyof S as S[K] extends false | undefined | null ? never : K]: S[K] extends object
+            ? P extends Field<K, (infer O)[]>
+              ? O extends Payload
+                ? Found<O, S[K]>[]
+                : never
+              : P extends Field<K, infer O | null>
+                ? O extends Payload
+                  ? Found<O, S[K]> | (Nullable<P, K> & null)
+                  : never
+                : K extends "_count"
+                  ? Count<Found<P, S[K]>>
+                  : never
+            : P extends Field<K, (infer O)[]>
+              ? O extends Payload
+                ? Defaults<O>[]
+                : never
+              : P extends Field<K, infer O | null>
+                ? O extends Payload
+                  ? Defaults<O> | (Nullable<P, K> & null)
+                  : never
+                : P extends { scalars: { [k in K]: infer O } }
+                  ? O
+                  : K extends "_count"
+                    ? Count<P["objects"]>
+                    : never;
+        } & (A extends { include: any } & Record<string, unknown> ? Defaults<P> : unknown)
+    : Defaults<P>;
+
+type Operation = "findUnique" | "findMany" | "create";
+
+export type Result<P extends Payload, A, O extends Operation, Null = null> = {
+  findUnique: Found<P, A> | Null;
+  findMany: Found<P, A>[];
+  create: Found<P, A>;
+}[O];
+
+type Subset<T, U> = { [key in keyof T]: key extends keyof U ? T[key] : never } & (T extends {
+  select: any;
+  include: any;
+}
+  ? "choose select or include"
+  : {});
+
+export type UserPayload<E extends Extensions = NoExtensions> = {
+  objects: {
+    articles: ArticlePayload<E>[];
+    comments: CommentPayload<E>[];
+    favorites: ArticlePayload<E>[];
+    followers: UserPayload<E>[];
+    following: UserPayload<E>[];
+  };
+  scalars: Computed<{ id: number; email: string; name: string }, E["result"]["user"]>;
+  composites: {};
+};
+
+export type ArticlePayload<E extends Extensions = NoExtensions> = {
+  objects: {
+    tags: TagPayload<E>[];
+    author: UserPayload<E>;
+    favoritedBy: UserPayload<E>[];
+    comments: CommentPayload<E>[];
+  };
+  scalars: Computed<
+    { id: number; slug: string; title: string; body: string; authorId: number },
+    E["result"]["article"]
+  >;
+  composites: {};
+};
+
+export type CommentPayload<E extends Extensions = NoExtensions> = {
+  objects: { author: UserPayload<E>; article: ArticlePayload<E> };
+  scalars: Computed<{ id: number; body: string }, E["result"]["comment"]>;
+  composites: {};
+};
+
+export type TagPayload<E extends Extensions = NoExtensions> = {
+  objects: { articles: ArticlePayload<E>[] };
+  scalars: Computed<{ id: number; name: string }, E["result"]["tag"]>;
+  composites: {};
+};
+
+type FindManyArgs = { where?: Row; select?: Row | null; include?: Row | null };
+
+export declare class Client<E extends Extensions = NoExtensions> {
+  article: {
+    findMany<T extends FindManyArgs>(
+      args?: Subset<T, FindManyArgs>
+    ): Promise<Result<ArticlePayload<E>, T, "findMany">>;
+  };
+  user: {
+    findUnique<T extends FindManyArgs>(
+      args: Subset<T, FindManyArgs>
+    ): Promise<Result<UserPayload<E>, T, "findUnique">>;
+  };
+}
+
+export declare const db: Client;

test-fixtures/projects/recursive-payload/service.ts

@@ -0,0 +1,12 @@
+import { db } from "./client";
+
+export const recent = async (author: string) => {
+  const articles = await db.article.findMany({
+    where: { author: { name: author } },
+    include: { author: true, tags: true }
+  });
+  const user = await db.user.findUnique({ where: { name: author } });
+  return { articles, user, count: articles.length };
+};
+
+export const title = (name: string) => name.trim();