📚

TypeScript 本体のコードを読んでみよう

2023/06/13に公開

みんなお世話になっている TypeScript のコードを読みたいと思ったことはないだろうか。読んだ。

https://github.com/microsoft/TypeScript

一週間ぐらいかかった。完全に読み切ったとは言えないが、概要は掴んだ。

なかなかに複雑でドメイン知識を得るのが難しかったので、これから読む人向けに、登場人物や概念を整理して紹介したい。

読んだのは 2023/6/8 時点で git clone したコード。

最初に: 自分のゴール設定

複数ファイルにまたがった参照を、 TypeScript の Language Service が提供する findReferences()findRenameLocations(), goToDefinitions() を使って、インクリメンタルに書き換えたかった。

Terser を使うと、今触ってるオブジェクトが何で、何のメンバを書き換えたかの情報が残らない。これを TypeScript のレイヤーでやれば、より積極的な minify ができるはず。

ただ、自分が簡単に使ってみた範囲では、あやふやな知識はあっても内部のライフサイクルや実装に自信が持てなかった。

TypeScript Compiler API の基本知識

エントリポイントから考えることとする。TypeScript は tsc や vscode 経由に使う以外にも、モジュールとして Compiler API を提供している。

例えば次のようなコードが基本形になる。

import ts from "typescript";

// このディレクトリの tsconfig.json を元に config や解析対象のファイルをロードする
const tsconfig = ts.readConfigFile("./tsconfig.json", ts.sys.readFile);
const options = ts.parseJsonConfigFileContent(tsconfig.config, ts.sys, "./");

// Compiler Host を取得
const host = ts.createCompilerHost(options.options);
const program = ts.createProgram(options.fileNames, options.options, host);

// 型の解析結果を取得
for (const diagnostic of program.getSemanticDiagnostics()) {
  console.log(
    "[semantic diagnostics]",
    diagnostic.file?.fileName.replace(process.cwd(), ""),
    `#${diagnostic.start}:${diagnostic.length}`,
    "-",
    diagnostic.messageText,
    indexSourceFile.getLineAndCharacterOfPosition(diagnostic.start!),
  );
}

ここを基準に登場人物の解説していく。

