[diffs/edit] tokenizer cache optimization

packages/diffs/scripts/benchmarkEditorTokenizer.ts

@@ -469,7 +469,7 @@ function createBenchmarkCases(
     {
       name: 'bracket-cache-invalidation',
       description:
-        'Edit one character at line 0 after caching bracket ranges for every line.',
+        'Edit one character at line 0, then query an unchanged cached line at the document tail.',
       run() {
         resetMessages();
         const counters = { grammarCalls: 0, setThemeCalls: 0 };
@@ -501,9 +501,12 @@ function createBenchmarkCases(
           ...fullRange,
           totalLines: 1,
         });
+        const editGrammarCalls = counters.grammarCalls;
+        tokenizer.getStringCommentRegexpRangesInLine(config.lines - 1);
         const elapsedMs = performance.now() - started;
         const operations = {
-          grammarCalls: counters.grammarCalls,
+          editGrammarCalls,
+          downstreamGrammarCalls: counters.grammarCalls - editGrammarCalls,
           dirtyLines: dirtyLines.size,
         };
         tokenizer.cleanUp();

packages/diffs/src/editor/tokenizer.ts

@@ -395,14 +395,6 @@ export class EditorTokenizer {
       );
     }
 
-    if (this.#matchBrackets) {
-      // Clear ignored token ranges for lines invalidated by the edit.
-      this.#bracketIgnoredRanges.length = Math.min(
-        this.#bracketIgnoredRanges.length,
-        change.startLine
-      );
-    }
-
     const { lineCount } = this.#textDocument;
     const { startingLine = 0, totalLines = Infinity } = renderRange ?? {};
     const renderRangeEndLine =
@@ -422,6 +414,15 @@ export class EditorTokenizer {
       change.lineDelta === 0 &&
       (change.changedLineChanges?.every(([, , lineDelta]) => lineDelta === 0) ??
         true);
+    if (this.#matchBrackets && !canReuseCachedStates) {
+      // Structural edits shift cache indexes, so only the untouched prefix is
+      // safe. Same-line edits overwrite every range they re-tokenize and can
+      // retain the untouched suffix once grammar state reconverges.
+      this.#bracketIgnoredRanges.length = Math.min(
+        this.#bracketIgnoredRanges.length,
+        change.startLine
+      );
+    }
     const canReuseShiftedStates =
       hostRealignsRows && change.lineDelta !== 0 && dirtyStart >= startingLine;
     const canCacheTokenizedStates =
@@ -587,6 +588,12 @@ export class EditorTokenizer {
     }
 
     if (backgroundStartLine !== undefined) {
+      if (this.#matchBrackets && canReuseCachedStates) {
+        this.#bracketIgnoredRanges.length = Math.min(
+          this.#bracketIgnoredRanges.length,
+          backgroundStartLine
+        );
+      }
       this.#scheduleBackgroundTokenize(
         backgroundStartLine,
         changedLineRanges,
@@ -599,6 +606,12 @@ export class EditorTokenizer {
           : dirtyStart < viewStart && !canReuseCachedStates
             ? dirtyStart
             : line;
+      if (this.#matchBrackets && canReuseCachedStates) {
+        this.#bracketIgnoredRanges.length = Math.min(
+          this.#bracketIgnoredRanges.length,
+          backgroundLine
+        );
+      }
       this.#scheduleBackgroundTokenize(
         backgroundLine,
         changedLineRanges,

packages/diffs/test/editorTokenizer.test.ts

@@ -348,6 +348,71 @@ describe('EditorTokenizer', () => {
     );
   });
 
+  test('keeps downstream bracket ranges after tokenizer state reconverges', () => {
+    const stringTokenMetadata = 2 << 8;
+    let tokenizeLineCount = 0;
+    const grammar = {
+      tokenizeLine2(lineText: string, ruleStack: StateStack) {
+        tokenizeLineCount++;
+        return {
+          tokens: new Uint32Array([0, stringTokenMetadata]),
+          ruleStack,
+          stoppedEarly: false,
+          lineText,
+        };
+      },
+    } as unknown as IGrammar;
+    const textDocument = new TextDocument(
+      'test.ts',
+      ['first[', 'second[', 'third['].join('\n'),
+      'typescript'
+    );
+    const tokenizer = new EditorTokenizer({
+      highlighter: createTestHighlighter({ getLanguage: () => grammar }),
+      textDocument,
+      codeOptions: { theme: 'test-theme', themeType: 'dark' },
+      setStyle: noopSetStyle,
+      onDeferTokenize: () => {},
+    });
+    tokenizer.tokenize(
+      {
+        startLine: 0,
+        startCharacter: 0,
+        endCharacter: 0,
+        endLine: 2,
+        endedAtDocumentEnd: false,
+        previousLineCount: 3,
+        lineCount: 3,
+        lineDelta: 0,
+        changes: [],
+        changedLineRanges: [[0, 2]],
+      },
+      { startingLine: 0, totalLines: 3, bufferBefore: 0, bufferAfter: 0 }
+    );
+
+    const change = textDocument.applyEdits([
+      {
+        range: {
+          start: { line: 0, character: 0 },
+          end: { line: 0, character: 1 },
+        },
+        newText: 'F',
+      },
+    ])!;
+    tokenizeLineCount = 0;
+    tokenizer.tokenize(change, {
+      startingLine: 0,
+      totalLines: 1,
+      bufferBefore: 0,
+      bufferAfter: 0,
+    });
+
+    expect(tokenizeLineCount).toBe(1);
+    expect(tokenizer.getStringCommentRegexpRangesInLine(2)).toEqual([[0, 6]]);
+    expect(tokenizeLineCount).toBe(1);
+    tokenizer.cleanUp();
+  });
+
   test('limits foreground tokenization to the render range after prepending lines', () => {
     const originalAddEventListener = globalThis.addEventListener;
     const originalPostMessage = globalThis.postMessage;