AtCoder で TypeScript の型レベルプログラムを動かす

## はじめに

まずはこの提出をご覧ください。

https://atcoder.jp/contests/language-test-202301/submissions/40508634

TypeScript の型レベルプログラミングを使って [PracticeA - Welcome to AtCoder](https://atcoder.jp/contests/language-test-202301/tasks/practice_1) を解くことができています。

本記事ではこのように、TypeScript の型レベルプログラムを AtCoder にジャッジしてもらう方法を紹介します。

## 2023 年の AtCoder の言語アップデート

2023 年 1 月から 4 月にかけて、AtCoder で大規模な言語アップデートが行われました。

https://atcoder.jp/posts/966

この言語アップデートで、TypeScript のバージョンが 3.8 系から 5.0 系に上がりました。(さらに処理系として Node.js に加えて Deno も追加されました。)

TypeScript といえば、4.1 で追加された [Template Literal Types](https://devblogs.microsoft.com/typescript/announcing-typescript-4-1/#template-literal-types) を使った型レベルプログラミングが賑わいを見せています。

今回のアップデートでジャッジ環境に TypeScript 5.0 系が組み込まれたことにより、TypeScript をライブラリとして呼び出すことで、型レベルプログラムを動かすことができるようになります。

## TypeScript Compiler API

TypeScript は src/compiler/checker.ts に巨大な型チェッカを備えています。こうした型チェッカをはじめとする TypeScript Compiler API は `typescript` パッケージを介して利用可能であることが知られています。

https://github.com/microsoft/TypeScript/wiki/Using-the-Compiler-API

今回の言語アップデートにおいて、TypeScript (Node.js) の実行環境は以下のようになっています。

- `/usr/lib/node_modules` (と、そこへのシンボリックリンクである `/home/runner/node_modules`) に `node_modules` がある
- カレントディレクトリは `/judge` で、`/judge/Main.ts` が実行されている

よって、この `node_modules` から `typescript` を抜き出し、`/judge/Main.ts` そのものを解析すればよさそうです。

## PracticeA - Welcome to AtCoder を解く

あらためて提出コードを見てみます。

https://atcoder.jp/contests/language-test-202301/submissions/40508634

このコードは 3 つの部分に分けられます。

- 1 ～ 6 行目: メイン部分
- 8 ～ 154 行目: 数値計算用ヘルパー部分
- 156 ～ 205 行目: 型レベル実行部分

### メイン部分

問題を解くための型レベル計算の部分です。

ここでは、「入力文字列を型パラメータとして受け取って文字列リテラル型を作るジェネリクス型」をデフォルトエクスポートします。

```ts
type Main<Input extends string> =
  Input extends `${infer A}\n${infer B} ${infer C}\n${infer S}\n`
    ? `${Decode<Add<Encode<A>, Add<Encode<B>, Encode<C>>>>} ${S}\n`
    : never;

export default Main;
```

### 数値計算用ヘルパー部分

上のコードに出てきている `Encode` `Decode` `Add` などのヘルパー型の中身です。

ここは今回の本質ではないので省略します。

<details><summary>数値計算用ヘルパー</summary>

```ts
type Natural = Digit[];

type Digit =
  | Digit0
  | Digit1
  | Digit2
  | Digit3
  | Digit4
  | Digit5
  | Digit6
  | Digit7
  | Digit8
  | Digit9;

type Digit0 = [];
type Digit1 = [unknown];
type Digit2 = [unknown, unknown];
type Digit3 = [unknown, unknown, unknown];
type Digit4 = [unknown, unknown, unknown, unknown];
type Digit5 = [unknown, unknown, unknown, unknown, unknown];
type Digit6 = [unknown, unknown, unknown, unknown, unknown, unknown];
type Digit7 = [unknown, unknown, unknown, unknown, unknown, unknown, unknown];
type Digit8 = [
  unknown,
  unknown,
  unknown,
  unknown,
  unknown,
  unknown,
  unknown,
  unknown
];
type Digit9 = [
  unknown,
  unknown,
  unknown,
  unknown,
  unknown,
  unknown,
  unknown,
  unknown,
  unknown
];

type EncodeNatural<N extends string> = N extends `${infer D}${infer DS}`
  ? [...EncodeNatural<DS>, EncodeDigit<D>]
  : [];

type DecodeNatural<N extends Natural> =
  DecodeNaturalRec<N> extends infer Result extends string
    ? Result extends ``
      ? `0`
      : Result
    : never;

type DecodeNaturalRec<N extends Natural> = N extends [
  infer D extends Digit,
  ...infer DS extends Natural
]
  ? `${DecodeNaturalRec<DS>}${DecodeDigit<D>}`
  : ``;

type EncodeDigit<D extends string> = D extends `0`
  ? Digit0
  : D extends `1`
  ? Digit1
  : D extends `2`
  ? Digit2
  : D extends `3`
  ? Digit3
  : D extends `4`
  ? Digit4
  : D extends `5`
  ? Digit5
  : D extends `6`
  ? Digit6
  : D extends `7`
  ? Digit7
  : D extends `8`
  ? Digit8
  : D extends `9`
  ? Digit9
  : never;

type DecodeDigit<D extends Digit> = D extends Digit0
  ? `0`
  : D extends Digit1
  ? `1`
  : D extends Digit2
  ? `2`
  : D extends Digit3
  ? `3`
  : D extends Digit4
  ? `4`
  : D extends Digit5
  ? `5`
  : D extends Digit6
  ? `6`
  : D extends Digit7
  ? `7`
  : D extends Digit8
  ? `8`
  : D extends Digit9
  ? `9`
  : never;

type AddDigits<
  D0 extends Digit,
  D1 extends Digit,
  Carry extends Digit0 | Digit1 = Digit0
> = [...D0, ...D1, ...Carry] extends infer Sum
  ? Sum extends [
      unknown,
      unknown,
      unknown,
      unknown,
      unknown,
      unknown,
      unknown,
      unknown,
      unknown,
      unknown,
      ...infer Rest
    ]
    ? [Rest, Digit1]
    : [Sum, Digit0]
  : never;

type AddNaturals<
  N0 extends Natural,
  N1 extends Natural,
  C0 extends Digit0 | Digit1 = Digit0
> = N0 extends [infer D0 extends Digit, ...infer DS0 extends Natural]
  ? N1 extends [infer D1 extends Digit, ...infer DS1 extends Natural]
    ? AddDigits<D0, D1, C0> extends [infer S0, infer C1 extends Digit0 | Digit1]
      ? [S0, ...AddNaturals<DS0, DS1, C1>]
      : never
    : C0 extends Digit1
    ? AddNaturals<N0, [C0]>
    : N0
  : C0 extends Digit1
  ? AddNaturals<N1, [C0]>
  : N1;

type Encode<N extends string> = EncodeNatural<N>;
type Decode<N extends Natural> = DecodeNatural<N>;
type Add<N0 extends Natural, N1 extends Natural> = AddNaturals<N0, N1>;
```
</details>

### 型レベル実行部分

本記事で紹介したいのが以下の部分です。

```ts
import * as fs from "node:fs";
import * as ts from "/usr/lib/node_modules/typescript";

function main() {
  const input = fs.readFileSync("/dev/stdin", "utf8");
  const entryFileName = "entry.ts";
  const entrySourceFile = ts.createSourceFile(
    entryFileName,
    [
      `import type Main from "/judge/Main";\n`,
      `type Input = ${JSON.stringify(input)};\n`,
      `type Output = Main<Input>;\n`,
    ].join(""),
    ts.ScriptTarget.Latest
  );
  const defaultCompilerHost = ts.createCompilerHost({});
  const customCompilerHost: ts.CompilerHost = {
    ...defaultCompilerHost,
    getSourceFile(fileName, ...rest) {
      if (fileName === entryFileName) {
        return entrySourceFile;
      } else {
        return defaultCompilerHost.getSourceFile.call(this, fileName, ...rest);
      }
    },
  };
  const program = ts.createProgram({
    rootNames: [entryFileName],
    options: {
      strict: true,
    },
    host: customCompilerHost,
  });
  const checker = program.getTypeChecker();
  function visit(node: ts.Node) {
    const symbol = checker.getSymbolAtLocation(node);
    if (symbol && symbol.getName() === `Output`) {
      const type = checker.getDeclaredTypeOfSymbol(symbol);
      if (type.isStringLiteral()) {
        process.stdout.write(type.value);
      } else {
        throw new Error(`type error: ${checker.typeToString(type)}`);
      }
    }
    node.forEachChild(visit);
  }
  entrySourceFile.forEachChild(visit);
}

main();
```

流れとしては、以下のようになります。

- ジャッジが実行されたら、入力をもとに、以下のような仮想の `SourceFile`、`entry.ts` を作る
  ```ts
  import type Main from "/judge/Main";
  type Input = "【ここに入力が来る】";
  type Output = Main<Input>;
  ```
- `entry.ts` をルートとして `Program` を作る
- `TypeCheker` を使って `entry.ts` 中の型を走査し、`Output` 型の情報を抜き出す

## 関連項目

TypeScript Compiler API を用いて型情報を抜き出す方法は、別のジャッジシステムで使うために既に用意してありました。経緯や使用例は以下にまとめています。

https://qiita.com/n4o847/items/5fb8e1cfe0344eee599f

以下はその際に作った CLI です。

https://github.com/n4o847/compile-time-typescript

## おわりに

本当はこれを使って精選 10 問を解きたいところでしたが、大変なのでできませんでした。

本記事中のスニペットは CC0 ライセンス (https://creativecommons.org/publicdomain/zero/1.0/deed.ja) のもとで公開としますので、誰かやってください。

実際にジャッジで使うとなると実行時間や再帰制限がシビアになりそうですね。