基本的に src/compiler/*.ts を読むことになる。特にお世話になるのは src/compiler/types.ts で、だいたいすべての公開インターフェースの型がここに書いてある。

ts.System

ts.System で表現されるその環境を表現する API の型が定義されている。 readFilewriteFile など。 node 環境でロードした ts.sys には node 用のts.System が実装されている。node 以外で読みだした場合(ブラウザやdenoなど)は ts.sys が空になるが、型を満たす形で自分で実装することでどの環境でもコンパイラを動かすことができる。 最悪 FileSystem のバックエンドがなくてもよく、例えば TS Playground なんかはインメモリのモックで動いている。

https://www.typescriptlang.org/play

ts.CompilerHost

一番基本的なコンパイラ環境のホスト抽象。 ts.sys を引数に取ったり要求したりする。

この CompilerHost を引数にもつことで、 後述する ProgramBuilderProgram, WatcherProgram, LanguageService などを動かすことができる。

create~Host がいくつかあり、だいたい create~Program と対になっている。目標を持って読む場合、だいたいは自分が使いたい~Program~Host を生成することを目指すことになる。

ts.SourceFile

ファイルと 1:1 に紐づく抽象。

const newSource = ts.createSourceFile(
    'index.ts',
    'export const x: number = 1;',
    ts.ScriptTarget.ESNext,
);

ts.createSourceFile で生成したファイルはこれだけでは複数ファイルにまたがって解析することができないので、複数ファイルを操作したい場合、基本的に ts.Program()を使う。

(ソースコードを読めばわかるが、これに解析された情報が大量に書き込まれている)

ただ、これだけでも基本的な AST 操作なら可能。

例えば次のような transformer によるコード変形には型が不要なので、AST遊びだけならこれだけでもいい。

import ts from "typescript";

const source = ts.createSourceFile(
  "test.ts",
  `export const x = 1`,
  ts.ScriptTarget.ESNext,
);

const transformerFactory = (context: ts.TransformationContext) => {
  const visit: ts.Visitor = (node) => {
    if (ts.isVariableDeclaration(node)) {
      return ts.factory.createVariableDeclaration(
        node.name,
        undefined,
        undefined,
        ts.factory.createNumericLiteral("2"),
      );
    }
    return ts.visitEachChild(node, visit, context);
  };
  return (node: ts.SourceFile) => ts.visitNode(node, visit);
};

const result = ts.transform(source, [
  transformerFactory as ts.TransformerFactory<ts.SourceFile>,
]);

const transformedSource = result.transformed[0];
const printer = ts.createPrinter({
  newLine: ts.NewLineKind.LineFeed,
});

console.log(printer.printFile(transformedSource));

ts.Program

複数の SourceFile で構成されるプログラム抽象。これを経由することで複数ファイルにまたがる依存解析や、型の解析が可能になる。

簡単な使い方。

const program = ts.createProgram(...);
const diagnostics = program.getSemanticDiganostics();
const soureFiles = program.getSourceFiles();
// 出力
program.emit();

内部的には host の IO 抽象を経由して、 src/index.ts のようなファイルを読み出し、それがどのファイルから呼ばれているか、呼んでいるか、型は壊れてないか、という解析をする。

emit は与えられた CompilerHost の writeFile が実装されていればそれを使う。

tsc -p . 相当

ts.TypeChecker

ts.Program の内部で使われる型の検査の実体だが、表面的には使えるAPIが限られており、 ts.Program ごしに使うことが意図されている。

直接使うとしたら typeChecker.getSymbolsInScope() だと思われる。そのファイルの特定のノードを引っ貼りだすことができる。

const typeChecker = program.getTypeChecker();
const symobls = typeChecker.getSymbolsInScope(
    sourceFile,
    ts.Symbolchecks.Type,
);

とはいえTS本体の実装を理解をしていないと使うのが難しい。

中を読んでて面白かったのが、 node.checks で、 node は AST Node の基本形である (ts.Node) で、 解析した結果、取りうることが可能な状態(例えば string | number)をビットフラグで書き込んでいる。

つまりある程度プリミティブ型で構成されるとき、ユニオン型 A | B の計算は、 実際の計算上でも nodeA.checks | nodeB.checks であり、インターセクション型 A & BnodeA.checks & nodeB.checks になっている。実際には Object 型と組み合わせてもっと複雑だが、基本はメンバー同士でそれをやるだけとも言える。

また node.links には node 同士の接続情報が書き込まれていて、それを辿ることでコントロールフロー解析、例えば if (x != null) {/* x が null でない scope */}が実現されている。

ちょっと自分が ts.Nodets.Symbol の区別が曖昧なのだが、基本は一緒で、AST Node に対してそれを行うか、宣言された Symbol に対しての違いでしかない。

ts.BuliderProgram

実は Program 自体は現在の SourceFiles を対象に解析を掛けた1回きりの状態を抽象しており、 SourceFile の中身を書き換えて再検査することができない。

複数回に渡って SourceFile を書き換える内部抽象が BuilderProgram になる。

const builder = ts.createSemanticDiagnosticsBuilderProgram(
  options.fileNames,
  options.options,
  host,
);
const program1 = builder.getProgram();
// ...
const program2 = builder.getProgram();

// 一致しない可能性がある
console.log(program1 === program2)

簡単に作れるように見えるが、実は BuilderProgram 自体は書き換えたことを内部に伝えるロジックを持っていない。インクリメンタルな操作のために BuilderProgram に渡す CompilerHost をいい感じにラップして生成するのが、WatchProgramLanguageService になる。(というのを理解するのに膨大な時間がかかった)

program の再構築はコストが掛かるように思えるが、 creatProgram は引数で別の program を取って、その内部キャッシュの SourceFiles と解析結果を引き継ぐことができるようになっている。

ts.WatchProgram

tsc --watch 相当。

