[lexical] Feature: error checking for node key re-use with type mismatch in __DEV__

jest.config.js

@@ -15,12 +15,17 @@ const common = {
 };
 
 // Use tsconfig's paths to configure jest's module name mapper
-const moduleNameMapper = Object.fromEntries(
-  Object.entries(tsconfig.compilerOptions.paths).map(([name, [firstPath]]) => [
-    `^${name}$`,
-    firstPath.replace(/^\./, '<rootDir>'),
-  ]),
-);
+const moduleNameMapper = {
+  ...Object.fromEntries(
+    Object.entries(tsconfig.compilerOptions.paths).map(
+      ([name, [firstPath]]) => [
+        `^${name}$`,
+        firstPath.replace(/^\./, '<rootDir>'),
+      ],
+    ),
+  ),
+  '^shared/invariant$': '<rootDir>/packages/shared/src/__mocks__/invariant.ts',
+};
 
 module.exports = {
   projects: [

packages/lexical/src/LexicalUpdates.ts

@@ -121,6 +121,10 @@ export function internalGetActiveEditor(): LexicalEditor | null {
   return activeEditor;
 }
 
+export function internalGetActiveEditorState(): EditorState | null {
+  return activeEditorState;
+}
+
 export function $applyTransforms(
   editor: LexicalEditor,
   node: LexicalNode,

packages/lexical/src/LexicalUtils.ts

@@ -64,6 +64,7 @@ import {
   errorOnReadOnly,
   getActiveEditor,
   getActiveEditorState,
+  internalGetActiveEditorState,
   isCurrentlyReadOnlyMode,
   triggerCommandListeners,
   updateEditor,
@@ -223,6 +224,9 @@ export function $setNodeKey(
   existingKey: NodeKey | null | undefined,
 ): void {
   if (existingKey != null) {
+    if (__DEV__) {
+      errorOnNodeKeyConstructorMismatch(node, existingKey);
+    }
     node.__key = existingKey;
     return;
   }
@@ -243,6 +247,35 @@ export function $setNodeKey(
   node.__key = key;
 }
 
+function errorOnNodeKeyConstructorMismatch(
+  node: LexicalNode,
+  existingKey: NodeKey,
+) {
+  const editorState = internalGetActiveEditorState();
+  if (!editorState) {
+    // tests expect to be able to do this kind of clone without an active editor state
+    return;
+  }
+  const existingNode = editorState._nodeMap.get(existingKey);
+  if (existingNode && existingNode.constructor !== node.constructor) {
+    // Lifted condition to if statement because the inverted logic is a bit confusing
+    if (node.constructor.name !== existingNode.constructor.name) {
+      invariant(
+        false,
+        'Lexical node with constructor %s attempted to re-use key from node in active editor state with constructor %s. Keys must not be re-used when the type is changed.',
+        node.constructor.name,
+        existingNode.constructor.name,
+      );
+    } else {
+      invariant(
+        false,
+        'Lexical node with constructor %s attempted to re-use key from node in active editor state with different constructor with the same name (possibly due to invalid Hot Module Replacement). Keys must not be re-used when the type is changed.',
+        node.constructor.name,
+      );
+    }
+  }
+}
+
 type IntentionallyMarkedAsDirtyElement = boolean;
 
 function internalMarkParentElementsAsDirty(

packages/lexical/src/__tests__/unit/LexicalEditor.test.tsx

@@ -2053,6 +2053,44 @@ describe('LexicalEditor tests', () => {
       expect(onError).not.toHaveBeenCalled();
     });
 
+    it('should fail if node keys are re-used', async () => {
+      const onError = jest.fn();
+
+      const newEditor = createTestEditor({
+        nodes: [
+          TestTextNode,
+          {
+            replace: TextNode,
+            // @ts-ignore
+            with: (node: TextNode) =>
+              new TestTextNode(node.getTextContent(), node.getKey()),
+          },
+        ],
+        onError: onError,
+        theme: {
+          text: {
+            bold: 'editor-text-bold',
+            italic: 'editor-text-italic',
+            underline: 'editor-text-underline',
+          },
+        },
+      });
+
+      newEditor.setRootElement(container);
+
+      await newEditor.update(() => {
+        // this will throw
+        $createTextNode('123');
+        expect(false).toBe('unreachable');
+      });
+
+      expect(onError).toHaveBeenCalledWith(
+        expect.objectContaining({
+          message: expect.stringMatching(/TestTextNode.*re-use key.*TextNode/),
+        }),
+      );
+    });
+
     it('node transform to the nodes specified by "replace" should not be applied to the nodes specified by "with" when "withKlass" is not specified', async () => {
       const onError = jest.fn();
 

packages/lexical/src/__tests__/unit/LexicalNode.test.ts

@@ -87,6 +87,19 @@ describe('LexicalNode tests', () => {
         });
       });
 
+      test('LexicalNode.constructor: type change detected', async () => {
+        const {editor} = testEnv;
+
+        await editor.update(() => {
+          const validNode = new TextNode(textNode.__text, textNode.__key);
+          expect(textNode.getLatest()).toBe(textNode);
+          expect(validNode.getLatest()).toBe(textNode);
+          expect(() => new TestNode(textNode.__key)).toThrowError(
+            /TestNode.*re-use key.*TextNode/,
+          );
+        });
+      });
+
       test('LexicalNode.clone()', async () => {
         const {editor} = testEnv;
 

packages/shared/src/__mocks__/invariant.ts

@@ -0,0 +1,24 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ */
+
+// invariant(condition, message) will refine types based on "condition", and
+// if "condition" is false will throw an error. This function is special-cased
+// in flow itself, so we can't name it anything else.
+export default function invariant(
+  cond?: boolean,
+  message?: string,
+  ...args: string[]
+): asserts cond {
+  if (cond) {
+    return;
+  }
+
+  throw new Error(
+    args.reduce((msg, arg) => msg.replace('%s', String(arg)), message || ''),
+  );
+}