vocs 2.0.6 → 2.0.7

package/src/internal/twoslash/file-patcher.ts

@@ -0,0 +1,60 @@
+import * as node_fs from 'node:fs'
+import MagicString from 'magic-string'
+
+/**
+ * Accumulates character-range patches per file and flushes them all at once.
+ *
+ * Used by the twoslash inline cache to write `// @twoslash-cache: ...` comments
+ * back into the original markdown source after a code block has been processed.
+ *
+ * Ported from `@shikijs/vitepress-twoslash`'s `FilePatcher`.
+ */
+export class FilePatcher {
+  private files = new Map<string, { content: string; patches: Map<string, string> } | null>()
+
+  /** Build a patch key from a character range. Omitting `to` denotes an insertion. */
+  static key(from: number, to?: number): string {
+    return `${from}${to ? `:${to}` : ''}`
+  }
+
+  /**
+   * Load a file's contents (cached for the lifetime of the patcher, until
+   * flushed via {@link patch}). Returns `null` if the file does not exist.
+   */
+  load(path: string): { content: string; patches: Map<string, string> } | null {
+    let file = this.files.get(path)
+    if (file === undefined) {
+      if (node_fs.existsSync(path)) {
+        const content = node_fs.readFileSync(path, { encoding: 'utf-8' })
+        file = { content, patches: new Map() }
+      } else {
+        file = null
+      }
+      this.files.set(path, file)
+    }
+    return file
+  }
+
+  /** Apply all queued patches for `path` and write the result back to disk. */
+  patch(path: string): void {
+    const file = this.files.get(path)
+    if (!file) return
+
+    if (file.patches.size) {
+      const s = new MagicString(file.content)
+
+      for (const [key, value] of file.patches) {
+        const [from, to] = key.split(':').map((x) => (x !== '' ? Number(x) : undefined))
+        if (from === undefined) continue
+
+        if (to !== undefined) s.update(from, to, value)
+        else s.appendRight(from, value)
+      }
+
+      const content = s.toString()
+      if (content !== file.content) node_fs.writeFileSync(path, content, { encoding: 'utf-8' })
+    }
+
+    this.files.delete(path)
+  }
+}

package/src/internal/twoslash/index.ts

@@ -1,3 +1,4 @@
 export * from './experimental_rust.js'
+export * as InlineCache from './inline-cache.js'
 export * as Renderer from './renderer.js'
 export * as TypesCache from './types-cache.js'

package/src/internal/twoslash/inline-cache.test.ts

