Core: use local fork of Shiki rehype integration

.changeset/dry-showers-do.md

@@ -0,0 +1,5 @@
+---
+'fumadocs-core': patch
+---
+
+use local fork of Shiki rehype integration

apps/docs/package.json

@@ -25,7 +25,6 @@
     "@radix-ui/react-hover-card": "^1.1.15",
     "@radix-ui/react-presence": "^1.1.5",
     "@radix-ui/react-tooltip": "^1.2.8",
-    "@shikijs/rehype": "^4.0.2",
     "@shikijs/transformers": "^4.0.2",
     "@takumi-rs/image-response": "^1.0.4",
     "ai": "^6.0.158",

packages/core/package.json

@@ -99,7 +99,6 @@
   },
   "dependencies": {
     "@orama/orama": "^3.1.18",
-    "@shikijs/rehype": "^4.0.2",
     "@shikijs/transformers": "^4.0.2",
     "estree-util-value-to-estree": "^3.5.0",
     "github-slugger": "^2.0.0",

packages/core/src/highlight/index.ts

@@ -10,10 +10,10 @@ import type {
 import type { ReactNode } from 'react';
 import type { Root } from 'hast';
 import { defaultShikiFactory, wasmShikiFactory } from './shiki/full';
-import { loadMissingLanguage, loadMissingTheme, applyDefaultThemes } from './utils';
 import { type Components, toJsxRuntime } from 'hast-util-to-jsx-runtime';
 import * as JsxRuntime from 'react/jsx-runtime';
 import { highlightHast as highlightHastBase } from './shiki';
+import { applyDefaultThemes, loadMissingLanguage, loadMissingTheme } from './utils';
 
 export type HighlightOptions = HighlightHastOptions & {
   components?: Partial<Components>;
@@ -53,6 +53,7 @@ export async function getHighlighter(
 ) {
   const factory = engineType === 'js' ? defaultShikiFactory : wasmShikiFactory;
   const instance = await factory.getOrInit();
+
   await Promise.all([
     options.langs && loadMissingLanguage(instance, options.langs),
     options.themes && loadMissingTheme(instance, options.themes),

packages/core/src/highlight/shiki/index.ts

@@ -7,8 +7,8 @@ import type {
   CodeToHastOptionsCommon,
   HighlighterCore,
 } from 'shiki';
-import { loadMissingTheme, getRequiredThemes } from '../utils';
 import type { Awaitable } from '@/types';
+import { loadMissingTheme, getRequiredThemes } from '../utils';
 
 interface InitOptions {
   langAlias?: Record<string, string>;

packages/core/src/highlight/utils.ts

@@ -10,10 +10,10 @@ export async function loadMissingTheme(
   highlighter: HighlighterCore,
   themes: (ThemeRegistrationAny | string)[],
 ) {
-  const { isSpecialTheme } = await import('shiki/core');
-
+  const bundled = highlighter.getBundledThemes();
   const missingThemes = themes.filter((theme) => {
-    if (isSpecialTheme(theme)) return false;
+    if (typeof theme === 'string' && !(theme in bundled)) return false;
+
     try {
       highlighter.getTheme(theme);
       return false;
@@ -29,10 +29,10 @@ export async function loadMissingLanguage(
   highlighter: HighlighterCore,
   langs: (LanguageRegistration | string)[],
 ) {
-  const { isSpecialLang } = await import('shiki/core');
-
+  const bundled = highlighter.getBundledLanguages();
   const missingLangs = langs.filter((lang) => {
-    if (isSpecialLang(lang)) return false;
+    if (typeof lang === 'string' && !(lang in bundled)) return false;
+
     try {
       highlighter.getLanguage(lang);
       return false;

packages/core/src/mdx-plugins/rehype-code.core.ts

@@ -1,6 +1,5 @@
 import type { Root } from 'hast';
-import type { RehypeShikiOptions } from '@shikijs/rehype';
-import rehypeShikiFromHighlighter from '@shikijs/rehype/core';
+import rehypeShikiFromHighlighter, { type RehypeShikiCoreOptions } from './rehype-code/shiki';
 import {
   transformerNotationDiff,
   transformerNotationFocus,
@@ -8,12 +7,12 @@ import {
   transformerNotationWordHighlight,
 } from '@shikijs/transformers';
 import type { Processor, Transformer } from 'unified';
-import type { Highlighter, HighlighterCore, ShikiTransformer } from 'shiki';
+import type { BuiltinLanguage, HighlighterCore, LanguageInput, ShikiTransformer } from 'shiki';
 import type { MdxJsxFlowElement } from 'mdast-util-mdx';
 import type { CodeBlockIcon, IconOptions } from './transformer-icon';
 import { transformerIcon } from './transformer-icon';
 import { parseCodeBlockAttributes } from '@/mdx-plugins/codeblock-utils';
-import type { Awaitable, DistributiveOmit } from '@/types';
+import type { Awaitable } from '@/types';
 import type { ShikiFactory } from '@/highlight/shiki';
 import { defaultThemes, getRequiredThemes } from '@/highlight/utils';
 
@@ -41,7 +40,7 @@ export function rehypeCodeDefaultOptions(): RehypeCodeOptionsCommon {
       const data: Record<string, unknown> = parsed.attributes;
       parsed.rest = parseLineNumber(parsed.rest, data);
 
-      data.__parsed_raw = parsed.rest;
+      data.__raw = parsed.rest;
       return data;
     },
   };
@@ -57,17 +56,18 @@ function parseLineNumber(str: string, data: Record<string, unknown>) {
   });
 }
 
-export type RehypeCodeOptionsCommon = DistributiveOmit<RehypeShikiOptions, 'lazy'> & {
+export type RehypeCodeOptionsCommon = RehypeShikiCoreOptions & {
   /**
-   * Load languages and themes on-demand.
-   * @defaultValue true
+   * Language names to include.
+   *
+   * @default Object.keys(bundledLanguages)
    */
-  lazy?: boolean;
-
+  langs?: Array<LanguageInput | BuiltinLanguage>;
   /**
-   * Filter meta string before processing
+   * Alias of languages
+   * @example { 'my-lang': 'javascript' }
    */
-  filterMetaString?: (metaString: string) => string;
+  langAlias?: Record<string, string>;
 
   /**
    * Add icon to code blocks
@@ -108,22 +108,6 @@ export function createRehypeCode<
     }
 
     const transformers = options.transformers ? [...options.transformers] : [];
-    transformers.unshift({
-      name: 'rehype-code:pre-process',
-      preprocess(code, { meta }) {
-        if (meta && '__parsed_raw' in meta) {
-          meta.__raw = meta.__parsed_raw;
-          delete meta.__parsed_raw;
-        }
-
-        if (meta && options.filterMetaString) {
-          meta.__raw = options.filterMetaString(meta.__raw ?? '');
-        }
-
-        // Remove empty line at end
-        return code.replace(/\n$/, '');
-      },
-    });
 
     if (options.icon !== false) {
       transformers.push(transformerIcon(options.icon));
@@ -141,7 +125,7 @@ export function createRehypeCode<
       highlighter.loadTheme(...(getRequiredThemes(options) as never[])),
       highlighter.loadLanguage(...(langs as never[])),
     ]);
-    return rehypeShikiFromHighlighter(highlighter as Highlighter, {
+    return rehypeShikiFromHighlighter(highlighter, {
       ...options,
       transformers,
     });

packages/core/src/mdx-plugins/rehype-code/LICENSE

@@ -0,0 +1,22 @@
+MIT License
+
+Copyright (c) 2021 Pine Wu
+Copyright (c) 2023 Anthony Fu <https://github.com/antfu>
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

packages/core/src/mdx-plugins/rehype-code/parsers.ts

@@ -0,0 +1,52 @@
+import type { Element, Root } from 'hast';
+import { flattenNodeHast } from '../utils';
+
+const RE_TAILING_CURLY_COLON = /(.+)\{:([\w-]+)\}$/;
+
+export interface ShikiParsedData {
+  type: 'inline' | 'pre';
+  meta?: string;
+  lang?: string;
+  code: string;
+}
+
+export type ShikiParser = (tree: Root, node: Element) => ShikiParsedData | undefined;
+
+export type InlineCodeParser = 'tailing-curly-colon';
+
+export const InlineCodeParsers: Record<InlineCodeParser, ShikiParser> = {
+  'tailing-curly-colon': (_tree, node) => {
+    const raw = flattenNodeHast(node);
+    const match = raw.match(RE_TAILING_CURLY_COLON);
+    if (!match) return;
+
+    return {
+      type: 'inline',
+      code: match[1] ?? raw,
+      lang: match.at(2),
+    };
+  },
+};
+
+const languagePrefix = 'language-';
+
+export const PreParser: ShikiParser = (_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: flattenNodeHast(head),
+    meta: head.data?.meta ?? head.properties.metastring?.toString() ?? '',
+  };
+};

packages/core/src/mdx-plugins/rehype-code/shiki.ts

@@ -0,0 +1,245 @@
+import type { Root, Element } from 'hast';
+import type { Transformer } from 'unified';
+import { isSpecialLang } from 'shiki/core';
+import { visit } from 'unist-util-visit';
+import {
+  type InlineCodeParser,
+  InlineCodeParsers,
+  PreParser,
+  type ShikiParsedData,
+} from './parsers';
+import type {
+  BuiltinTheme,
+  CodeOptionsMeta,
+  CodeOptionsThemes,
+  CodeToHastOptions,
+  CodeToHastOptionsCommon,
+  HighlighterCore,
+  TransformerOptions,
+} from 'shiki';
+
+export interface MapLike<K = unknown, V = unknown> {
+  get: (key: K) => V | undefined;
+  set: (key: K, value: V) => this;
+}
+
+export interface RehypeShikiExtraOptions {
+  /**
+   * Add `language-*` class to code element
+   *
+   * @default false
+   */
+  addLanguageClass?: boolean;
+
+  /**
+   * The default language to use when is not specified
+   */
+  defaultLanguage?: string;
+
+  /**
+   * The fallback language to use when specified language is not loaded, or not included in the bundle
+   */
+  fallbackLanguage?: string;
+
+  /**
+   * Load languages and themes on-demand.
+   *
+   * @default true
+   */
+  lazy?: boolean;
+
+  /**
+   * `mdast-util-to-hast` adds a newline to the end of code blocks
+   *
+   * This option strips that newline from the code block
+   *
+   * @default true
+   * @see https://github.com/syntax-tree/mdast-util-to-hast/blob/f511a93817b131fb73419bf7d24d73a5b8b0f0c2/lib/handlers/code.js#L22
+   */
+  stripEndNewline?: boolean;
+
+  /**
+   * Custom meta string parser
+   * Return an object to merge with `meta`
+   */
+  parseMetaString?: (
+    metaString: string,
+    node: Element,
+    tree: Root,
+  ) => Record<string, unknown> | undefined | null;
+
+  /**
+   * Filter meta string before processing
+   */
+  filterMetaString?: (metaString: string) => string;
+
+  /**
+   * Highlight inline code blocks
+   *
+   * - `false`: disable inline code block highlighting
+   * - `tailing-curly-colon`: highlight with `\`code{:lang}\``
+   *
+   * @see https://shiki.style/packages/rehype#inline-code
+   * @default false
+   */
+  inline?: InlineCodeParser | false;
+
+  /**
+   * Custom map to cache transformed codeToHast result
+   *
+   * @default undefined
+   */
+  cache?: MapLike<string, Root>;
+
+  /**
+   * Chance to handle the error
+   * If not provided, the error will be thrown
+   */
+  onError?: (error: unknown) => void;
+}
+
+export type RehypeShikiCoreOptions = CodeOptionsThemes<BuiltinTheme> &
+  TransformerOptions &
+  CodeOptionsMeta &
+  RehypeShikiExtraOptions &
+  Omit<CodeToHastOptionsCommon, 'lang'>;
+
+function rehypeShikiFromHighlighter(
+  highlighter: HighlighterCore,
+  options: RehypeShikiCoreOptions,
+): Transformer<Root, Root> {
+  const {
+    addLanguageClass = false,
+    parseMetaString,
+    filterMetaString,
+    cache,
+    defaultLanguage,
+    fallbackLanguage,
+    onError,
+    stripEndNewline = true,
+    inline = false,
+    lazy = false,
+    ...rest
+  } = options;
+
+  function highlight(
+    lang: string,
+    code: string,
+    metaString = '',
+    tree: Root,
+    node: Element,
+  ): Root | undefined {
+    if (filterMetaString) {
+      metaString = filterMetaString(metaString);
+    }
+
+    const cacheKey = `${lang}:${metaString}:${code}`;
+    const cachedValue = cache?.get(cacheKey);
+
+    if (cachedValue) {
+      return cachedValue;
+    }
+
+    const codeOptions: CodeToHastOptions = {
+      ...rest,
+      lang,
+      meta: {
+        ...rest.meta,
+        __raw: metaString,
+        ...parseMetaString?.(metaString, node, tree),
+      },
+    };
+
+    if (addLanguageClass) {
+      // always construct a new array, avoid adding the transformer repeatedly
+      codeOptions.transformers = [
+        ...(codeOptions.transformers ?? []),
+        {
+          name: 'rehype-shiki:code-language-class',
+          code(node) {
+            this.addClassToHast(node, `language-${lang}`);
+            return node;
+          },
+        },
+      ];
+    }
+
+    if (stripEndNewline && code.endsWith('\n')) code = code.slice(0, -1);
+
+    try {
+      const fragment = highlighter.codeToHast(code, codeOptions);
+      cache?.set(cacheKey, fragment);
+      return fragment;
+    } catch (error) {
+      if (onError) onError(error);
+      else throw error;
+    }
+  }
+
+  async function onNode(tree: Root, node: Element, parsed: ShikiParsedData) {
+    let lang = parsed.lang ?? defaultLanguage;
+    if (!lang) return;
+
+    const isLoaded = isSpecialLang(lang) || highlighter.getLoadedLanguages().includes(lang);
+
+    if (lazy && !isLoaded && lang in highlighter.getBundledLanguages()) {
+      await highlighter.loadLanguage(lang as never);
+    } else if (!isLoaded && fallbackLanguage) {
+      lang = fallbackLanguage;
+      await highlighter.loadLanguage(fallbackLanguage as never);
+    }
+
+    const fragment = highlight(lang, parsed.code, parsed.meta, tree, node);
+    if (!fragment) return;
+
+    if (parsed.type === 'inline') {
+      const head = fragment.children[0];
+      if (head.type === 'element' && head.tagName === 'pre') {
+        head.tagName = 'span';
+      }
+    }
+
+    return fragment;
+  }
+
+  return async (tree) => {
+    const queue: Promise<void>[] = [];
+
+    visit(tree, 'element', (node, index, parent) => {
+      let parsed: ShikiParsedData | undefined;
+
+      // needed for hast node replacement
+      if (!parent || index == null) return;
+
+      if (node.tagName === 'pre') {
+        parsed = PreParser(tree, node);
+      } else if (node.tagName === 'code' && inline) {
+        parsed = InlineCodeParsers[inline](tree, node);
+      }
+
+      if (!parsed) return 'skip';
+
+      queue.push(
+        onNode(tree, node, parsed)
+          .then((fragment) => {
+            if (fragment) parent.children[index] = fragment as never;
+          })
+          .catch((error) => {
+            if (onError) {
+              onError(error);
+            } else {
+              throw error;
+            }
+          }),
+      );
+
+      return 'skip';
+    });
+
+    if (queue.length > 0) {
+      await Promise.all(queue);
+    }
+  };
+}
+
+export default rehypeShikiFromHighlighter;

packages/core/src/mdx-plugins/utils.ts

@@ -2,14 +2,22 @@ import type { RootContent } from 'mdast';
 import { valueToEstree } from 'estree-util-value-to-estree';
 import type { Expression } from 'estree-jsx';
 import type { MdxjsEsm } from 'mdast-util-mdx';
+import type Hast from 'hast';
 
 export function flattenNode(node: RootContent): string {
-  if ('children' in node) return node.children.map((child) => flattenNode(child)).join('');
+  if ('children' in node) return node.children.map(flattenNode).join('');
 
   if ('value' in node && typeof node.value === 'string') return node.value;
 
   return '';
 }
+export function flattenNodeHast(node: Hast.RootContent): string {
+  if ('children' in node) {
+    return node.children.map(flattenNodeHast).join('');
+  }
+
+  return 'value' in node && typeof node.value === 'string' ? node.value : '';
+}
 
 export function toMdxExport(name: string, value: unknown): MdxjsEsm {
   return toMdxExportRaw(name, valueToEstree(value));

pnpm-lock.yaml

@@ -74,9 +74,6 @@ importers:
       '@radix-ui/react-tooltip':
         specifier: ^1.2.8
         version: 1.2.8(@types/[email protected](@types/[email protected]))(@types/[email protected])([email protected]([email protected]))([email protected])
-      '@shikijs/rehype':
-        specifier: ^4.0.2
-        version: 4.0.2
       '@shikijs/transformers':
         specifier: ^4.0.2
         version: 4.0.2
@@ -1843,9 +1840,6 @@ importers:
       '@orama/orama':
         specifier: ^3.1.18
         version: 3.1.18
-      '@shikijs/rehype':
-        specifier: ^4.0.2
-        version: 4.0.2
       '@shikijs/transformers':
         specifier: ^4.0.2
         version: 4.0.2
@@ -6629,10 +6623,6 @@ packages:
     resolution: {integrity: sha512-M6UMPrSa3fN5ayeJwFVl9qWofl273wtK1VG8ySDZ1mQBfhCpdd8nEx7nPZ/tk7k+TYcpqBZzj/AnwxT9lO+HJw==}
     engines: {node: '>=20'}
 
-  '@shikijs/[email protected]':
-    resolution: {integrity: sha512-cmPlKLD8JeojasNFoY64162ScpEdEdQUMuVodPCrv1nx1z3bjmGwoKWDruQWa/ejSznImlaeB0Ty6Q3zPaVQAA==}
-    engines: {node: '>=20'}
-
   '@shikijs/[email protected]':
     resolution: {integrity: sha512-mjCafwt8lJJaVSsQvNVrJumbnnj1RI8jbUKrPKgE6E3OvQKxnuRoBaYC51H4IGHePsGN/QtALglWBU7DoKDFnA==}
     engines: {node: '>=20'}
@@ -9405,9 +9395,6 @@ packages:
   [email protected]:
     resolution: {integrity: sha512-MlWT6Pjt4CG9lFCjiz4BH7l9wmrMkfkJYCxFwKQic8+RTZgWPuWxwAfjJElsXkex7DJjfSJsQIt931ilUgmwdA==}
 
-  [email protected]:
-    resolution: {integrity: sha512-XelQVTDWvqcl3axRfI0xSeoVKzyIFPwsAGSLIsKdJKQMXDYJS4WYrBNF/8J7RdhIcFI2BOHgAifggsvsxp/3+A==}
-
   [email protected]:
     resolution: {integrity: sha512-KK6y/BN8lbaq654j7JgBydev7wuNMcID54lkRav1P0CaE1e47P72AWWPiGKXTJU271ooYzcvTAn/Zt0REnvc7A==}
 
@@ -16138,15 +16125,6 @@ snapshots:
       '@shikijs/vscode-textmate': 10.0.2
       '@types/hast': 3.0.4
 
-  '@shikijs/[email protected]':
-    dependencies:
-      '@shikijs/types': 4.0.2
-      '@types/hast': 3.0.4
-      hast-util-to-string: 3.0.1
-      shiki: 4.0.2
-      unified: 11.0.5
-      unist-util-visit: 5.1.0
-
   '@shikijs/[email protected]':
     dependencies:
       '@shikijs/types': 4.0.2
@@ -19580,10 +19558,6 @@ snapshots:
       web-namespaces: 2.0.1
       zwitch: 2.0.4
 
-  [email protected]:
-    dependencies:
-      '@types/hast': 3.0.4
-
   [email protected]:
     dependencies:
       '@types/hast': 3.0.4