From be37c26c884f8abf44afd6042d029076965c7f7e Mon Sep 17 00:00:00 2001 From: "N. Taylor Mullen" Date: Fri, 16 Jan 2026 09:33:13 -0800 Subject: [PATCH] perf(ui): optimize text buffer and highlighting for large inputs (#16782) Co-authored-by: Jacob Richman --- .../ui/components/shared/performance.test.ts | 93 +++++++++++ .../ui/components/shared/text-buffer.test.ts | 140 +++++++++++++++- .../src/ui/components/shared/text-buffer.ts | 150 +++++++++++++----- packages/cli/src/ui/constants.ts | 1 + packages/cli/src/ui/utils/highlight.test.ts | 4 +- packages/cli/src/ui/utils/highlight.ts | 25 ++- packages/cli/src/ui/utils/textUtils.ts | 33 ++-- packages/core/src/index.ts | 1 + packages/core/src/utils/LruCache.test.ts | 70 ++++++++ packages/core/src/utils/LruCache.ts | 28 ++-- 10 files changed, 469 insertions(+), 76 deletions(-) create mode 100644 packages/cli/src/ui/components/shared/performance.test.ts create mode 100644 packages/core/src/utils/LruCache.test.ts diff --git a/packages/cli/src/ui/components/shared/performance.test.ts b/packages/cli/src/ui/components/shared/performance.test.ts new file mode 100644 index 0000000000..683995745b --- /dev/null +++ b/packages/cli/src/ui/components/shared/performance.test.ts @@ -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); + }); +}); diff --git a/packages/cli/src/ui/components/shared/text-buffer.test.ts b/packages/cli/src/ui/components/shared/text-buffer.test.ts index 0cd7c7bfd2..4a264c2cd8 100644 --- a/packages/cli/src/ui/components/shared/text-buffer.test.ts +++ b/packages/cli/src/ui/components/shared/text-buffer.test.ts @@ -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'); + }); + }); }); diff --git a/packages/cli/src/ui/components/shared/text-buffer.ts b/packages/cli/src/ui/components/shared/text-buffer.ts index a80d088bf2..61e6b98e55 100644 --- a/packages/cli/src/ui/components/shared/text-buffer.ts +++ b/packages/cli/src/ui/components/shared/text-buffer.ts @@ -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( + 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( + 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 diff --git a/packages/cli/src/ui/constants.ts b/packages/cli/src/ui/constants.ts index b211c9ce8d..f6c867f56f 100644 --- a/packages/cli/src/ui/constants.ts +++ b/packages/cli/src/ui/constants.ts @@ -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; diff --git a/packages/cli/src/ui/utils/highlight.test.ts b/packages/cli/src/ui/utils/highlight.test.ts index ba2cc09eea..db98ae90a6 100644 --- a/packages/cli/src/ui/utils/highlight.test.ts +++ b/packages/cli/src/ui/utils/highlight.test.ts @@ -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' }, ]); }); diff --git a/packages/cli/src/ui/utils/highlight.ts b/packages/cli/src/ui/utils/highlight.ts index 475854a3dd..3d182f2451 100644 --- a/packages/cli/src/ui/utils/highlight.ts +++ b/packages/cli/src/ui/utils/highlight.ts @@ -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( + 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; } diff --git a/packages/cli/src/ui/utils/textUtils.ts b/packages/cli/src/ui/utils/textUtils.ts index 13c009a136..86d345047d 100644 --- a/packages/cli/src/ui/utils/textUtils.ts +++ b/packages/cli/src/ui/utils/textUtils.ts @@ -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(); +// Cache for code points const MAX_STRING_LENGTH_TO_CACHE = 1000; +const codePointsCache = new LruCache( + 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(); +const stringWidthCache = new LruCache( + 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; diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index 506e602ebf..1a387b5828 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -92,6 +92,7 @@ export * from './utils/version.js'; export * from './utils/checkpointUtils.js'; export * from './utils/apiConversionUtils.js'; export * from './utils/channel.js'; +export * from './utils/LruCache.js'; // Export services export * from './services/fileDiscoveryService.js'; diff --git a/packages/core/src/utils/LruCache.test.ts b/packages/core/src/utils/LruCache.test.ts new file mode 100644 index 0000000000..10cdbb600c --- /dev/null +++ b/packages/core/src/utils/LruCache.test.ts @@ -0,0 +1,70 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect } from 'vitest'; +import { LruCache } from './LruCache.js'; + +describe('LruCache', () => { + it('should store and retrieve values', () => { + const cache = new LruCache(10); + cache.set('a', 1); + expect(cache.get('a')).toBe(1); + expect(cache.get('b')).toBeUndefined(); + }); + + it('should evict oldest items when limit is reached', () => { + const cache = new LruCache(2); + cache.set('a', 1); + cache.set('b', 2); + cache.set('c', 3); + + expect(cache.get('a')).toBeUndefined(); + expect(cache.get('b')).toBe(2); + expect(cache.get('c')).toBe(3); + }); + + it('should update LRU order on get', () => { + const cache = new LruCache(2); + cache.set('a', 1); + cache.set('b', 2); + cache.get('a'); // Mark 'a' as recently used + cache.set('c', 3); // Should evict 'b' + + expect(cache.get('b')).toBeUndefined(); + expect(cache.get('a')).toBe(1); + expect(cache.get('c')).toBe(3); + }); + + it('should correctly handle falsy values (truthiness bug fix)', () => { + const cache = new LruCache(2); + cache.set('zero', 0); + cache.set('empty', ''); + + // Both should be in cache + expect(cache.get('zero')).toBe(0); + expect(cache.get('empty')).toBe(''); + + // Access 'zero' to move it to end + cache.get('zero'); + + // Add new item, should evict 'empty' (because 'zero' was just accessed) + cache.set('new', 'item'); + + expect(cache.get('empty')).toBeUndefined(); + expect(cache.get('zero')).toBe(0); + expect(cache.get('new')).toBe('item'); + }); + + it('should correctly handle undefined as a value', () => { + // NOTE: mnemonist LRUMap returns undefined if not found. + // If we store undefined, we can't distinguish between "not found" and "found undefined". + // But we should at least check its behavior. + const cache = new LruCache(10); + cache.set('key', undefined); + expect(cache.has('key')).toBe(true); + expect(cache.get('key')).toBeUndefined(); + }); +}); diff --git a/packages/core/src/utils/LruCache.ts b/packages/core/src/utils/LruCache.ts index 076828c459..7c917a8d30 100644 --- a/packages/core/src/utils/LruCache.ts +++ b/packages/core/src/utils/LruCache.ts @@ -4,34 +4,24 @@ * SPDX-License-Identifier: Apache-2.0 */ +import { LRUMap } from 'mnemonist'; + export class LruCache { - private cache: Map; - private maxSize: number; + private cache: LRUMap; constructor(maxSize: number) { - this.cache = new Map(); - this.maxSize = maxSize; + this.cache = new LRUMap(maxSize); } get(key: K): V | undefined { - const value = this.cache.get(key); - if (value) { - // Move to end to mark as recently used - this.cache.delete(key); - this.cache.set(key, value); - } - return value; + return this.cache.get(key); + } + + has(key: K): boolean { + return this.cache.has(key); } set(key: K, value: V): void { - if (this.cache.has(key)) { - this.cache.delete(key); - } else if (this.cache.size >= this.maxSize) { - const firstKey = this.cache.keys().next().value; - if (firstKey !== undefined) { - this.cache.delete(firstKey); - } - } this.cache.set(key, value); }