From 1aa2cb4a032061643e237933b0c908a9d2dc2b84 Mon Sep 17 00:00:00 2001
From: Bob Ippolito <bob@redivi.com>
Date: Fri, 3 Jan 2025 08:49:03 -0800
Subject: [PATCH] Fix registerNodeTransform regression introduced in #6894

---
 .../__tests__/unit/LexicalCodeNode.test.ts    | 29 +++++++++++++++++++
 packages/lexical/src/LexicalUtils.ts          |  9 ++++--
 2 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/packages/lexical-code/src/__tests__/unit/LexicalCodeNode.test.ts b/packages/lexical-code/src/__tests__/unit/LexicalCodeNode.test.ts
index c4f8a9bf6e7..fede2b48818 100644
--- a/packages/lexical-code/src/__tests__/unit/LexicalCodeNode.test.ts
+++ b/packages/lexical-code/src/__tests__/unit/LexicalCodeNode.test.ts
@@ -7,6 +7,7 @@
  */
 
 import {
+  $createCodeHighlightNode,
   $createCodeNode,
   $isCodeHighlightNode,
   registerCodeHighlighting,
@@ -856,5 +857,33 @@ describe('LexicalCodeNode tests', () => {
         }
       }
     });
+    describe('initial editor state before transforms', () => {
+      test('can be registered after initial editor state (regression #7014)', async () => {
+        const {editor} = testEnv;
+        await editor.update(
+          () => {
+            const root = $getRoot();
+            const codeBlock = $createCodeNode('javascript');
+            codeBlock.append(
+              $createCodeHighlightNode('const lexical = "awesome"'),
+            );
+            root.append(codeBlock);
+          },
+          {tag: 'history-merge'},
+        );
+        // before transform
+        expect(testEnv.innerHTML).toBe(
+          '<code spellcheck="false" data-language="javascript" data-highlight-language="javascript" dir="ltr"><span data-lexical-text="true">const lexical = "awesome"</span></code>',
+        );
+        registerRichText(editor);
+        registerTabIndentation(editor);
+        registerCodeHighlighting(editor);
+        await Promise.resolve(undefined);
+        // after transforms
+        expect(testEnv.innerHTML).toBe(
+          '<code spellcheck="false" data-language="javascript" data-highlight-language="javascript" dir="ltr" data-gutter="1"><span data-lexical-text="true">const</span><span data-lexical-text="true"> lexical </span><span data-lexical-text="true">=</span><span data-lexical-text="true"> </span><span data-lexical-text="true">"awesome"</span></code>',
+        );
+      });
+    });
   });
 });
diff --git a/packages/lexical/src/LexicalUtils.ts b/packages/lexical/src/LexicalUtils.ts
index 1cedfae5d78..40ee5de0695 100644
--- a/packages/lexical/src/LexicalUtils.ts
+++ b/packages/lexical/src/LexicalUtils.ts
@@ -543,8 +543,13 @@ export function markNodesWithTypesAsDirty(
   editor.update(
     () => {
       for (const nodeMap of dirtyNodeMaps) {
-        for (const node of nodeMap.values()) {
-          node.markDirty();
+        for (const nodeKey of nodeMap.keys()) {
+          // We are only concerned with nodes that are still in the latest NodeMap,
+          // if they no longer exist then markDirty would raise an exception
+          const latest = $getNodeByKey(nodeKey);
+          if (latest) {
+            latest.markDirty();
+          }
         }
       }
     },