perf(editor): reuse unchanged review block line counts (#34934)

* perf(editor): reuse unchanged review block line counts

* test(editor): drop serializer call-count assertions for review block reuse

---------

Co-authored-by: OrcaWin <293788423+OrcaWin@users.noreply.github.com>

src/renderer/src/components/editor/rich-markdown-comment-blocks-incremental.test.ts

@@ -0,0 +1,60 @@
+// @vitest-environment happy-dom
+import { Editor } from '@tiptap/core'
+import { afterEach, describe, expect, it, vi } from 'vitest'
+import {
+  buildRichMarkdownCommentBlocks,
+  getRichMarkdownCommentBlocks
+} from './rich-markdown-comment-blocks'
+import { createRichMarkdownExtensions } from './rich-markdown-extensions'
+import { createRichMarkdownEditorCodec } from './rich-markdown-source-transport'
+
+const editors: Editor[] = []
+
+function createEditor(source: string) {
+  const editor = new Editor({
+    element: document.createElement('div'),
+    extensions: createRichMarkdownExtensions({
+      codec: createRichMarkdownEditorCodec()
+    }),
+    content: source,
+    contentType: 'markdown'
+  })
+  editors.push(editor)
+  return editor
+}
+
+afterEach(() => {
+  for (const editor of editors.splice(0)) {
+    editor.destroy()
+  }
+  vi.restoreAllMocks()
+})
+
+describe('review source block serialization after edits', () => {
+  it('preserves exact block positions through multiline edits, moves, insertions, deletion and undo', () => {
+    const editor = createEditor(
+      'First\n\n```text\none\ntwo\n```\n\n> quote\n>\n> second\n\n- item one\n- item two\n\nLast'
+    )
+    const check = () =>
+      expect(getRichMarkdownCommentBlocks(editor)).toEqual(buildRichMarkdownCommentBlocks(editor))
+    check()
+    editor.view.dispatch(
+      editor.state.tr.insertText('extra\n', editor.state.doc.child(0).nodeSize + 1)
+    )
+    check()
+    const node = editor.state.doc.child(0)
+    const move = editor.state.tr.delete(0, node.nodeSize)
+    move.insert(move.doc.content.size, node)
+    editor.view.dispatch(move)
+    check()
+    const paragraph = editor.schema.nodes.paragraph.create(null, editor.schema.text('inserted'))
+    editor.view.dispatch(editor.state.tr.insert(0, paragraph))
+    check()
+    editor.view.dispatch(editor.state.tr.delete(0, paragraph.nodeSize))
+    check()
+    expect(editor.commands.undo()).toBe(true)
+    check()
+    expect(editor.commands.redo()).toBe(true)
+    check()
+  })
+})

src/renderer/src/components/editor/rich-markdown-comment-blocks.test.ts

@@ -130,15 +130,12 @@ describe('Markdown review source block reuse', () => {
     const from = selectLastParagraph(editor)
     vi.spyOn(Range.prototype, 'getBoundingClientRect').mockReturnValue(new DOMRect(10, 20, 80, 20))
     const serialize = vi.spyOn(editor.markdown!, 'serialize')
-    const json = vi.spyOn(editor, 'getJSON')
     const note = comment()
     const expected = [{ from, to: from + 4 }]
 
     expect(getRichMarkdownAnnotationHighlightRanges(editor, [note], 20)).toEqual(expected)
     expect(serialize).toHaveBeenCalledTimes(2 * editor.state.doc.childCount - 1)
-    expect(json).toHaveBeenCalledTimes(1)
     serialize.mockClear()
-    json.mockClear()
 
     for (let index = 0; index < 20; index++) {
       expect(getRichMarkdownAnnotationHighlightRanges(editor, [note], 20)).toEqual(expected)
@@ -153,7 +150,6 @@ describe('Markdown review source block reuse', () => {
       expect(getRichMarkdownReviewRailBlocks(editor)).toBe(getRichMarkdownCommentBlocks(editor))
     }
     expect(serialize).not.toHaveBeenCalled()
-    expect(json).not.toHaveBeenCalled()
   })
 
   it('updates selected text and viewport geometry without rebuilding source lines', () => {
@@ -198,17 +194,13 @@ describe('Markdown review source block reuse', () => {
     editor.view.dispatch(editor.state.tr.insertText('new\n', codeFrom))
     selectLastParagraph(editor)
     expect(getRichMarkdownAnnotationTarget(editor, root)?.lineNumber).toBe(10)
-    expect(serialize).toHaveBeenCalledTimes(2 * editor.state.doc.childCount - 1)
     expect(
       getRichMarkdownCommentAtPos(editor, [comment(30)], 20, editor.state.selection.from)
     ).toEqual(comment(30))
-    expect(serialize).toHaveBeenCalledTimes(2 * editor.state.doc.childCount - 1)
-    serialize.mockClear()
 
     expect(editor.commands.undo()).toBe(true)
     selectLastParagraph(editor)
     expect(getRichMarkdownAnnotationTarget(editor, root)?.lineNumber).toBe(9)
-    expect(serialize).toHaveBeenCalledTimes(2 * editor.state.doc.childCount - 1)
     serialize.mockClear()
     expect(getRichMarkdownAnnotationTarget(editor, root)?.lineNumber).toBe(9)
     expect(serialize).not.toHaveBeenCalled()

src/renderer/src/components/editor/rich-markdown-comment-blocks.ts

@@ -1,4 +1,5 @@
 import type { Editor, JSONContent } from '@tiptap/core'
+import type { Node as ProseMirrorNode } from '@tiptap/pm/model'
 import { countRichMarkdownReviewMarkdownLines } from './rich-markdown-review-line-count'
 
 export type RichMarkdownCommentBlock = {
@@ -9,30 +10,64 @@ export type RichMarkdownCommentBlock = {
   to: number
 }
 
+function richMarkdownNodeJson(node: ProseMirrorNode): JSONContent {
+  return node.toJSON()
+}
+
 function serializeRichMarkdownJson(editor: Editor, content: JSONContent[]): string {
   return (editor.markdown?.serialize({ type: 'doc', content }) ?? '').trimEnd()
 }
 
+type RichMarkdownCommentBlockCounts = {
+  nodes: WeakMap<ProseMirrorNode, number>
+  separators: WeakMap<ProseMirrorNode, WeakMap<ProseMirrorNode, number>>
+}
+
+function createCommentBlockCounts(): RichMarkdownCommentBlockCounts {
+  return { nodes: new WeakMap(), separators: new WeakMap() }
+}
+
 export function buildRichMarkdownCommentBlocks(editor: Editor): RichMarkdownCommentBlock[] {
-  const jsonContent = editor.getJSON().content ?? []
+  return buildCommentBlocksWithCounts(editor, createCommentBlockCounts())
+}
+
+function buildCommentBlocksWithCounts(
+  editor: Editor,
+  counts: RichMarkdownCommentBlockCounts
+): RichMarkdownCommentBlock[] {
   const blocks: RichMarkdownCommentBlock[] = []
   let nextLine = 1
-  let previousNodeJson: JSONContent | null = null
+  let previousNode: ProseMirrorNode | null = null
+  let previousNodeJson: JSONContent | undefined
   let previousNodeLineCount = 0
 
   editor.state.doc.forEach((node, nodeOffset, index) => {
-    const nodeJson = jsonContent[index]
-    if (!nodeJson) {
-      return
-    }
-    const nodeMarkdown = serializeRichMarkdownJson(editor, [nodeJson])
-    const nodeLineCount = countRichMarkdownReviewMarkdownLines(nodeMarkdown)
-    if (previousNodeJson) {
-      const pairMarkdown = serializeRichMarkdownJson(editor, [previousNodeJson, nodeJson])
-      const separatorLineCount = Math.max(
-        0,
-        countRichMarkdownReviewMarkdownLines(pairMarkdown) - previousNodeLineCount - nodeLineCount
+    let nodeJson: JSONContent | undefined
+    let nodeLineCount = counts.nodes.get(node)
+    if (nodeLineCount === undefined) {
+      nodeJson = richMarkdownNodeJson(node)
+      nodeLineCount = countRichMarkdownReviewMarkdownLines(
+        serializeRichMarkdownJson(editor, [nodeJson])
       )
+      counts.nodes.set(node, nodeLineCount)
+    }
+    if (previousNode) {
+      let separatorCounts = counts.separators.get(previousNode)
+      if (!separatorCounts) {
+        separatorCounts = new WeakMap()
+        counts.separators.set(previousNode, separatorCounts)
+      }
+      let separatorLineCount = separatorCounts.get(node)
+      if (separatorLineCount === undefined) {
+        previousNodeJson ??= richMarkdownNodeJson(previousNode)
+        nodeJson ??= richMarkdownNodeJson(node)
+        const pairMarkdown = serializeRichMarkdownJson(editor, [previousNodeJson, nodeJson])
+        separatorLineCount = Math.max(
+          0,
+          countRichMarkdownReviewMarkdownLines(pairMarkdown) - previousNodeLineCount - nodeLineCount
+        )
+        separatorCounts.set(node, separatorLineCount)
+      }
       nextLine += separatorLineCount
     }
     const startLine = nextLine
@@ -46,6 +81,7 @@ export function buildRichMarkdownCommentBlocks(editor: Editor): RichMarkdownComm
       to: from + Math.max(0, node.nodeSize - 1)
     })
     nextLine = endLine + 1
+    previousNode = node
     previousNodeJson = nodeJson
     previousNodeLineCount = nodeLineCount
   })
@@ -62,6 +98,7 @@ type RichMarkdownCommentBlocksSnapshot = {
   markdown: Editor['markdown']
   serialize: NonNullable<Editor['markdown']>['serialize'] | undefined
   blocks: readonly RichMarkdownCommentBlock[]
+  counts: RichMarkdownCommentBlockCounts
 }
 
 // Keep only the current document per editor; scrolling changes geometry, not source lines.
@@ -76,7 +113,11 @@ export function getRichMarkdownCommentBlocks(editor: Editor): readonly RichMarkd
     return cached.blocks
   }
 
-  const blocks = buildRichMarkdownCommentBlocks(editor)
-  blocksByEditor.set(editor, { doc, markdown, serialize, blocks })
+  const counts =
+    cached && cached.markdown === markdown && cached.serialize === serialize
+      ? cached.counts
+      : createCommentBlockCounts()
+  const blocks = buildCommentBlocksWithCounts(editor, counts)
+  blocksByEditor.set(editor, { doc, markdown, serialize, blocks, counts })
   return blocks
 }

src/renderer/src/components/editor/rich-markdown-review-annotations.test.ts

@@ -62,7 +62,9 @@ describe('rich markdown annotation block reuse', () => {
     }))
     const doc = {
       forEach(callback: (node: unknown, offset: number, index: number) => void): void {
-        content.forEach((node, index) => callback(node, index * 10, index))
+        content.forEach((node, index) =>
+          callback({ ...node, toJSON: () => node }, index * 10, index)
+        )
       },
       // These fixtures never match selected text; traversal only needs to be supported.
       nodesBetween(): void {},