first draft for fine grained bundle support

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(() => <div data-testid="mocked-hast">Mocked HAST</div>),
+}));
+
+// 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 <div data-testid="test-component">{highlighted}</div>;
+};
+
+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(
+      <TestComponent
+        code={code}
+        language="typescript"
+        theme="github-dark"
+        options={{ fineGrainedBundle }}
+      />
+    );
+
+    // 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(
+      <TestComponent
+        code={code}
+        language="typescript"
+        theme="github-dark"
+        options={{ fineGrainedBundle }}
+      />
+    );
+
+    await waitFor(() => {
+      expect(createFineGrainedBundle).toHaveBeenCalledTimes(1);
+    });
+
+    // Re-render with same config
+    rerender(
+      <TestComponent
+        code={code + "// new comment"}
+        language="typescript"
+        theme="github-dark"
+        options={{ fineGrainedBundle }}
+      />
+    );
+
+    // 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(
+      <TestComponent
+        code={code}
+        language="typescript"
+        theme="github-dark"
+        options={{ fineGrainedBundle: fineGrainedBundle1 }}
+      />
+    );
+
+    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(
+      <TestComponent
+        code={code}
+        language="typescript"
+        theme="github-dark"
+        options={{ fineGrainedBundle: fineGrainedBundle2 }}
+      />
+    );
+
+    // 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(
+      <TestComponent
+        code={code}
+        language="typescript"
+        theme="github-dark"
+        options={{ fineGrainedBundle: fineGrainedBundle1 }}
+      />
+    );
+
+    await waitFor(() => {
+      expect(createFineGrainedBundle).toHaveBeenCalledTimes(1);
+    });
+
+    // Re-render with same config but different order
+    rerender(
+      <TestComponent
+        code={code}
+        language="typescript"
+        theme="github-dark"
+        options={{ fineGrainedBundle: fineGrainedBundle2 }}
+      />
+    );
+
+    // Should not create a new bundle
+    await waitFor(() => {
+      expect(createFineGrainedBundle).toHaveBeenCalledTimes(1);
+    });
+  });
+});

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,
+  };
+}