diff --git a/package/src/__tests__/fine-grained-bundle.test.tsx b/package/src/__tests__/fine-grained-bundle.test.tsx new file mode 100644 index 0000000..3043bc5 --- /dev/null +++ b/package/src/__tests__/fine-grained-bundle.test.tsx @@ -0,0 +1,220 @@ +// @ts-nocheck +// ai written test +import { render, waitFor } from "@testing-library/react"; +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { useShikiHighlighter } from "../hook"; +import { createFineGrainedBundle } from "../bundle"; +import React from "react"; +import type { HighlighterOptions } from "../types"; + +// Mock the bundle module +vi.mock("../bundle", () => ({ + createFineGrainedBundle: vi.fn(), +})); + +// Mock the toJsxRuntime function to avoid full rendering +vi.mock("hast-util-to-jsx-runtime", () => ({ + toJsxRuntime: vi + .fn() + .mockImplementation(() =>
Mocked HAST
), +})); + +// Helper component for testing useShikiHighlighter hook +const TestComponent = ({ + code, + language, + theme, + options, +}: { + code: string; + language: any; + theme: any; + options?: HighlighterOptions; +}) => { + const highlighted = useShikiHighlighter(code, language, theme, options); + return
{highlighted}
; +}; + +describe("Fine-grained bundle", () => { + const mockCodeToHast = vi + .fn() + .mockResolvedValue({ type: "element", tagName: "div" }); + + beforeEach(() => { + vi.resetAllMocks(); + // Mock the bundle with a predefined codeToHast function + (createFineGrainedBundle as any).mockResolvedValue({ + codeToHast: mockCodeToHast, + bundledLanguages: {}, + bundledThemes: {}, + }); + }); + + it("should use createCodegenBundle when fineGrainedBundle is provided", async () => { + const code = 'const test = "Hello";'; + const fineGrainedBundle = { + langs: ["typescript"], + themes: ["github-dark"], + engine: "javascript" as const, + }; + + render( + + ); + + // Wait for the asynchronous operations to complete + await waitFor(() => { + // Include precompiled: false when checking call arguments + expect(createFineGrainedBundle).toHaveBeenCalledWith({ + ...fineGrainedBundle, + precompiled: false, + }); + expect(mockCodeToHast).toHaveBeenCalled(); + }); + }); + + it("should cache the bundle for the same configuration", async () => { + const code = 'const test = "Hello";'; + const fineGrainedBundle = { + langs: ["typescript"], + themes: ["github-dark"], + engine: "javascript" as const, + }; + + // Render with the same config twice + const { rerender } = render( + + ); + + await waitFor(() => { + expect(createFineGrainedBundle).toHaveBeenCalledTimes(1); + }); + + // Re-render with same config + rerender( + + ); + + // Should not create a new bundle + await waitFor(() => { + expect(createFineGrainedBundle).toHaveBeenCalledTimes(1); + }); + }); + + it("should create a new bundle when configuration changes", async () => { + const code = 'const test = "Hello";'; + const fineGrainedBundle1 = { + langs: ["typescript"], + themes: ["github-dark"], + engine: "javascript" as const, + }; + + const fineGrainedBundle2 = { + langs: ["typescript", "javascript"], + themes: ["github-dark", "github-light"], + engine: "javascript" as const, + }; + + // Render with first config + const { rerender } = render( + + ); + + await waitFor(() => { + // Include precompiled: false when checking call arguments + expect(createFineGrainedBundle).toHaveBeenCalledWith({ + ...fineGrainedBundle1, + precompiled: false, + }); + expect(createFineGrainedBundle).toHaveBeenCalledTimes(1); + }); + + // Re-render with different config + rerender( + + ); + + // Should create a new bundle + await waitFor(() => { + // Include precompiled: false when checking call arguments + expect(createFineGrainedBundle).toHaveBeenCalledWith({ + ...fineGrainedBundle2, + precompiled: false, + }); + expect(createFineGrainedBundle).toHaveBeenCalledTimes(2); + }); + }); + + it("should create deterministic cache keys regardless of order", async () => { + const code = 'const test = "Hello";'; + + // Same configuration with different order + const fineGrainedBundle1 = { + langs: ["typescript", "javascript"], + themes: ["github-dark", "github-light"], + engine: "javascript" as const, + precompiled: false, + }; + + const fineGrainedBundle2 = { + langs: ["javascript", "typescript"], + themes: ["github-light", "github-dark"], + engine: "javascript" as const, + precompiled: false, + }; + + // Render with first config + const { rerender } = render( + + ); + + await waitFor(() => { + expect(createFineGrainedBundle).toHaveBeenCalledTimes(1); + }); + + // Re-render with same config but different order + rerender( + + ); + + // Should not create a new bundle + await waitFor(() => { + expect(createFineGrainedBundle).toHaveBeenCalledTimes(1); + }); + }); +}); diff --git a/package/src/bundle.ts b/package/src/bundle.ts new file mode 100644 index 0000000..ea812c4 --- /dev/null +++ b/package/src/bundle.ts @@ -0,0 +1,137 @@ +import type { BundledLanguage, BundledTheme, RegexEngine } from 'shiki'; +import { + bundledLanguagesInfo, + bundledThemesInfo, +} from 'shiki/bundle/full'; +import { createOnigurumaEngine } from 'shiki/engine/oniguruma'; +import { + createJavaScriptRegexEngine, + createJavaScriptRawEngine, +} from 'shiki/engine/javascript'; +import { + createSingletonShorthands, + createdBundledHighlighter, +} from 'shiki/core'; + +export interface CodegenBundleOptions { + /** + * The languages to bundle, specified by their string identifiers. + * @example ['typescript', 'javascript', 'vue'] + */ + langs: readonly BundledLanguage[]; + + /** + * The themes to bundle, specified by their string identifiers. + * @example ['github-dark', 'github-light'] + */ + themes: readonly BundledTheme[]; + + /** + * The engine to use for syntax highlighting. + * @default 'oniguruma' + */ + engine: 'oniguruma' | 'javascript' | 'javascript-raw'; + + /** + * Use precompiled grammars. + * Only available when `engine` is set to `javascript` or `javascript-raw`. + * @default false + */ + precompiled?: boolean; +} + +/** + * Create a fine grained bundle based on simple string identifiers of languages, themes, and engines. + * This is a runtime adaptation of shiki-codegen's approach. + */ +export async function createFineGrainedBundle({ + langs, + themes, + engine = 'oniguruma', + precompiled = false, +}: CodegenBundleOptions) { + if ( + precompiled && + engine !== 'javascript' && + engine !== 'javascript-raw' + ) { + throw new Error( + 'Precompiled grammars are only available when using the JavaScript engine' + ); + } + + const langImports = langs.map(async (lang) => { + const info = bundledLanguagesInfo.find( + (i) => i.id === lang || i.aliases?.includes(lang as string) + ); + if (!info) { + throw new Error(`Language ${lang} not found`); + } + + const module = await import( + `@shikijs/${precompiled ? 'langs-precompiled' : 'langs'}/${info.id}` + ); + return { id: lang, module: module.default || module }; + }); + + const themeImports = themes.map(async (theme) => { + const info = bundledThemesInfo.find((i) => i.id === theme); + if (!info) { + throw new Error(`Theme ${theme} not found`); + } + + const module = await import(`@shikijs/themes/${info.id}`); + return { id: theme, module: module.default || module }; + }); + + const resolvedLangs = await Promise.all(langImports); + const resolvedThemes = await Promise.all(themeImports); + + let engineInstance: RegexEngine; + switch (engine) { + case 'javascript': + engineInstance = createJavaScriptRegexEngine(); + break; + case 'javascript-raw': + engineInstance = createJavaScriptRawEngine(); + break; + case 'oniguruma': + engineInstance = await createOnigurumaEngine(import('shiki/wasm')); + break; + } + + const createHighlighter = createdBundledHighlighter({ + langs: Object.fromEntries( + resolvedLangs.map(({ id, module }) => [ + id, + () => Promise.resolve(module), + ]) + ), + themes: Object.fromEntries( + resolvedThemes.map(({ id, module }) => [ + id, + () => Promise.resolve(module), + ]) + ), + engine: () => engineInstance, + }); + + const shorthands = createSingletonShorthands(createHighlighter); + + return { + bundledLanguages: Object.fromEntries( + resolvedLangs.map(({ id, module }) => [id, module]) + ), + bundledThemes: Object.fromEntries( + resolvedThemes.map(({ id, module }) => [id, module]) + ), + // createHighlighter, + // codeToHtml: shorthands.codeToHtml, + codeToHast: shorthands.codeToHast, + // codeToTokens: shorthands.codeToTokens, + // codeToTokensBase: shorthands.codeToTokensBase, + // codeToTokensWithThemes: shorthands.codeToTokensWithThemes, + // getSingletonHighlighter: shorthands.getSingletonHighlighter, + // getLastGrammarState: shorthands.getLastGrammarState, + }; +} diff --git a/package/src/component.tsx b/package/src/component.tsx index fefb3c1..3f4bf24 100644 --- a/package/src/component.tsx +++ b/package/src/component.tsx @@ -1,14 +1,9 @@ -import './styles.css'; -import { clsx } from 'clsx'; -import { useShikiHighlighter } from './hook'; -import { resolveLanguage } from './resolvers'; +import "./styles.css"; +import { clsx } from "clsx"; +import { useShikiHighlighter } from "./hook"; +import { resolveLanguage } from "./resolvers"; -import type { - HighlighterOptions, - Language, - Theme, - Themes, -} from './types'; +import type { HighlighterOptions, Language, Theme, Themes } from "./types"; /** * Props for the ShikiHighlighter component @@ -115,8 +110,9 @@ export const ShikiHighlighter = ({ langClassName, showLanguage = true, children: code, - as: Element = 'pre', + as: Element = "pre", customLanguages, + fineGrainedBundle, ...shikiOptions }: ShikiHighlighterProps): React.ReactElement => { const options: HighlighterOptions = { @@ -125,29 +121,21 @@ export const ShikiHighlighter = ({ customLanguages, defaultColor, cssVariablePrefix, + fineGrainedBundle, ...shikiOptions, }; - // Use resolveLanguage to get displayLanguageId directly - const { displayLanguageId } = resolveLanguage( - language, - customLanguages - ); + const { displayLanguageId } = resolveLanguage(language, customLanguages); - const highlightedCode = useShikiHighlighter( - code, - language, - theme, - options - ); + const highlightedCode = useShikiHighlighter(code, language, theme, options); return ( {showLanguage && displayLanguageId ? ( diff --git a/package/src/hook.ts b/package/src/hook.ts index 7fadbba..059bf5a 100644 --- a/package/src/hook.ts +++ b/package/src/hook.ts @@ -19,7 +19,9 @@ import { } from 'shiki'; import type { ShikiLanguageRegistration } from './extended-types'; +import { createFineGrainedBundle } from './bundle'; +import type { Root } from 'hast'; import type { Language, Theme, @@ -130,9 +132,48 @@ export const useShikiHighlighter = ( timeoutId: undefined, }); + const useFineGrainedBundle = !!stableOpts.fineGrainedBundle; + + const bundleCache = useRef>>(new Map()); + + const bundle = useMemo(() => { + if (!useFineGrainedBundle) return null; + + const { + langs, + themes, + engine = 'javascript', + precompiled = false, + // biome-ignore lint/style/noNonNullAssertion: will fix + } = stableOpts.fineGrainedBundle!; // TODO: Fix types + + const cacheKey = `${[...langs].sort().join(',')}|${[...themes] + .sort() + .join(',')}|${engine}|${precompiled}`; + + if (bundleCache.current.has(cacheKey)) { + return bundleCache.current.get(cacheKey); + } + + const bundlePromise = createFineGrainedBundle({ + langs, + themes, + engine, + precompiled, + }); + + bundleCache.current.set(cacheKey, bundlePromise); + + return bundlePromise; + }, [useFineGrainedBundle, stableOpts.fineGrainedBundle]); + const shikiOptions = useMemo(() => { - const { defaultColor, cssVariablePrefix, ...restOptions } = - stableOpts; + const { + defaultColor, + cssVariablePrefix, + fineGrainedBundle, + ...restOptions + } = stableOpts; const languageOption = { lang: languageId }; const themeOptions = isMultiTheme @@ -153,12 +194,17 @@ export const useShikiHighlighter = ( const highlightCode = async () => { if (!languageId) return; - const highlighter = await getSingletonHighlighter({ - langs: [langsToLoad as ShikiLanguageRegistration], - themes: themesToLoad, - }); - - const hast = highlighter.codeToHast(code, shikiOptions); + let hast: Root; + if (useFineGrainedBundle && bundle) { + const resolvedBundle = await bundle; + hast = await resolvedBundle.codeToHast(code, shikiOptions); + } else { + const highlighter = await getSingletonHighlighter({ + langs: [langsToLoad as ShikiLanguageRegistration], + themes: themesToLoad, + }); + hast = highlighter.codeToHast(code, shikiOptions); + } if (isMounted) { setHighlightedCode(toJsxRuntime(hast, { jsx, jsxs, Fragment })); @@ -177,7 +223,13 @@ export const useShikiHighlighter = ( isMounted = false; clearTimeout(timeoutControl.current.timeoutId); }; - }, [code, shikiOptions, stableOpts.delay]); + }, [ + code, + shikiOptions, + stableOpts.delay, + stableOpts.fineGrainedBundle, + useFineGrainedBundle, + ]); return highlightedCode; }; diff --git a/package/src/index.ts b/package/src/index.ts index f57b2cf..9088379 100644 --- a/package/src/index.ts +++ b/package/src/index.ts @@ -1,10 +1,11 @@ -export { useShikiHighlighter } from './hook'; -export { isInlineCode, rehypeInlineCodeProperty } from './utils'; +export { useShikiHighlighter } from "./hook"; +export { isInlineCode, rehypeInlineCodeProperty } from "./utils"; +export { createFineGrainedBundle } from "./bundle"; export { ShikiHighlighter as default, type ShikiHighlighterProps, -} from './component'; +} from "./component"; export type { Language, @@ -12,4 +13,5 @@ export type { Themes, Element, HighlighterOptions, -} from './types'; + FineGrainedBundleOptions, +} from "./types"; diff --git a/package/src/types.ts b/package/src/types.ts index d34f7c6..29254a0 100644 --- a/package/src/types.ts +++ b/package/src/types.ts @@ -6,11 +6,11 @@ import type { ThemeRegistrationAny, StringLiteralUnion, CodeToHastOptions, -} from 'shiki'; +} from "shiki"; -import type { LanguageRegistration } from './extended-types'; +import type { LanguageRegistration } from "./extended-types"; -import type { Element as HastElement } from 'hast'; +import type { Element as HastElement } from "hast"; /** * HTML Element, use to type `node` from react-markdown @@ -32,6 +32,36 @@ type Language = */ type Theme = ThemeRegistrationAny | StringLiteralUnion; +/** + * Fine-grained bundle options. + */ +interface FineGrainedBundleOptions { + /** + * The languages to bundle, specified by their string identifiers. + * @example ['typescript', 'javascript', 'vue'] + */ + langs: readonly BundledLanguage[]; + + /** + * The themes to bundle, specified by their string identifiers. + * @example ['github-dark', 'github-light'] + */ + themes: readonly BundledTheme[]; + + /** + * The engine to use for syntax highlighting. + * @default 'javascript' + */ + engine?: "oniguruma" | "javascript" | "javascript-raw"; + + /** + * Use precompiled grammars. + * Only available when `engine` is set to `javascript` or `javascript-raw`. + * @default false + */ + precompiled?: boolean; +} + /** * A map of color names to themes. * This allows you to specify multiple themes for the generated code. @@ -76,9 +106,15 @@ interface HighlighterOptions extends ReactShikiOptions, Pick< CodeOptionsMultipleThemes, - 'defaultColor' | 'cssVariablePrefix' + "defaultColor" | "cssVariablePrefix" >, - Omit {} + Omit { + /** + * Fine-grained bundle options to reduce bundle size. + * When provided, only specified languages and themes will be included. + */ + fineGrainedBundle?: FineGrainedBundleOptions; +} /** * State for the throttling logic @@ -101,4 +137,5 @@ export type { Element, TimeoutState, HighlighterOptions, + FineGrainedBundleOptions, };