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