feat(vitepress-twoslash): cache type informations to improve performance (#798)

Co-authored-by: Anthony Fu <github@antfu.me>

Author: nakasyou

docs/.vitepress/config.ts

@@ -6,6 +6,7 @@ import { withMermaid } from 'vitepress-plugin-mermaid'
 import { version } from '../../package.json'
 import { transformerColorizedBrackets } from '../../packages/colorized-brackets/src'
 import { transformerMetaWordHighlight, transformerNotationWordHighlight, transformerRemoveNotationEscape } from '../../packages/transformers/src'
+import { createFileSystemTypesCache } from '../../packages/vitepress-twoslash/src/cache-fs'
 import { defaultHoverInfoProcessor, transformerTwoslash } from '../../packages/vitepress-twoslash/src/index'
 import vite from './vite.config'
 
@@ -125,6 +126,7 @@ export default withMermaid(defineConfig({
             // Remove shiki_core namespace
             .replace(/_shikijs_core\w*\./g, '')
         },
+        typesCache: createFileSystemTypesCache(),
       }),
       transformerRemoveNotationEscape(),
       transformerColorizedBrackets({ explicitTrigger: true }),

docs/packages/vitepress.md

@@ -117,3 +117,26 @@ onMounted(() => {
   </button>
 </template>
 ```
+
+### File System Cache
+
+To speed up the build process, you can enable the file system cache for the generated types, that shares across multiple builds. By default the cache is stored in the `.vitepress/cache/twoslash` along with other VitePress caches.
+
+In your [`.vitepress/config.ts`](https://vitepress.dev/reference/site-config):
+
+```ts
+// .vitepress/config.ts
+import { transformerTwoslash } from '@shikijs/vitepress-twoslash'
+import { createFileSystemTypesCache } from '@shikijs/vitepress-twoslash/cache-fs' // [!code hl]
+import { defineConfig } from 'vitepress'
+
+export default defineConfig({
+  markdown: {
+    codeTransformers: [
+      transformerTwoslash({
+        typesCache: createFileSystemTypesCache() // [!code hl]
+      })
+    ]
+  }
+})
+```

package.json

@@ -26,6 +26,8 @@
     "@rollup/plugin-node-resolve": "catalog:",
     "@rollup/plugin-replace": "catalog:",
     "@rollup/plugin-terser": "catalog:",
+    "@shikijs/engine-javascript": "workspace:*",
+    "@shikijs/engine-oniguruma": "workspace:*",
     "@shikijs/markdown-it": "workspace:*",
     "@shikijs/monaco": "workspace:*",
     "@shikijs/rehype": "workspace:*",
@@ -75,6 +77,8 @@
   "resolutions": {
     "@shikijs/compat": "workspace:*",
     "@shikijs/core": "workspace:*",
+    "@shikijs/engine-javascript": "workspace:*",
+    "@shikijs/engine-oniguruma": "workspace:*",
     "@shikijs/markdown-it": "workspace:*",
     "@shikijs/rehype": "workspace:*",
     "@shikijs/transformers": "workspace:*",

packages/vitepress-twoslash/build.config.ts

@@ -5,6 +5,7 @@ export default defineBuildConfig({
   entries: [
     'src/index.ts',
     'src/client.ts',
+    'src/cache-fs.ts',
   ],
   declaration: true,
   rollup: {

packages/vitepress-twoslash/package.json

@@ -27,6 +27,10 @@
       "types": "./dist/client.d.mts",
       "default": "./dist/client.mjs"
     },
+    "./cache-fs": {
+      "types": "./dist/cache-fs.d.mts",
+      "default": "./dist/cache-fs.mjs"
+    },
     "./style.css": "./style.css",
     "./style-core.css": "./style-core.css"
   },

packages/vitepress-twoslash/src/cache-fs.ts

@@ -0,0 +1,38 @@
+import type { TwoslashTypesCache } from './types'
+import { createHash } from 'node:crypto'
+import { existsSync, mkdirSync, readFileSync, writeFileSync } from 'node:fs'
+import { join, resolve } from 'node:path'
+import process from 'node:process'
+
+export interface FileSystemTypeResultCacheOptions {
+  /**
+   * The directory to store the cache files.
+   *
+   * @default '.vitepress/cache/twoslash'
+   */
+  dir?: string
+}
+
+export function createFileSystemTypesCache(options: FileSystemTypeResultCacheOptions = {}): TwoslashTypesCache {
+  const dir = resolve(process.cwd(), options.dir ?? '.vitepress/cache/twoslash')
+
+  return {
+    init() {
+      mkdirSync(dir, { recursive: true })
+    },
+    read(code) {
+      const hash = createHash('SHA256').update(code).digest('hex').slice(0, 12)
+      const filePath = join(dir, `${hash}.json`)
+      if (!existsSync(filePath)) {
+        return null
+      }
+      return JSON.parse(readFileSync(filePath, { encoding: 'utf-8' }))
+    },
+    write(code, data) {
+      const hash = createHash('SHA256').update(code).digest('hex').slice(0, 12)
+      const filePath = join(dir, `${hash}.json`)
+      const json = JSON.stringify(data)
+      writeFileSync(filePath, json, { encoding: 'utf-8' })
+    },
+  }
+}

packages/vitepress-twoslash/src/index.ts

@@ -1,26 +1,14 @@
 /* eslint-disable node/prefer-global/process */
-import type { TransformerTwoslashOptions } from '@shikijs/twoslash/core'
 import type { ShikiTransformer } from 'shiki'
-import type { VueSpecificOptions } from 'twoslash-vue'
-import type { TwoslashFloatingVueRendererOptions } from './renderer-floating-vue'
+import type { TwoslashExecuteOptions, TwoslashReturn } from 'twoslash'
+import type { VitePressPluginTwoslashOptions } from './types'
 import { createTransformerFactory } from '@shikijs/twoslash/core'
 import { removeTwoslashNotations } from 'twoslash'
 import { createTwoslasher } from 'twoslash-vue'
 import { rendererFloatingVue } from './renderer-floating-vue'
 
 export * from './renderer-floating-vue'
-
-interface TransformerTwoslashVueOptions extends TransformerTwoslashOptions {
-  twoslashOptions?: TransformerTwoslashOptions['twoslashOptions'] & VueSpecificOptions
-}
-
-export interface VitePressPluginTwoslashOptions extends TransformerTwoslashVueOptions, TwoslashFloatingVueRendererOptions {
-  /**
-   * Requires adding `twoslash` to the code block explicitly to run twoslash
-   * @default true
-   */
-  explicitTrigger?: TransformerTwoslashOptions['explicitTrigger']
-}
+export * from './types'
 
 /**
  * Create a Shiki transformer for VitePress to enable twoslash integration
@@ -30,6 +18,7 @@ export interface VitePressPluginTwoslashOptions extends TransformerTwoslashVueOp
 export function transformerTwoslash(options: VitePressPluginTwoslashOptions = {}): ShikiTransformer {
   const {
     explicitTrigger = true,
+    typesCache,
   } = options
 
   const onError = (error: any, code: string): void => {
@@ -44,9 +33,24 @@ export function transformerTwoslash(options: VitePressPluginTwoslashOptions = {}
     removeTwoslashNotations(code)
   }
 
-  const twoslash = createTransformerFactory(
-    createTwoslasher(options.twoslashOptions),
-  )({
+  const defaultTwoslasher = createTwoslasher(options.twoslashOptions)
+
+  let twoslasher = defaultTwoslasher
+  // Wrap twoslasher with cache when `resultCache` is provided
+  if (typesCache) {
+    twoslasher = ((code: string, extension?: string, options?: TwoslashExecuteOptions): TwoslashReturn => {
+      const cached = typesCache.read(code) // Restore cache
+      if (cached)
+        return cached
+
+      const twoslashResult = defaultTwoslasher(code, extension, options)
+      typesCache.write(code, twoslashResult)
+      return twoslashResult
+    }) as typeof defaultTwoslasher
+    twoslasher.getCacheMap = defaultTwoslasher.getCacheMap
+  }
+
+  const twoslash = createTransformerFactory(twoslasher)({
     langs: ['ts', 'tsx', 'js', 'jsx', 'json', 'vue'],
     renderer: rendererFloatingVue(options),
     onTwoslashError: onError,
@@ -59,6 +63,8 @@ export function transformerTwoslash(options: VitePressPluginTwoslashOptions = {}
     ? explicitTrigger
     : /\btwoslash\b/
 
+  typesCache?.init?.()
+
   return {
     ...twoslash,
     name: '@shikijs/vitepress-twoslash',

packages/vitepress-twoslash/src/types.ts

@@ -0,0 +1,55 @@
+import type { TransformerTwoslashOptions } from '@shikijs/twoslash/core'
+import type { TwoslashReturn } from 'twoslash'
+import type { VueSpecificOptions } from 'twoslash-vue'
+import type { TwoslashFloatingVueRendererOptions } from './renderer-floating-vue'
+
+interface TransformerTwoslashVueOptions extends TransformerTwoslashOptions {
+  twoslashOptions?: TransformerTwoslashOptions['twoslashOptions'] & VueSpecificOptions
+}
+
+export interface VitePressPluginTwoslashOptions extends TransformerTwoslashVueOptions, TwoslashFloatingVueRendererOptions {
+  /**
+   * Requires adding `twoslash` to the code block explicitly to run twoslash
+   * @default true
+   */
+  explicitTrigger?: TransformerTwoslashOptions['explicitTrigger']
+
+  /**
+   * The options for caching resolved types
+   *
+   * @example
+   * ```ts
+   * import { transformerTwoslash } from '@shikijs/vitepress-twoslash'
+   * import { createFileSystemTypesCache } from '@shikijs/vitepress-twoslash/cache-fs'
+   *
+   * transformerTwoslash({
+   *   typesCache: createFileSystemTypesCache({
+   *     dir: './my-cache-dir'
+   *   })
+   * })
+   * ```
+   */
+  typesCache?: TwoslashTypesCache
+}
+
+export interface TwoslashTypesCache {
+  /**
+   * Read cached result
+   *
+   * @param code Source code
+   */
+  read: (code: string) => TwoslashReturn | null
+
+  /**
+   * Save result to cache
+   *
+   * @param code Source code
+   * @param data Twoslash data
+   */
+  write: (code: string, data: TwoslashReturn) => void
+
+  /**
+   * On initialization
+   */
+  init?: () => void
+}

pnpm-lock.yaml

@@ -12,6 +12,7 @@
       "@shikijs/transformers": ["./packages/transformers/src/index.ts"],
       "@shikijs/twoslash/core": ["./packages/twoslash/src/core.ts"],
       "@shikijs/twoslash": ["./packages/twoslash/src/index.ts"],
+      "@shikijs/vitepress-twoslash/cache-fs": ["./packages/vitepress-twoslash/src/cache-fs.ts"],
       "@shikijs/vitepress-twoslash/client": ["./packages/vitepress-twoslash/src/client.ts"],
       "@shikijs/vitepress-twoslash": ["./packages/vitepress-twoslash/src/index.ts"],
       "@shikijs/markdown-it": ["./packages/markdown-it/src/index.ts"],