Skip to content

Commit cdc579f

Browse files
AVGVSTVS96AVGVSTVS96
committed
first draft for fine grained bundle support
1 parent 91a4fa2 commit cdc579f

File tree

6 files changed

+480
-44
lines changed

6 files changed

+480
-44
lines changed

‎package/src/__tests__/fine-grained-bundle.test.tsx‎

Lines changed: 220 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,220 @@
1+
// @ts-nocheck
2+
// ai written test
3+
import { render, waitFor } from "@testing-library/react";
4+
import { describe, it, expect, vi, beforeEach } from "vitest";
5+
import { useShikiHighlighter } from "../hook";
6+
import { createFineGrainedBundle } from "../bundle";
7+
import React from "react";
8+
import type { HighlighterOptions } from "../types";
9+
10+
// Mock the bundle module
11+
vi.mock("../bundle", () => ({
12+
createFineGrainedBundle: vi.fn(),
13+
}));
14+
15+
// Mock the toJsxRuntime function to avoid full rendering
16+
vi.mock("hast-util-to-jsx-runtime", () => ({
17+
toJsxRuntime: vi
18+
.fn()
19+
.mockImplementation(() => <div data-testid="mocked-hast">Mocked HAST</div>),
20+
}));
21+
22+
// Helper component for testing useShikiHighlighter hook
23+
const TestComponent = ({
24+
code,
25+
language,
26+
theme,
27+
options,
28+
}: {
29+
code: string;
30+
language: any;
31+
theme: any;
32+
options?: HighlighterOptions;
33+
}) => {
34+
const highlighted = useShikiHighlighter(code, language, theme, options);
35+
return <div data-testid="test-component">{highlighted}</div>;
36+
};
37+
38+
describe("Fine-grained bundle", () => {
39+
const mockCodeToHast = vi
40+
.fn()
41+
.mockResolvedValue({ type: "element", tagName: "div" });
42+
43+
beforeEach(() => {
44+
vi.resetAllMocks();
45+
// Mock the bundle with a predefined codeToHast function
46+
(createFineGrainedBundle as any).mockResolvedValue({
47+
codeToHast: mockCodeToHast,
48+
bundledLanguages: {},
49+
bundledThemes: {},
50+
});
51+
});
52+
53+
it("should use createCodegenBundle when fineGrainedBundle is provided", async () => {
54+
const code = 'const test = "Hello";';
55+
const fineGrainedBundle = {
56+
langs: ["typescript"],
57+
themes: ["github-dark"],
58+
engine: "javascript" as const,
59+
};
60+
61+
render(
62+
<TestComponent
63+
code={code}
64+
language="typescript"
65+
theme="github-dark"
66+
options={{ fineGrainedBundle }}
67+
/>
68+
);
69+
70+
// Wait for the asynchronous operations to complete
71+
await waitFor(() => {
72+
// Include precompiled: false when checking call arguments
73+
expect(createFineGrainedBundle).toHaveBeenCalledWith({
74+
...fineGrainedBundle,
75+
precompiled: false,
76+
});
77+
expect(mockCodeToHast).toHaveBeenCalled();
78+
});
79+
});
80+
81+
it("should cache the bundle for the same configuration", async () => {
82+
const code = 'const test = "Hello";';
83+
const fineGrainedBundle = {
84+
langs: ["typescript"],
85+
themes: ["github-dark"],
86+
engine: "javascript" as const,
87+
};
88+
89+
// Render with the same config twice
90+
const { rerender } = render(
91+
<TestComponent
92+
code={code}
93+
language="typescript"
94+
theme="github-dark"
95+
options={{ fineGrainedBundle }}
96+
/>
97+
);
98+
99+
await waitFor(() => {
100+
expect(createFineGrainedBundle).toHaveBeenCalledTimes(1);
101+
});
102+
103+
// Re-render with same config
104+
rerender(
105+
<TestComponent
106+
code={code + "// new comment"}
107+
language="typescript"
108+
theme="github-dark"
109+
options={{ fineGrainedBundle }}
110+
/>
111+
);
112+
113+
// Should not create a new bundle
114+
await waitFor(() => {
115+
expect(createFineGrainedBundle).toHaveBeenCalledTimes(1);
116+
});
117+
});
118+
119+
it("should create a new bundle when configuration changes", async () => {
120+
const code = 'const test = "Hello";';
121+
const fineGrainedBundle1 = {
122+
langs: ["typescript"],
123+
themes: ["github-dark"],
124+
engine: "javascript" as const,
125+
};
126+
127+
const fineGrainedBundle2 = {
128+
langs: ["typescript", "javascript"],
129+
themes: ["github-dark", "github-light"],
130+
engine: "javascript" as const,
131+
};
132+
133+
// Render with first config
134+
const { rerender } = render(
135+
<TestComponent
136+
code={code}
137+
language="typescript"
138+
theme="github-dark"
139+
options={{ fineGrainedBundle: fineGrainedBundle1 }}
140+
/>
141+
);
142+
143+
await waitFor(() => {
144+
// Include precompiled: false when checking call arguments
145+
expect(createFineGrainedBundle).toHaveBeenCalledWith({
146+
...fineGrainedBundle1,
147+
precompiled: false,
148+
});
149+
expect(createFineGrainedBundle).toHaveBeenCalledTimes(1);
150+
});
151+
152+
// Re-render with different config
153+
rerender(
154+
<TestComponent
155+
code={code}
156+
language="typescript"
157+
theme="github-dark"
158+
options={{ fineGrainedBundle: fineGrainedBundle2 }}
159+
/>
160+
);
161+
162+
// Should create a new bundle
163+
await waitFor(() => {
164+
// Include precompiled: false when checking call arguments
165+
expect(createFineGrainedBundle).toHaveBeenCalledWith({
166+
...fineGrainedBundle2,
167+
precompiled: false,
168+
});
169+
expect(createFineGrainedBundle).toHaveBeenCalledTimes(2);
170+
});
171+
});
172+
173+
it("should create deterministic cache keys regardless of order", async () => {
174+
const code = 'const test = "Hello";';
175+
176+
// Same configuration with different order
177+
const fineGrainedBundle1 = {
178+
langs: ["typescript", "javascript"],
179+
themes: ["github-dark", "github-light"],
180+
engine: "javascript" as const,
181+
precompiled: false,
182+
};
183+
184+
const fineGrainedBundle2 = {
185+
langs: ["javascript", "typescript"],
186+
themes: ["github-light", "github-dark"],
187+
engine: "javascript" as const,
188+
precompiled: false,
189+
};
190+
191+
// Render with first config
192+
const { rerender } = render(
193+
<TestComponent
194+
code={code}
195+
language="typescript"
196+
theme="github-dark"
197+
options={{ fineGrainedBundle: fineGrainedBundle1 }}
198+
/>
199+
);
200+
201+
await waitFor(() => {
202+
expect(createFineGrainedBundle).toHaveBeenCalledTimes(1);
203+
});
204+
205+
// Re-render with same config but different order
206+
rerender(
207+
<TestComponent
208+
code={code}
209+
language="typescript"
210+
theme="github-dark"
211+
options={{ fineGrainedBundle: fineGrainedBundle2 }}
212+
/>
213+
);
214+
215+
// Should not create a new bundle
216+
await waitFor(() => {
217+
expect(createFineGrainedBundle).toHaveBeenCalledTimes(1);
218+
});
219+
});
220+
});