主に node 経由で FileSystem を監視し、そのコールバックで内部の BuilderProgram に変更を通知する。
ファイル変更通知は 250ms 単位でバッファリングされており、内部キャッシュを破棄して再ロードする。

内部キャッシュは ts.ModuleResolution で表現されているが、表面的には直接触る方法がない。

ts.LanguageService

LSP の LanguageService の部分。厳密に LSP のインターフェースを喋るわけではなく、それは ts.Server にある。というか ts.LanugageService のあとに LSP を仕様化したらしいので、細かい違いがあるらしい。

例えば vscode の findReferencesgoToDefinition のロジックは、この上で TypeChecker の解析結果を使って実装されている。補完インターフェースも持っていて、このトークンの続きはどう補完する? みたいなのもこいつが持っている。

これを使うのに、ちょっと特殊な CompilerHost である LanguageServiceHost を要求される。次のドキュメントが参考になる。というかこれしか表向きのドキュメントがない。

https://github.com/microsoft/TypeScript/wiki/Using-the-Compiler-API#incremental-build-support-using-the-language-services

  // Create the language service host to allow the LS to communicate with the host
  const servicesHost: ts.LanguageServiceHost = {
    getScriptFileNames: () => rootFileNames,
    getScriptVersion: fileName =>
      files[fileName] && files[fileName].version.toString(),
    getScriptSnapshot: fileName => {
      if (!fs.existsSync(fileName)) {
        return undefined;
      }

      return ts.ScriptSnapshot.fromString(fs.readFileSync(fileName).toString());
    },
    getCurrentDirectory: () => process.cwd(),
    getCompilationSettings: () => options,
    getDefaultLibFileName: options => ts.getDefaultLibFilePath(options),
    fileExists: ts.sys.fileExists,
    readFile: ts.sys.readFile,
    readDirectory: ts.sys.readDirectory,
    directoryExists: ts.sys.directoryExists,
    getDirectories: ts.sys.getDirectories,
  };

getScriptVersion と getScriptSnapshot がキモで、内部的にドキュメントを書き換えたらそのファイルに対応するバージョンを変更する。そうすることで getScriptSnapshot が再び呼ばれるので、そこで新しい文字列を返す。

自分はこれをちゃんとした理解で扱う方法が知りたくて、頑張って読み解いた。その結果、こういう理解をした。

  • LanguageService はソースコードの変更に応じて内部の program を再生成する
  • SourceFile は documentRegistry で管理され、 scriptSnapshot や scirptVersion と同時にこれを更新することで変更を通知できる
  • Program や TypeChecker はソースコードを変更するごとに作り直される
  • IScriptSnapshot の getChangeRange を定義することで変更レンジを返すことができ、 IncrementalParser はこれを元に部分的にASTを作り直す

これを自分でいい感じに使うのは実践編で後述。

ts.Server

主に vscode と話すためのインターフェースを持っている。ここのソースコードは読んでない。 Compiler API には含まれないが、そこで実装されてる API を使いたければ、 直接 typescript/lib/tsserver.js から呼べなくもない。

import ts from "typescript/lib/tsserver.js";
console.log(ts.server); // 通常の Compiler API より多くの機能がある

ts.LanugageServer 実践編

ここまでに得た知識を使って、 LanguageService 上にインメモリキャッシュを構築し、インクリメンタルに書き換えを行う LanguageService を実装してみた。

import ts from "typescript";
import fs from "node:fs";
import path from "node:path";
import { DocumentRegistry } from "typescript";

const tsconfig = ts.readConfigFile("./tsconfig.json", ts.sys.readFile);
const options = ts.parseJsonConfigFileContent(tsconfig.config, ts.sys, "./");
const defaultHost = ts.createCompilerHost(options.options);

const expandPath = (fname: string) => {
  if (fname.startsWith("/")) {
    return fname;
  }
  const root = process.cwd();
  return path.join(root, fname);
};

function applyRenameLocations(
  code: string,
  toName: string,
  renameLocations: readonly ts.RenameLocation[],
) {
  let current = code;
  let offset = 0;
  for (const loc of renameLocations) {
    const start = loc.textSpan.start;
    const end = loc.textSpan.start + loc.textSpan.length;
    current = current.slice(0, start + offset) + toName +
      current.slice(end + offset);
    offset += toName.length - (end - start);
  }
  return current;
}

