feat(rehype): support lazy load languages (#787)

Co-authored-by: Anthony Fu <[email protected]>

packages/rehype/src/core.ts

@@ -2,12 +2,12 @@ import type {
   CodeToHastOptions,
   HighlighterGeneric,
 } from '@shikijs/types'
-import type { Element, Root } from 'hast'
+import type { Root } from 'hast'
 import type { Transformer } from 'unified'
+import type { RehypeShikiHandler } from './handlers'
 import type { RehypeShikiCoreOptions } from './types'
-import { toString } from 'hast-util-to-string'
 import { visit } from 'unist-util-visit'
-import { InlineCodeProcessors } from './inline'
+import { InlineCodeHandlers, PreHandler } from './handlers'
 
 export * from './types'
 
@@ -17,7 +17,6 @@ function rehypeShikiFromHighlighter(
   highlighter: HighlighterGeneric<any, any>,
   options: RehypeShikiCoreOptions,
 ): Transformer<Root, Root> {
-  const langs = highlighter.getLoadedLanguages()
   const {
     addLanguageClass = false,
     parseMetaString,
@@ -27,18 +26,10 @@ function rehypeShikiFromHighlighter(
     onError,
     stripEndNewline = true,
     inline = false,
+    lazy = false,
     ...rest
   } = options
 
-  /**
-   * Get the determined language of code block (with default language & fallbacks)
-   */
-  function getLanguage(lang = defaultLanguage): string | undefined {
-    if (lang && fallbackLanguage && !langs.includes(lang))
-      return fallbackLanguage
-    return lang
-  }
-
   function highlight(
     lang: string,
     code: string,
@@ -92,65 +83,86 @@ function rehypeShikiFromHighlighter(
     }
   }
 
-  function processPre(tree: Root, node: Element): Root | undefined {
-    const head = node.children[0]
-
-    if (
-      !head
-      || head.type !== 'element'
-      || head.tagName !== 'code'
-      || !head.properties
-    ) {
-      return
-    }
+  return (tree) => {
+    // use this queue if lazy is enabled
+    const languageQueue: string[] = []
+    const queue: (() => void)[] = []
 
-    const classes = head.properties.className
-    const languageClass = Array.isArray(classes)
-      ? classes.find(
-        d => typeof d === 'string' && d.startsWith(languagePrefix),
-      )
-      : undefined
+    function getLanguage(lang: string | undefined): string | undefined {
+      if (!lang)
+        return defaultLanguage
 
-    const lang = getLanguage(
-      typeof languageClass === 'string'
-        ? languageClass.slice(languagePrefix.length)
-        : undefined,
-    )
+      if (highlighter.getLoadedLanguages().includes(lang))
+        return lang
 
-    if (!lang)
-      return
-
-    const code = toString(head)
-    const metaString = head.data?.meta ?? head.properties.metastring?.toString() ?? ''
-    const meta = parseMetaString?.(metaString, node, tree) || {}
+      if (lazy) {
+        languageQueue.push(lang)
+        return lang
+      }
 
-    return highlight(lang, code, metaString, meta)
-  }
+      if (fallbackLanguage)
+        return fallbackLanguage
+    }
 
-  return function (tree) {
     visit(tree, 'element', (node, index, parent) => {
+      let handler: RehypeShikiHandler | undefined
+
       // needed for hast node replacement
       if (!parent || index == null)
         return
 
       if (node.tagName === 'pre') {
-        const result = processPre(tree, node)
-
-        if (result) {
-          parent.children.splice(index, 1, ...result.children)
-        }
-
-        // don't look for the `code` node inside
-        return 'skip'
+        handler = PreHandler
       }
 
       if (node.tagName === 'code' && inline) {
-        const result = InlineCodeProcessors[inline]?.({ node, getLanguage, highlight })
-        if (result) {
-          parent.children.splice(index, 1, ...result.children)
+        handler = InlineCodeHandlers[inline]
+      }
+
+      if (!handler)
+        return
+
+      const res = handler(tree, node)
+      if (!res)
+        return
+
+      const lang = getLanguage(res.lang)
+      if (!lang)
+        return
+
+      const processNode = (): void => {
+        const meta = res.meta ? parseMetaString?.(res.meta, node, tree) : undefined
+
+        const fragment = highlight(lang, res.code, res.meta, meta ?? {})
+        if (!fragment)
+          return
+
+        if (res.type === 'inline') {
+          const head = fragment.children[0]
+          if (head.type === 'element' && head.tagName === 'pre') {
+            head.tagName = 'span'
+          }
         }
+
+        parent.children.splice(index, 1, ...fragment.children)
       }
+
+      if (lazy)
+        queue.push(processNode)
+      else
+        processNode()
+
+      // don't visit processed nodes
+      return 'skip'
     })
+
+    if (lazy) {
+      return highlighter
+        .loadLanguage(...languageQueue)
+        .then(() => {
+          queue.forEach(fn => fn())
+        })
+    }
   }
 }
 

packages/rehype/src/handlers.ts

@@ -0,0 +1,61 @@
+import type { Element, Root } from 'hast'
+import type { RehypeShikiCoreOptions } from './types'
+import { toString } from 'hast-util-to-string'
+
+type Truthy<T> = T extends false | '' | 0 | null | undefined ? never : T
+
+export type RehypeShikiHandler = (
+  tree: Root,
+  node: Element
+) => {
+  type: 'inline' | 'pre'
+  meta?: string
+  lang?: string
+  code: string
+} | undefined
+
+export const InlineCodeHandlers: Record<Truthy<RehypeShikiCoreOptions['inline']>, RehypeShikiHandler> = {
+  'tailing-curly-colon': (_tree, node) => {
+    const raw = toString(node)
+    const match = raw.match(/(.+)\{:([\w-]+)\}$/)
+    if (!match)
+      return
+
+    return {
+      type: 'inline',
+      code: match[1] ?? raw,
+      lang: match.at(2),
+    }
+  },
+}
+
+const languagePrefix = 'language-'
+
+export const PreHandler: RehypeShikiHandler = (_tree, node) => {
+  const head = node.children[0]
+
+  if (
+    !head
+    || head.type !== 'element'
+    || head.tagName !== 'code'
+    || !head.properties
+  ) {
+    return
+  }
+
+  const classes = head.properties.className
+  const languageClass = Array.isArray(classes)
+    ? classes.find(
+      d => typeof d === 'string' && d.startsWith(languagePrefix),
+    )
+    : undefined
+
+  return {
+    type: 'pre',
+    lang: typeof languageClass === 'string'
+      ? languageClass.slice(languagePrefix.length)
+      : undefined,
+    code: toString(head),
+    meta: head.data?.meta ?? head.properties.metastring?.toString() ?? '',
+  }
+}

packages/rehype/src/types.ts

@@ -27,9 +27,19 @@ export interface RehypeShikiExtraOptions {
 
   /**
    * The fallback language to use when specified language is not loaded
+   *
+   * Ignored if `lazy` is enabled
    */
   fallbackLanguage?: string
 
+  /**
+   * Load languages and themes on-demand.
+   * When enable, this would make requires the unified pipeline to be async.
+   *
+   * @default false
+   */
+  lazy?: boolean
+
   /**
    * `mdast-util-to-hast` adds a newline to the end of code blocks
    *

packages/rehype/test/core.test.ts

@@ -38,6 +38,31 @@ it('run', async () => {
   expect(file.toString()).toMatchFileSnapshot('./fixtures/a.core.out.html')
 })
 
+it('run with lazy', async () => {
+  const highlighter = await createHighlighter({
+    themes: [
+      'vitesse-light',
+    ],
+    langs: [],
+  })
+
+  const file = await unified()
+    .use(remarkParse)
+    .use(remarkRehype)
+    .use(rehypeShikiFromHighlighter, highlighter, {
+      lazy: true,
+      theme: 'vitesse-light',
+      defaultLanguage: 'text',
+      transformers: [
+        transformerMetaHighlight(),
+      ],
+    })
+    .use(rehypeStringify)
+    .process(await fs.readFile(new URL('./fixtures/a.md', import.meta.url)))
+
+  expect(file.toString()).toMatchFileSnapshot('./fixtures/a.core.out.html')
+})
+
 it('run with rehype-raw', async () => {
   const highlighter = await createHighlighter({
     themes: [