1import {
2 renderTokenLines,
3 type DiffsHighlighter,
4 type DiffsThemeNames,
5 type ThemedToken,
6 type TokensResult,
7} from "@pierre/diffs";
8
9function highlightedDocument(
10 result: TokensResult,
11 lines = renderTokenLines(result.tokens, {}),
12): { type: "root"; children: ReturnType<typeof renderTokenLines> } {
13 return {
14 type: "root" as const,
15 children: [
16 {
17 type: "element" as const,
18 tagName: "pre",
19 properties: {
20 class: "diffs",
21 style:
22 result.rootStyle === false
23 ? undefined
24 : (result.rootStyle ??
25 [
26 result.fg != null ? `color:${result.fg}` : "",
27 result.bg != null ? `background-color:${result.bg}` : "",
28 ]
29 .filter(Boolean)
30 .join(";")),
31 tabIndex: 0,
32 },
33 children: [
34 {
35 type: "element" as const,
36 tagName: "code",
37 properties: {},
38 children: lines.flatMap((line, index) =>
39 index === 0 ? [line] : [{ type: "text" as const, value: "\n" }, line],
40 ),
41 },
42 ],
43 },
44 ],
45 };
46}
47
48export function highlightCodeToHast(
49 highlighter: DiffsHighlighter,
50 code: string,
51 language: string,
52 theme: DiffsThemeNames,
53) {
54 return highlightedDocument(highlighter.codeToTokens(code, { lang: language, theme }));
55}
56
57export function createIncrementalHighlightedDocument(
58 highlighter: DiffsHighlighter,
59 language: string,
60 theme: DiffsThemeNames,
61) {
62 const options = { lang: language, theme };
63 let shell: TokensResult | undefined;
64 let tokenizer: ReturnType<DiffsHighlighter["createStreamTokenizer"]> | undefined;
65 let previous = "";
66 const completedLines: ReturnType<typeof renderTokenLines> = [];
67
68 function dispose() {
69 tokenizer?.dispose();
70 tokenizer = undefined;
71 previous = "";
72 completedLines.length = 0;
73 }
74
75 return Object.assign(
76 (code: string) => {
77 // ANSI and partial CRLF sequences need the full-source normalization path.
78 if (language === "ansi" || code.includes("\r")) {
79 dispose();
80 return highlightCodeToHast(highlighter, code, language, theme);
81 }
82 if (!code.startsWith(previous)) dispose();
83 // Created on first use so an unknown language throws inside the caller's fallback.
84 shell ??= highlighter.codeToTokens("", options);
85 tokenizer ??= highlighter.createStreamTokenizer(options);
86 const update = tokenizer.enqueue(code.slice(previous.length));
87 if (update instanceof Promise) {
88 dispose();
89 void update.catch(() => {});
90 return highlightCodeToHast(highlighter, code, language, theme);
91 }
92 previous = code;
93 let line: ThemedToken[] = [];
94 for (const token of update.stable) {
95 if (token.content === "\n") {
96 completedLines.push(...renderTokenLines([line], {}));
97 line = [];
98 } else line.push(token);
99 }
100 // Completed line nodes retain their identity while the final line changes.
101 return highlightedDocument(shell, [
102 ...completedLines,
103 ...renderTokenLines([update.unstable], {}),
104 ]);
105 },
106 { dispose },
107 );
108}