From 82dafbebe043f14b548adfd68d5492f872de9f56 Mon Sep 17 00:00:00 2001
From: Dominic Gannaway <trueadm@users.noreply.github.com>
Date: Sat, 17 Dec 2022 14:07:00 +0000
Subject: [PATCH] Use the actual node count rather than the dirty node count
 (#3569)

---
 .../lexical-react/src/LexicalTreeView.tsx     | 25 +++++++++++--------
 packages/lexical/src/LexicalSelection.ts      |  4 +--
 packages/lexical/src/LexicalUpdates.ts        |  4 +--
 3 files changed, 18 insertions(+), 15 deletions(-)

diff --git a/packages/lexical-react/src/LexicalTreeView.tsx b/packages/lexical-react/src/LexicalTreeView.tsx
index 17da5c892dc..2a74122a3f5 100644
--- a/packages/lexical-react/src/LexicalTreeView.tsx
+++ b/packages/lexical-react/src/LexicalTreeView.tsx
@@ -98,20 +98,23 @@ export function TreeView({
   );
 
   useEffect(() => {
-    setContent(
-      generateContent(
-        editor.getEditorState(),
-        editor._config,
-        editor._compositionKey,
-        editor._editable,
-      ),
-    );
-  }, [editor]);
+    const editorState = editor.getEditorState();
+    if (!showLimited && editorState._nodeMap.size > 1000) {
+      setContent(
+        generateContent(
+          editorState,
+          editor._config,
+          editor._compositionKey,
+          editor._editable,
+        ),
+      );
+    }
+  }, [editor, showLimited]);
 
   useEffect(() => {
     return mergeRegister(
-      editor.registerUpdateListener(({editorState, dirtyLeaves}) => {
-        if (!showLimited && dirtyLeaves.size > 1000) {
+      editor.registerUpdateListener(({editorState}) => {
+        if (!showLimited && editorState._nodeMap.size > 1000) {
           lastEditorStateRef.current = editorState;
           setIsLimited(true);
           if (!showLimited) {
diff --git a/packages/lexical/src/LexicalSelection.ts b/packages/lexical/src/LexicalSelection.ts
index 0d907d284c0..25ebac38091 100644
--- a/packages/lexical/src/LexicalSelection.ts
+++ b/packages/lexical/src/LexicalSelection.ts
@@ -2743,7 +2743,7 @@ export function updateDOMSelection(
   domSelection: Selection,
   tags: Set<string>,
   rootElement: HTMLElement,
-  dirtyLeavesCount: number,
+  nodeCount: number,
 ): void {
   const anchorDOMNode = domSelection.anchorNode;
   const focusDOMNode = domSelection.focusNode;
@@ -2852,7 +2852,7 @@ export function updateDOMSelection(
       // The downside is that is makes the computation within Lexical more
       // complex, as now, we've sync update the DOM, but selection no longer
       // matches.
-      if (IS_CHROME && dirtyLeavesCount > 1000) {
+      if (IS_CHROME && nodeCount > 1000) {
         window.requestAnimationFrame(() =>
           domSelection.setBaseAndExtent(
             nextAnchorNode as Node,
diff --git a/packages/lexical/src/LexicalUpdates.ts b/packages/lexical/src/LexicalUpdates.ts
index 76ed8de9ac7..b1d54c2ebcb 100644
--- a/packages/lexical/src/LexicalUpdates.ts
+++ b/packages/lexical/src/LexicalUpdates.ts
@@ -521,7 +521,7 @@ export function commitPendingUpdates(editor: LexicalEditor): void {
   const normalizedNodes = editor._normalizedNodes;
   const tags = editor._updateTags;
   const deferred = editor._deferred;
-  const dirtyLeavesCount = dirtyLeaves.size;
+  const nodeCount = pendingEditorState._nodeMap.size;
 
   if (needsUpdate) {
     editor._dirtyType = NO_DIRTY_NODES;
@@ -569,7 +569,7 @@ export function commitPendingUpdates(editor: LexicalEditor): void {
           domSelection,
           tags,
           rootElement as HTMLElement,
-          dirtyLeavesCount,
+          nodeCount,
         );
       }
       updateDOMBlockCursorElement(