feat(perf): lazy embed languages bundle for SFCs and Docs (#791)

Author: antfu

packages/markdown-it/test/index.test.ts

@@ -1,12 +1,18 @@
 import fs from 'node:fs/promises'
 import { transformerMetaHighlight } from '@shikijs/transformers'
 import MarkdownIt from 'markdown-it'
+import { createHighlighter } from 'shiki'
 import { expect, it } from 'vitest'
 import Shiki from '../src'
+import { fromHighlighter } from '../src/core'
 
-it('run for base', async () => {
+it('run for base', { timeout: 10_000 }, async () => {
   const md = MarkdownIt()
-  md.use(await Shiki({
+  const shiki = await createHighlighter({
+    langs: ['js'],
+    themes: ['vitesse-light', 'vitesse-dark'],
+  })
+  md.use(fromHighlighter(shiki, {
     themes: {
       light: 'vitesse-light',
       dark: 'vitesse-dark',
@@ -19,11 +25,12 @@ it('run for base', async () => {
   const result = md.render(await fs.readFile(new URL('./fixtures/a.md', import.meta.url), 'utf-8'))
 
   expect(result).toMatchFileSnapshot('./fixtures/a.out.html')
-}, { timeout: 10_000 })
+})
 
-it('run for fallback language', async () => {
+it('run for fallback language', { timeout: 10_000 }, async () => {
   const md = MarkdownIt()
   md.use(await Shiki({
+    langs: ['js'],
     themes: {
       light: 'vitesse-light',
       dark: 'vitesse-dark',
@@ -37,15 +44,16 @@ it('run for fallback language', async () => {
   const result = md.render(await fs.readFile(new URL('./fixtures/b.md', import.meta.url), 'utf-8'))
 
   expect(result).toMatchFileSnapshot('./fixtures/b.out.html')
-}, { timeout: 10_000 })
+})
 
-it('run for default language', async () => {
+it('run for default language', { timeout: 10_000 }, async () => {
   const md = MarkdownIt()
   md.use(await Shiki({
     themes: {
       light: 'vitesse-light',
       dark: 'vitesse-dark',
     },
+    langs: ['js', 'ts'],
     defaultLanguage: 'js',
     transformers: [
       transformerMetaHighlight(),
@@ -55,4 +63,4 @@ it('run for default language', async () => {
   const result = md.render(await fs.readFile(new URL('./fixtures/c.md', import.meta.url), 'utf-8'))
 
   expect(result).toMatchFileSnapshot('./fixtures/c.out.html')
-}, { timeout: 10_000 })
+})

packages/shiki/scripts/prepare/langs.ts

@@ -5,12 +5,49 @@ import { grammars, injections } from 'tm-grammars'
 import { COMMENT_HEAD } from './constants'
 
 /**
- * Languages that includes a lot of embedded langs,
- * We only load on-demand for these langs.
+ * Document-like languages that have embedded langs
  */
-const LANGS_LAZY_EMBEDDED = [
+const LANGS_LAZY_EMBEDDED_ALL = {
+  markdown: [],
+  mdx: [],
+  wikitext: [],
+  asciidoc: [],
+  latex: ['tex'],
+} as Record<string, string[]>
+
+/**
+ * Single-file-component-like languages that have embedded langs
+ * For these langs, we exclude the standalone embedded langs from the main bundle
+ */
+const LANGS_LAZY_EMBEDDED_PARTIAL = [
+  'vue',
+  'vue-html',
+  'svelte',
+  'pug',
+  'haml',
+  'astro',
+]
+
+/**
+ * Languages to be excluded from SFC langs
+ */
+const STANDALONG_LANGS_EMBEDDED = [
+  'pug',
+  'stylus',
+  'sass',
+  'scss',
+  'coffee',
+  'jsonc',
+  'json5',
+  'yaml',
+  'toml',
+  'scss',
+  'graphql',
   'markdown',
-  'mdx',
+  'less',
+  'jsx',
+  'tsx',
+  'ruby',
 ]
 
 export async function prepareLangs() {
@@ -22,6 +59,8 @@ export async function prepareLangs() {
 
   allLangFiles.sort()
 
+  const resolvedLangs: LanguageRegistration[] = []
+
   for (const file of allLangFiles) {
     const content = await fs.readJSON(file)
     const lang = grammars.find(i => i.name === content.name) || injections.find(i => i.name === content.name)
@@ -40,12 +79,21 @@ export async function prepareLangs() {
     }
 
     // We don't load all the embedded langs for markdown
-    if (LANGS_LAZY_EMBEDDED.includes(lang.name)) {
-      json.embeddedLangsLazy = json.embeddedLangs
-      json.embeddedLangs = []
+    if (LANGS_LAZY_EMBEDDED_ALL[lang.name]) {
+      const includes = LANGS_LAZY_EMBEDDED_ALL[lang.name]
+      json.embeddedLangsLazy = (json.embeddedLangs || []).filter(i => !includes.includes(i)) || []
+      json.embeddedLangs = includes
+    }
+    else if (LANGS_LAZY_EMBEDDED_PARTIAL.includes(lang.name)) {
+      json.embeddedLangsLazy = (json.embeddedLangs || []).filter(i => STANDALONG_LANGS_EMBEDDED.includes(i)) || []
+      json.embeddedLangs = (json.embeddedLangs || []).filter(i => !STANDALONG_LANGS_EMBEDDED.includes(i)) || []
     }
 
     const deps: string[] = json.embeddedLangs || []
+    resolvedLangs.push(json)
+
+    if (deps.length > 10)
+      console.log(json.name, json.embeddedLangs)
 
     await fs.writeFile(
       `./src/assets/langs/${lang.name}.js`,
@@ -102,12 +150,10 @@ export default langs
     while (changed) {
       changed = false
       for (const id of bundledIds) {
-        if (LANGS_LAZY_EMBEDDED.includes(id))
-          continue
-        const lang = grammars.find(i => i.name === id)
+        const lang = resolvedLangs.find(i => i.name === id)
         if (!lang)
           continue
-        for (const e of lang.embedded || []) {
+        for (const e of lang.embeddedLangs || []) {
           if (!bundledIds.has(e)) {
             bundledIds.add(e)
             changed = true

packages/shiki/src/assets/langs-bundle-web.ts

@@ -122,11 +122,6 @@ export const bundledLanguagesInfo: BundledLanguageInfo[] = [
     'name': 'JSON',
     'import': (() => import('./langs/json')) as DynamicImportLanguageRegistration
   },
-  {
-    'id': 'json5',
-    'name': 'JSON5',
-    'import': (() => import('./langs/json5')) as DynamicImportLanguageRegistration
-  },
   {
     'id': 'jsonc',
     'name': 'JSON with Comments',
@@ -155,11 +150,6 @@ export const bundledLanguagesInfo: BundledLanguageInfo[] = [
     'name': 'Less',
     'import': (() => import('./langs/less')) as DynamicImportLanguageRegistration
   },
-  {
-    'id': 'lua',
-    'name': 'Lua',
-    'import': (() => import('./langs/lua')) as DynamicImportLanguageRegistration
-  },
   {
     'id': 'markdown',
     'name': 'Markdown',
@@ -222,14 +212,6 @@ export const bundledLanguagesInfo: BundledLanguageInfo[] = [
     ],
     'import': (() => import('./langs/regexp')) as DynamicImportLanguageRegistration
   },
-  {
-    'id': 'ruby',
-    'name': 'Ruby',
-    'aliases': [
-      'rb'
-    ],
-    'import': (() => import('./langs/ruby')) as DynamicImportLanguageRegistration
-  },
   {
     'id': 'sass',
     'name': 'Sass',
@@ -269,11 +251,6 @@ export const bundledLanguagesInfo: BundledLanguageInfo[] = [
     'name': 'Svelte',
     'import': (() => import('./langs/svelte')) as DynamicImportLanguageRegistration
   },
-  {
-    'id': 'toml',
-    'name': 'TOML',
-    'import': (() => import('./langs/toml')) as DynamicImportLanguageRegistration
-  },
   {
     'id': 'ts-tags',
     'name': 'TypeScript with Tags',
@@ -364,14 +341,12 @@ export type BundledLanguage =
   | 'jl'
   | 'js'
   | 'json'
-  | 'json5'
   | 'jsonc'
   | 'jsonl'
   | 'jsx'
   | 'julia'
   | 'less'
   | 'lit'
-  | 'lua'
   | 'markdown'
   | 'marko'
   | 'md'
@@ -383,10 +358,8 @@ export type BundledLanguage =
   | 'py'
   | 'python'
   | 'r'
-  | 'rb'
   | 'regex'
   | 'regexp'
-  | 'ruby'
   | 'sass'
   | 'scss'
   | 'sh'
@@ -396,7 +369,6 @@ export type BundledLanguage =
   | 'styl'
   | 'stylus'
   | 'svelte'
-  | 'toml'
   | 'ts'
   | 'ts-tags'
   | 'tsx'

packages/shiki/test/bundle.test.ts

@@ -17,5 +17,5 @@ it('bundle-web', async () => {
   }))
 
   expect(highlighter.getLoadedLanguages().length)
-    .toMatchInlineSnapshot(`91`)
+    .toMatchInlineSnapshot(`86`)
 })

packages/shiki/test/general.test.ts

@@ -55,39 +55,19 @@ describe('should', () => {
     expect(shiki.getLoadedLanguages().sort())
       .toMatchInlineSnapshot(`
         [
-          "coffee",
-          "coffeescript",
           "css",
-          "gql",
-          "graphql",
           "html",
           "html-derivative",
-          "jade",
           "javascript",
           "js",
           "json",
-          "json5",
-          "jsonc",
-          "jsx",
-          "less",
-          "markdown",
           "markdown-vue",
-          "md",
-          "pug",
-          "sass",
-          "scss",
-          "styl",
-          "stylus",
-          "toml",
           "ts",
-          "tsx",
           "typescript",
           "vue",
           "vue-directives",
           "vue-interpolations",
           "vue-sfc-style-variable-injection",
-          "yaml",
-          "yml",
         ]
       `)
   })
@@ -150,6 +130,49 @@ describe('should', () => {
       `)
   })
 
+  it('dynamic load lang with vue', async () => {
+    const shiki = await createHighlighter({
+      langs: [],
+      themes: [],
+    })
+
+    await shiki.loadTheme('vitesse-dark')
+    await shiki.loadLanguage('vue')
+
+    expect(shiki.getLoadedLanguages())
+      .not
+      .includes('scss')
+
+    const code = `
+      <template>
+        <h1>Hello</h1>
+      </template>
+
+      <script setup lang="ts">
+      const a: number = 1
+      </script>
+
+      <style lang="scss">
+      h1 {
+        span {
+          color: red;
+        }
+      }
+      </style>
+    `
+
+    const html1 = shiki.codeToHtml(code, { lang: 'vue', theme: 'vitesse-dark' })
+
+    await shiki.loadLanguage('scss')
+
+    expect(shiki.getLoadedLanguages())
+      .includes('scss')
+
+    const html2 = shiki.codeToHtml(code, { lang: 'vue', theme: 'vitesse-dark' })
+
+    expect(html1).not.toEqual(html2)
+  })
+
   it('monokai underline', async () => {
     expect(await codeToHtml('type Foo = { bar: string }', {
       theme: 'monokai',