Upgrade shiki and optimize shiki rehype plugin (#223)

Co-authored-by: Luc van Kampen <[email protected]>

app/app/globals.css

@@ -1,6 +1,6 @@
 :root.dark {
-    --shiki-color-text: theme('colors.ens.dark.text.primary');
-    --shiki-color-background: theme('colors.ens.dark.background.primary');
+    --shiki-foreground: theme('colors.ens.dark.text.primary');
+    --shiki-background: theme('colors.ens.dark.background.primary');
     --shiki-token-constant: theme('colors.ens.dark.blue.primary');
     --shiki-token-string: theme('colors.ens.dark.green.primary');
     --shiki-token-comment: theme('colors.ens.dark.text.secondary');
@@ -12,8 +12,8 @@
 }
 
 :root {
-    --shiki-color-text: #0969da;
-    --shiki-color-background: #afb8c133;
+    --shiki-foreground: #0969da;
+    --shiki-background: #afb8c133;
     --shiki-token-constant: #0044ff;
     --shiki-token-string: #0a3069;
     --shiki-token-comment: #6e7781;

app/mdx/rehype.mjs

@@ -3,9 +3,12 @@ import * as acorn from 'acorn';
 import { toString } from 'mdast-util-to-string';
 import { mdxAnnotations } from 'mdx-annotations';
 import rehypeMdxTitle from 'rehype-mdx-title';
-import shiki from 'shiki';
+import {
+    bundledLanguages,
+    createCssVariablesTheme,
+    getHighlighter,
+} from 'shiki';
 import { visit } from 'unist-util-visit';
-
 /**
  * @type {import('unified').Plugin[]}
  */
@@ -42,48 +45,128 @@ function rehypeParseCodeBlocks() {
     };
 }
 
-let highlighter;
+const rehypeShikiFromHighlighter = function (highlighter, options) {
+    const {
+        addLanguageClass = false,
+        parseMetaString,
+        cache,
+        ...rest
+    } = options;
+    const prefix = 'language-';
 
-function rehypeShiki() {
-    return async (tree) => {
-        // lightHighlighter =
-        //     lightHighlighter ??
-        //     (await shiki.getHighlighter({ theme: 'min-light' }));
-        highlighter =
-            highlighter ??
-            (await shiki.getHighlighter({ theme: 'css-variables' }));
+    return function (tree) {
+        visit(tree, 'element', (node, index, parent) => {
+            if (!parent || index == undefined || node.tagName !== 'pre') return;
+
+            const [head] = node.children;
 
-        visit(tree, 'element', (node) => {
             if (
-                node.tagName === 'pre' &&
-                node.children[0]?.tagName === 'code'
-            ) {
-                const codeNode = node.children[0];
-                const textNode = codeNode.children[0];
-
-                node.properties.code = textNode.value;
-                node.properties.meta = codeNode?.data?.meta;
-
-                if (node.properties.language) {
-                    const lightTokens = highlighter.codeToThemedTokens(
-                        textNode.value,
-                        node.properties.language
-                    );
-
-                    const renderLight = shiki.renderToHtml(lightTokens, {
-                        elements: {
-                            pre: ({ children }) => children,
-                            code: ({ children }) => children,
-                            line: ({ children }) => `<span>${children}</span>`,
-                        },
-                    });
-
-                    textNode.value = renderLight;
-                }
+                !head ||
+                head.type !== 'element' ||
+                head.tagName !== 'code' ||
+                !head.properties
+            )
+                return;
+
+            const [textNode] = head.children;
+
+            const classes = head.properties.className;
+
+            if (!Array.isArray(classes)) return;
+
+            const language = classes.find(
+                (d) => typeof d === 'string' && d.startsWith(prefix)
+            );
+
+            if (typeof language !== 'string') return;
+
+            const code = toString(textNode);
+
+            const cachedValue = cache?.get(code);
+
+            if (cachedValue) {
+                textNode.value = cachedValue;
+
+                return;
+            }
+
+            const metaString =
+                head.data?.meta ?? head.properties.metastring ?? '';
+
+            const meta = parseMetaString?.(metaString, node, tree) || {};
+            const codeOptions = {
+                ...rest,
+                lang: language.slice(prefix.length),
+                meta: {
+                    ...rest.meta,
+                    ...meta,
+                    __raw: metaString,
+                },
+                transformers: [
+                    {
+                        root: ({ children }) => children.at(0),
+                        pre: ({ children }) => children.at(0),
+                        code: ({ children }) => children,
+                    },
+                ],
+            };
+
+            if (addLanguageClass) {
+                codeOptions.transformers || (codeOptions.transformers = []);
+                codeOptions.transformers.push({
+                    name: 'rehype-shiki:code-language-class',
+                    code(node2) {
+                        this.addClassToHast(node2, language);
+
+                        return node2;
+                    },
+                });
+            }
+
+            try {
+                const fragment = highlighter.codeToHtml(code, codeOptions);
+
+                cache?.set(code, fragment);
+                textNode.value = fragment;
+            } catch (error) {
+                if (options.onError) options.onError(error);
+                else throw error;
             }
         });
     };
-}
+};
+
+const rehypeShiki = function (options = {}) {
+    const themeNames = (
+        'themes' in options ? Object.values(options.themes) : [options.theme]
+    ).filter(Boolean);
+    const langs = options.langs || Object.keys(bundledLanguages);
+    // eslint-disable-next-line unicorn/no-this-assignment
+    const context = this;
+    let promise;
+
+    return async function (tree) {
+        if (!promise) {
+            promise = getHighlighter({
+                themes: themeNames,
+                langs,
+            }).then((highlighter) =>
+                rehypeShikiFromHighlighter.call(context, highlighter, options)
+            );
+        }
+
+        const handler = await promise;
+
+        return handler(tree);
+    };
+};
+
+const shikiCssTheme = createCssVariablesTheme({
+    name: 'css-variables',
+    variablePrefix: '--shiki-',
+    variableDefaults: {},
+    fontStyle: true,
+});
 
 function rehypeSlugify() {
     return (tree) => {
@@ -208,7 +291,13 @@ export const rehypePlugins = [
     /**
      * Add syntax highlighting to code blocks
      */
-    rehypeShiki,
+    [
+        rehypeShiki,
+        {
+            theme: shikiCssTheme,
+            addLanguageClass: true,
+        },
+    ],
     /**
      * Add an id to h2 elements
      */

app/package.json

@@ -17,7 +17,7 @@
         "@headlessui/react": "^1.7.18",
         "@mdx-js/loader": "^3.0.0",
         "@mdx-js/react": "^3.0.0",
-        "@next/mdx": "^14.1.0",
+        "@next/mdx": "^14.1.4",
         "@sindresorhus/slugify": "^2.2.1",
         "@tailwindcss/typography": "^0.5.10",
         "@tanstack/react-query": "^5.18.1",
@@ -49,7 +49,7 @@
         "remark-gfm": "^4.0.0",
         "schema-dts": "^1.1.2",
         "seedrandom": "^3.0.5",
-        "shiki": "^0.11.1",
+        "shiki": "^1.2.4",
         "simple-git": "^3.22.0",
         "siwe": "^2.1.4",
         "swr": "^2.2.4",
@@ -73,7 +73,7 @@
         "eslint-config-next": "14.1.0",
         "eslint-plugin-tailwindcss": "^3.14.2",
         "eslint-plugin-v3xlabs": "^1.6.2",
-        "next": "14.1.0",
+        "next": "14.1.4",
         "prettier": "^3.2.5",
         "react": "18.2.0",
         "react-dom": "18.2.0",