From 4dac35fc9ec3eecb4371befa619626ee4db1ba58 Mon Sep 17 00:00:00 2001
From: Divyansh Singh <40380293+brc-dd@users.noreply.github.com>
Date: Sun, 3 Nov 2024 17:09:29 +0530
Subject: [PATCH] fix: (temporary patch) lang lazy load not working with
 twoslash

---
 src/node/cli.ts                        |  6 ++-
 src/node/markdown/plugins/highlight.ts | 52 ++++++++++++++++----------
 src/node/worker_shikiResolveLang.ts    |  8 +++-
 3 files changed, 43 insertions(+), 23 deletions(-)

diff --git a/src/node/cli.ts b/src/node/cli.ts
index fb60c6f612e2..6b7f977cdce2 100644
--- a/src/node/cli.ts
+++ b/src/node/cli.ts
@@ -2,10 +2,14 @@ import minimist from 'minimist'
 import c from 'picocolors'
 import { createLogger } from 'vite'
 import { build, createServer, serve } from '.'
-import { init } from './init/init'
 import { version } from '../../package.json'
+import { init } from './init/init'
 import { bindShortcuts } from './shortcuts'
 
+if (process.env.DEBUG) {
+  Error.stackTraceLimit = Infinity
+}
+
 const argv: any = minimist(process.argv.slice(2))
 
 const logVersion = (logger = createLogger()) => {
diff --git a/src/node/markdown/plugins/highlight.ts b/src/node/markdown/plugins/highlight.ts
index 0ae9ed751ef3..5e0abf2862dc 100644
--- a/src/node/markdown/plugins/highlight.ts
+++ b/src/node/markdown/plugins/highlight.ts
@@ -9,7 +9,7 @@ import {
 import { customAlphabet } from 'nanoid'
 import { createRequire } from 'node:module'
 import c from 'picocolors'
-import type { ShikiTransformer } from 'shiki'
+import type { LanguageRegistration, ShikiTransformer } from 'shiki'
 import { createHighlighter, isSpecialLang } from 'shiki'
 import { createSyncFn } from 'synckit'
 import type { Logger } from 'vite'
@@ -61,10 +61,14 @@ export async function highlight(
   logger: Pick<Logger, 'warn'> = console
 ): Promise<[(str: string, lang: string, attrs: string) => string, () => void]> {
   const {
-    defaultHighlightLang: defaultLang = '',
+    defaultHighlightLang: defaultLang = 'txt',
     codeTransformers: userTransformers = []
   } = options
 
+  const usingTwoslash = userTransformers.some(
+    ({ name }) => name === '@shikijs/vitepress-twoslash'
+  )
+
   const highlighter = await createHighlighter({
     themes:
       typeof theme === 'object' && 'light' in theme && 'dark' in theme
@@ -72,11 +76,29 @@ export async function highlight(
         : [theme],
     langs: [
       ...(options.languages || []),
-      ...Object.values(options.languageAlias || {})
+      ...Object.values(options.languageAlias || {}),
+
+      // patch for twoslash - https://github.com/vuejs/vitepress/issues/4334
+      ...(usingTwoslash
+        ? Object.keys((await import('shiki')).bundledLanguages)
+        : [])
     ],
     langAlias: options.languageAlias
   })
 
+  function loadLanguage(name: string | LanguageRegistration) {
+    const lang = typeof name === 'string' ? name : name.name
+    if (
+      !isSpecialLang(lang) &&
+      !highlighter.getLoadedLanguages().includes(lang)
+    ) {
+      const resolvedLang = resolveLangSync(lang)
+      if (resolvedLang.length) highlighter.loadLanguageSync(resolvedLang)
+      else return false
+    }
+    return true
+  }
+
   await options?.shikiSetup?.(highlighter)
 
   const transformers: ShikiTransformer[] = [
@@ -116,23 +138,13 @@ export async function highlight(
           .replace(vueRE, '')
           .toLowerCase() || defaultLang
 
-      if (lang) {
-        const langLoaded = highlighter.getLoadedLanguages().includes(lang)
-        if (!langLoaded && !isSpecialLang(lang)) {
-          const resolvedLang = resolveLangSync(lang)
-          if (!resolvedLang.length) {
-            logger.warn(
-              c.yellow(
-                `\nThe language '${lang}' is not loaded, falling back to '${
-                  defaultLang || 'txt'
-                }' for syntax highlighting.`
-              )
-            )
-            lang = defaultLang
-          } else {
-            highlighter.loadLanguageSync(resolvedLang)
-          }
-        }
+      if (!loadLanguage(lang)) {
+        logger.warn(
+          c.yellow(
+            `\nThe language '${lang}' is not loaded, falling back to '${defaultLang}' for syntax highlighting.`
+          )
+        )
+        lang = defaultLang
       }
 
       const lineOptions = attrsToLines(attrs)
diff --git a/src/node/worker_shikiResolveLang.ts b/src/node/worker_shikiResolveLang.ts
index 22a2c2287bab..644caf9ef3c9 100644
--- a/src/node/worker_shikiResolveLang.ts
+++ b/src/node/worker_shikiResolveLang.ts
@@ -1,4 +1,8 @@
-import { bundledLanguages, type DynamicImportLanguageRegistration } from 'shiki'
+import {
+  bundledLanguages,
+  type DynamicImportLanguageRegistration,
+  type LanguageRegistration
+} from 'shiki'
 import { runAsWorker } from 'synckit'
 
 async function resolveLang(lang: string) {
@@ -10,7 +14,7 @@ async function resolveLang(lang: string) {
       >
     )
       [lang]?.()
-      .then((m) => m.default) || []
+      .then((m) => m.default) || ([] as LanguageRegistration[])
   )
 }