‎package/src/bundle.ts‎

Lines changed: 137 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,137 @@
1+
import type { BundledLanguage, BundledTheme, RegexEngine } from 'shiki';
2+
import {
3+
bundledLanguagesInfo,
4+
bundledThemesInfo,
5+
} from 'shiki/bundle/full';
6+
import { createOnigurumaEngine } from 'shiki/engine/oniguruma';
7+
import {
8+
createJavaScriptRegexEngine,
9+
createJavaScriptRawEngine,
10+
} from 'shiki/engine/javascript';
11+
import {
12+
createSingletonShorthands,
13+
createdBundledHighlighter,
14+
} from 'shiki/core';
15+
16+
export interface CodegenBundleOptions {
17+
/**
18+
* The languages to bundle, specified by their string identifiers.
19+
* @example ['typescript', 'javascript', 'vue']
20+
*/
21+
langs: readonly BundledLanguage[];
22+
23+
/**
24+
* The themes to bundle, specified by their string identifiers.
25+
* @example ['github-dark', 'github-light']
26+
*/
27+
themes: readonly BundledTheme[];
28+
29+
/**
30+
* The engine to use for syntax highlighting.
31+
* @default 'oniguruma'
32+
*/
33+
engine: 'oniguruma' | 'javascript' | 'javascript-raw';
34+
35+
/**
36+
* Use precompiled grammars.
37+
* Only available when `engine` is set to `javascript` or `javascript-raw`.
38+
* @default false
39+
*/
40+
precompiled?: boolean;
41+
}
42+
43+
/**
44+
* Create a fine grained bundle based on simple string identifiers of languages, themes, and engines.
45+
* This is a runtime adaptation of shiki-codegen's approach.
46+
*/
47+
export async function createFineGrainedBundle({
48+
langs,
49+
themes,
50+
engine = 'oniguruma',
51+
precompiled = false,
52+
}: CodegenBundleOptions) {
53+
if (
54+
precompiled &&
55+
engine !== 'javascript' &&
56+
engine !== 'javascript-raw'
57+
) {
58+
throw new Error(
59+
'Precompiled grammars are only available when using the JavaScript engine'
60+
);
61+
}
62+
63+
const langImports = langs.map(async (lang) => {
64+
const info = bundledLanguagesInfo.find(
65+
(i) => i.id === lang || i.aliases?.includes(lang as string)
66+
);
67+
if (!info) {
68+
throw new Error(`Language ${lang} not found`);
69+
}
70+
71+
const module = await import(
72+
`@shikijs/${precompiled ? 'langs-precompiled' : 'langs'}/${info.id}`
73+
);
74+
return { id: lang, module: module.default || module };
75+
});
76+
77+
const themeImports = themes.map(async (theme) => {
78+
const info = bundledThemesInfo.find((i) => i.id === theme);
79+
if (!info) {
80+
throw new Error(`Theme ${theme} not found`);
81+
}
82+
83+
const module = await import(`@shikijs/themes/${info.id}`);
84+
return { id: theme, module: module.default || module };
85+
});
86+
87+
const resolvedLangs = await Promise.all(langImports);
88+
const resolvedThemes = await Promise.all(themeImports);
89+
90+
let engineInstance: RegexEngine;
91+
switch (engine) {
92+
case 'javascript':
93+
engineInstance = createJavaScriptRegexEngine();
94+
break;
95+
case 'javascript-raw':
96+
engineInstance = createJavaScriptRawEngine();
97+
break;
98+
case 'oniguruma':
99+
engineInstance = await createOnigurumaEngine(import('shiki/wasm'));
100+
break;
101+
}
102+
103+
const createHighlighter = createdBundledHighlighter({
104+
langs: Object.fromEntries(
105+
resolvedLangs.map(({ id, module }) => [
106+
id,
107+
() => Promise.resolve(module),
108+
])
109+
),
110+
themes: Object.fromEntries(
111+
resolvedThemes.map(({ id, module }) => [
112+
id,
113+
() => Promise.resolve(module),
114+
])
115+
),
116+
engine: () => engineInstance,
117+
});
118+
119+
const shorthands = createSingletonShorthands(createHighlighter);
120+
121+
return {
122+
bundledLanguages: Object.fromEntries(
123+
resolvedLangs.map(({ id, module }) => [id, module])
124+
),
125+
bundledThemes: Object.fromEntries(
126+
resolvedThemes.map(({ id, module }) => [id, module])
127+
),
128+
// createHighlighter,
129+
// codeToHtml: shorthands.codeToHtml,
130+
codeToHast: shorthands.codeToHast,
131+
// codeToTokens: shorthands.codeToTokens,
132+
// codeToTokensBase: shorthands.codeToTokensBase,
133+
// codeToTokensWithThemes: shorthands.codeToTokensWithThemes,
134+
// getSingletonHighlighter: shorthands.getSingletonHighlighter,
135+
// getLastGrammarState: shorthands.getLastGrammarState,
136+
};
137+
}

0 commit comments

Comments
 (0)