Open2024/03/19にコメント追加3

typesafe な tagged template のテンプレートエンジン

AI 用のプロンプトを大量に生成してると、いい感じに推論されるテンプレートエンジンがほしくなります。

というわけで TypeScript の Tagged Template Literal で作ってみました。

function hash(str: string) {
  let hash = 5381;
  let i = str.length;
  while (i) {
    hash = (hash * 33) ^ str.charCodeAt(--i);
  }
  return hash >>> 0;
}

type Tmpl<T extends Record<string, string>> = string & { _types: T };

type TupleToUniqueUnion<T> =
  T extends [infer First, ...infer Rest]
  ? string extends First
  ? TupleToUniqueUnion<Rest>
  : First | TupleToUniqueUnion<Rest>
  : never;

export function tpl<
  In extends string[],
  K extends string = TupleToUniqueUnion<In>,
  P extends string = K extends `@${infer V}` ? V : never,
  Final extends {} = { [k in P]: string }
>(input: TemplateStringsArray, ...params: In): (values: Final) => string {
  const prebuilt = build(input, ...params);
  const fn: any = (values: Final) => {
    return format(prebuilt, values as any);
  }
  return fn;
}

function build<
  In extends string[],
  K extends string = TupleToUniqueUnion<In>,
  P extends string = K extends `@${infer V}` ? V : never
>(input: TemplateStringsArray, ...params: In): Tmpl<{ [k in P]: string }> {
  let out = '';
  for (let i = 0; i < input.length; i++) {
    const t = input[i];
    const param = params[i];
    if (param === undefined) {
      continue;
    }
    if (param.startsWith('@')) {
      const insertion = `_@_{${hash(param.slice(1) as string)}}`;
      out += t + insertion;
    } else {
      out += t + param;
    }
  }
  return out as Tmpl<any>;
}

function format<
  T extends Tmpl<Record<string, string>>,
>(
  t: T,
  values: T['_types'],
): string {
  const hashedValues = Object.fromEntries(
    Object.entries(values).map(([k, v]) => [hash(k), v])
  );

  const built = dedent(t);
  const minIndent = getMinIndent(built);
  const indented = (content: string, replacer: string) => {
    const at = built.indexOf(replacer);
    const lastNewline = built.lastIndexOf('\n', at);
    const indentSize = built.slice(lastNewline + 1, at).match(/^[ \t]*/)?.[0]?.length ?? 0;
    return indent(content, Math.max(indentSize - minIndent, 0));
  }

  return built.replaceAll(/\_\@\_{(\d+)}/g, (_matched, hash) => {
    const v = hashedValues[hash];
    return indented(v, `_@_{${hash}}`);
  }).trim();
}

function getMinIndent(str: string) {
  const match = str.match(/^[ \t]*(?=\S)/gm);
  if (!match) {
    return 0;
  }
  return Math.min(...match.map((x) => x.length));
}

export function dedent(str: string) {
  const matched = str.match(/^[ \t]*(?=\S)/gm);
  if (!matched) {
    return str;
  }
  const indent = Math.min(...matched.map((x) => x.length));
  const re = new RegExp(`^[ \\t]{${indent}}`, 'gm');
  const result = indent > 0 ? str.replaceAll(re, '') : str;
  return result.replaceAll(/^\n/g, '');
}

export function indent(str: string, count: number) {
  if (count === 0) return str;
  const prefix = ' '.repeat(count);
  return str.split('\n').map((x, idx) => {
    if (idx === 0) return x;
    return prefix + x
  }).join('\n');
}

使い方

const n = 1;

const template = build`
  ${'@a'} = ${'@b'} | ${'raw insert value'} n: ${n.toString()}
    ${'@ml'}
`;

const result = format(template, {
  a: 'x',
  b: 'y',
  ml: 'd\nx',
});

console.log(result);

const t = tpl`${'@a'} = ${'@b'} | ${'raw insert value'} n: ${n.toString()}`;
console.log(`template`, tpl`${'@x'}`({ x: 'hello' }), t({ a: 'x', b: 'y' }));

先に型を宣言してから展開するのではなく、tagged template literal で${'@foo'} みたいなプロパティ名を宣言したら、逆にそれが推論されて {foo: string} のような型が決まります。

インデントされている場所にプロパティを挿入するときは、改行してもインデント状態を保持します。なのでこういう出力になります。

