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