authorgravatar for julius0216@outlook.comJulius Marminge <julius0216@outlook.com> 2026-09-11 12:04:44-07:00
committergravatar for noreply@github.comGitHub <noreply@github.com> 2026-09-11 12:04:44-07:00
log7bd7f99e6cab940892ff9240549c507bfd066aee
tree33db37c83b4d03fd9e9326d15439aa856ebf2368
parent2b7d3a45e6ee9454e4dd46aaf0276ca05d8555d3
signature Signed by PGP key B5690EEEBB952194

perf(mobile): reuse completed code lines while streaming (#11211)

11 files changed, 783 insertions(+), 184 deletions(-)

apps/mobile/modules/t3-markdown-text/src/NativeMarkdownBlock.tsx+66-166
......@@ -1,10 +1,14 @@
1import { createContext, useContext, useEffect, useState } from "react";
1import { createContext, memo, useContext, useMemo } from "react";
22import { Image, Platform, ScrollView, Text, useColorScheme, View } from "react-native";
33import type { MarkdownNode } from "react-native-nitro-markdown/headless";
44
55import { CopyTextButton } from "./CopyTextButton";
66import { MarkdownTextPrimitive } from "./MarkdownTextPrimitive";
7import { nativeMarkdownDocumentRuns, nativeMarkdownListItemBlocks } from "./nativeMarkdownText";
7import {
8 nativeMarkdownDocumentRuns,
9 nativeMarkdownListItemBlocks,
10 nativeMarkdownNodePosition,
11} from "./nativeMarkdownText";
812import { NativeMarkdownSelectableText } from "./NativeMarkdownSelectableText";
913import type {
1014 MarkdownCodeHighlighter,
......@@ -13,15 +17,11 @@ import type {
1317 NativeMarkdownTextStyle,
1418 SelectableMarkdownSkill,
1519} from "./SelectableMarkdownText.types";
20import { useHighlightedCode, type HighlightedCode } from "./useHighlightedCode";
1621
1722/** Set by SelectableMarkdownText so images anywhere in the block tree can use it. */
1823export const MarkdownImageRendererContext = createContext<MarkdownImageRenderer | null>(null);
1924
20type HighlightedCode = ReadonlyArray<ReadonlyArray<MarkdownHighlightedToken>>;
21
22const highlightedCodeCache = new Map<string, HighlightedCode>();
23const highlightedCodePromiseCache = new Map<string, Promise<HighlightedCode>>();
24const HIGHLIGHTED_CODE_CACHE_LIMIT = 64;
2525const MONO_FONT_FAMILY = Platform.select({
2626 ios: "ui-monospace",
2727 android: "monospace",
......@@ -29,7 +29,7 @@ const MONO_FONT_FAMILY = Platform.select({
2929});
3030
3131function nodeKey(node: MarkdownNode, index: number): string {
32 return `${node.type}:${node.beg ?? index}:${node.end ?? index}`;
32 return `${node.type}:${nativeMarkdownNodePosition(node, index)}`;
3333}
3434
3535/** Code inside markdown scales with the base text size (12pt at the default 15pt body). */
......@@ -67,174 +67,74 @@ function SelectableNode(props: {
6767 );
6868}
6969
70function codeHighlightCacheKey(
71 code: string,
72 language: string | undefined,
73 theme: "light" | "dark",
74): string {
75 return `${theme}:${language ?? "text"}:${code}`;
76}
77
78function cacheHighlightedCode(key: string, tokens: HighlightedCode): void {
79 highlightedCodeCache.delete(key);
80 highlightedCodeCache.set(key, tokens);
81
82 while (highlightedCodeCache.size > HIGHLIGHTED_CODE_CACHE_LIMIT) {
83 const oldestKey = highlightedCodeCache.keys().next().value;
84 if (oldestKey === undefined) {
85 break;
86 }
87 highlightedCodeCache.delete(oldestKey);
88 }
89}
90
91function loadHighlightedCode(
92 code: string,
93 language: string | undefined,
94 theme: "light" | "dark",
95 highlightCode: MarkdownCodeHighlighter,
96): Promise<HighlightedCode> {
97 const key = codeHighlightCacheKey(code, language, theme);
98 const cached = highlightedCodeCache.get(key);
99 if (cached) {
100 return Promise.resolve(cached);
101 }
102
103 const pending = highlightedCodePromiseCache.get(key);
104 if (pending) {
105 return pending;
106 }
107
108 const promise = highlightCode({ code, language, theme })
109 .then((tokens) => {
110 cacheHighlightedCode(key, tokens);
111 highlightedCodePromiseCache.delete(key);
112 return tokens;
113 })
114 .catch((error) => {
115 highlightedCodePromiseCache.delete(key);
116 throw error;
117 });
118 highlightedCodePromiseCache.set(key, promise);
119 return promise;
120}
121
122function useHighlightedCode(
123 code: string,
124 language: string | undefined,
125 theme: "light" | "dark",
126 highlightCode: MarkdownCodeHighlighter,
127): HighlightedCode | null {
128 const key = codeHighlightCacheKey(code, language, theme);
129 const [highlighted, setHighlighted] = useState<{
130 readonly key: string;
131 readonly tokens: HighlightedCode | null;
132 }>(() => ({
133 key,
134 tokens: highlightedCodeCache.get(key) ?? null,
135 }));
136
137 useEffect(() => {
138 let active = true;
139 const cached = highlightedCodeCache.get(key);
140 if (cached) {
141 cacheHighlightedCode(key, cached);
142 setHighlighted({ key, tokens: cached });
143 return () => {
144 active = false;
145 };
146 }
147
148 void loadHighlightedCode(code, language, theme, highlightCode)
149 .then((tokens) => {
150 if (active) {
151 setHighlighted({ key, tokens });
152 }
153 })
154 .catch(() => {
155 if (active) {
156 setHighlighted({ key, tokens: null });
157 }
158 });
159 return () => {
160 active = false;
161 };
162 }, [code, highlightCode, key, language, theme]);
163
164 return highlighted.key === key ? highlighted.tokens : null;
165}
166
167function HighlightedCodeText(props: {
168 readonly content: string;
169 readonly highlighted: HighlightedCode | null;
170 readonly textStyle: NativeMarkdownTextStyle;
70const HighlightedCodeLine = memo(function HighlightedCodeLine(props: {
71 readonly tokens: ReadonlyArray<MarkdownHighlightedToken>;
72 readonly color: string;
73 readonly newline: boolean;
17174}) {
172 if (!props.highlighted) {
173 return (
75 let offset = 0;
76 const children = [];
77 for (const token of props.tokens) {
78 if (!token.content) continue;
79 children.push(
17480 <MarkdownTextPrimitive
175 uiTextView
176 selectable
81 key={offset}
17782 style={{
178 color: props.textStyle.codeColor,
83 color: token.color ?? props.color,
17984 fontFamily: MONO_FONT_FAMILY,
180 fontSize: codeBlockFontSize(props.textStyle),
181 lineHeight: codeBlockLineHeight(props.textStyle),
85 fontStyle: token.fontStyle !== null && (token.fontStyle & 1) === 1 ? "italic" : "normal",
86 fontWeight: token.fontStyle !== null && (token.fontStyle & 2) === 2 ? "700" : "400",
18287 }}
18388 >
184 {props.content}
185 </MarkdownTextPrimitive>
89 {token.content}
90 </MarkdownTextPrimitive>,
18691 );
92 offset += token.content.length;
18793 }
188 const highlighted = props.highlighted;
189 let sourceOffset = 0;
190 const keyOccurrences = new Map<string, number>();
191 const keyedLines = highlighted.map((line) => {
192 const lineStart = sourceOffset;
193 const tokens = line.map((token) => {
194 const start = sourceOffset;
195 sourceOffset += token.content.length;
196 const signature = `${start}:${token.content}:${token.color ?? ""}:${token.fontStyle ?? ""}`;
197 const occurrence = keyOccurrences.get(signature) ?? 0;
198 keyOccurrences.set(signature, occurrence + 1);
199 return { key: `${signature}:${occurrence}`, token };
200 });
201 sourceOffset += 1;
202 return {
203 key: `line:${lineStart}:${line.map((token) => token.content).join("")}`,
204 tokens,
205 };
206 });
94 return (
95 <MarkdownTextPrimitive>
96 {children}
97 {props.newline ? "\n" : ""}
98 </MarkdownTextPrimitive>
99 );
100});
207101
102function HighlightedCodeText(props: {
103 readonly content: string;
104 readonly highlighted: HighlightedCode | null;
105 readonly textStyle: NativeMarkdownTextStyle;
106}) {
107 // The text root provides inherited styles through context. A new style object
108 // would rerender every token even when its completed line is unchanged.
109 const fontSize = codeBlockFontSize(props.textStyle);
110 const lineHeight = codeBlockLineHeight(props.textStyle);
111 const style = useMemo(
112 () => ({
113 color: props.textStyle.codeColor,
114 fontFamily: MONO_FONT_FAMILY,
115 fontSize,
116 lineHeight,
117 }),
118 [props.textStyle.codeColor, fontSize, lineHeight],
119 );
120 let offset = 0;
121 const lines = [];
122 if (props.highlighted) {
123 for (const tokens of props.highlighted) {
124 lines.push(
125 <HighlightedCodeLine
126 key={offset}
127 tokens={tokens}
128 color={props.textStyle.codeColor}
129 newline={lines.length + 1 < props.highlighted.length}
130 />,
131 );
132 offset += tokens.reduce((length, token) => length + token.content.length, 0) + 1;
133 }
134 }
208135 return (
209 <MarkdownTextPrimitive
210 uiTextView
211 selectable
212 style={{
213 color: props.textStyle.codeColor,
214 fontFamily: MONO_FONT_FAMILY,
215 fontSize: codeBlockFontSize(props.textStyle),
216 lineHeight: codeBlockLineHeight(props.textStyle),
217 }}
218 >
219 {keyedLines.map((line, lineIndex) => (
220 <MarkdownTextPrimitive key={line.key}>
221 {line.tokens.map(({ key, token }) => (
222 <MarkdownTextPrimitive
223 key={key}
224 style={{
225 color: token.color ?? props.textStyle.codeColor,
226 fontFamily: MONO_FONT_FAMILY,
227 fontStyle:
228 token.fontStyle !== null && (token.fontStyle & 1) === 1 ? "italic" : "normal",
229 fontWeight: token.fontStyle !== null && (token.fontStyle & 2) === 2 ? "700" : "400",
230 }}
231 >
232 {token.content}
233 </MarkdownTextPrimitive>
234 ))}
235 {lineIndex + 1 < keyedLines.length ? "\n" : ""}
236 </MarkdownTextPrimitive>
237 ))}
136 <MarkdownTextPrimitive uiTextView selectable style={style}>
137 {props.highlighted ? lines : props.content}
238138 </MarkdownTextPrimitive>
239139 );
240140}
apps/mobile/modules/t3-markdown-text/src/SelectableMarkdownText.types.ts+13-2
......@@ -25,11 +25,22 @@ export interface MarkdownHighlightedToken {
2525 readonly fontStyle: number | null;
2626}
2727
28export type MarkdownCodeHighlighter = (input: {
28export interface MarkdownCodeHighlightInput {
29 /** Identity of the mounted code block, for incremental highlighting. */
30 readonly session?: object;
2931 readonly code: string;
3032 readonly language?: string | null;
3133 readonly theme: "light" | "dark";
32}) => Promise<ReadonlyArray<ReadonlyArray<MarkdownHighlightedToken>>>;
34}
35export interface MarkdownCodeHighlighter {
36 (
37 input: MarkdownCodeHighlightInput,
38 ): Promise<ReadonlyArray<ReadonlyArray<MarkdownHighlightedToken>>>;
39 /** Optional synchronous result for a small append to an already warm block. */
40 read?: (
41 input: MarkdownCodeHighlightInput,
42 ) => ReadonlyArray<ReadonlyArray<MarkdownHighlightedToken>> | undefined;
43}
3344
3445export interface SelectableMarkdownSkill {
3546 readonly name: string;
apps/mobile/modules/t3-markdown-text/src/nativeMarkdownText.ts+18-5
......@@ -694,21 +694,31 @@ function containsRichBlock(node: MarkdownNode): boolean {
694694 return (node.children ?? []).some(containsRichBlock);
695695}
696696
697/**
698 * Sibling identity for React keys. A source offset survives appends while the
699 * document streams; the child index is the fallback for offset-free nodes. The
700 * two never share a namespace, so a positioned node cannot collide with an
701 * offset-free sibling whose index happens to equal its offset.
702 */
703export function nativeMarkdownNodePosition(node: MarkdownNode, index: number): string {
704 return node.beg === undefined ? `index:${index}` : `offset:${node.beg}`;
705}
706
697707export function nativeMarkdownDocumentChunks(
698708 document: MarkdownNode,
699709): ReadonlyArray<NativeMarkdownDocumentChunk> {
700710 const chunks: NativeMarkdownDocumentChunk[] = [];
701711 let selectableNodes: MarkdownNode[] = [];
712 let selectableStart = 0;
702713
703714 const flushSelectable = () => {
704 if (selectableNodes.length === 0) {
715 const first = selectableNodes[0];
716 if (!first) {
705717 return;
706718 }
707 const first = selectableNodes[0];
708 const last = selectableNodes.at(-1);
709719 chunks.push({
710720 kind: "selectable",
711 key: `selectable:${first?.beg ?? "start"}:${last?.end ?? "end"}`,
721 key: `selectable:${nativeMarkdownNodePosition(first, selectableStart)}`,
712722 node: {
713723 type: "document",
714724 children: selectableNodes,
......@@ -719,6 +729,9 @@ export function nativeMarkdownDocumentChunks(
719729
720730 for (const [index, child] of (document.children ?? []).entries()) {
721731 if (!containsRichBlock(child)) {
732 if (selectableNodes.length === 0) {
733 selectableStart = index;
734 }
722735 selectableNodes.push(child);
723736 continue;
724737 }
......@@ -726,7 +739,7 @@ export function nativeMarkdownDocumentChunks(
726739 flushSelectable();
727740 chunks.push({
728741 kind: "rich",
729 key: `rich:${child.type}:${child.beg ?? index}:${child.end ?? index}`,
742 key: `rich:${child.type}:${nativeMarkdownNodePosition(child, index)}`,
730743 node: child,
731744 });
732745 }
apps/mobile/modules/t3-markdown-text/src/pendingCodeHighlight.test.ts created+28
......@@ -0,0 +1,28 @@
1import { expect, it } from "vite-plus/test";
2import { pendingCodeHighlight } from "./pendingCodeHighlight";
3
4it("keeps completed colors and exact current text without retaining an edited tail", () => {
5 const colored = [
6 [{ content: "const n = 1;", color: "red", fontStyle: 0 }],
7 [{ content: "partial", color: "blue", fontStyle: 0 }],
8 ];
9 const result = pendingCodeHighlight(
10 "const n = 1;\npartial",
11 "const n = 1;\nchanged\nnext",
12 colored,
13 )!;
14 expect(result[0]).toBe(colored[0]);
15 expect(result.map((line) => line.map((token) => token.content).join("")).join("\n")).toBe(
16 "const n = 1;\nchanged\nnext",
17 );
18 expect(
19 result
20 .slice(1)
21 .flat()
22 .every((token) => token.color === null),
23 ).toBe(true);
24 expect(
25 pendingCodeHighlight("const n = 1;\npartial", "const n = 2;\npartial", colored),
26 ).toBeNull();
27 expect(pendingCodeHighlight("partial", "other", colored)).toBeNull();
28});
apps/mobile/modules/t3-markdown-text/src/pendingCodeHighlight.ts created+19
......@@ -0,0 +1,19 @@
1import type { MarkdownHighlightedToken } from "./SelectableMarkdownText.types";
2
3/** Keep finished lines colored while the current line awaits highlighting. */
4export function pendingCodeHighlight(
5 previousCode: string,
6 code: string,
7 tokens: ReadonlyArray<ReadonlyArray<MarkdownHighlightedToken>>,
8): ReadonlyArray<ReadonlyArray<MarkdownHighlightedToken>> | null {
9 const end = previousCode.lastIndexOf("\n") + 1;
10 if (!end || !code.startsWith(previousCode.slice(0, end))) return null;
11 const completedLines = previousCode.slice(0, end).split("\n").length - 1;
12 return [
13 ...tokens.slice(0, completedLines),
14 ...code
15 .slice(end)
16 .split("\n")
17 .map((content) => [{ content, color: null, fontStyle: null }]),
18 ];
19}
apps/mobile/modules/t3-markdown-text/src/useHighlightedCode.test.tsx created+189
......@@ -0,0 +1,189 @@
1import { act, useLayoutEffect } from "react";
2import { createRoot, type Root } from "react-dom/client";
3import { afterEach, beforeEach, describe, expect, it, vi } from "vite-plus/test";
4
5import type { MarkdownCodeHighlighter } from "./SelectableMarkdownText.types";
6import { useHighlightedCode, type HighlightedCode } from "./useHighlightedCode";
7
8type Deferred = {
9 readonly code: string;
10 readonly resolve: (tokens: HighlightedCode) => void;
11};
12
13/** One colored line per newline-separated line, so a line's color proves where it came from. */
14function tokensFor(code: string, color: string): HighlightedCode {
15 return code.split("\n").map((content) => [{ content, color, fontStyle: 0 }]);
16}
17
18function lineText(tokens: HighlightedCode | null): string | null {
19 return tokens?.map((line) => line.map((token) => token.content).join("")).join("\n") ?? null;
20}
21
22function lineColors(tokens: HighlightedCode | null): ReadonlyArray<string | null> {
23 return tokens?.map((line) => line[0]?.color ?? null) ?? [];
24}
25
26const deferred: Deferred[] = [];
27const readResults = new Map<string, HighlightedCode>();
28let renders = 0;
29let result: HighlightedCode | null = null;
30let root: Root;
31
32const highlightCode: MarkdownCodeHighlighter = Object.assign(
33 vi.fn(
34 (input: { code: string }) =>
35 new Promise<HighlightedCode>((resolve) => {
36 deferred.push({ code: input.code, resolve });
37 }),
38 ),
39 {
40 read: vi.fn((input: { code: string }) => readResults.get(input.code)),
41 },
42);
43
44function Probe(props: {
45 readonly code: string;
46 readonly language?: string;
47 readonly theme?: "light" | "dark";
48}) {
49 const tokens = useHighlightedCode(
50 props.code,
51 props.language ?? "ts",
52 props.theme ?? "dark",
53 highlightCode,
54 );
55 useLayoutEffect(() => {
56 renders += 1;
57 result = tokens;
58 });
59 return null;
60}
61
62async function render(props: Parameters<typeof Probe>[0]) {
63 await act(() => root.render(<Probe {...props} />));
64}
65
66async function resolveAsync(index: number, color: string) {
67 const pending = deferred[index];
68 if (!pending) throw new Error(`no pending highlight at ${index}`);
69 await act(async () => {
70 pending.resolve(tokensFor(pending.code, color));
71 });
72}
73
74beforeEach(async () => {
75 // The probe has no DOM output, but ReactDOM needs an event target.
76 const document = {
77 nodeType: 9,
78 addEventListener() {},
79 removeEventListener() {},
80 };
81 const container = {
82 nodeType: 1,
83 tagName: "DIV",
84 namespaceURI: "http://www.w3.org/1999/xhtml",
85 ownerDocument: document,
86 addEventListener() {},
87 removeEventListener() {},
88 };
89 vi.stubGlobal("document", document);
90 vi.stubGlobal("window", { document, HTMLIFrameElement: EventTarget });
91 vi.stubGlobal("IS_REACT_ACT_ENVIRONMENT", true);
92 deferred.length = 0;
93 readResults.clear();
94 renders = 0;
95 result = null;
96 vi.mocked(highlightCode).mockClear();
97 vi.mocked(highlightCode.read!).mockClear();
98 root = createRoot(container as unknown as HTMLElement);
99});
100
101afterEach(async () => {
102 await act(() => root.unmount());
103 vi.unstubAllGlobals();
104});
105
106// Every test uses a distinct language so the module-level token cache from an
107// earlier test can never satisfy a later one.
108describe("useHighlightedCode", () => {
109 it("keeps colors from the latest synchronous read when a later read misses", async () => {
110 const language = "sync-baseline";
111 await render({ code: "a", language });
112 expect(result).toBeNull();
113 await resolveAsync(0, "async");
114 expect(lineColors(result)).toEqual(["async"]);
115
116 readResults.set("a\nb", tokensFor("a\nb", "sync-1"));
117 readResults.set("a\nb\nc", tokensFor("a\nb\nc", "sync-2"));
118 const rendersBeforeAppends = renders;
119 await render({ code: "a\nb", language });
120 await render({ code: "a\nb\nc", language });
121 expect(lineColors(result)).toEqual(["sync-2", "sync-2", "sync-2"]);
122 // A synchronous hit is one render per append: no state update, no highlighter call.
123 expect(renders - rendersBeforeAppends).toBe(2);
124 expect(highlightCode).toHaveBeenCalledTimes(1);
125
126 await render({ code: "a\nb\nc\nd", language });
127 expect(highlightCode).toHaveBeenCalledTimes(2);
128 expect(lineText(result)).toBe("a\nb\nc\nd");
129 // Lines before the previous code's final newline are the completed ones.
130 expect(lineColors(result)).toEqual(["sync-2", "sync-2", null, null]);
131
132 await resolveAsync(1, "async-2");
133 expect(lineColors(result)).toEqual(["async-2", "async-2", "async-2", "async-2"]);
134 });
135
136 it("does not show obsolete text after a non-append edit", async () => {
137 const language = "non-append";
138 await render({ code: "a\nb", language });
139 await resolveAsync(0, "async");
140 readResults.set("a\nb\nc", tokensFor("a\nb\nc", "sync"));
141 await render({ code: "a\nb\nc", language });
142 expect(lineColors(result)).toEqual(["sync", "sync", "sync"]);
143
144 await render({ code: "x\nb\nc", language });
145 expect(result).toBeNull();
146 await resolveAsync(1, "async-2");
147 expect(lineText(result)).toBe("x\nb\nc");
148 });
149
150 it("never reuses colors across a language or theme change", async () => {
151 const language = "theme-change";
152 await render({ code: "a\nb", language });
153 await resolveAsync(0, "async");
154 readResults.set("a\nb\nc", tokensFor("a\nb\nc", "sync"));
155 await render({ code: "a\nb\nc", language });
156 expect(lineColors(result)).toEqual(["sync", "sync", "sync"]);
157
158 await render({ code: "a\nb\nc\nd", language, theme: "light" });
159 expect(result).toBeNull();
160 await render({ code: "a\nb\nc\nd", language: `${language}-other` });
161 expect(result).toBeNull();
162 });
163
164 it("ignores an asynchronous completion that is older than the current code", async () => {
165 const language = "stale-async";
166 await render({ code: "a\nb", language });
167 await resolveAsync(0, "async");
168
169 await render({ code: "a\nb\nc", language });
170 expect(deferred).toHaveLength(2);
171 expect(lineColors(result)).toEqual(["async", null, null]);
172
173 readResults.set("a\nb\nc\nd", tokensFor("a\nb\nc\nd", "sync"));
174 await render({ code: "a\nb\nc\nd", language });
175 expect(lineColors(result)).toEqual(["sync", "sync", "sync", "sync"]);
176
177 await render({ code: "a\nb\nc\nd\ne", language });
178 expect(deferred).toHaveLength(3);
179 expect(lineColors(result)).toEqual(["sync", "sync", "sync", null, null]);
180
181 // The older request finishes after newer synchronous results exist.
182 await resolveAsync(1, "stale");
183 expect(lineText(result)).toBe("a\nb\nc\nd\ne");
184 expect(lineColors(result)).toEqual(["sync", "sync", "sync", null, null]);
185
186 await resolveAsync(2, "async-3");
187 expect(lineColors(result)).toEqual(["async-3", "async-3", "async-3", "async-3", "async-3"]);
188 });
189});
apps/mobile/modules/t3-markdown-text/src/useHighlightedCode.ts created+149
......@@ -0,0 +1,149 @@
1import { useEffect, useMemo, useRef, useState } from "react";
2
3import { pendingCodeHighlight } from "./pendingCodeHighlight";
4import type {
5 MarkdownCodeHighlighter,
6 MarkdownHighlightedToken,
7} from "./SelectableMarkdownText.types";
8
9export type HighlightedCode = ReadonlyArray<ReadonlyArray<MarkdownHighlightedToken>>;
10
11interface HighlightedCodeResult {
12 readonly key: string;
13 readonly code: string;
14 readonly language: string | undefined;
15 readonly theme: "light" | "dark";
16 readonly tokens: HighlightedCode | null;
17}
18
19const highlightedCodeCache = new Map<string, HighlightedCode>();
20const highlightedCodePromiseCache = new Map<string, Promise<HighlightedCode>>();
21const HIGHLIGHTED_CODE_CACHE_LIMIT = 64;
22
23function codeHighlightCacheKey(
24 code: string,
25 language: string | undefined,
26 theme: "light" | "dark",
27): string {
28 return `${theme}:${language ?? "text"}:${code}`;
29}
30
31function cacheHighlightedCode(key: string, tokens: HighlightedCode): void {
32 highlightedCodeCache.delete(key);
33 highlightedCodeCache.set(key, tokens);
34
35 while (highlightedCodeCache.size > HIGHLIGHTED_CODE_CACHE_LIMIT) {
36 const oldestKey = highlightedCodeCache.keys().next().value;
37 if (oldestKey === undefined) {
38 break;
39 }
40 highlightedCodeCache.delete(oldestKey);
41 }
42}
43
44function loadHighlightedCode(
45 code: string,
46 language: string | undefined,
47 theme: "light" | "dark",
48 highlightCode: MarkdownCodeHighlighter,
49 session: object,
50): Promise<HighlightedCode> {
51 const key = codeHighlightCacheKey(code, language, theme);
52 const cached = highlightedCodeCache.get(key);
53 if (cached) {
54 return Promise.resolve(cached);
55 }
56
57 const pending = highlightedCodePromiseCache.get(key);
58 if (pending) {
59 return pending;
60 }
61
62 const promise = highlightCode({ code, language, theme, session })
63 .then((tokens) => {
64 cacheHighlightedCode(key, tokens);
65 highlightedCodePromiseCache.delete(key);
66 return tokens;
67 })
68 .catch((error) => {
69 highlightedCodePromiseCache.delete(key);
70 throw error;
71 });
72 highlightedCodePromiseCache.set(key, promise);
73 return promise;
74}
75
76/**
77 * Tokens for a code block, or null while nothing usable exists yet. A cached or
78 * synchronous result renders in the same pass. Otherwise the most recent
79 * completed result seeds `pendingCodeHighlight` so finished lines keep their
80 * colors while the asynchronous highlighter catches up.
81 */
82export function useHighlightedCode(
83 code: string,
84 language: string | undefined,
85 theme: "light" | "dark",
86 highlightCode: MarkdownCodeHighlighter,
87): HighlightedCode | null {
88 const [session] = useState(() => ({}));
89 const key = codeHighlightCacheKey(code, language, theme);
90 const ready = useMemo(
91 () => highlightedCodeCache.get(key) ?? highlightCode.read?.({ code, language, theme, session }),
92 [code, language, theme, key, highlightCode, session],
93 );
94 const [highlighted, setHighlighted] = useState<HighlightedCodeResult>(() => ({
95 code,
96 language,
97 theme,
98 key,
99 tokens: highlightedCodeCache.get(key) ?? null,
100 }));
101 // Synchronous reads never touch state, so each append costs one render. The
102 // ref remembers the newest of them, and is cleared whenever an asynchronous
103 // result commits so `highlighted` is the baseline again from then on.
104 const latestRead = useRef<HighlightedCodeResult | null>(null);
105
106 useEffect(() => {
107 if (ready) {
108 latestRead.current = { code, language, theme, key, tokens: ready };
109 return;
110 }
111 let active = true;
112 const commit = (tokens: HighlightedCode | null) => {
113 latestRead.current = null;
114 setHighlighted({ code, language, theme, key, tokens });
115 };
116 const cached = highlightedCodeCache.get(key);
117 if (cached) {
118 cacheHighlightedCode(key, cached);
119 commit(cached);
120 return () => {
121 active = false;
122 };
123 }
124
125 void loadHighlightedCode(code, language, theme, highlightCode, session)
126 .then((tokens) => {
127 if (active) {
128 commit(tokens);
129 }
130 })
131 .catch(() => {
132 if (active) {
133 commit(null);
134 }
135 });
136 return () => {
137 active = false;
138 };
139 }, [code, highlightCode, key, language, theme, ready, session]);
140
141 if (ready) return ready;
142 // oxlint-disable-next-line react/refs -- Written only after commit; a discarded render never advances it.
143 const baseline = latestRead.current ?? highlighted;
144 if (baseline.key === key) return baseline.tokens;
145 if (baseline.tokens && baseline.language === language && baseline.theme === theme) {
146 return pendingCodeHighlight(baseline.code, code, baseline.tokens);
147 }
148 return null;
149}
apps/mobile/src/features/review/incrementalSnippet.test.ts created+85
......@@ -0,0 +1,85 @@
1import { createHighlighterCore } from "@shikijs/core";
2import { createJavaScriptRegexEngine } from "@shikijs/engine-javascript";
3import typescript from "@shikijs/langs/typescript";
4import dark from "@shikijs/themes/github-dark-default";
5import { describe, expect, it } from "vite-plus/test";
6
7import { createIncrementalSnippet } from "./incrementalSnippet";
8import { highlightCodeSnippet } from "./shikiReviewHighlighter";
9
10const highlighter = await createHighlighterCore({
11 langs: [typescript],
12 themes: [dark],
13 engine: createJavaScriptRegexEngine(),
14});
15const rendered = (lines: ReturnType<typeof highlighter.codeToTokensBase>) =>
16 lines.map((line) => line.map(({ content, color, fontStyle }) => ({ content, color, fontStyle })));
17const options = { lang: "typescript", theme: "github-dark-default" };
18
19describe("incremental snippet highlighting", () => {
20 it("matches full highlighting through partial lines, multiline syntax, edits and truncation", async () => {
21 const highlight = createIncrementalSnippet(highlighter, options.lang, options.theme);
22 const source = "/* comment\nstill comment */\nconst text = `first\nsecond ${42}`;\n";
23 for (let end = 1; end <= source.length; end++) {
24 const code = source.slice(0, end);
25 expect(rendered(highlight.read(code) ?? (await highlight(code)))).toEqual(
26 rendered(highlighter.codeToTokensBase(code, options)),
27 );
28 }
29 for (const code of ['const other = "new";\n', "const other", source, source.slice(0, 20)]) {
30 expect(rendered(highlight.read(code) ?? (await highlight(code)))).toEqual(
31 rendered(highlighter.codeToTokensBase(code, options)),
32 );
33 }
34 });
35
36 it("carries grammar state across batches and overlapping updates", async () => {
37 const highlight = createIncrementalSnippet(highlighter, options.lang, options.theme);
38 const code = "/*\n" + "still a comment\n".repeat(410) + "*/\nconst value = 42;";
39 const first = highlight(code);
40 const next = highlight("const changed = true;\n");
41 expect(rendered(await first)).toEqual(rendered(highlighter.codeToTokensBase(code, options)));
42 expect(rendered(await next)).toEqual(
43 rendered(highlighter.codeToTokensBase("const changed = true;\n", options)),
44 );
45 const last = 'const changed = true;\nconst tail = "ok";';
46 expect(rendered(await highlight(last))).toEqual(
47 rendered(highlighter.codeToTokensBase(last, options)),
48 );
49 });
50
51 it("reuses completed token rows on warm reads and falls back for large updates", async () => {
52 const session = {};
53 const input = { language: "ts", theme: "dark" as const, session };
54 const code = "const a = 1;\nconst b = 2;";
55 expect(highlightCodeSnippet.read({ ...input, code })).toBeUndefined();
56 const first = await highlightCodeSnippet({ ...input, code });
57 const nextCode = code + "\nconst c = 3;";
58 const next = highlightCodeSnippet.read({ ...input, code: nextCode })!;
59 expect(next[0]).toBe(first[0]);
60 expect(next).toEqual(
61 await highlightCodeSnippet({ code: nextCode, language: "ts", theme: "dark" }),
62 );
63 expect(
64 highlightCodeSnippet.read({ ...input, code: nextCode + "\n" + "const x = 1;\n".repeat(210) }),
65 ).toBeUndefined();
66 expect(highlightCodeSnippet.read({ ...input, code: nextCode, theme: "light" })).toBeUndefined();
67 });
68
69 it("resets session on language, theme, CRLF, long-line and empty input changes", async () => {
70 const session = {};
71 const cases = [
72 { code: "/* hello\nworld */", language: "ts", theme: "dark" as const },
73 { code: "/* hello\nworld */\nconst n = 3;", language: "ts", theme: "light" as const },
74 { code: "hello\nworld", language: "text", theme: "light" as const },
75 { code: "const n = 2;\r\nconst m = 3;", language: "ts", theme: "dark" as const },
76 { code: "a".repeat(1100), language: "ts", theme: "dark" as const },
77 { code: "", language: "ts", theme: "dark" as const },
78 { code: "const fresh = true;\n", language: "ts", theme: "dark" as const },
79 ];
80 for (const input of cases)
81 expect(await highlightCodeSnippet({ ...input, session })).toEqual(
82 await highlightCodeSnippet(input),
83 );
84 });
85});
apps/mobile/src/features/review/incrementalSnippet.ts created+68
......@@ -0,0 +1,68 @@
1import type { HighlighterCore } from "@shikijs/core";
2
3/** A code block owns this session. Only completed lines survive the next update. */
4export function createIncrementalSnippet(
5 highlighter: HighlighterCore,
6 language: string,
7 theme: string,
8) {
9 const options = { lang: language, theme };
10 let cached:
11 | {
12 prefix: string;
13 tokens: ReturnType<HighlighterCore["codeToTokensBase"]>;
14 state: ReturnType<HighlighterCore["getLastGrammarState"]>;
15 }
16 | undefined;
17 let revision = 0;
18
19 function* tokenize(code: string) {
20 const currentRevision = ++revision;
21 const previous = cached && code.startsWith(cached.prefix) ? cached : undefined;
22 const end = code.lastIndexOf("\n") + 1;
23 let state = previous?.state;
24 const tokens = [...(previous?.tokens ?? [])];
25 const completed = code.slice(previous?.prefix.length ?? 0, end);
26 if (completed) {
27 const lines = completed.slice(0, -1).split("\n");
28 for (let offset = 0; offset < lines.length; offset += 200) {
29 const batch = highlighter.codeToTokensBase(lines.slice(offset, offset + 200).join("\n"), {
30 ...options,
31 ...(state ? { grammarState: state } : {}),
32 });
33 state = highlighter.getLastGrammarState(batch);
34 tokens.push(...batch);
35 if (offset + 200 < lines.length) yield;
36 }
37 }
38 if (currentRevision === revision) {
39 cached = state ? { prefix: code.slice(0, end), tokens, state } : undefined;
40 }
41 return [
42 ...tokens,
43 ...highlighter.codeToTokensBase(code.slice(end), {
44 ...options,
45 ...(state ? { grammarState: state } : {}),
46 }),
47 ];
48 }
49 const highlight = async (code: string) => {
50 const work = tokenize(code);
51 let next = work.next();
52 while (!next.done) {
53 await new Promise<void>((resolve) => setTimeout(resolve, 0));
54 next = work.next();
55 }
56 return next.value;
57 };
58 // A small append can finish during render, avoiding a plain-text commit before
59 // the asynchronous effect supplies colors. Cold/large changes stay asynchronous.
60 highlight.read = (code: string) => {
61 if (!cached || !code.startsWith(cached.prefix)) return undefined;
62 const tail = code.slice(cached.prefix.length);
63 if (tail.length > 2_000 || tail.split("\n").length > 200) return undefined;
64 const next = tokenize(code).next();
65 return next.done ? next.value : undefined;
66 };
67 return highlight;
68}
apps/mobile/src/features/review/shikiReviewHighlighter.ts+60-6
......@@ -17,6 +17,7 @@ import {
1717 resolveReviewHighlighterEnginePreference,
1818 type ReviewHighlighterEngine,
1919} from "./reviewHighlighterEngine";
20import { createIncrementalSnippet } from "./incrementalSnippet";
2021import type { ReviewRenderableLineRow } from "./reviewModel";
2122import { applyDiffRangesToTokens, computeWordAltDiffRanges } from "./reviewWordDiffs";
2223
......@@ -453,16 +454,23 @@ async function resolveLanguageFromPath(
453454 return candidate;
454455}
455456
457type RawHighlightedLine = ReadonlyArray<{ content: string; color?: string; fontStyle?: number }>;
458const normalizedLines = new WeakMap<RawHighlightedLine, ReadonlyArray<ReviewHighlightedToken>>();
459
456460function normalizeHighlightedLines(
457 tokenLines: ReadonlyArray<ReadonlyArray<{ content: string; color?: string; fontStyle?: number }>>,
461 tokenLines: ReadonlyArray<RawHighlightedLine>,
458462): ReadonlyArray<ReadonlyArray<ReviewHighlightedToken>> {
459 return tokenLines.map((line) =>
460 line.map((token) => ({
463 return tokenLines.map((line) => {
464 const cached = normalizedLines.get(line);
465 if (cached) return cached;
466 const normalized = line.map((token) => ({
461467 content: token.content,
462468 color: token.color ?? null,
463469 fontStyle: token.fontStyle ?? null,
464 })),
465 );
470 }));
471 normalizedLines.set(line, normalized);
472 return normalized;
473 });
466474}
467475
468476function applyWordAltDiffHighlightsToSelectedLines(input: {
......@@ -594,16 +602,62 @@ async function highlightLines(
594602 return highlightedLines;
595603}
596604
605const snippetSessions = new WeakMap<
606 object,
607 { language: string; theme: string; highlight: ReturnType<typeof createIncrementalSnippet> }
608>();
609
597610export async function highlightCodeSnippet(input: {
611 readonly session?: object;
598612 readonly code: string;
599613 readonly language?: string | null;
600614 readonly theme: ReviewDiffTheme;
601615}): Promise<ReadonlyArray<ReadonlyArray<ReviewHighlightedToken>>> {
602616 const languageHint = input.language?.trim() || "text";
603617 const language = await resolveLanguageFromPath(`snippet.${languageHint}`, languageHint);
604 return highlightLines(input.code, language, SHIKI_THEME_NAME_BY_SCHEME[input.theme]);
618 const theme = SHIKI_THEME_NAME_BY_SCHEME[input.theme];
619 // Bound retained text and preserve the existing plain-text/long-line fallback.
620 if (
621 !input.session ||
622 language === "text" ||
623 input.code.length === 0 ||
624 input.code.length > 100_000 ||
625 input.code.includes("\r") ||
626 input.code.split("\n").some((line) => line.length > REVIEW_TOKENIZE_MAX_LINE_LENGTH)
627 ) {
628 if (input.session) snippetSessions.delete(input.session);
629 return highlightLines(input.code, language, theme);
630 }
631 const highlighter = await getHighlighter();
632 let session = snippetSessions.get(input.session);
633 if (!session || session.language !== language || session.theme !== theme) {
634 session = {
635 language,
636 theme,
637 highlight: createIncrementalSnippet(highlighter, language, theme),
638 };
639 snippetSessions.set(input.session, session);
640 }
641 return normalizeHighlightedLines(await session.highlight(input.code));
605642}
606643
644highlightCodeSnippet.read = (input: Parameters<typeof highlightCodeSnippet>[0]) => {
645 if (!input.session || !input.code || input.code.length > 100_000 || input.code.includes("\r"))
646 return undefined;
647 const session = snippetSessions.get(input.session);
648 const hint = input.language?.trim() || "text";
649 const language = resolveLoadedLanguageFromPath(`snippet.${hint}`, hint);
650 if (
651 !session ||
652 session.language !== language ||
653 session.theme !== SHIKI_THEME_NAME_BY_SCHEME[input.theme] ||
654 input.code.split("\n").some((line) => line.length > REVIEW_TOKENIZE_MAX_LINE_LENGTH)
655 )
656 return undefined;
657 const tokens = session.highlight.read(input.code);
658 return tokens ? normalizeHighlightedLines(tokens) : undefined;
659};
660
607661export async function highlightSourceFile(input: {
608662 readonly path: string;
609663 readonly contents: string;
apps/mobile/src/lib/nativeMarkdownText.test.ts+88-5
......@@ -536,7 +536,7 @@ describe("nativeMarkdownDocumentChunks", () => {
536536 ).toEqual([
537537 {
538538 kind: "rich",
539 key: "rich:blockquote:0:120",
539 key: "rich:blockquote:offset:0",
540540 node: blockquote,
541541 },
542542 ]);
......@@ -663,7 +663,7 @@ describe("nativeMarkdownDocumentChunks", () => {
663663 expect(chunks[0]).toMatchObject({ kind: "selectable" });
664664 expect(chunks[1]).toEqual({
665665 kind: "rich",
666 key: "rich:code_block:11:35",
666 key: "rich:code_block:offset:11",
667667 node: document.children?.[1],
668668 });
669669 expect(chunks[2]).toMatchObject({ kind: "selectable" });
......@@ -700,7 +700,7 @@ describe("nativeMarkdownDocumentChunks", () => {
700700 expect(nativeMarkdownDocumentChunks(document)).toEqual([
701701 {
702702 kind: "rich",
703 key: "rich:list:0:45",
703 key: "rich:list:offset:0",
704704 node: document.children?.[0],
705705 },
706706 ]);
......@@ -728,7 +728,7 @@ describe("nativeMarkdownDocumentChunks", () => {
728728 expect(chunks[0]).toMatchObject({ kind: "selectable" });
729729 expect(chunks[1]).toEqual({
730730 kind: "rich",
731 key: "rich:horizontal_rule:1:1",
731 key: "rich:horizontal_rule:index:1",
732732 node: document.children?.[1],
733733 });
734734 expect(chunks[2]).toMatchObject({ kind: "selectable" });
......@@ -774,7 +774,7 @@ describe("nativeMarkdownDocumentChunks", () => {
774774 expect(chunks[0]).toMatchObject({ kind: "selectable" });
775775 expect(chunks[1]).toEqual({
776776 kind: "rich",
777 key: "rich:list:1:1",
777 key: "rich:list:index:1",
778778 node: document.children?.[1],
779779 });
780780 expect(chunks[2]).toMatchObject({ kind: "selectable" });
......@@ -824,6 +824,89 @@ describe("nativeMarkdownDocumentChunks", () => {
824824 });
825825 });
826826
827 it("keys positioned and offset-free siblings in separate namespaces", () => {
828 const positioned: MarkdownNode = {
829 type: "blockquote",
830 beg: 1,
831 end: 10,
832 children: [{ type: "paragraph", children: [{ type: "text", content: "Positioned" }] }],
833 };
834 const offsetFree: MarkdownNode = {
835 type: "blockquote",
836 children: [{ type: "paragraph", children: [{ type: "text", content: "Generated" }] }],
837 };
838 const chunks = nativeMarkdownDocumentChunks({
839 type: "document",
840 children: [
841 { type: "paragraph", beg: 0, end: 0, children: [] },
842 offsetFree,
843 positioned,
844 { type: "paragraph", children: [{ type: "text", content: "Tail" }] },
845 ],
846 });
847
848 expect(chunks.map((chunk) => chunk.key)).toEqual([
849 "selectable:offset:0",
850 "rich:blockquote:index:1",
851 "rich:blockquote:offset:1",
852 "selectable:index:3",
853 ]);
854 });
855
856 it("gives every offset-free selectable group its own key", () => {
857 const chunks = nativeMarkdownDocumentChunks({
858 type: "document",
859 children: [
860 { type: "paragraph", children: [{ type: "text", content: "One" }] },
861 { type: "horizontal_rule" },
862 { type: "paragraph", children: [{ type: "text", content: "Two" }] },
863 { type: "horizontal_rule" },
864 { type: "paragraph", children: [{ type: "text", content: "Three" }] },
865 ],
866 });
867
868 expect(chunks.map((chunk) => chunk.key)).toEqual([
869 "selectable:index:0",
870 "rich:horizontal_rule:index:1",
871 "selectable:index:2",
872 "rich:horizontal_rule:index:3",
873 "selectable:index:4",
874 ]);
875 });
876
877 it("keeps positioned chunk keys stable while text streams in", () => {
878 const before: MarkdownNode = {
879 type: "document",
880 children: [
881 { type: "paragraph", beg: 0, end: 5, children: [{ type: "text", content: "Intro" }] },
882 {
883 type: "code_block",
884 language: "ts",
885 beg: 7,
886 end: 20,
887 children: [{ type: "text", content: "const a" }],
888 },
889 ],
890 };
891 const after: MarkdownNode = {
892 type: "document",
893 children: [
894 { type: "paragraph", beg: 0, end: 5, children: [{ type: "text", content: "Intro" }] },
895 {
896 type: "code_block",
897 language: "ts",
898 beg: 7,
899 end: 40,
900 children: [{ type: "text", content: "const a = 1;\nconst b" }],
901 },
902 ],
903 };
904
905 expect(nativeMarkdownDocumentChunks(after).map((chunk) => chunk.key)).toEqual(
906 nativeMarkdownDocumentChunks(before).map((chunk) => chunk.key),
907 );
908 });
909
827910 it("keeps a plain list in one selectable native text container", () => {
828911 const list: MarkdownNode = {
829912 type: "list",