x = y | raw insert value n: 1
  d
  x
template hello x = y | raw insert value n: 1

mizchi

今のままでも便利ですが、このままではリスト展開や分岐を行うのに、テンプレート外でのロジック処理が必要になってしまっています。

ある程度内にロジックを記述できるように頑張ってみました。具体的にはテンプレートの型合成、 each(key, tpl) と when(key, tpl) を実装します。

function hash(str: string) {
  let hash = 5381;
  let i = str.length;
  while (i) {
    hash = (hash * 33) ^ str.charCodeAt(--i);
  }
  return hash >>> 0;
}

type Identity<T> = T extends object ? {} & {
  [P in keyof T]: T[P]
} : T;

type ItemRenderer<T> = {
  kind: 'item',
  (values: T, hashed?: Record<number, any>): string;
};

type WhenRenderer<K extends string, T, U> = {
  kind: 'when',
  key: K;
  _elseType: U;
  (values: T, hashed?: Record<number, any>): string;
  else(values: U, hashed?: Record<number, any>): string;
};

type ListRenderer<K extends string, T> = {
  kind: 'list';
  key: K;
  joiner: string;
  _type: T;
  (values: T, hashed?: Record<number, any>): string;
};

type TupleToUniqueUnion<T> =
  T extends [infer First, ...infer Rest]
  ? First extends (...args: any[]) => any
  ? TupleToUniqueUnion<Rest>
  : string extends First
  ? TupleToUniqueUnion<Rest>
  : First | TupleToUniqueUnion<Rest>
  : never;

type ToSubValues<T> =
  T extends [infer First, ...infer Rest]
  ? First extends ItemRenderer<infer P>
  ? P & ToSubValues<Rest>
  : ToSubValues<Rest>
  : {};

type ToWhenValues<T> =
  T extends [infer First, ...infer Rest]
  ? First extends WhenRenderer<infer K, infer P, infer U>
  ? { [k in K]: boolean } & P & U & ToWhenValues<Rest>
  : ToWhenValues<Rest>
  : {};

type ToListValues<T> =
  T extends [infer First, ...infer Rest]
  ? First extends ListRenderer<infer K, infer P>
  ? { [k in K]: P[] } & ToListValues<Rest>
  : ToListValues<Rest>
  : {};

type RawKey<T extends string> = T extends `${'@' | '!'}${infer R}` ? R : never;

export function each<K extends `${'@' | '!'}${string}`, T>(key: K, tpl: ItemRenderer<T>, joiner?: string): ListRenderer<RawKey<K>, T> {
  const fn: any = (values: T) => {
    return tpl(values);
  }
  fn.kind = 'list';
  fn.key = key.slice(1);
  fn.joiner = joiner ?? '\n';
  return fn;
}

export function when<K extends `${'@' | '!'}${string}`, T, U>(key: K, tpl: ItemRenderer<T>, else_?: ItemRenderer<U>): WhenRenderer<RawKey<K>, T, U> {
  const fn: any = (values: T) => {
    return tpl(values);
  }
  fn.kind = 'when';
  fn.key = key.slice(1);
  fn.else = else_;
  return fn;
}

export function tpl<
  I extends Array<string | ItemRenderer<any> | ListRenderer<any, any> | WhenRenderer<any, any, any>>,
  K extends string = TupleToUniqueUnion<I>,
  P extends string = RawKey<K>,
  Final = Identity<{ [k in P | keyof ToSubValues<I>]: string } & ToListValues<I> & ToWhenValues<I>>
