[diffs/edit] avoid full retokenization on newline inserts

packages/diffs/scripts/benchmarkEditorTokenizer.ts

@@ -237,7 +237,8 @@ function createTokenizer(
 
 function primeTokenizer(
   tokenizer: EditorTokenizer,
-  textDocument: TextDocument<unknown>
+  textDocument: TextDocument<unknown>,
+  renderRange?: RenderRange
 ): void {
   const lineCount = textDocument.lineCount;
   tokenizer.tokenize(
@@ -252,7 +253,7 @@ function primeTokenizer(
       lineDelta: 0,
       changedLineRanges: [[0, lineCount - 1]],
     },
-    {
+    renderRange ?? {
       startingLine: 0,
       totalLines: lineCount,
       bufferBefore: 0,
@@ -319,8 +320,70 @@ function createBenchmarkCases(
     bufferBefore: 0,
     bufferAfter: 0,
   };
+  const visibleRange: RenderRange = {
+    startingLine: 0,
+    totalLines: Math.min(100, config.lines),
+    bufferBefore: 0,
+    bufferAfter: 0,
+  };
+  const createVisibleNewlineCase = (
+    hostRealignsRows: boolean
+  ): BenchmarkCase => ({
+    name: hostRealignsRows
+      ? 'visible-newline-shifted'
+      : 'visible-newline-dense',
+    description: hostRealignsRows
+      ? 'Insert a newline and settle against shifted grammar state for a host that realigns rows.'
+      : 'Insert a newline and tokenize the full visible range for a host that cannot realign rows.',
+    run() {
+      resetMessages();
+      const counters = { grammarCalls: 0, setThemeCalls: 0 };
+      const textDocument = new TextDocument(
+        'visible-newline.ts',
+        sourceText,
+        'typescript'
+      );
+      const tokenizer = createTokenizer(textDocument, counters, {
+        matchBrackets: false,
+      });
+      primeTokenizer(tokenizer, textDocument, visibleRange);
+      tokenizer.stopBackgroundTokenize();
+      const change = textDocument.applyEdits([
+        {
+          range: {
+            start: { line: 0, character: 0 },
+            end: { line: 0, character: 0 },
+          },
+          newText: 'inserted\n',
+        },
+      ]);
+      if (change === undefined) {
+        throw new Error('Expected the newline edit to change the document');
+      }
+      counters.grammarCalls = 0;
+      postedMessages.length = 0;
+      collectGarbage();
+      const started = performance.now();
+      const dirtyLines = tokenizer.tokenize(
+        change,
+        visibleRange,
+        hostRealignsRows
+      );
+      const elapsedMs = performance.now() - started;
+      const operations = {
+        visibleLines: visibleRange.totalLines,
+        grammarCalls: counters.grammarCalls,
+        dirtyLines: dirtyLines.size,
+        backgroundMessages: postedMessages.length,
+      };
+      tokenizer.cleanUp();
+      return { elapsedMs, operations };
+    },
+  });
 
   return [
+    createVisibleNewlineCase(false),
+    createVisibleNewlineCase(true),
     {
       name: 'unbounded-structural-edit',
       description:

packages/diffs/src/editor/editor.ts

@@ -3111,7 +3111,7 @@ export class Editor<LAnnotation> implements DiffsEditor<LAnnotation> {
     tokenizer.stopBackgroundTokenize();
 
     const t = performance.now();
-    const dirtyLines = tokenizer.tokenize(change, renderRange);
+    const dirtyLines = tokenizer.tokenize(change, renderRange, !this.#isDiff);
     const t2 = performance.now();
 
     if (dirtyLines.size > 0) {

packages/diffs/src/editor/tokenizer.ts

@@ -381,7 +381,8 @@ export class EditorTokenizer {
   // the state stack map for the given render range.
   tokenize(
     change: TextDocumentChange,
-    renderRange?: RenderRange
+    renderRange?: RenderRange,
+    hostRealignsRows = false
   ): Map<number, Array<HighlightedToken>> {
     this.#ensureGrammar();
     this.#ensureActiveTheme();
@@ -421,6 +422,8 @@ export class EditorTokenizer {
       change.lineDelta === 0 &&
       (change.changedLineChanges?.every(([, , lineDelta]) => lineDelta === 0) ??
         true);
+    const canReuseShiftedStates =
+      hostRealignsRows && change.lineDelta !== 0 && dirtyStart >= startingLine;
     const canCacheTokenizedStates =
       canReuseCachedStates ||
       renderRange === undefined ||
@@ -492,7 +495,9 @@ export class EditorTokenizer {
     for (; line < renderRangeEndLine; ) {
       const previousNextState = canReuseCachedStates
         ? this.#stateStack[line + 1]
-        : undefined;
+        : canReuseShiftedStates
+          ? this.#getPreviousEndState(line + 1)
+          : undefined;
       if (canCacheTokenizedStates) {
         this.#stateStack[line] = state;
       }
@@ -514,7 +519,7 @@ export class EditorTokenizer {
       }
       settled =
         line >= currentChangedRangeEnd &&
-        canReuseCachedStates &&
+        (canReuseCachedStates || canReuseShiftedStates) &&
         previousNextState !== undefined &&
         state.equals(previousNextState);
       if (settled) {
@@ -528,12 +533,26 @@ export class EditorTokenizer {
           backgroundChangedRangeIndex = changedRangeIndex;
           break;
         }
-        if (this.#stateStack[nextRange[0]] === undefined) {
+        let nextState: StateStack | undefined = this.#stateStack[nextRange[0]];
+        if (canReuseShiftedStates) {
+          for (
+            let stateLine = line + 2;
+            stateLine <= nextRange[0];
+            stateLine++
+          ) {
+            nextState = this.#getPreviousEndState(stateLine);
+            if (nextState === undefined) {
+              break;
+            }
+            this.#stateStack[stateLine] = nextState;
+          }
+        }
+        if (nextState === undefined) {
           currentChangedRangeEnd = nextRange[1];
           line++;
         } else {
           line = nextRange[0];
-          state = this.#stateStack[line] ?? state;
+          state = nextState;
           currentChangedRangeEnd = nextRange[1];
         }
         settled = false;
@@ -550,6 +569,19 @@ export class EditorTokenizer {
       }
     }
 
+    if (settled && canReuseShiftedStates && backgroundStartLine === undefined) {
+      for (let stateLine = line + 2; stateLine <= lineCount; stateLine++) {
+        const previousState = this.#getPreviousEndState(stateLine);
+        if (previousState === undefined) {
+          break;
+        }
+        this.#stateStack[stateLine] = previousState;
+      }
+      this.#comparisonStateStack = [];
+      this.#comparisonStateStackStart = 0;
+      this.#comparisonLineChanges = [];
+    }
+
     if (offscreenDirtyLines !== undefined && offscreenDirtyLines.size > 0) {
       this.#onDeferTokenize(offscreenDirtyLines, this.#themeType);
     }

packages/diffs/test/editorTokenizer.test.ts

@@ -434,6 +434,115 @@ describe('EditorTokenizer', () => {
     }
   });
 
+  test('settles only net line-count changes for rerendering hosts', () => {
+    let tokenizeLineCount = 0;
+    const grammar = {
+      tokenizeLine2(lineText: string, ruleStack: StateStack) {
+        tokenizeLineCount++;
+        return {
+          tokens: new Uint32Array([0, 0]),
+          ruleStack,
+          stoppedEarly: false,
+          lineText,
+        };
+      },
+    } as unknown as IGrammar;
+    const textDocument = new TextDocument(
+      'test.ts',
+      Array.from({ length: 200 }, (_, i) => `line ${i}`).join('\n'),
+      'typescript'
+    );
+    const tokenizer = new EditorTokenizer({
+      highlighter: createTestHighlighter({
+        getLanguage: () => grammar,
+      }),
+      textDocument,
+      codeOptions: { theme: 'test-theme', themeType: 'dark' },
+      setStyle: noopSetStyle,
+      onDeferTokenize: () => {},
+    });
+    const renderRange = {
+      startingLine: 0,
+      totalLines: 100,
+      bufferBefore: 0,
+      bufferAfter: 0,
+    };
+
+    tokenizer.tokenize(
+      {
+        startLine: 0,
+        startCharacter: 0,
+        endCharacter: 0,
+        endLine: 199,
+        endedAtDocumentEnd: false,
+        previousLineCount: textDocument.lineCount,
+        lineCount: textDocument.lineCount,
+        lineDelta: 0,
+        changes: [],
+        changedLineRanges: [[0, 199]],
+      },
+      renderRange
+    );
+    tokenizer.stopBackgroundTokenize();
+    tokenizeLineCount = 0;
+
+    const insertLine = textDocument.applyEdits([
+      {
+        range: {
+          start: { line: 50, character: 0 },
+          end: { line: 50, character: 0 },
+        },
+        newText: '\n',
+      },
+    ])!;
+    const insertedLines = tokenizer.tokenize(insertLine, renderRange, true);
+
+    expect([...insertedLines.keys()]).toEqual([50, 51]);
+    expect(tokenizeLineCount).toBe(1);
+    tokenizeLineCount = 0;
+
+    const editLowerLine = textDocument.applyEdits([
+      {
+        range: {
+          start: { line: 80, character: 0 },
+          end: { line: 80, character: 0 },
+        },
+        newText: 'changed ',
+      },
+    ])!;
+    const editedLines = tokenizer.tokenize(editLowerLine, renderRange, true);
+
+    expect([...editedLines.keys()]).toEqual([80]);
+    expect(tokenizeLineCount).toBe(1);
+
+    const netZeroChange = textDocument.applyEdits([
+      {
+        range: {
+          start: { line: 10, character: 0 },
+          end: { line: 10, character: 0 },
+        },
+        newText: 'inserted\n',
+      },
+      {
+        range: {
+          start: { line: 20, character: textDocument.getLineText(20).length },
+          end: { line: 21, character: 0 },
+        },
+        newText: '',
+      },
+    ])!;
+    expect(netZeroChange.lineDelta).toBe(0);
+    expect(
+      netZeroChange.changedLineChanges?.map((change) => change[2])
+    ).toEqual([1, -1]);
+
+    const netZeroLines = tokenizer.tokenize(netZeroChange, renderRange, true);
+    expect([...netZeroLines.keys()]).toEqual(
+      Array.from({ length: 90 }, (_, index) => index + 10)
+    );
+    tokenizer.cleanUp();
+  });
+
   test('flushes offscreen line 0 when select-all delete shrinks the document', () => {
     const grammar = {
       tokenizeLine2(lineText: string, ruleStack: StateStack) {