type SnapshotManager = {
  readFileSnapshot(fileName: string): string | undefined;
  writeFileSnapshot(fileName: string, content: string): ts.SourceFile;
};

export interface InMemoryLanguageServiceHost extends ts.LanguageServiceHost {
  getSnapshotManager: (
    registory: DocumentRegistry,
  ) => SnapshotManager;
}

export function createInMemoryLanguageServiceHost(): InMemoryLanguageServiceHost {
  // read once, write on memory
  const fileContents = new Map<string, string>();
  const fileSnapshots = new Map<string, ts.IScriptSnapshot>();
  const fileVersions = new Map<string, number>();
  const fileDirtySet = new Set<string>();

  const getSnapshotManagerInternal: (
    registory: DocumentRegistry,
  ) => SnapshotManager = (registory: ts.DocumentRegistry) => {
    return {
      readFileSnapshot(fileName: string) {
        fileName = expandPath(fileName);
        console.log("[readFileSnapshot]", fileName);
        if (fileContents.has(fileName)) {
          return fileContents.get(fileName) as string;
        }
        return defaultHost.readFile(fileName);
      },
      writeFileSnapshot(fileName: string, content: string) {
        fileName = expandPath(fileName);
        const nextVersion = (fileVersions.get(fileName) || 0) + 1;
        // fileVersions.set(fileName, nextVersion);
        fileContents.set(fileName, content);
        console.log(
          "[writeFileSnapshot]",
          fileName,
          nextVersion,
          content.length,
        );
        fileDirtySet.add(fileName);
        const newSource = registory.updateDocument(
          fileName,
          serviceHost,
          ts.ScriptSnapshot.fromString(content),
          String(nextVersion),
        );
        return newSource;
      },
    };
  };

  const serviceHost: InMemoryLanguageServiceHost = {
    getDefaultLibFileName: defaultHost.getDefaultLibFileName,
    fileExists: ts.sys.fileExists,
    readDirectory: ts.sys.readDirectory,
    directoryExists: ts.sys.directoryExists,
    getDirectories: ts.sys.getDirectories,
    getCurrentDirectory: defaultHost.getCurrentDirectory,
    getScriptFileNames: () => options.fileNames,
    getCompilationSettings: () => options.options,
    readFile: (fname, encode) => {
      fname = expandPath(fname);
      // console.log("[readFile]", fname);
      if (fileContents.has(fname)) {
        return fileContents.get(fname) as string;
      }
      const rawFileResult = ts.sys.readFile(fname, encode);
      if (rawFileResult) {
        fileContents.set(fname, rawFileResult);
        fileVersions.set(
          fname,
          (fileVersions.get(fname) || 0) + 1,
        );
      }
      return rawFileResult;
    },
    writeFile: (fileName, content) => {
      fileName = expandPath(fileName);
      console.log("[writeFile:mock]", fileName, content.length);
    },
    getScriptSnapshot: (fileName) => {
      fileName = expandPath(fileName);
      if (fileName.includes("src/index.ts")) {
        console.log("[getScriptSnapshot]", fileName);
      }
      if (fileSnapshots.has(fileName)) {
        return fileSnapshots.get(fileName)!;
      }
      const contentCache = fileContents.get(fileName);
      if (contentCache) {
        const newSnapshot = ts.ScriptSnapshot.fromString(contentCache);
        fileSnapshots.set(fileName, newSnapshot);
        return newSnapshot;
      }
      if (!fs.existsSync(fileName)) return;
      const raw = ts.sys.readFile(fileName, "utf8")!;
      const snopshot = ts.ScriptSnapshot.fromString(raw);
      fileSnapshots.set(fileName, snopshot);
      return snopshot;
    },
    getScriptVersion: (fileName) => {
      fileName = expandPath(fileName);
      const isDirty = fileDirtySet.has(fileName);
      if (isDirty) {
        const current = fileVersions.get(fileName) || 0;
        fileDirtySet.delete(fileName);
        fileVersions.set(fileName, current + 1);
      }
      return (fileVersions.get(fileName) || 0).toString();
    },
    getSnapshotManager: getSnapshotManagerInternal,
  };
  return serviceHost;
}

