1import { Extension } from "@tiptap/core";
2import { type EditorState, Plugin, PluginKey } from "@tiptap/pm/state";
3import { Decoration, DecorationSet } from "@tiptap/pm/view";
4import type { Node as ProseMirrorNode } from "@tiptap/pm/model";
5
6import type { DiffThemeName } from "~/lib/diffRendering";
7import { languageOfInfoString } from "~/composer-code-languages";
8import { getSyntaxHighlighterPromise } from "~/lib/syntaxHighlighting";
9
10interface HighlightedBlock {
11 readonly signature: string;
12 readonly decorations: ReadonlyArray<{ from: number; to: number; color: string }>;
13}
14
15const composerCodeBlockHighlightKey = new PluginKey<DecorationSet>("composerCodeBlockHighlight");
16
17// Tokenizing larger fences can stall the main thread.
18const MAX_HIGHLIGHTED_CODE_BLOCK_LENGTH = 20_000;
19
20// Keep mapped colors while typing in long fences, then refresh after a pause.
21const DEFERRED_HIGHLIGHT_CODE_BLOCK_LENGTH = 2_000;
22const DEFERRED_HIGHLIGHT_DELAY_MS = 300;
23
24function shouldHighlightCodeBlock(code: string): boolean {
25 return code.length <= MAX_HIGHLIGHTED_CODE_BLOCK_LENGTH;
26}
27function blockSignature(node: ProseMirrorNode, theme: DiffThemeName): string {
28 return [theme, String(node.attrs.language ?? ""), node.textContent].join("\u0000");
29}
30
31/** Code blocks are always top-level; avoid walking the whole draft per keystroke. */
32function collectCodeBlocks(state: EditorState): Array<{ node: ProseMirrorNode; pos: number }> {
33 const blocks: Array<{ node: ProseMirrorNode; pos: number }> = [];
34 state.doc.forEach((node, pos) => {
35 if (node.type.name === "codeBlock") blocks.push({ node, pos });
36 });
37 return blocks;
38}
39
40export function composerCodeBlockHighlight(options: {
41 resolveTheme: () => DiffThemeName;
42}): Extension {
43 return Extension.create({
44 name: "composerCodeBlockHighlight",
45
46 addProseMirrorPlugins() {
47 // Prune obsolete signatures each scan so the cache is bounded by the document without evicting live
48 // blocks.
49 const cache = new Map<string, HighlightedBlock>();
50 // Track live signatures so late tokenization cannot restore edited-away blocks.
51 let live = new Set<string>();
52
53 return [
54 new Plugin<DecorationSet>({
55 key: composerCodeBlockHighlightKey,
56
57 state: {
58 init: () => DecorationSet.empty,
59 apply(transaction, value, _oldState, newState) {
60 const mapped = value.map(transaction.mapping, transaction.doc);
61 if (!transaction.docChanged && !transaction.getMeta(composerCodeBlockHighlightKey)) {
62 return mapped;
63 }
64 return buildDecorations(newState, cache, options.resolveTheme(), mapped);
65 },
66 },
67
68 props: {
69 decorations(state) {
70 return composerCodeBlockHighlightKey.getState(state);
71 },
72 },
73
74 view(view) {
75 let disposed = false;
76 let paintedTheme = options.resolveTheme();
77 let scannedDoc: ProseMirrorNode | null = null;
78 let deferredRefresh: ReturnType<typeof setTimeout> | undefined;
79
80 // Repaint dispatches and selection changes must not restart the debounce.
81 const refresh = (deferLong: boolean) => {
82 const theme = options.resolveTheme();
83 const themeChanged = theme !== paintedTheme;
84 paintedTheme = theme;
85 // Selection changes cannot affect tokens; avoid rebuilding every block's text signature.
86 if (!themeChanged && view.state.doc === scannedDoc) return;
87 clearTimeout(deferredRefresh);
88 scannedDoc = view.state.doc;
89 const blocks = collectCodeBlocks(view.state);
90 const signatures = new Set(blocks.map(({ node }) => blockSignature(node, theme)));
91 live = signatures;
92 for (const signature of cache.keys()) {
93 if (!live.has(signature)) cache.delete(signature);
94 }
95 const missing = blocks.filter(
96 ({ node }) =>
97 shouldHighlightCodeBlock(node.textContent) &&
98 !cache.has(blockSignature(node, theme)),
99 );
100 const pending =
101 deferLong && !themeChanged
102 ? missing.filter(
103 ({ node }) => node.textContent.length <= DEFERRED_HIGHLIGHT_CODE_BLOCK_LENGTH,
104 )
105 : missing;
106 if (pending.length < missing.length) {
107 deferredRefresh = setTimeout(() => {
108 scannedDoc = null;
109 refresh(false);
110 }, DEFERRED_HIGHLIGHT_DELAY_MS);
111 }
112 if (pending.length === 0) {
113 if (themeChanged) {
114 view.dispatch(view.state.tr.setMeta(composerCodeBlockHighlightKey, true));
115 }
116 return;
117 }
118
119 void Promise.allSettled(
120 pending.map(async ({ node }) => {
121 // Stored info includes metadata after the language; Shiki needs only the name.
122 const language =
123 languageOfInfoString(String(node.attrs.language ?? "")) || "text";
124 const signature = blockSignature(node, theme);
125 // Cache failure as plain text to clear retained colors.
126 const highlighter = await getSyntaxHighlighterPromise(language).catch(() => null);
127 // An older request must not bypass the debounce if its text returns.
128 if (disposed || live !== signatures || cache.has(signature)) return;
129 cache.set(signature, {
130 signature,
131 decorations: highlighter
132 ? tokenizeBlock(highlighter, node.textContent, language, theme)
133 : [],
134 });
135 }),
136 ).then(() => {
137 if (disposed || live !== signatures) return;
138 view.dispatch(view.state.tr.setMeta(composerCodeBlockHighlightKey, true));
139 });
140 };
141
142 refresh(false);
143 // The theme lives on <html>, outside the editor, so nothing would
144 // otherwise tell the view a repaint is due.
145 const themeObserver = new MutationObserver(() => refresh(false));
146 themeObserver.observe(document.documentElement, {
147 attributeFilter: ["class"],
148 });
149
150 return {
151 update: () => refresh(true),
152 destroy() {
153 disposed = true;
154 clearTimeout(deferredRefresh);
155 themeObserver.disconnect();
156 },
157 };
158 },
159 }),
160 ];
161 },
162 });
163}
164
165type Highlighter = Awaited<ReturnType<typeof getSyntaxHighlighterPromise>>;
166
167export function tokenizeBlock(
168 highlighter: Highlighter,
169 code: string,
170 language: string,
171 theme: DiffThemeName,
172): ReadonlyArray<{ from: number; to: number; color: string }> {
173 let tokens;
174 try {
175 tokens = highlighter.codeToTokens(code, { lang: language, theme }).tokens;
176 } catch {
177 return [];
178 }
179
180 const decorations: Array<{ from: number; to: number; color: string }> = [];
181 let offset = 0;
182 for (const [lineIndex, line] of tokens.entries()) {
183 // The separator between lines, which Shiki drops; CRLF is two characters.
184 if (lineIndex > 0) offset += code.startsWith("\r\n", offset) ? 2 : 1;
185 for (const token of line) {
186 const length = token.content.length;
187 if (token.color && token.content.trim()) {
188 decorations.push({ from: offset, to: offset + length, color: token.color });
189 }
190 offset += length;
191 }
192 }
193 return decorations;
194}
195
196function buildDecorations(
197 state: EditorState,
198 cache: Map<string, HighlightedBlock>,
199 theme: DiffThemeName,
200 previous: DecorationSet,
201): DecorationSet {
202 const decorations: Decoration[] = [];
203 for (const { node, pos } of collectCodeBlocks(state)) {
204 if (!shouldHighlightCodeBlock(node.textContent)) continue;
205 const highlighted = cache.get(blockSignature(node, theme));
206 if (!highlighted) {
207 // A mapped token can span split fences and the paragraph between them.
208 const start = pos + 1;
209 const end = start + node.content.size;
210 for (const decoration of previous.find(start, end)) {
211 if (decoration.from >= start && decoration.to <= end) {
212 decorations.push(decoration);
213 continue;
214 }
215 const from = Math.max(decoration.from, start);
216 const to = Math.min(decoration.to, end);
217 if (from < to) decorations.push(colorDecoration(from, to, decoration.spec.color));
218 }
219 continue;
220 }
221 // Token offsets index the block's text, which lines up with document
222 // positions only because `codeBlock` is `text*` with no marks — one
223 // unmarked text node, so `content.size === textContent.length`. Allowing
224 // marks or hard breaks in the code block would desync these silently.
225 if (node.content.size !== node.textContent.length) continue;
226 const start = pos + 1;
227 for (const decoration of highlighted.decorations) {
228 decorations.push(
229 colorDecoration(start + decoration.from, start + decoration.to, decoration.color),
230 );
231 }
232 }
233 return DecorationSet.create(state.doc, decorations);
234}
235
236// Keep the color available when clipping mapped decorations to a fence.
237function colorDecoration(from: number, to: number, color: string): Decoration {
238 return Decoration.inline(from, to, { style: `color:${color}` }, { color });
239}