@@ -0,0 +1,158 @@
+import * as fs from 'node:fs'
+import * as os from 'node:os'
+import * as path from 'node:path'
+import { afterEach, describe, expect, it } from 'vitest'
+import {
+  createInlineTypesCache,
+  extractSourceMapComment,
+  injectSourceMapComment,
+} from './inline-cache.js'
+
+describe('source map codec', () => {
+  it('round-trips a source map through inject/extract', () => {
+    const sourceMap = { path: '/docs/a.md', from: 15, to: 42 }
+    const injected = injectSourceMapComment('const a = 1', sourceMap)
+    expect(injected.startsWith('// @vocs-twoslash-source:')).toBe(true)
+
+    const result = extractSourceMapComment(injected)
+    expect(result.code).toBe('const a = 1')
+    expect(result.sourceMap).toEqual(sourceMap)
+  })
+
+  it('returns null source map when none present', () => {
+    const result = extractSourceMapComment('const a = 1')
+    expect(result.code).toBe('const a = 1')
+    expect(result.sourceMap).toBeNull()
+  })
+})
+
+describe('inline types cache', () => {
+  const tmpFiles: string[] = []
+
+  afterEach(() => {
+    for (const file of tmpFiles.splice(0)) fs.rmSync(file, { force: true })
+  })
+
+  function createMarkdown(body: string) {
+    const file = path.join(fs.mkdtempSync(path.join(os.tmpdir(), 'vocs-inline-cache-')), 'page.md')
+    const md = `\`\`\`ts twoslash\n${body}\n\`\`\`\n`
+    fs.writeFileSync(file, md, 'utf-8')
+    tmpFiles.push(file)
+    // body starts right after the first newline (the opening fence line).
+    const from = md.indexOf('\n') + 1
+    return { file, md, from, to: md.length }
+  }
+
+  // Pull the on-disk fence body the way shiki would see it (everything between
+  // the opening fence line and the closing fence).
+  function readBody(file: string, from: number) {
+    const content = fs.readFileSync(file, 'utf-8')
+    return content.slice(from, content.lastIndexOf('\n```'))
+  }
+
+  it('writes a cache comment back to the source and reads it on the next pass', () => {
+    const code = 'const a: string = "x"'
+    const { file, from, to } = createMarkdown(code)
+
+    const data = { nodes: [], code: 'compiled-output' }
+
+    // --- Pass 1: cache miss -> twoslash runs -> write back ---
+    {
+      const { typesCache, patcher } = createInlineTypesCache()
+      const meta = { sourceMap: { path: file, from, to } } as never
+
+      const preprocessed = typesCache.preprocess?.(code, 'ts', {}, meta)
+      expect(preprocessed).toBe(code)
+      expect(typesCache.read(code, 'ts', {}, meta)).toBeNull()
+
+      typesCache.write(code, data as never, 'ts', {}, meta)
+      patcher.patch(file)
+    }
+
+    const written = fs.readFileSync(file, 'utf-8')
+    expect(written).toContain('// @twoslash-cache:')
+
+    // --- Pass 2: cache hit -> twoslash skipped ---
+    {
+      const { typesCache } = createInlineTypesCache()
+      const body = readBody(file, from)
+      const meta = { sourceMap: { path: file, from, to } } as never
+
+      const preprocessed = typesCache.preprocess?.(body, 'ts', {}, meta)
+      // the cache comment is stripped before twoslash would run
+      expect(preprocessed).toBe(code)
+
+      const cached = typesCache.read(body, 'ts', {}, meta)
+      expect(cached).toEqual(data)
+    }
+  })
+
+  it('invalidates the cache when the hash no longer matches', () => {
+    const code = 'const a: string = "x"'
+    const { file, from, to } = createMarkdown(code)
+
+    // Seed a cache entry.
+    {
+      const { typesCache, patcher } = createInlineTypesCache()
+      const meta = { sourceMap: { path: file, from, to } } as never
+      typesCache.preprocess?.(code, 'ts', {}, meta)
+      typesCache.write(code, { nodes: [], code: 'x' } as never, 'ts', {}, meta)
+      patcher.patch(file)
+    }
+
+    // Read with different code -> hash mismatch -> miss.
+    {
+      const { typesCache } = createInlineTypesCache()
+      const body = readBody(file, from)
+      const meta = { sourceMap: { path: file, from, to } } as never
+      // The body still has the old cache line, but we validate against new code.
+      typesCache.preprocess?.(`${body}\nconst b = 2`, 'ts', {}, meta)
+      expect(typesCache.read(body, 'ts', {}, meta)).toBeNull()
+    }
+  })
+
+  it('ignoreCache skips reading existing cache', () => {
+    const code = 'const a: string = "x"'
+    const { file, from, to } = createMarkdown(code)
+
+    {
+      const { typesCache, patcher } = createInlineTypesCache()
+      const meta = { sourceMap: { path: file, from, to } } as never
+      typesCache.preprocess?.(code, 'ts', {}, meta)
+      typesCache.write(code, { nodes: [], code: 'x' } as never, 'ts', {}, meta)
+      patcher.patch(file)
+    }
+
+    {
+      const { typesCache } = createInlineTypesCache({ ignoreCache: true })
+      const body = readBody(file, from)
+      const meta = { sourceMap: { path: file, from, to } } as never
+      typesCache.preprocess?.(body, 'ts', {}, meta)
+      expect(typesCache.read(body, 'ts', {}, meta)).toBeNull()
+    }
+  })
+
+  it('remove mode strips the cache comment from the source', () => {
+    const code = 'const a: string = "x"'
+    const { file, from, to } = createMarkdown(code)
+
+    {
+      const { typesCache, patcher } = createInlineTypesCache()
+      const meta = { sourceMap: { path: file, from, to } } as never
+      typesCache.preprocess?.(code, 'ts', {}, meta)
+      typesCache.write(code, { nodes: [], code: 'x' } as never, 'ts', {}, meta)
+      patcher.patch(file)
+    }
+    expect(fs.readFileSync(file, 'utf-8')).toContain('// @twoslash-cache:')
+
+    {
+      const { typesCache, patcher } = createInlineTypesCache({ remove: true })
+      const body = readBody(file, from)
+      const meta = { sourceMap: { path: file, from, to } } as never
+      typesCache.preprocess?.(body, 'ts', {}, meta)
+      typesCache.write(body, { nodes: [], code: 'x' } as never, 'ts', {}, meta)
+      patcher.patch(file)
+    }
+    expect(fs.readFileSync(file, 'utf-8')).not.toContain('// @twoslash-cache:')
+  })
+})