これは最初の読み込みだけ FileSystem から返し、以降は専用の SnapshotManager でインメモリキャッシュを書き換える実装。色々と試行錯誤したのでもっと単純化はできると思う。

TypeScript 本体の実装パターンを踏襲して、新しいインターフェースを切って create~Host の結果を食う create~ とした。

手元の src/index.ts はこんな感じのものを用意した (わざと型違反している)

const x: number = "";

これを使うコード例。

// ...続き
{
  // usage
  const prefs: ts.UserPreferences = {};
  const registory = ts.createDocumentRegistry();
  const serviceHost = createInMemoryLanguageServiceHost();
  const languageService = ts.createLanguageService(
    serviceHost,
    registory,
  );

  // languageService.
  const snapshotManager = serviceHost.getSnapshotManager(registory);

  // write src/index.ts and check types
  const raw = snapshotManager.readFileSnapshot("src/index.ts");
  const newSource = snapshotManager.writeFileSnapshot(
    "src/index.ts",
    raw + "\nconst y: number = x;",
  );

  // find scoped variables

  // languageService.getSemanticDiagnostics("src/index.ts");
  const program = languageService.getProgram()!;
  const checker = program.getTypeChecker();
  const localVariables = checker.getSymbolsInScope(
    newSource,
    ts.SymbolFlags.BlockScopedVariable,
  );

  // rename x to x_?
  const symbol = localVariables.find((s) => s.name === "x")!;
  const renameLocations = languageService.findRenameLocations(
    "src/index.ts",
    symbol.valueDeclaration!.getStart(),
    false,
    false,
    prefs,
  );
  const targets = new Set(renameLocations!.map((loc) => loc.fileName));

  let current = snapshotManager.readFileSnapshot("src/index.ts")!;
  for (const target of targets) {
    const renameLocationsToTarget = renameLocations!.filter(
      (loc) => expandPath(target) === expandPath(loc.fileName),
    );
    const newSymbolName = `${symbol.name}_${
      Math.random().toString(36).slice(2)
    }`;
    current = applyRenameLocations(
      current,
      newSymbolName,
      renameLocationsToTarget,
    );
  }
  snapshotManager.writeFileSnapshot("src/index.ts", current);
  const result = languageService.getSemanticDiagnostics("src/index.ts");
  console.log("post error", result.length);
  console.log(snapshotManager.readFileSnapshot("src/index.ts"));

  const oldProgram = program;
  {
    // rename y to y_?
    const program = languageService.getProgram()!;
    const program2 = languageService.getProgram()!;
    console.log(
      "------- program updated",
      program !== oldProgram,
      program2 === program,
    );
    const checker = program.getTypeChecker();
    const newSource = program.getSourceFile("src/index.ts")!;
    const localVariables = checker.getSymbolsInScope(
      newSource,
      ts.SymbolFlags.BlockScopedVariable,
    );
    const symbol = localVariables.find((s) => s.name === "y")!;
    const renameLocations = languageService.findRenameLocations(
      "src/index.ts",
      symbol.valueDeclaration!.getStart(),
      false,
      false,
      prefs,
    );
    const targets = new Set(renameLocations!.map((loc) => loc.fileName));
    let current = snapshotManager.readFileSnapshot("src/index.ts")!;
    for (const target of targets) {
      const renameLocationsToTarget = renameLocations!.filter(
        (loc) => expandPath(target) === expandPath(loc.fileName),
      );
      const newSymbolName = `${symbol.name}_${
        Math.random().toString(36).slice(2)
      }`;
      current = applyRenameLocations(
        current,
        newSymbolName,
        renameLocationsToTarget,
      );
    }
    snapshotManager.writeFileSnapshot("src/index.ts", current);
    const result = languageService.getSemanticDiagnostics("src/index.ts");
    console.log("post error", result.length);
    console.log(snapshotManager.readFileSnapshot("src/index.ts"));
  }
}

