perf(mobile): reuse completed code lines while streaming

apps/mobile/modules/t3-markdown-text/src/NativeMarkdownBlock.tsx

@@ -1,10 +1,11 @@
-import { createContext, useContext, useEffect, useState } from "react";
+import { createContext, memo, useContext, useEffect, useMemo, useState } from "react";
 import { Image, Platform, ScrollView, Text, useColorScheme, View } from "react-native";
 import type { MarkdownNode } from "react-native-nitro-markdown/headless";
 
 import { CopyTextButton } from "./CopyTextButton";
 import { MarkdownTextPrimitive } from "./MarkdownTextPrimitive";
 import { nativeMarkdownDocumentRuns, nativeMarkdownListItemBlocks } from "./nativeMarkdownText";
+import { pendingCodeHighlight } from "./pendingCodeHighlight";
 import { NativeMarkdownSelectableText } from "./NativeMarkdownSelectableText";
 import type {
   MarkdownCodeHighlighter,
@@ -29,7 +30,7 @@ const MONO_FONT_FAMILY = Platform.select({
 });
 
 function nodeKey(node: MarkdownNode, index: number): string {
-  return `${node.type}:${node.beg ?? index}:${node.end ?? index}`;
+  return `${node.type}:${node.beg ?? index}`;
 }
 
 /** Code inside markdown scales with the base text size (12pt at the default 15pt body). */
@@ -93,6 +94,7 @@ function loadHighlightedCode(
   language: string | undefined,
   theme: "light" | "dark",
   highlightCode: MarkdownCodeHighlighter,
+  session: object,
 ): Promise<HighlightedCode> {
   const key = codeHighlightCacheKey(code, language, theme);
   const cached = highlightedCodeCache.get(key);
@@ -105,7 +107,7 @@ function loadHighlightedCode(
     return pending;
   }
 
-  const promise = highlightCode({ code, language, theme })
+  const promise = highlightCode({ code, language, theme, session })
     .then((tokens) => {
       cacheHighlightedCode(key, tokens);
       highlightedCodePromiseCache.delete(key);
@@ -125,116 +127,130 @@ function useHighlightedCode(
   theme: "light" | "dark",
   highlightCode: MarkdownCodeHighlighter,
 ): HighlightedCode | null {
+  const [session] = useState(() => ({}));
   const key = codeHighlightCacheKey(code, language, theme);
+  const ready = useMemo(
+    () => highlightedCodeCache.get(key) ?? highlightCode.read?.({ code, language, theme, session }),
+    [code, language, theme, key, highlightCode, session],
+  );
   const [highlighted, setHighlighted] = useState<{
     readonly key: string;
+    readonly code: string;
+    readonly language: string | undefined;
+    readonly theme: "light" | "dark";
     readonly tokens: HighlightedCode | null;
   }>(() => ({
+    code,
+    language,
+    theme,
     key,
     tokens: highlightedCodeCache.get(key) ?? null,
   }));
 
   useEffect(() => {
+    if (ready) return;
     let active = true;
     const cached = highlightedCodeCache.get(key);
     if (cached) {
       cacheHighlightedCode(key, cached);
-      setHighlighted({ key, tokens: cached });
+      setHighlighted({ code, language, theme, key, tokens: cached });
       return () => {
         active = false;
       };
     }
 
-    void loadHighlightedCode(code, language, theme, highlightCode)
+    void loadHighlightedCode(code, language, theme, highlightCode, session)
       .then((tokens) => {
         if (active) {
-          setHighlighted({ key, tokens });
+          setHighlighted({ code, language, theme, key, tokens });
         }
       })
       .catch(() => {
         if (active) {
-          setHighlighted({ key, tokens: null });
+          setHighlighted({ code, language, theme, key, tokens: null });
         }
       });
     return () => {
       active = false;
     };
-  }, [code, highlightCode, key, language, theme]);
+  }, [code, highlightCode, key, language, theme, ready, session]);
 
-  return highlighted.key === key ? highlighted.tokens : null;
+  if (ready) return ready;
+  if (highlighted.key === key) return highlighted.tokens;
+  if (highlighted.tokens && highlighted.language === language && highlighted.theme === theme) {
+    return pendingCodeHighlight(highlighted.code, code, highlighted.tokens);
+  }
+  return null;
 }
 
-function HighlightedCodeText(props: {
-  readonly content: string;
-  readonly highlighted: HighlightedCode | null;
-  readonly textStyle: NativeMarkdownTextStyle;
+const HighlightedCodeLine = memo(function HighlightedCodeLine(props: {
+  readonly tokens: ReadonlyArray<MarkdownHighlightedToken>;
+  readonly color: string;
+  readonly newline: boolean;
 }) {
-  if (!props.highlighted) {
-    return (
+  let offset = 0;
+  const children = [];
+  for (const token of props.tokens) {
+    if (!token.content) continue;
+    children.push(
       <MarkdownTextPrimitive
-        uiTextView
-        selectable
+        key={offset}
         style={{
-          color: props.textStyle.codeColor,
+          color: token.color ?? props.color,
           fontFamily: MONO_FONT_FAMILY,
-          fontSize: codeBlockFontSize(props.textStyle),
-          lineHeight: codeBlockLineHeight(props.textStyle),
+          fontStyle: token.fontStyle !== null && (token.fontStyle & 1) === 1 ? "italic" : "normal",
+          fontWeight: token.fontStyle !== null && (token.fontStyle & 2) === 2 ? "700" : "400",
         }}
       >
-        {props.content}
-      </MarkdownTextPrimitive>
+        {token.content}
+      </MarkdownTextPrimitive>,
     );
+    offset += token.content.length;
   }
-  const highlighted = props.highlighted;
-  let sourceOffset = 0;
-  const keyOccurrences = new Map<string, number>();
-  const keyedLines = highlighted.map((line) => {
-    const lineStart = sourceOffset;
-    const tokens = line.map((token) => {
-      const start = sourceOffset;
-      sourceOffset += token.content.length;
-      const signature = `${start}:${token.content}:${token.color ?? ""}:${token.fontStyle ?? ""}`;
-      const occurrence = keyOccurrences.get(signature) ?? 0;
-      keyOccurrences.set(signature, occurrence + 1);
-      return { key: `${signature}:${occurrence}`, token };
-    });
-    sourceOffset += 1;
-    return {
-      key: `line:${lineStart}:${line.map((token) => token.content).join("")}`,
-      tokens,
-    };
-  });
+  return (
+    <MarkdownTextPrimitive>
+      {children}
+      {props.newline ? "\n" : ""}
+    </MarkdownTextPrimitive>
+  );
+});
 
+function HighlightedCodeText(props: {
+  readonly content: string;
+  readonly highlighted: HighlightedCode | null;
+  readonly textStyle: NativeMarkdownTextStyle;
+}) {
+  // The text root provides inherited styles through context. A new style object
+  // would rerender every token even when its completed line is unchanged.
+  const fontSize = codeBlockFontSize(props.textStyle);
+  const lineHeight = codeBlockLineHeight(props.textStyle);
+  const style = useMemo(
+    () => ({
+      color: props.textStyle.codeColor,
+      fontFamily: MONO_FONT_FAMILY,
+      fontSize,
+      lineHeight,
+    }),
+    [props.textStyle.codeColor, fontSize, lineHeight],
+  );
+  let offset = 0;
+  const lines = [];
+  if (props.highlighted) {
+    for (const tokens of props.highlighted) {
+      lines.push(
+        <HighlightedCodeLine
+          key={offset}
+          tokens={tokens}
+          color={props.textStyle.codeColor}
+          newline={lines.length + 1 < props.highlighted.length}
+        />,
+      );
+      offset += tokens.reduce((length, token) => length + token.content.length, 0) + 1;
+    }
+  }
   return (
-    <MarkdownTextPrimitive
-      uiTextView
-      selectable
-      style={{
-        color: props.textStyle.codeColor,
-        fontFamily: MONO_FONT_FAMILY,
-        fontSize: codeBlockFontSize(props.textStyle),
-        lineHeight: codeBlockLineHeight(props.textStyle),
-      }}
-    >
-      {keyedLines.map((line, lineIndex) => (
-        <MarkdownTextPrimitive key={line.key}>
-          {line.tokens.map(({ key, token }) => (
-            <MarkdownTextPrimitive
-              key={key}
-              style={{
-                color: token.color ?? props.textStyle.codeColor,
-                fontFamily: MONO_FONT_FAMILY,
-                fontStyle:
-                  token.fontStyle !== null && (token.fontStyle & 1) === 1 ? "italic" : "normal",
-                fontWeight: token.fontStyle !== null && (token.fontStyle & 2) === 2 ? "700" : "400",
-              }}
-            >
-              {token.content}
-            </MarkdownTextPrimitive>
-          ))}
-          {lineIndex + 1 < keyedLines.length ? "\n" : ""}
-        </MarkdownTextPrimitive>
-      ))}
+    <MarkdownTextPrimitive uiTextView selectable style={style}>
+      {props.highlighted ? lines : props.content}
     </MarkdownTextPrimitive>
   );
 }

apps/mobile/modules/t3-markdown-text/src/SelectableMarkdownText.types.ts

@@ -25,11 +25,22 @@ export interface MarkdownHighlightedToken {
   readonly fontStyle: number | null;
 }
 
-export type MarkdownCodeHighlighter = (input: {
+export interface MarkdownCodeHighlightInput {
+  /** Identity of the mounted code block, for incremental highlighting. */
+  readonly session?: object;
   readonly code: string;
   readonly language?: string | null;
   readonly theme: "light" | "dark";
-}) => Promise<ReadonlyArray<ReadonlyArray<MarkdownHighlightedToken>>>;
+}
+export interface MarkdownCodeHighlighter {
+  (
+    input: MarkdownCodeHighlightInput,
+  ): Promise<ReadonlyArray<ReadonlyArray<MarkdownHighlightedToken>>>;
+  /** Optional synchronous result for a small append to an already warm block. */
+  read?: (
+    input: MarkdownCodeHighlightInput,
+  ) => ReadonlyArray<ReadonlyArray<MarkdownHighlightedToken>> | undefined;
+}
 
 export interface SelectableMarkdownSkill {
   readonly name: string;

apps/mobile/modules/t3-markdown-text/src/nativeMarkdownText.ts

@@ -705,10 +705,9 @@ export function nativeMarkdownDocumentChunks(
       return;
     }
     const first = selectableNodes[0];
-    const last = selectableNodes.at(-1);
     chunks.push({
       kind: "selectable",
-      key: `selectable:${first?.beg ?? "start"}:${last?.end ?? "end"}`,
+      key: `selectable:${first?.beg ?? "start"}`,
       node: {
         type: "document",
         children: selectableNodes,
@@ -726,7 +725,7 @@ export function nativeMarkdownDocumentChunks(
     flushSelectable();
     chunks.push({
       kind: "rich",
-      key: `rich:${child.type}:${child.beg ?? index}:${child.end ?? index}`,
+      key: `rich:${child.type}:${child.beg ?? index}`,
       node: child,
     });
   }

apps/mobile/modules/t3-markdown-text/src/pendingCodeHighlight.test.ts

@@ -0,0 +1,28 @@
+import { expect, it } from "vite-plus/test";
+import { pendingCodeHighlight } from "./pendingCodeHighlight";
+
+it("keeps completed colors and exact current text without retaining an edited tail", () => {
+  const colored = [
+    [{ content: "const n = 1;", color: "red", fontStyle: 0 }],
+    [{ content: "partial", color: "blue", fontStyle: 0 }],
+  ];
+  const result = pendingCodeHighlight(
+    "const n = 1;\npartial",
+    "const n = 1;\nchanged\nnext",
+    colored,
+  )!;
+  expect(result[0]).toBe(colored[0]);
+  expect(result.map((line) => line.map((token) => token.content).join("")).join("\n")).toBe(
+    "const n = 1;\nchanged\nnext",
+  );
+  expect(
+    result
+      .slice(1)
+      .flat()
+      .every((token) => token.color === null),
+  ).toBe(true);
+  expect(
+    pendingCodeHighlight("const n = 1;\npartial", "const n = 2;\npartial", colored),
+  ).toBeNull();
+  expect(pendingCodeHighlight("partial", "other", colored)).toBeNull();
+});

apps/mobile/modules/t3-markdown-text/src/pendingCodeHighlight.ts

@@ -0,0 +1,19 @@
+import type { MarkdownHighlightedToken } from "./SelectableMarkdownText.types";
+
+/** Keep finished lines colored while the current line awaits highlighting. */
+export function pendingCodeHighlight(
+  previousCode: string,
+  code: string,
+  tokens: ReadonlyArray<ReadonlyArray<MarkdownHighlightedToken>>,
+): ReadonlyArray<ReadonlyArray<MarkdownHighlightedToken>> | null {
+  const end = previousCode.lastIndexOf("\n") + 1;
+  if (!end || !code.startsWith(previousCode.slice(0, end))) return null;
+  const completedLines = previousCode.slice(0, end).split("\n").length - 1;
+  return [
+    ...tokens.slice(0, completedLines),
+    ...code
+      .slice(end)
+      .split("\n")
+      .map((content) => [{ content, color: null, fontStyle: null }]),
+  ];
+}