package/src/internal/twoslash/inline-cache.ts

@@ -0,0 +1,282 @@
+import * as crypto from 'node:crypto'
+import type { TwoslashShikiReturn, TwoslashTypesCache } from '@shikijs/twoslash'
+import LZString from 'lz-string'
+import { getObjectHash, type TwoslashExecuteOptions } from 'twoslash/core'
+import { FilePatcher } from './file-patcher.js'
+
+/**
+ * Inline twoslash cache.
+ *
+ * Persists the serialized twoslash result directly into the markdown source as a
+ * `// @twoslash-cache: ...` comment inside the fenced code block. On the next
+ * build the comment is read, validated against a hash of the source, and used
+ * directly — skipping the TypeScript compiler entirely.
+ *
+ * Unlike the filesystem types cache (which lives in a separate, usually
+ * gitignored directory), the inline cache travels with the source files, so it
+ * stays warm across fresh clones and CI runs.
+ *
+ * Ported from `@shikijs/vitepress-twoslash`'s inline cache, adapted to vocs's
+ * MDX/remark pipeline:
+ * - source-map injection happens via a remark plugin (see `remarkInlineCache` in
+ *   `mdx.ts`) instead of a Vite `load` hook + markdown-it.
+ * - the queued patches are flushed by the `vocs:mdx` Vite plugin after each file
+ *   is transformed.
+ */
+
+/** Source position of a fenced code block's body within its markdown file. */
+export type SourceMap = {
+  path: string
+  from: number
+  to: number
+}
+
+declare module '@shikijs/types' {
+  interface ShikiTransformerContextMeta {
+    sourceMap?: SourceMap | null
+    __cache?: TwoslashShikiReturn
+    __patch?: (newCache: string) => void
+  }
+}
+
+/**
+ * Tag used for the ephemeral source-map comment. This comment is injected into
+ * the in-memory code value during compilation and stripped again before
+ * twoslash/shiki run — it is never written to disk.
+ */
+const SOURCE_MAP_KEY = '@vocs-twoslash-source'
+const SOURCE_MAP_REGEX = new RegExp(`// ${SOURCE_MAP_KEY}:(.*)(?:\n|$)`)
+
+export function injectSourceMapComment(value: string, sourceMap: SourceMap): string {
+  return `// ${SOURCE_MAP_KEY}:${JSON.stringify(sourceMap)}\n${value}`
+}
+
+export function extractSourceMapComment(code: string): {
+  code: string
+  sourceMap: SourceMap | null
+} {
+  let sourceMap: SourceMap | null = null
+  try {
+    code = code.replace(SOURCE_MAP_REGEX, (_, p1: string) => {
+      sourceMap = JSON.parse(p1) as SourceMap
+      return ''
+    })
+  } catch {
+    // ignore malformed source map
+  }
+  return { code, sourceMap }
+}
+
+type CachePayload = {
+  v: number
+  hash: string
+  data: string
+}
+
+const CODE_INLINE_CACHE_KEY = '@twoslash-cache'
+const CODE_INLINE_CACHE_REGEX = new RegExp(`// ${CODE_INLINE_CACHE_KEY}: (.*)(?:\n|$)`, 'g')
+
+export function createInlineTypesCache(
+  options: { remove?: boolean | undefined; ignoreCache?: boolean | undefined } = {},
+): { typesCache: TwoslashTypesCache; patcher: FilePatcher } {
+  const { remove, ignoreCache } = options
+  const patcher = new FilePatcher()
+
+  const optionsHashCache = new WeakMap<TwoslashExecuteOptions, string>()
+  function getOptionsHash(options: TwoslashExecuteOptions = {}): string {
+    let hash = optionsHashCache.get(options)
+    if (!hash) {
+      hash = getObjectHash(options)
+      optionsHashCache.set(options, hash)
+    }
+    return hash
+  }
+
+  function cacheHash(code: string, lang?: string, options?: TwoslashExecuteOptions): string {
+    return crypto
+      .createHash('sha256')
+      .update(`${getOptionsHash(options)}:${lang ?? ''}:${code}`)
+      .digest('hex')
+  }
+
+  function stringifyCachePayload(
+    data: TwoslashShikiReturn,
+    code: string,
+    lang?: string,
+    options?: TwoslashExecuteOptions,
+  ): string {
+    const payload: CachePayload = {
+      v: 1,
+      hash: cacheHash(code, lang, options),
+      data: LZString.compressToBase64(JSON.stringify(data)),
+    }
+    return JSON.stringify(payload)
+  }
+
+  function resolveCachePayload(cache: string): {
+    payload: CachePayload
+    twoslash: () => TwoslashShikiReturn | null
+  } | null {
+    if (!cache) return null
+    try {
+      const payload = JSON.parse(cache) as CachePayload
+      if (payload.v === 1) {
+        return {
+          payload,
+          twoslash: () => {
+            try {
+              return JSON.parse(LZString.decompressFromBase64(payload.data))
+            } catch {
+              return null
+            }
+          },
+        }
+      }
+    } catch {
+      // ignore malformed payload
+    }
+    return null
+  }
+
+  function resolveSourcePatcher(
+    source: SourceMap,
+    search?: string,
+  ): ((newCache: string) => void) | undefined {
+    const file = patcher.load(source.path)
+    if (file === null) return undefined
+
+    const range: { from: number; to?: number } = { from: source.from }
+    let linebreak = true
+
+    if (search) {
+      const cachePos = file.content.indexOf(search, source.from)
+      if (cachePos !== -1 && cachePos < source.to) {
+        range.from = cachePos
+        range.to = cachePos + search.length
+        linebreak = search.endsWith('\n')
+      }
+    }
+
+    const patchKey = FilePatcher.key(range.from, range.to)
+    return (newCache: string) => {
+      if (newCache === '') {
+        // remove the existing cache comment if one was found
+        if (range.to !== undefined) file.patches.set(patchKey, '')
+        return
+      }
+      file.patches.set(patchKey, newCache + (linebreak ? '\n' : ''))
+    }
+  }
+
+  const typesCache: TwoslashTypesCache = {
+    preprocess(code, lang, options, meta) {
+      if (!meta) return
+
+      let rawCache = ''
+      let cacheString = ''
+
+      code = code.replaceAll(CODE_INLINE_CACHE_REGEX, (full, p1: string) => {
+        // keep only the first occurrence (duplicates may appear via @include)
+        if (!rawCache.length) {
+          cacheString = p1
+          rawCache = full
+        }
+        return ''
+      })
+
+      const shouldLoadCache = !ignoreCache && !remove
+      if (shouldLoadCache) {
+        const cache = resolveCachePayload(cacheString)
+        if (cache?.payload.hash === cacheHash(code, lang, options)) {
+          const twoslash = cache.twoslash()
+          if (twoslash) meta.__cache = twoslash
+        }
+      }
+
+      if (meta.sourceMap) {
+        const patch = resolveSourcePatcher(meta.sourceMap, rawCache)
+        if (patch) meta.__patch = patch
+      }
+
+      return code
+    },
+    read(_code, _lang, _options, meta) {
+      return meta?.__cache ?? null
+    },
+    write(code, data, lang, options, meta) {
+      if (remove) {
+        meta?.__patch?.('')
+        return
+      }
+      const twoslashShiki = simplifyTwoslashReturn(data)
+      const cacheStr = `// ${CODE_INLINE_CACHE_KEY}: ${stringifyCachePayload(twoslashShiki, code, lang, options)}`
+      meta?.__patch?.(cacheStr)
+    },
+  }
+
+  return { typesCache, patcher }
+}
+
+/** Keep only the fields shiki needs when serializing a twoslash result. */
+function simplifyTwoslashReturn(ret: TwoslashShikiReturn): TwoslashShikiReturn {
+  return {
+    nodes: ret.nodes,
+    code: ret.code,
+    ...(ret.meta?.extension !== undefined ? { meta: { extension: ret.meta.extension } } : {}),
+  }
+}
+
+function isEnabledEnv(key: string): boolean | null {
+  const val = process.env?.[key]?.toLowerCase()
+  if (val)
+    return (
+      (
+        {
+          true: true,
+          false: false,
+          1: true,
+          0: false,
+          yes: true,
+          no: false,
+          y: true,
+          n: false,
+        } as Record<string, boolean>
+      )[val] ?? null
+    )
+  return null
+}
+
+/**
+ * Resolve whether the inline cache is enabled. The `TWOSLASH_INLINE_CACHE`
+ * environment variable takes precedence over the config flag.
+ */
+export function enabled(configFlag?: boolean | undefined): boolean {
+  const env = isEnabledEnv('TWOSLASH_INLINE_CACHE')
+  if (env !== null) return env
+  return Boolean(configFlag)
+}
+
+function envOptions(): { remove: boolean; ignoreCache: boolean } {
+  return {
+    remove: isEnabledEnv('TWOSLASH_INLINE_CACHE_REMOVE') === true,
+    ignoreCache: isEnabledEnv('TWOSLASH_INLINE_CACHE_IGNORE') === true,
+  }
+}
+
+let current: { typesCache: TwoslashTypesCache; patcher: FilePatcher } | undefined
+
+/** Get (or lazily create) the process-wide inline cache instance. */
+export function getOrCreate(): { typesCache: TwoslashTypesCache; patcher: FilePatcher } {
+  if (!current) current = createInlineTypesCache(envOptions())
+  return current
+}
+
+/** Flush any queued write-backs for `path` to disk. No-op if not initialized. */
+export function flush(path: string): void {
+  current?.patcher.patch(path)
+}
+
+/** Reset the singleton (primarily for tests). */
+export function reset(): void {
+  current = undefined
+}

package/src/internal/vite-plugins.ts

@@ -15,6 +15,7 @@ import * as Mdx from './mdx.js'
 import { SearchDocuments, SearchIndex } from './search.js'
 import * as ShikiTransformers from './shiki-transformers.js'
 import * as TaskRunner from './task-runner.js'
+import * as InlineCache from './twoslash/inline-cache.js'
 
 export { default as icons } from 'unplugin-icons/vite'
 export { default as arraybuffer } from 'vite-plugin-arraybuffer'
@@ -282,6 +283,9 @@ export function mdx(config: Config.Config): PluginOption {
 
       try {
         const result = normalizeTransformResult(await plugin.transform.call(this, code, id))
+        // Flush any inline twoslash cache write-backs queued during this
+        // file's transform (no-op when the inline cache is disabled).
+        InlineCache.flush(id)
         if (result) cache.set(id, { code, result })
         return result
       } catch (error) {