>(input: TemplateStringsArray, ...params: I): ItemRenderer<Final> {
  let built = '';
  const subs: Array<ItemRenderer<any> | ListRenderer<any, any> | WhenRenderer<any, any, any>> = [];
  for (let i = 0; i < input.length; i++) {
    const t = input[i];
    const param = params[i];
    built += t;
    if (param === undefined) {
      continue;
    }
    if (typeof param === 'function') {
      // TODO: sub template
      built += `_fn_{${subs.length}}`;
      subs.push(param);
    } else if (param.startsWith('@') || param.startsWith('!')) {
      const insertion = `_@_{${hash(param.slice(1) as string)}}`;
      built += insertion;
    } else {
      built += param;
    }
  }

  built = dedent(built);
  const minIndent = getMinIndent(built);

  const indented = (content: string, replacer: string) => {
    const at = built.indexOf(replacer);
    const lastNewline = built.lastIndexOf('\n', at);
    const indentSize = built.slice(lastNewline + 1, at).match(/^[ \t]*/)?.[0]?.length ?? 0;
    return indent(content, Math.max(indentSize - minIndent, 0));
  }
  const fn: any = (values: Final) => {
    const hashed = Object.fromEntries(
      Object.entries(values as any).map(([k, v]) => [hash(k), v])
    ) as any;
    return built
      .replaceAll(/_[@\!]_{(\d+)}/g, (matched, hash) => {
        const content = hashed[hash];
        const isRaw = matched.startsWith('_!_');
        if (isRaw) {
          return content;
        }
        return indented(content, matched);
      })
      .replaceAll(/\_fn\_{(\d+)}/g, (matched, fnIdx) => {
        const r = subs[Number(fnIdx)];
        if (r.kind === 'list') {
          const childValues = hashed[hash(r.key)];
          return childValues.map((v: any) => r(v)).map((v: string) => indented(v, matched)).join(r.joiner);
        } else if (r.kind === 'when') {
          const val = hashed[hash(r.key)];
          if (val) {
            const content = r(values, hashed);
            return indented(content, matched);
          } else if (r.else) {
            return r.else(values, hashed);
          } else {
            return ''
          }
        } else {
          const content = r(values, hashed);
          return indented(content, matched);
        }
      })
      .trim();
  };
  fn.kind = 'item';
  return fn;
}

export function dedent(str: string) {
  const matched = str.match(/^[ \t]*(?=\S)/gm);
  if (!matched) {
    return str;
  }
  const indent = Math.min(...matched.map((x) => x.length));
  const re = new RegExp(`^[ \\t]{${indent}}`, 'gm');
  const result = indent > 0 ? str.replaceAll(re, '') : str;
  return result.replaceAll(/^\n/g, '');
}

export function indent(str: string, count: number) {
  if (count === 0) return str;
  const prefix = ' '.repeat(count);
  return str.split('\n').map((x, idx) => {
    if (idx === 0) return x;
    return prefix + x
  }).join('\n');
}

function getMinIndent(str: string) {
  const match = str.match(/^[ \t]*(?=\S)/gm);
  if (!match) {
    return 0;
  }
  return Math.min(...match.map((x) => x.length));
}

tagged template literal で宣言済みのテンプレートを挿入すると、テンプレート同士で型が合成される。
each('@items', Renderer<T>) で {items: T[]} な型を推論する
when('@flag', Renderer<T>) で { flag: boolean } & T の型を推論する

使い方

{
  const template = tpl`
    ${each('@items', tpl`
      ${'@key'} = ${'@value'}
    `)}
    ${each('@xs', tpl`
      ${'@v'}
    `, ' ')}
    ${'@key'} on top
    ${when('@flag', tpl`
      ${'@key'} is active
    `)}
    ${when('@flag2',
    tpl`${'@thenKey'} is active`,
    tpl`else: ${'@elseKey'}`
  )}
  `;

  const result = template({
    flag: true,
    flag2: false,
    thenKey: 'then',
    elseKey: 'else',
    key: 'key0',
    xs: [{ v: 'x' }, { v: 'y' }],
    items: [
      { key: 'x', value: 'y' },
      { key: 'z', value: 'w' },
    ]
  });

  console.log(result);
}

出力

x = y | raw
  d
  nl
x = y
z = w
x y
key0 on top
key0 is active
else: else

自分がほしい機能を積み込んでみただけですが、こういう機能をサクッと作れるのが TypeScript のいいところですね。

mizchi

今回このテンプレートエンジンを実装するのにあたって苦労したのが、 ['a', 'b', string] な Tuple を 'a' | 'b' のユニオン型に展開するための型ユーティリティを実装するところでした。

type R1 = ['a', 'b', string][number] // 'a' | 'b' | string 型を簡約した段階で string になってしまう

type TupleToUniqueUnion<T> =
  T extends [infer First, ...infer Rest]
  ? string extends First
  ? TupleToUniqueUnion<Rest>
  : First | TupleToUniqueUnion<Rest>
  : never;

// Tuple を一つずつ分解し string 型を除外することで string literal 型のユニオンが実現できる
type R2 = TupleToUniqueUnion<['a', 'b', string]> // => 'a' | 'b'

'string extends ${string}' の否定で string 型を除外できることは、結構使うテクニックです。