1import { use, useMemo, type CSSProperties } from "react";
2
3import { resolveDiffThemeName } from "../../lib/diffRendering";
4import type { EmbeddedScript } from "../../lib/embeddedScripts";
5import { getSyntaxHighlighterPromise } from "../../lib/syntaxHighlighting";
6
7interface SyntaxToken {
8 readonly content: string;
9 readonly offset: number;
10 readonly color?: string;
11 readonly fontStyle?: number;
12}
13
14type TokenLines = ReadonlyArray<ReadonlyArray<SyntaxToken>>;
15
16const NO_EMBEDDED_SCRIPTS: ReadonlyArray<EmbeddedScript> = [];
17
18function syntaxTokenStyle(token: SyntaxToken): CSSProperties {
19 const fontStyle = token.fontStyle ?? 0;
20 return {
21 ...(token.color ? { color: token.color } : {}),
22 ...(fontStyle & 1 ? { fontStyle: "italic" } : {}),
23 ...(fontStyle & 2 ? { fontWeight: 700 } : {}),
24 ...(fontStyle & 4 ? { textDecoration: "underline" } : {}),
25 };
26}
27
28/**
29 * Colors `code` inside the caller's `<pre>` without changing its text, so
30 * selection and copy match the plain version. Suspends while the grammar
31 * loads; wrap it in Suspense with the plain text as the fallback.
32 * `wordClassName` goes on each whitespace-separated word. `embedded`
33 * scripts are colored with their own grammar inside `code`.
34 */
35export function HighlightedTokens({
36 code,
37 language,
38 embedded = NO_EMBEDDED_SCRIPTS,
39 theme,
40 wordClassName,
41}: {
42 code: string;
43 language: string;
44 embedded?: ReadonlyArray<EmbeddedScript>;
45 theme: "light" | "dark";
46 wordClassName?: string;
47}) {
48 const highlighter = use(getSyntaxHighlighterPromise(language));
49 // Every grammar loads into the same shared highlighter.
50 for (const script of embedded) use(getSyntaxHighlighterPromise(script.language));
51 const lines = useMemo(() => {
52 const themeName = resolveDiffThemeName(theme);
53 const tokenize = (text: string, lang: string) =>
54 highlighter.codeToTokens(text, { lang, theme: themeName }).tokens;
55 return keyedLines(
56 code,
57 withEmbeddedScripts(code, tokenize(code, language), embedded, tokenize),
58 );
59 }, [code, embedded, highlighter, language, theme]);
60
61 return lines.map(({ key, tokens, ending }) => (
62 <span key={key}>
63 {wordClassName
64 ? wordsOf(tokens).map((part) =>
65 typeof part === "string" ? (
66 part
67 ) : (
68 <span key={part.key} className={wordClassName}>
69 {part.pieces.map((piece) => (
70 <span key={piece.key} style={syntaxTokenStyle(piece.token)}>
71 {piece.text}
72 </span>
73 ))}
74 </span>
75 ),
76 )
77 : tokens.map((token) => (
78 <span key={`${token.offset}:${token.content}`} style={syntaxTokenStyle(token)}>
79 {token.content}
80 </span>
81 ))}
82 {ending}
83 </span>
84 ));
85}
86
87/**
88 * Recolors each embedded script's span of `code` with its own grammar,
89 * outermost first so nested scripts win. A script whose grammar is missing
90 * keeps the colors around it.
91 */
92export function withEmbeddedScripts(
93 code: string,
94 lines: TokenLines,
95 embedded: ReadonlyArray<EmbeddedScript>,
96 tokenize: (text: string, language: string) => TokenLines,
97): TokenLines {
98 if (embedded.length === 0) return lines;
99 const styles = Array.from<SyntaxToken | undefined>({ length: code.length });
100 for (const token of lines.flat()) {
101 styles.fill(token, token.offset, token.offset + token.content.length);
102 }
103 for (const script of embedded) {
104 let scriptLines: TokenLines;
105 try {
106 scriptLines = tokenize(script.text, script.language);
107 } catch {
108 continue;
109 }
110 for (const token of scriptLines.flat()) {
111 if (token.content === "") continue;
112 const last = token.offset + token.content.length - 1;
113 styles.fill(token, script.starts[token.offset], script.ends[last]);
114 }
115 }
116 return restyledLines(code, styles);
117}
118
119/** Splits `code` into lines of tokens, one per run of characters sharing a style. */
120function restyledLines(
121 code: string,
122 styles: ReadonlyArray<SyntaxToken | undefined>,
123): SyntaxToken[][] {
124 const lines: SyntaxToken[][] = [];
125 let start = 0;
126 for (const lineBreak of [...code.matchAll(/\r?\n/gu), undefined]) {
127 const end = lineBreak?.index ?? code.length;
128 const tokens: SyntaxToken[] = [];
129 for (let index = start; index < end;) {
130 const style = styles[index];
131 let next = index + 1;
132 while (
133 next < end &&
134 styles[next]?.color === style?.color &&
135 styles[next]?.fontStyle === style?.fontStyle
136 ) {
137 next += 1;
138 }
139 tokens.push({
140 content: code.slice(index, next),
141 offset: index,
142 ...(style?.color ? { color: style.color } : {}),
143 ...(style?.fontStyle ? { fontStyle: style.fontStyle } : {}),
144 });
145 index = next;
146 }
147 lines.push(tokens);
148 if (lineBreak) start = lineBreak.index + lineBreak[0].length;
149 }
150 return lines;
151}
152
153/**
154 * Keys each line by its start offset, which stays unique when several lines
155 * are empty, and pairs it with the line ending `code` has after it. Shiki drops
156 * `\n` and `\r\n` from tokens alike, so CRLF has to come back from the source.
157 */
158export function keyedLines<Token extends SyntaxToken>(
159 code: string,
160 lines: ReadonlyArray<ReadonlyArray<Token>>,
161) {
162 const endings = code.match(/\r?\n/gu) ?? [];
163 let lineStart = 0;
164 return lines.map((tokens, index) => {
165 const text = tokens.map((token) => token.content).join("");
166 const key = `${lineStart}:${text}`;
167 const ending = endings[index] ?? "";
168 lineStart += text.length + ending.length;
169 return { key, tokens, ending };
170 });
171}
172
173interface Word {
174 readonly key: string;
175 readonly pieces: Array<{ key: string; text: string; token: SyntaxToken }>;
176}
177
178/**
179 * Regroups a line's tokens into whitespace runs and whole words, matching
180 * `text.split(/(\s+)/)`, so a word made of several tokens (`"$HOME/x"`) wraps
181 * exactly like its plain-text fallback.
182 */
183export function wordsOf(tokens: ReadonlyArray<SyntaxToken>): Array<string | Word> {
184 const parts: Array<string | Word> = [];
185 let word: Word | null = null;
186 for (const token of tokens) {
187 let offset = token.offset;
188 for (const text of token.content.split(/(\s+)/u)) {
189 if (text === "") continue;
190 if (/^\s/u.test(text)) {
191 parts.push(text);
192 word = null;
193 } else {
194 if (!word) {
195 word = { key: String(offset), pieces: [] };
196 parts.push(word);
197 }
198 word.pieces.push({ key: String(offset), text, token });
199 }
200 offset += text.length;
201 }
202 }
203 return parts;
204}