From b64f5a655153f4d64c4e77b0b6fb300a2e0d9991 Mon Sep 17 00:00:00 2001
From: Steven Hall <steven@gitbook.io>
Date: Wed, 17 Apr 2024 13:02:03 +0100
Subject: [PATCH] Fully dynamic imports of languages

---
 .../DocumentView/CodeBlock/highlight.ts       | 68 +++++++++++--------
 src/components/Header/Header.tsx              |  8 ++-
 2 files changed, 46 insertions(+), 30 deletions(-)

diff --git a/src/components/DocumentView/CodeBlock/highlight.ts b/src/components/DocumentView/CodeBlock/highlight.ts
index 7d39955e2..7c1a6ff15 100644
--- a/src/components/DocumentView/CodeBlock/highlight.ts
+++ b/src/components/DocumentView/CodeBlock/highlight.ts
@@ -1,13 +1,11 @@
 import { DocumentBlockCode, DocumentBlockCodeLine, DocumentInlineAnnotation } from '@gitbook/api';
 import {
-    loadWasm,
     ThemedToken,
-    getHighlighter,
     createCssVariablesTheme,
-    HighlighterGeneric,
-    bundledLanguages,
-    bundledThemes,
+    HighlighterCore,
+    LanguageInput,
 } from 'shiki';
+import { getHighlighterCore } from 'shiki/core';
 // @ts-ignore - onigWasm is a Wasm module
 import onigWasm from 'shiki/onig.wasm?module';
 
@@ -15,8 +13,6 @@ import { asyncMutexFunction, singleton } from '@/lib/async';
 import { getNodeText } from '@/lib/document';
 import { trace } from '@/lib/tracing';
 
-import { DocumentContext } from '../DocumentView';
-
 export type HighlightLine = {
     highlighted: boolean;
     tokens: HighlightToken[];
@@ -35,8 +31,8 @@ type PositionedToken = ThemedToken & { start: number; end: number };
  * Highlight a code block while preserving inline elements.
  */
 export async function highlight(block: DocumentBlockCode): Promise<HighlightLine[]> {
-    const langName = block.data.syntax ? getLanguageForSyntax(block.data.syntax) : null;
-    if (!langName) {
+    const langName = block.data.syntax; // ? getLanguageForSyntax(block.data.syntax) : null;
+    if (!langName || !loaders[langName]) {
         // Language not found, fallback to plain highlighting
         return plainHighlighting(block);
     }
@@ -87,16 +83,16 @@ export async function highlight(block: DocumentBlockCode): Promise<HighlightLine
 /**
  * Validate a language name.
  */
-function getLanguageForSyntax(syntax: string): keyof typeof bundledLanguages | null {
-    // @ts-ignore
-    const lang = bundledLanguages[syntax];
-    if (!lang) {
-        return null;
-    }
+// function getLanguageForSyntax(syntax: string): keyof typeof bundledLanguages | null {
+//     // @ts-ignore
+//     const lang = bundledLanguages[syntax];
+//     if (!lang) {
+//         return null;
+//     }
 
-    // @ts-ignore
-    return syntax;
-}
+//     // @ts-ignore
+//     return syntax;
+// }
 
 /**
  * Parse a code block without highlighting it.
@@ -305,16 +301,26 @@ function cleanupLine(line: string): string {
  */
 const loadHighlighter = singleton(async () => {
     return await trace('highlighting.loadHighlighter', async () => {
-        if (typeof onigWasm !== 'string') {
-            // When running bun test, the import is a string, we ignore it and let the module
-            // loads it on its own.
-            //
-            // Otherwise for Vercel/Cloudflare, we need to load it ourselves.
-            await loadWasm((obj) => WebAssembly.instantiate(onigWasm, obj));
-        }
-        const highlighter = await getHighlighter({
+        // if (typeof onigWasm !== 'string') {
+        //     // When running bun test, the import is a string, we ignore it and let the module
+        //     // loads it on its own.
+        //     //
+        //     // Otherwise for Vercel/Cloudflare, we need to load it ourselves.
+        //     await loadWasm((obj) => WebAssembly.instantiate(onigWasm, obj));
+        // }
+        const highlighter = await getHighlighterCore({
             themes: [createCssVariablesTheme()],
             langs: [],
+            loadWasm:
+                typeof onigWasm !== 'string'
+                    ? async (obj) => {
+                          // When running bun test, the import is a string, we ignore it and let the module
+                          // loads it on its own.
+                          //
+                          // Otherwise for Vercel/Cloudflare, we need to load it ourselves.
+                          return WebAssembly.instantiate(onigWasm, obj);
+                      }
+                    : undefined,
         });
         return highlighter;
     });
@@ -322,7 +328,7 @@ const loadHighlighter = singleton(async () => {
 
 const loadLanguagesMutex = asyncMutexFunction();
 async function loadHighlighterLanguage(
-    highlighter: HighlighterGeneric<keyof typeof bundledLanguages, keyof typeof bundledThemes>,
+    highlighter: HighlighterCore,
     lang: keyof typeof bundledLanguages,
 ) {
     await loadLanguagesMutex.runBlocking(async () => {
@@ -332,7 +338,13 @@ async function loadHighlighterLanguage(
 
         await trace(
             `highlighting.loadLanguage(${lang})`,
-            async () => await highlighter.loadLanguage(lang),
+            async () => await highlighter.loadLanguage(loaders[lang]),
         );
     });
 }
+
+const loaders: { [key: string]: LanguageInput } = {
+    javascript: () => import('shiki/langs/javascript.mjs'),
+    typescript: () => import('shiki/langs/typescript.mjs'),
+    tsx: () => import('shiki/langs/tsx.mjs'),
+}
\ No newline at end of file
diff --git a/src/components/Header/Header.tsx b/src/components/Header/Header.tsx
index 0cf62fa51..7d57f141b 100644
--- a/src/components/Header/Header.tsx
+++ b/src/components/Header/Header.tsx
@@ -33,7 +33,9 @@ export function Header(props: {
 
     const isCustomizationDefault =
         customization.header.preset === CustomizationHeaderPreset.Default;
-    const isMultiVariants = parent?.object === 'collection' || (parent && parent.object === 'site' && spaces.length > 1)
+    const isMultiVariants =
+        parent?.object === 'collection' ||
+        (parent && parent.object === 'site' && spaces.length > 1);
 
     return (
         <header
@@ -74,7 +76,9 @@ export function Header(props: {
                     )}
                 >
                     <HeaderLogo parent={parent} space={space} customization={customization} />
-                    <span>{isMultiVariants ? <SpacesDropdown space={space} spaces={spaces} /> : null}</span>
+                    <span>
+                        {isMultiVariants ? <SpacesDropdown space={space} spaces={spaces} /> : null}
+                    </span>
                     <HeaderLinks>
                         {customization.header.links.map((link, index) => {
                             return (