1import type { DiffsHighlighter } from "@pierre/diffs";
2
3import type { DiffThemeName } from "./diffRendering";
4
5function codeChildren(root: ReturnType<DiffsHighlighter["codeToHast"]>) {
6 const pre = root.children.find((node) => node.type === "element" && node.tagName === "pre");
7 if (pre?.type !== "element") throw new Error("Missing highlighted pre element");
8 const code = pre.children.find((node) => node.type === "element" && node.tagName === "code");
9 if (code?.type !== "element") throw new Error("Missing highlighted code element");
10 return code.children;
11}
12
13/** Resume tokenization after the last completed line. Keep its grammar state so
14 * multiline strings, comments, and embedded languages continue to highlight as
15 * they do in a full pass. The current line is always highlighted again.
16 */
17export function createIncrementalHighlightedDocument(
18 highlighter: DiffsHighlighter,
19 language: string,
20 theme: DiffThemeName,
21): (code: string) => ReturnType<DiffsHighlighter["codeToHast"]> {
22 const options = { lang: language, theme };
23 const newline = { type: "text" as const, value: "\n" };
24 let cached:
25 | {
26 prefix: string;
27 state: ReturnType<DiffsHighlighter["getLastGrammarState"]>;
28 children: ReturnType<typeof codeChildren>;
29 }
30 | undefined;
31
32 return (code: string) => {
33 // Plain text and ANSI do not have a TextMate grammar state. A CR at the end
34 // of a chunk can still become a CRLF, so keep that input on the full path.
35 if (
36 !language ||
37 ["text", "plaintext", "plain", "txt", "ansi"].includes(language) ||
38 code.includes("\r")
39 ) {
40 return highlighter.codeToHast(code, options);
41 }
42 if (cached && !code.startsWith(cached.prefix)) cached = undefined;
43 const end = code.lastIndexOf("\n") + 1;
44 if (end > (cached?.prefix.length ?? 0)) {
45 // Omit the final newline: Shiki would tokenize an extra empty line and
46 // advance the grammar state twice before we process the following line.
47 const root = highlighter.codeToHast(code.slice(cached?.prefix.length ?? 0, end - 1), {
48 ...options,
49 ...(cached ? { grammarState: cached.state } : {}),
50 });
51 const state = highlighter.getLastGrammarState(root);
52 if (!state) {
53 cached = undefined;
54 return highlighter.codeToHast(code, options);
55 }
56 cached = {
57 prefix: code.slice(0, end),
58 state,
59 children: [...(cached ? [...cached.children, newline] : []), ...codeChildren(root)],
60 };
61 }
62 const prefix = cached;
63 if (!prefix) return highlighter.codeToHast(code, options);
64 return highlighter.codeToHast(code.slice(prefix.prefix.length), {
65 ...options,
66 grammarState: prefix.state,
67 transformers: [
68 {
69 code: (node) => ({ ...node, children: [...prefix.children, newline, ...node.children] }),
70 },
71 ],
72 });
73 };
74}