perf(ui): optimize text buffer and highlighting for large inputs (#16782)

Co-authored-by: Jacob Richman <jacob314@gmail.com>
2026-04-21 18:44:30 -07:00 · 2026-01-16 09:33:13 -08:00
parent fcd860e1b0
commit be37c26c88
10 changed files with 469 additions and 76 deletions
@@ -0,0 +1,93 @@
+/**
+ * @license
+ * Copyright 2026 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+import { describe, it, expect, afterEach, vi } from 'vitest';
+import { act } from 'react';
+import { renderHook } from '../../../test-utils/render.js';
+import { useTextBuffer } from './text-buffer.js';
+import { parseInputForHighlighting } from '../../utils/highlight.js';
+
+describe('text-buffer performance', () => {
+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+
+  it('should handle pasting large amounts of text efficiently', () => {
+    const viewport = { width: 80, height: 24 };
+    const { result } = renderHook(() =>
+      useTextBuffer({
+        viewport,
+        isValidPath: () => false,
+      }),
+    );
+
+    const lines = 5000;
+    const largeText = Array.from(
+      { length: lines },
+      (_, i) =>
+        `Line ${i}: some sample text with many @path/to/image${i}.png and maybe some more @path/to/another/image.png references to trigger regex. This line is much longer than the previous one to test wrapping.`,
+    ).join('\n');
+
+    const start = Date.now();
+    act(() => {
+      result.current.insert(largeText, { paste: true });
+    });
+    const end = Date.now();
+
+    const duration = end - start;
+    expect(duration).toBeLessThan(5000);
+  });
+
+  it('should handle character-by-character insertion in a large buffer efficiently', () => {
+    const lines = 5000;
+    const initialText = Array.from(
+      { length: lines },
+      (_, i) => `Line ${i}: some sample text with @path/to/image.png`,
+    ).join('\n');
+    const viewport = { width: 80, height: 24 };
+
+    const { result } = renderHook(() =>
+      useTextBuffer({
+        initialText,
+        viewport,
+        isValidPath: () => false,
+      }),
+    );
+
+    const start = Date.now();
+    const charsToInsert = 100;
+    for (let i = 0; i < charsToInsert; i++) {
+      act(() => {
+        result.current.insert('a');
+      });
+    }
+    const end = Date.now();
+
+    const duration = end - start;
+    expect(duration).toBeLessThan(5000);
+  });
+
+  it('should highlight many lines efficiently', () => {
+    const lines = 5000;
+    const sampleLines = Array.from(
+      { length: lines },
+      (_, i) =>
+        `Line ${i}: some sample text with @path/to/image${i}.png /command and more @file.txt`,
+    );
+
+    const start = Date.now();
+    for (let i = 0; i < 100; i++) {
+      // Simulate 100 renders
+      for (const line of sampleLines.slice(0, 20)) {
+        // 20 visible lines
+        parseInputForHighlighting(line, 1, []);
+      }
+    }
+    const end = Date.now();
+
+    const duration = end - start;
+    expect(duration).toBeLessThan(500);
+  });
+});
@@ -4,10 +4,13 @@
 * SPDX-License-Identifier: Apache-2.0
 */

-import { describe, it, expect, beforeEach } from 'vitest';
+import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest';
 import stripAnsi from 'strip-ansi';
 import { act } from 'react';
-import { renderHook } from '../../../test-utils/render.js';
+import {
+  renderHook,
+  renderHookWithProviders,
+} from '../../../test-utils/render.js';
 import type {
  Viewport,
  TextBuffer,
@@ -55,6 +58,10 @@ const initialState: TextBufferState = {
 };

 describe('textBufferReducer', () => {
+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+
  it('should return the initial state if state is undefined', () => {
    const action = { type: 'unknown_action' } as unknown as TextBufferAction;
    const state = textBufferReducer(initialState, action);
@@ -381,6 +388,10 @@ describe('useTextBuffer', () => {
    viewport = { width: 10, height: 3 }; // Default viewport for tests
  });

+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+
  describe('Initialization', () => {
    it('should initialize with empty text and cursor at (0,0) by default', () => {
      const { result } = renderHook(() =>
@@ -2371,6 +2382,10 @@ describe('Unicode helper functions', () => {
 });

 describe('Transformation Utilities', () => {
+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+
  describe('getTransformedImagePath', () => {
    it('should transform a simple image path', () => {
      expect(getTransformedImagePath('@test.png')).toBe('[Image test.png]');
@@ -2547,4 +2562,125 @@ describe('Transformation Utilities', () => {
      expect(result.transformedToLogMap).toEqual([0]); // Just the trailing position
    });
  });
+
+  describe('Layout Caching and Invalidation', () => {
+    it.each([
+      {
+        desc: 'via setText',
+        actFn: (result: { current: TextBuffer }) =>
+          result.current.setText('changed line'),
+        expected: 'changed line',
+      },
+      {
+        desc: 'via replaceRange',
+        actFn: (result: { current: TextBuffer }) =>
+          result.current.replaceRange(0, 0, 0, 13, 'changed line'),
+        expected: 'changed line',
+      },
+    ])(
+      'should invalidate cache when line content changes $desc',
+      ({ actFn, expected }) => {
+        const viewport = { width: 80, height: 24 };
+        const { result } = renderHookWithProviders(() =>
+          useTextBuffer({
+            initialText: 'original line',
+            viewport,
+            isValidPath: () => true,
+          }),
+        );
+
+        const originalLayout = result.current.visualLayout;
+
+        act(() => {
+          actFn(result);
+        });
+
+        expect(result.current.visualLayout).not.toBe(originalLayout);
+        expect(result.current.allVisualLines[0]).toBe(expected);
+      },
+    );
+
+    it('should invalidate cache when viewport width changes', () => {
+      const viewport = { width: 80, height: 24 };
+      const { result, rerender } = renderHookWithProviders(
+        ({ vp }) =>
+          useTextBuffer({
+            initialText:
+              'a very long line that will wrap when the viewport is small',
+            viewport: vp,
+            isValidPath: () => true,
+          }),
+        { initialProps: { vp: viewport } },
+      );
+
+      const originalLayout = result.current.visualLayout;
+
+      // Shrink viewport to force wrapping change
+      rerender({ vp: { width: 10, height: 24 } });
+
+      expect(result.current.visualLayout).not.toBe(originalLayout);
+      expect(result.current.allVisualLines.length).toBeGreaterThan(1);
+    });
+
+    it('should correctly handle cursor expansion/collapse in cached layout', () => {
+      const viewport = { width: 80, height: 24 };
+      const text = 'Check @image.png here';
+      const { result } = renderHookWithProviders(() =>
+        useTextBuffer({
+          initialText: text,
+          viewport,
+          isValidPath: () => true,
+        }),
+      );
+
+      // Cursor at start (collapsed)
+      act(() => {
+        result.current.moveToOffset(0);
+      });
+      expect(result.current.allVisualLines[0]).toContain('[Image image.png]');
+
+      // Move cursor onto the @path (expanded)
+      act(() => {
+        result.current.moveToOffset(7); // onto @
+      });
+      expect(result.current.allVisualLines[0]).toContain('@image.png');
+      expect(result.current.allVisualLines[0]).not.toContain(
+        '[Image image.png]',
+      );
+
+      // Move cursor away (collapsed again)
+      act(() => {
+        result.current.moveToOffset(0);
+      });
+      expect(result.current.allVisualLines[0]).toContain('[Image image.png]');
+    });
+
+    it('should reuse cache for unchanged lines during editing', () => {
+      const viewport = { width: 80, height: 24 };
+      const initialText = 'line 1\nline 2\nline 3';
+      const { result } = renderHookWithProviders(() =>
+        useTextBuffer({
+          initialText,
+          viewport,
+          isValidPath: () => true,
+        }),
+      );
+
+      const layout1 = result.current.visualLayout;
+
+      // Edit line 1
+      act(() => {
+        result.current.moveToOffset(0);
+        result.current.insert('X');
+      });
+
+      const layout2 = result.current.visualLayout;
+      expect(layout2).not.toBe(layout1);
+
+      // Verify that visual lines for line 2 and 3 (indices 1 and 2 in visualLines)
+      // are identical in content if not in object reference (the arrays are rebuilt, but contents are cached)
+      expect(result.current.allVisualLines[1]).toBe('line 2');
+      expect(result.current.allVisualLines[2]).toBe('line 3');
+    });
+  });
 });
@@ -18,6 +18,7 @@ import {
  type EditorType,
  getEditorCommand,
  isGuiEditor,
+  LruCache,
 } from '@google/gemini-cli-core';
 import {
  toCodePoints,
@@ -31,6 +32,7 @@ import type { Key } from '../../contexts/KeypressContext.js';
 import { keyMatchers, Command } from '../../keyMatchers.js';
 import type { VimAction } from './vim-buffer-actions.js';
 import { handleVimAction } from './vim-buffer-actions.js';
+import { LRU_BUFFER_PERF_CACHE_LIMIT } from '../../constants.js';

 export type Direction =
  | 'left'
@@ -726,9 +728,18 @@ export function getTransformedImagePath(filePath: string): string {
  return `[Image ${truncatedBase}${extension}]`;
 }

+const transformationsCache = new LruCache<string, Transformation[]>(
+  LRU_BUFFER_PERF_CACHE_LIMIT,
+);
+
 export function calculateTransformationsForLine(
  line: string,
 ): Transformation[] {
+  const cached = transformationsCache.get(line);
+  if (cached) {
+    return cached;
+  }
+
  const transformations: Transformation[] = [];
  let match: RegExpExecArray | null;

@@ -748,6 +759,8 @@ export function calculateTransformationsForLine(
    });
  }

+  transformationsCache.set(line, transformations);
+
  return transformations;
 }

@@ -845,6 +858,7 @@ export function calculateTransformedLine(

  return { transformedLine, transformedToLogMap };
 }
+
 export interface VisualLayout {
  visualLines: string[];
  // For each logical line, an array of [visualLineIndex, startColInLogical]
@@ -858,6 +872,33 @@ export interface VisualLayout {
  visualToTransformedMap: number[];
 }

+// Caches for layout calculation
+interface LineLayoutResult {
+  visualLines: string[];
+  logicalToVisualMap: Array<[number, number]>;
+  visualToLogicalMap: Array<[number, number]>;
+  transformedToLogMap: number[];
+  visualToTransformedMap: number[];
+}
+
+const lineLayoutCache = new LruCache<string, LineLayoutResult>(
+  LRU_BUFFER_PERF_CACHE_LIMIT,
+);
+
+function getLineLayoutCacheKey(
+  line: string,
+  viewportWidth: number,
+  isCursorOnLine: boolean,
+  cursorCol: number,
+): string {
+  // Most lines (99.9% in a large buffer) are not cursor lines.
+  // We use a simpler key for them to reduce string allocation overhead.
+  if (!isCursorOnLine) {
+    return `${viewportWidth}:N:${line}`;
+  }
+  return `${viewportWidth}:C:${cursorCol}:${line}`;
+}
+
 // Calculates the visual wrapping of lines and the mapping between logical and visual coordinates.
 // This is an expensive operation and should be memoized.
 function calculateLayout(
@@ -873,6 +914,34 @@ function calculateLayout(

  logicalLines.forEach((logLine, logIndex) => {
    logicalToVisualMap[logIndex] = [];
+
+    const isCursorOnLine = logIndex === logicalCursor[0];
+    const cacheKey = getLineLayoutCacheKey(
+      logLine,
+      viewportWidth,
+      isCursorOnLine,
+      logicalCursor[1],
+    );
+    const cached = lineLayoutCache.get(cacheKey);
+
+    if (cached) {
+      const visualLineOffset = visualLines.length;
+      visualLines.push(...cached.visualLines);
+      cached.logicalToVisualMap.forEach(([relVisualIdx, logCol]) => {
+        logicalToVisualMap[logIndex].push([
+          visualLineOffset + relVisualIdx,
+          logCol,
+        ]);
+      });
+      cached.visualToLogicalMap.forEach(([, logCol]) => {
+        visualToLogicalMap.push([logIndex, logCol]);
+      });
+      transformedToLogicalMaps[logIndex] = cached.transformedToLogMap;
+      visualToTransformedMap.push(...cached.visualToTransformedMap);
+      return;
+    }
+
+    // Not in cache, calculate
    const transformations = calculateTransformationsForLine(logLine);
    const { transformedLine, transformedToLogMap } = calculateTransformedLine(
      logLine,
@@ -880,13 +949,18 @@ function calculateLayout(
      logicalCursor,
      transformations,
    );
-    transformedToLogicalMaps[logIndex] = transformedToLogMap;
+
+    const lineVisualLines: string[] = [];
+    const lineLogicalToVisualMap: Array<[number, number]> = [];
+    const lineVisualToLogicalMap: Array<[number, number]> = [];
+    const lineVisualToTransformedMap: number[] = [];
+
    if (transformedLine.length === 0) {
      // Handle empty logical line
-      logicalToVisualMap[logIndex].push([visualLines.length, 0]);
-      visualToLogicalMap.push([logIndex, 0]);
-      visualToTransformedMap.push(0);
-      visualLines.push('');
+      lineLogicalToVisualMap.push([0, 0]);
+      lineVisualToLogicalMap.push([logIndex, 0]);
+      lineVisualToTransformedMap.push(0);
+      lineVisualLines.push('');
    } else {
      // Non-empty logical line
      let currentPosInLogLine = 0; // Tracks position within the current logical line (code point index)
@@ -929,15 +1003,6 @@ function calculateLayout(
                // Single character is wider than viewport, take it anyway
                currentChunk = char;
                numCodePointsInChunk = 1;
-              } else if (
-                numCodePointsInChunk === 0 &&
-                charVisualWidth <= viewportWidth
-              ) {
-                // This case should ideally be caught by the next iteration if the char fits.
-                // If it doesn't fit (because currentChunkVisualWidth was already > 0 from a previous char that filled the line),
-                // then numCodePointsInChunk would not be 0.
-                // This branch means the current char *itself* doesn't fit an empty line, which is handled by the above.
-                // If we are here, it means the loop should break and the current chunk (which is empty) is finalized.
              }
            }
            break; // Break from inner loop to finalize this chunk
@@ -955,55 +1020,57 @@ function calculateLayout(
          }
        }

-        // If the inner loop completed without breaking (i.e., remaining text fits)
-        // or if the loop broke but numCodePointsInChunk is still 0 (e.g. first char too wide for empty line)
        if (
          numCodePointsInChunk === 0 &&
          currentPosInLogLine < codePointsInLogLine.length
        ) {
-          // This can happen if the very first character considered for a new visual line is wider than the viewport.
-          // In this case, we take that single character.
          const firstChar = codePointsInLogLine[currentPosInLogLine];
          currentChunk = firstChar;
-          numCodePointsInChunk = 1; // Ensure we advance
-        }
-
-        // If after everything, numCodePointsInChunk is still 0 but we haven't processed the whole logical line,
-        // it implies an issue, like viewportWidth being 0 or less. Avoid infinite loop.
-        if (
-          numCodePointsInChunk === 0 &&
-          currentPosInLogLine < codePointsInLogLine.length
-        ) {
-          // Force advance by one character to prevent infinite loop if something went wrong
-          currentChunk = codePointsInLogLine[currentPosInLogLine];
          numCodePointsInChunk = 1;
        }

        const logicalStartCol = transformedToLogMap[currentPosInLogLine] ?? 0;
-        logicalToVisualMap[logIndex].push([
-          visualLines.length,
-          logicalStartCol,
-        ]);
-        visualToLogicalMap.push([logIndex, logicalStartCol]);
-        visualToTransformedMap.push(currentPosInLogLine);
-        visualLines.push(currentChunk);
+        lineLogicalToVisualMap.push([lineVisualLines.length, logicalStartCol]);
+        lineVisualToLogicalMap.push([logIndex, logicalStartCol]);
+        lineVisualToTransformedMap.push(currentPosInLogLine);
+        lineVisualLines.push(currentChunk);

        const logicalStartOfThisChunk = currentPosInLogLine;
        currentPosInLogLine += numCodePointsInChunk;

-        // If the chunk processed did not consume the entire logical line,
-        // and the character immediately following the chunk is a space,
-        // advance past this space as it acted as a delimiter for word wrapping.
        if (
          logicalStartOfThisChunk + numCodePointsInChunk <
            codePointsInLogLine.length &&
-          currentPosInLogLine < codePointsInLogLine.length && // Redundant if previous is true, but safe
+          currentPosInLogLine < codePointsInLogLine.length &&
          codePointsInLogLine[currentPosInLogLine] === ' '
        ) {
          currentPosInLogLine++;
        }
      }
    }
+
+    // Cache the result for this line
+    lineLayoutCache.set(cacheKey, {
+      visualLines: lineVisualLines,
+      logicalToVisualMap: lineLogicalToVisualMap,
+      visualToLogicalMap: lineVisualToLogicalMap,
+      transformedToLogMap,
+      visualToTransformedMap: lineVisualToTransformedMap,
+    });
+
+    const visualLineOffset = visualLines.length;
+    visualLines.push(...lineVisualLines);
+    lineLogicalToVisualMap.forEach(([relVisualIdx, logCol]) => {
+      logicalToVisualMap[logIndex].push([
+        visualLineOffset + relVisualIdx,
+        logCol,
+      ]);
+    });
+    lineVisualToLogicalMap.forEach(([, logCol]) => {
+      visualToLogicalMap.push([logIndex, logCol]);
+    });
+    transformedToLogicalMaps[logIndex] = transformedToLogMap;
+    visualToTransformedMap.push(...lineVisualToTransformedMap);
  });

  // If the entire logical text was empty, ensure there's one empty visual line.
@@ -2378,6 +2445,7 @@ export function useTextBuffer({
      transformedToLogicalMaps,
      visualToTransformedMap,
      transformationsByLine,
+      visualLayout,
      setText,
      insert,
      newline,
@@ -2447,6 +2515,7 @@ export function useTextBuffer({
      transformedToLogicalMaps,
      visualToTransformedMap,
      transformationsByLine,
+      visualLayout,
      setText,
      insert,
      newline,
@@ -2540,6 +2609,7 @@ export interface TextBuffer {
  visualToTransformedMap: number[];
  /** Cached transformations per logical line */
  transformationsByLine: Transformation[][];
+  visualLayout: VisualLayout;

  // Actions

@@ -32,3 +32,4 @@ export const MAX_MCP_RESOURCES_TO_SHOW = 10;
 export const WARNING_PROMPT_DURATION_MS = 1000;
 export const QUEUE_ERROR_DISPLAY_DURATION_MS = 3000;
 export const SHELL_ACTION_REQUIRED_TITLE_DELAY_MS = 30000;
+export const LRU_BUFFER_PERF_CACHE_LIMIT = 20000;
@@ -212,13 +212,13 @@ describe('parseInputForHighlighting with Transformations', () => {
  });

  it('should handle empty transformations array', () => {
-    const line = 'Check out @test.png';
+    const line = 'Check out @test_no_transform.png';
    const result = parseInputForHighlighting(line, 0, [], 0);

    // Should fall back to default highlighting
    expect(result).toEqual([
      { text: 'Check out ', type: 'default' },
-      { text: '@test.png', type: 'file' },
+      { text: '@test_no_transform.png', type: 'file' },
    ]);
  });

@@ -4,8 +4,10 @@
 * SPDX-License-Identifier: Apache-2.0
 */

+import { LruCache } from '@google/gemini-cli-core';
 import type { Transformation } from '../components/shared/text-buffer.js';
 import { cpLen, cpSlice } from './textUtils.js';
+import { LRU_BUFFER_PERF_CACHE_LIMIT } from '../constants.js';

 export type HighlightToken = {
  text: string;
@@ -18,12 +20,30 @@ export type HighlightToken = {
 // semicolon, common punctuation, and brackets.
 const HIGHLIGHT_REGEX = /(^\/[a-zA-Z0-9_-]+|@(?:\\ |[^,\s;!?()[\]{}])+)/g;

+const highlightCache = new LruCache<string, readonly HighlightToken[]>(
+  LRU_BUFFER_PERF_CACHE_LIMIT,
+);
+
 export function parseInputForHighlighting(
  text: string,
  index: number,
  transformations: Transformation[] = [],
  cursorCol?: number,
 ): readonly HighlightToken[] {
+  let isCursorInsideTransform = false;
+  if (cursorCol !== undefined) {
+    for (const transform of transformations) {
+      if (cursorCol >= transform.logStart && cursorCol <= transform.logEnd) {
+        isCursorInsideTransform = true;
+        break;
+      }
+    }
+  }
+
+  const cacheKey = `${index === 0 ? 'F' : 'N'}:${isCursorInsideTransform ? cursorCol : 'NC'}:${text}`;
+  const cached = highlightCache.get(cacheKey);
+  if (cached !== undefined) return cached;
+
  HIGHLIGHT_REGEX.lastIndex = 0;

  if (!text) {
@@ -79,7 +99,7 @@ export function parseInputForHighlighting(
    tokens.push(...parseUntransformedInput(textBeforeTransformation));

    const isCursorInside =
-      typeof cursorCol === 'number' &&
+      cursorCol !== undefined &&
      cursorCol >= transformation.logStart &&
      cursorCol <= transformation.logEnd;
    const transformationText = isCursorInside
@@ -92,6 +112,9 @@ export function parseInputForHighlighting(

  const textAfterFinalTransformation = cpSlice(text, column);
  tokens.push(...parseUntransformedInput(textAfterFinalTransformation));
+
+  highlightCache.set(cacheKey, tokens);
+
  return tokens;
 }

@@ -8,6 +8,8 @@ import stripAnsi from 'strip-ansi';
 import ansiRegex from 'ansi-regex';
 import { stripVTControlCharacters } from 'node:util';
 import stringWidth from 'string-width';
+import { LruCache } from '@google/gemini-cli-core';
+import { LRU_BUFFER_PERF_CACHE_LIMIT } from '../constants.js';

 /**
 * Calculates the maximum width of a multi-line ASCII art string.
@@ -28,9 +30,11 @@ export const getAsciiArtWidth = (asciiArt: string): number => {
 *  code units so that surrogate‑pair emoji count as one "column".)
 * ---------------------------------------------------------------------- */

-// Cache for code points to reduce GC pressure
-const codePointsCache = new Map<string, string[]>();
+// Cache for code points
 const MAX_STRING_LENGTH_TO_CACHE = 1000;
+const codePointsCache = new LruCache<string, string[]>(
+  LRU_BUFFER_PERF_CACHE_LIMIT,
+);

 export function toCodePoints(str: string): string[] {
  // ASCII fast path - check if all chars are ASCII (0-127)
@@ -48,14 +52,14 @@ export function toCodePoints(str: string): string[] {
  // Cache short strings
  if (str.length <= MAX_STRING_LENGTH_TO_CACHE) {
    const cached = codePointsCache.get(str);
-    if (cached) {
+    if (cached !== undefined) {
      return cached;
    }
  }

  const result = Array.from(str);

-  // Cache result (unlimited like Ink)
+  // Cache result
  if (str.length <= MAX_STRING_LENGTH_TO_CACHE) {
    codePointsCache.set(str, result);
  }
@@ -119,21 +123,26 @@ export function stripUnsafeCharacters(str: string): string {
    .join('');
 }

-// String width caching for performance optimization
-const stringWidthCache = new Map<string, number>();
+const stringWidthCache = new LruCache<string, number>(
+  LRU_BUFFER_PERF_CACHE_LIMIT,
+);

 /**
 * Cached version of stringWidth function for better performance
- * Follows Ink's approach with unlimited cache (no eviction)
 */
 export const getCachedStringWidth = (str: string): number => {
-  // ASCII printable chars have width 1
-  if (/^[\x20-\x7E]*$/.test(str)) {
-    return str.length;
+  // ASCII printable chars (32-126) have width 1.
+  // This is a very frequent path, so we use a fast numeric check.
+  if (str.length === 1) {
+    const code = str.charCodeAt(0);
+    if (code >= 0x20 && code <= 0x7e) {
+      return 1;
+    }
  }

-  if (stringWidthCache.has(str)) {
-    return stringWidthCache.get(str)!;
+  const cached = stringWidthCache.get(str);
+  if (cached !== undefined) {
+    return cached;
  }

  let width: number;