二回目の書き換えが冗長なのはコピペしてサボっているのだが、基本は3回コードを書き換える。

x を参照する y の行を追加する

const x: number = "";
const y: number = x;

x の変数名を乱数サフィックス付きで書き換える

const x_3fo8yrgzd8u: number = "";
const y: number = x_3fo8yrgzd8u;

ちゃんと二箇所書き換わっている。

ついでに y も書き換える。(このとき、renameInfo を更新するために program の再生成が必要だった)

const x_3fo8yrgzd8u: number = "xxxx";

const y_o7708up7yh8: number = x_3fo8yrgzd8u;

こうなる。

これらの操作を経てもローカルの src/index.ts は書き換わっていない。

これで自分が当初やりたいことと思っていたことが実現できた。

ソースコードリーディングの感想

TS 本体がESM へ移行しようとしている最中だったので、中途半端部分はあったのだが、一生物の知識が得られた気がする。

とくに IncrementalParsing や型チェッカーの内部実装、program 更新時のキャッシュ引き継ぎなど、宝の宝庫だと思う。

コンパイラのコードは難しいと思われるかもしれないが、TypeScript をある程度使った知識があれば、利用者としてのドメイン知識でゴリ押せるところがあった。

今書いてるコードが、型検査器の内部でどういう状態をとっているのか透けて見えるようになったのが一番の収穫。

実際のソースコードの読み方のコツ

  • 何か新しい型や概念に出会ったらまず src/compiler/types.ts で確認する。基本的な型はここに書いてあるので、自然とコードジャンプでここに戻ってくる。
  • 仮にあなたが TS の Compiler API に詳しくても実装上の振る舞いが予想に反することがある。 とくに /** @internal */ のコメントが付いてるインターフェースは、実装上は存在するが公開APIから落ちており、それが重要だったりする。
    • SourecFile と Node と Symbol は特に内部的なメタデータが多い。
  • findXxx という関数に対して findXxxWorker という関数が呼ばれているとき、その前後が tracing...Permance... でラップされている。このとき、 findXxxWorker の処理コストが重いことの合図で、腰を据えて覚悟をすることをおすすめする。また、再帰していることが多いが、その時は findXxx 側に戻る。
  • 大きなモジュールを生成する関数を読むときに上から読むとつらく感じがちだがコツがある
    • 前半部分は各種キャッシュ引き継ぎで、後半部分が新規生成、という順番になってるので、とくに createProgram などは後半から読み始めてもよいと思う。
    • とにかく Program が基本単位なのだが、 oldProgram の引き継ぎを読むのが辛かった。sourceFiles を引き継ぎます、ぐらいの雑な理解でもいいかもしれない。
  • TypeChecker の型推論の周りは基本的にビットマスクだと思って読む。難しいビット演算があるわけではないが、単純に手順が多い。
  • Parser は手書きの再帰下降パーサが愚直に書かれていて魔窟。よほど必要とされない限り(新しい文法を実装したいとか思わない限りは)読み飛ばしていい。
  • ほとんどの実装が関数とそのローカルスコープで内部状態を隠蔽しながら書かれている。class はほぼない。関数スコープで状態を隠蔽するので、大きな単位を生成する関数では、ローカルスコープに何がいるかを確認しながら読む
  • src/compiler/moduleResolver.ts で実装されているモジュール解決手順は node と ESM の関係を知っている人にはある程度直感的だが、とにかく仕様も実装も複雑で難しい。ある程度そういうものと思って読み飛ばした。

おわり

みなさんもぜひ読んでみよう! そして一つだけ大事なTIPS。 10年もののリポジトリで git clone が重すぎるので、必ず git clone --depth 1 --branch main するように!

読んでた時のスクラップ集

https://zenn.dev/mizchi/scraps/3c30ea6fa9f8e5
https://zenn.dev/mizchi/scraps/80d47cbc601a5f
https://zenn.dev/mizchi/scraps/9f0f6b3b08bff6

Discussion