1import { tokenizeCode, type HighlightScheme, type TokenizedCode } from "../../lib/highlights";
2import type { NativeReviewDiffFile, NativeReviewDiffSource } from "./nativeReviewDiffTypes";
3import type { NativeReviewDiffRow, NativeReviewDiffToken } from "./nativeReviewDiffSurface";
4import { diffHighlightTraceId, traceDiffHighlight } from "./nativeReviewDiffTrace";
5
6interface NativeReviewDiffLineRow extends NativeReviewDiffRow {
7 readonly kind: "line";
8 readonly fileId: string;
9 readonly content: string;
10}
11
12export interface HighlightNativeReviewDiffVisibleRowsInput {
13 readonly rows: ReadonlyArray<NativeReviewDiffRow>;
14 readonly files: ReadonlyArray<NativeReviewDiffFile>;
15 readonly scheme: HighlightScheme;
16 readonly firstRowIndex: number;
17 readonly lastRowIndex: number;
18 readonly overscanRows?: number;
19 readonly maxRows?: number;
20 readonly alreadyHighlightedRowIds?: ReadonlySet<string>;
21 readonly collapsedFileIds?: ReadonlySet<string>;
22 readonly signal?: AbortSignal;
23 readonly onChunk?: (chunk: NativeReviewDiffHighlightChunk) => void;
24}
25
26export interface NativeReviewDiffHighlightChunk {
27 readonly tokensByRowId: Record<string, ReadonlyArray<NativeReviewDiffToken>>;
28 readonly retryableRowIds: ReadonlyArray<string>;
29 readonly rowCount: number;
30 readonly durationMs: number;
31}
32
33const NATIVE_REVIEW_DIFF_VISIBLE_OVERSCAN_ROWS = 160;
34const NATIVE_REVIEW_DIFF_VISIBLE_MAX_ROWS = 360;
35function isHighlightableLineRow(row: NativeReviewDiffRow): row is NativeReviewDiffLineRow {
36 return row.kind === "line" && typeof row.fileId === "string" && typeof row.content === "string";
37}
38
39function hasConsecutiveLineNumbers(
40 previous: number | null | undefined,
41 next: number | null | undefined,
42): boolean {
43 return typeof previous === "number" && typeof next === "number" && next === previous + 1;
44}
45
46// Native lexing is non-incremental. Cache theme-independent raw records, never
47// expanded token objects for an entire large file. Bound both entry count and bytes.
48const MAX_TOKEN_CACHE_BYTES = 64 * 1024 * 1024;
49const MAX_TOKEN_CACHE_ENTRIES = 8;
50// LRU entries own only token data; source contexts and diff rows can collect independently.
51const tokenCacheKeys = new WeakMap<object, object>();
52const tokenCache = new Map<object, { language: string; tokens: TokenizedCode }>();
53let tokenCacheBytes = 0;
54
55function getTokenCacheKey(context: object): object {
56 let key = tokenCacheKeys.get(context);
57 if (!key) {
58 key = {};
59 tokenCacheKeys.set(context, key);
60 }
61 return key;
62}
63
64function getCachedTokens(
65 cacheKey: object,
66 contents: string | (() => string),
67 language: string,
68): TokenizedCode {
69 const cached = tokenCache.get(cacheKey);
70 if (cached) {
71 tokenCache.delete(cacheKey);
72 tokenCacheBytes -= cached.tokens.byteLength;
73 if (cached.language === language) {
74 tokenCache.set(cacheKey, cached);
75 tokenCacheBytes += cached.tokens.byteLength;
76 return cached.tokens;
77 }
78 }
79 const tokens = tokenizeCode(typeof contents === "string" ? contents : contents(), language);
80 if (tokens.byteLength <= MAX_TOKEN_CACHE_BYTES) {
81 tokenCache.set(cacheKey, { language, tokens });
82 tokenCacheBytes += tokens.byteLength;
83 }
84 while (tokenCacheBytes > MAX_TOKEN_CACHE_BYTES || tokenCache.size > MAX_TOKEN_CACHE_ENTRIES) {
85 const oldest = tokenCache.keys().next().value!;
86 tokenCacheBytes -= tokenCache.get(oldest)!.tokens.byteLength;
87 tokenCache.delete(oldest);
88 }
89 return tokens;
90}
91
92const fullSourceContexts = new WeakMap<
93 NativeReviewDiffSource,
94 { old: FullSourceContext; new: FullSourceContext }
95>();
96
97interface FullSourceContext {
98 readonly contents: string;
99 readonly comparisonLines: ReadonlyArray<string>;
100 readonly validation: WeakMap<ReadonlyArray<NativeReviewDiffLineRow>, boolean>;
101}
102
103const sourceRowsByFileCache = new WeakMap<
104 ReadonlyArray<NativeReviewDiffRow>,
105 Map<string, NativeReviewDiffLineRow[]>
106>();
107
108function getSourceRowsByFile(rows: ReadonlyArray<NativeReviewDiffRow>) {
109 let byFile = sourceRowsByFileCache.get(rows);
110 if (!byFile) {
111 byFile = new Map();
112 for (const row of rows) {
113 if (!isHighlightableLineRow(row)) continue;
114 const fileRows = byFile.get(row.fileId);
115 if (fileRows) fileRows.push(row);
116 else byFile.set(row.fileId, [row]);
117 }
118 sourceRowsByFileCache.set(rows, byFile);
119 }
120 return byFile;
121}
122
123function matchesSourceSide(
124 context: FullSourceContext,
125 rows: ReadonlyArray<NativeReviewDiffLineRow>,
126 side: DiffSide,
127) {
128 const cached = context.validation.get(rows);
129 if (cached !== undefined) return cached;
130 // Context participates on both sides, even though it displays NEW tokens.
131 // A known mismatch anywhere invalidates grammar from the entire source side.
132 const matches = rows.every((row) => {
133 if (row.change === (side === "old" ? "add" : "delete")) return true;
134 const lineNumber = side === "old" ? row.oldLineNumber : row.newLineNumber;
135 return (
136 typeof lineNumber === "number" && context.comparisonLines[lineNumber - 1] === row.content
137 );
138 });
139 context.validation.set(rows, matches);
140 return matches;
141}
142
143function createFullSourceContext(contents: string): FullSourceContext {
144 return {
145 contents,
146 // The diff parser retains CR on patch rows, unlike the token output.
147 comparisonLines: contents.split("\n"),
148 validation: new WeakMap(),
149 };
150}
151
152function getFullSourceContexts(source: NativeReviewDiffSource) {
153 let contexts = fullSourceContexts.get(source);
154 if (!contexts) {
155 contexts = {
156 old: createFullSourceContext(source.oldContents),
157 new: createFullSourceContext(source.newContents),
158 };
159 fullSourceContexts.set(source, contexts);
160 }
161 return contexts;
162}
163
164type DiffSide = "old" | "new";
165
166interface SourceSegment {
167 readonly fileId: string;
168 readonly side: DiffSide;
169 readonly rows: NativeReviewDiffLineRow[];
170}
171
172// Row arrays are immutable snapshots; comments and hidden hunks retain source context.
173const sourceSegmentsCache = new WeakMap<
174 ReadonlyArray<NativeReviewDiffRow>,
175 Map<string, { readonly segment: SourceSegment; readonly index: number }>
176>();
177
178function getSourceSegments(rows: ReadonlyArray<NativeReviewDiffRow>) {
179 const cached = sourceSegmentsCache.get(rows);
180 if (cached) return cached;
181
182 const segmentsByRowId = new Map<
183 string,
184 { readonly segment: SourceSegment; readonly index: number }
185 >();
186 for (const side of ["old", "new"] as const) {
187 let segment: SourceSegment | undefined;
188 for (const row of rows) {
189 if (row.kind === "comment" || row.kind === "hunk") continue;
190 if (!isHighlightableLineRow(row)) {
191 segment = undefined;
192 continue;
193 }
194 if (row.change === (side === "old" ? "add" : "delete")) continue;
195
196 const previous = segment?.rows.at(-1);
197 const lineNumber = side === "old" ? row.oldLineNumber : row.newLineNumber;
198 const previousLineNumber = side === "old" ? previous?.oldLineNumber : previous?.newLineNumber;
199 if (
200 !segment ||
201 segment.fileId !== row.fileId ||
202 !hasConsecutiveLineNumbers(previousLineNumber, lineNumber)
203 ) {
204 // A gap has no known incoming grammar state. Highlight the available
205 // fragment independently; never join it to an earlier, incomplete hunk.
206 segment = { fileId: row.fileId, side, rows: [] };
207 }
208 if (displaySide(row) === side) {
209 segmentsByRowId.set(row.id, { segment, index: segment.rows.length });
210 }
211 segment.rows.push(row);
212 }
213 }
214 sourceSegmentsCache.set(rows, segmentsByRowId);
215 return segmentsByRowId;
216}
217
218function displaySide(row: NativeReviewDiffLineRow): DiffSide {
219 // Unified context rows have one token array. Match the NEW side, while OLD
220 // context still advances deletion grammar independently.
221 return row.change === "delete" ? "old" : "new";
222}
223
224function createFileMap(files: ReadonlyArray<NativeReviewDiffFile>) {
225 return new Map(files.map((file) => [file.id, file]));
226}
227
228function clampRowIndex(index: number, rows: ReadonlyArray<NativeReviewDiffRow>) {
229 if (rows.length === 0) {
230 return 0;
231 }
232
233 return Math.min(rows.length - 1, Math.max(0, Math.floor(index)));
234}
235
236function makePlainTokenFallback(
237 row: NativeReviewDiffLineRow,
238): ReadonlyArray<NativeReviewDiffToken> {
239 return [{ content: row.content || " ", color: null, fontStyle: null }];
240}
241
242export async function highlightNativeReviewDiffVisibleRows(
243 input: HighlightNativeReviewDiffVisibleRowsInput,
244): Promise<{
245 readonly tokensByRowId: Record<string, ReadonlyArray<NativeReviewDiffToken>>;
246 // Deliver plain tokens for these rows without marking them permanently highlighted.
247 readonly retryableRowIds: ReadonlyArray<string>;
248 readonly rowCount: number;
249 readonly durationMs: number;
250}> {
251 const request = diffHighlightTraceId(input);
252 traceDiffHighlight("highlight-start", {
253 request,
254 first: input.firstRowIndex,
255 last: input.lastRowIndex,
256 overscan: input.overscanRows ?? NATIVE_REVIEW_DIFF_VISIBLE_OVERSCAN_ROWS,
257 });
258 if (input.signal?.aborted) {
259 return {
260 tokensByRowId: {},
261 retryableRowIds: [],
262 rowCount: 0,
263 durationMs: 0,
264 };
265 }
266
267 const startedAt = performance.now();
268 const fileMap = createFileMap(input.files);
269 const overscanRows = input.overscanRows ?? NATIVE_REVIEW_DIFF_VISIBLE_OVERSCAN_ROWS;
270 const maxRows = input.maxRows ?? NATIVE_REVIEW_DIFF_VISIBLE_MAX_ROWS;
271 const startIndex = clampRowIndex(input.firstRowIndex - overscanRows, input.rows);
272 const endIndex = clampRowIndex(input.lastRowIndex + overscanRows, input.rows);
273 const visibleStartIndex = clampRowIndex(input.firstRowIndex, input.rows);
274 const visibleEndIndex = clampRowIndex(input.lastRowIndex, input.rows);
275 const visibleFileIds = new Set<string>();
276 // Overscan warms visible files without loading unrelated neighboring sources.
277 // Headers and already-highlighted rows still identify files in the viewport.
278 for (let rowIndex = visibleStartIndex; rowIndex <= visibleEndIndex; rowIndex += 1) {
279 const fileId = input.rows[rowIndex]?.fileId;
280 if (typeof fileId === "string") visibleFileIds.add(fileId);
281 }
282 const selectedRows: NativeReviewDiffLineRow[] = [];
283
284 for (
285 let rowIndex = startIndex;
286 rowIndex <= endIndex && selectedRows.length < maxRows;
287 rowIndex += 1
288 ) {
289 const row = input.rows[rowIndex];
290 if (
291 row &&
292 isHighlightableLineRow(row) &&
293 visibleFileIds.has(row.fileId) &&
294 !input.collapsedFileIds?.has(row.fileId) &&
295 !input.alreadyHighlightedRowIds?.has(row.id) &&
296 fileMap.has(row.fileId)
297 ) {
298 selectedRows.push(row);
299 }
300 }
301
302 const tokensByRowId: Record<string, ReadonlyArray<NativeReviewDiffToken>> = {};
303 const retryableRowIds: string[] = [];
304 const selectedRowIds = new Set(selectedRows.map((row) => row.id));
305 const tokenize = (context: object, contents: string | (() => string), language: string) => {
306 const cacheKey = getTokenCacheKey(context);
307 const cacheHit = tokenCache.get(cacheKey)?.language === language;
308 const cached = getCachedTokens(cacheKey, contents, language);
309 traceDiffHighlight("tokenized", {
310 request,
311 characters: typeof contents === "string" ? contents.length : undefined,
312 cacheHit,
313 retainedBytes: cached.byteLength,
314 lines: cached.lineCount,
315 aborted: input.signal?.aborted ?? false,
316 });
317 return cached;
318 };
319
320 const sourceFiles = new Map<string, NativeReviewDiffLineRow[]>();
321 for (const row of selectedRows) {
322 if (!fileMap.get(row.fileId)?.loadContents) continue;
323 const fileRows = sourceFiles.get(row.fileId);
324 if (fileRows) fileRows.push(row);
325 else sourceFiles.set(row.fileId, [row]);
326 }
327 const emitFileChunk = (fileRows: ReadonlyArray<NativeReviewDiffLineRow>) => {
328 if (!input.onChunk || input.signal?.aborted) return;
329 const rowIds = new Set(fileRows.map((row) => row.id));
330 input.onChunk({
331 tokensByRowId: Object.fromEntries(fileRows.map((row) => [row.id, tokensByRowId[row.id]!])),
332 retryableRowIds: retryableRowIds.filter((id) => rowIds.has(id)),
333 rowCount: fileRows.length,
334 durationMs: Math.round(performance.now() - startedAt),
335 });
336 };
337 // The file filling the viewport should not wait for a sliver of its neighbor.
338 for (const [fileId, fileRows] of [...sourceFiles].sort((a, b) => b[1].length - a[1].length)) {
339 if (input.signal?.aborted) break;
340 const file = fileMap.get(fileId)!;
341 const fileTraceId = diffHighlightTraceId(fileId);
342 traceDiffHighlight("highlight-source", { request, file: fileTraceId, rows: fileRows.length });
343 let source: NativeReviewDiffSource | null;
344 try {
345 source = await file.loadContents!(input.signal);
346 } catch {
347 traceDiffHighlight("highlight-source-error", { request, file: fileTraceId });
348 for (const row of fileRows) {
349 selectedRowIds.delete(row.id);
350 tokensByRowId[row.id] = makePlainTokenFallback(row);
351 retryableRowIds.push(row.id);
352 }
353 emitFileChunk(fileRows);
354 continue;
355 }
356 if (input.signal?.aborted) break;
357 const contexts = source ? getFullSourceContexts(source) : null;
358 for (const side of ["old", "new"] as const) {
359 if (input.signal?.aborted) break;
360 const candidate = contexts?.[side];
361 const context =
362 candidate &&
363 matchesSourceSide(candidate, getSourceRowsByFile(input.rows).get(fileId)!, side)
364 ? candidate
365 : null;
366 traceDiffHighlight("source-validation", {
367 request,
368 file: fileTraceId,
369 side,
370 present: !!candidate,
371 matches: !!context,
372 knownLines: context?.comparisonLines.length,
373 });
374 const matchingRows: { row: NativeReviewDiffLineRow; lineIndex: number }[] = [];
375 for (const row of fileRows) {
376 if (displaySide(row) !== side) continue;
377 selectedRowIds.delete(row.id);
378 const lineNumber = side === "old" ? row.oldLineNumber : row.newLineNumber;
379 const lineIndex = typeof lineNumber === "number" ? lineNumber - 1 : -1;
380 if (
381 context &&
382 lineIndex >= 0 &&
383 lineIndex < context.comparisonLines.length &&
384 context.comparisonLines[lineIndex] === row.content
385 ) {
386 matchingRows.push({ row, lineIndex });
387 } else {
388 // A changed revision or missing source cannot
389 // supply trustworthy grammar. Do not cache guessed syntax colors.
390 tokensByRowId[row.id] = makePlainTokenFallback(row);
391 }
392 }
393 if (!context || matchingRows.length === 0) continue;
394 const tokens = tokenize(context, context.contents, file.language);
395 for (const { row, lineIndex } of matchingRows) {
396 tokensByRowId[row.id] =
397 tokens.highlightLine(lineIndex, input.scheme) ?? makePlainTokenFallback(row);
398 }
399 }
400 emitFileChunk(fileRows);
401 await new Promise<void>((resolve) => setTimeout(resolve, 0));
402 }
403
404 const selectedSegments = new Map<
405 SourceSegment,
406 { firstSelected: number; lastSelected: number }
407 >();
408 if (selectedRowIds.size > 0 && !input.signal?.aborted) {
409 const sourceSegments = getSourceSegments(input.rows);
410 for (const row of selectedRows) {
411 if (!selectedRowIds.has(row.id)) continue;
412 const source = sourceSegments.get(row.id);
413 if (!source) continue;
414 const selection = selectedSegments.get(source.segment);
415 if (selection) {
416 selection.lastSelected = source.index;
417 } else {
418 selectedSegments.set(source.segment, {
419 firstSelected: source.index,
420 lastSelected: source.index,
421 });
422 }
423 }
424 }
425
426 for (const [segment, { firstSelected, lastSelected }] of selectedSegments) {
427 if (input.signal?.aborted) break;
428 const file = fileMap.get(segment.fileId);
429 if (!file) continue;
430
431 const tokens = tokenize(
432 segment,
433 () => segment.rows.map((row) => row.content).join("\n"),
434 file.language,
435 );
436 for (let index = firstSelected; index <= lastSelected; index += 1) {
437 const row = segment.rows[index]!;
438 if (selectedRowIds.has(row.id) && displaySide(row) === segment.side) {
439 tokensByRowId[row.id] =
440 tokens.highlightLine(index, input.scheme) ?? makePlainTokenFallback(row);
441 }
442 }
443 }
444
445 if (input.signal?.aborted) {
446 traceDiffHighlight("highlight-aborted", { request });
447 return {
448 tokensByRowId: {},
449 retryableRowIds: [],
450 rowCount: 0,
451 durationMs: 0,
452 };
453 }
454
455 traceDiffHighlight("highlight-result", {
456 request,
457 selected: selectedRows.length,
458 delivered: Object.keys(tokensByRowId).length,
459 retryable: retryableRowIds.length,
460 colored: Object.values(tokensByRowId).filter((tokens) =>
461 tokens.some((token) => token.color !== null),
462 ).length,
463 });
464 return {
465 tokensByRowId,
466 retryableRowIds,
467 rowCount: selectedRows.length,
468 durationMs: Math.round(performance.now() - startedAt),
469 };
470}