[lexical] Bug Fix: Implement missing deserialization for flat $config NodeState (#7683)

Author: etrepum

packages/lexical/src/LexicalNode.ts

@@ -1147,7 +1147,7 @@ export class LexicalNode {
   updateFromJSON(
     serializedNode: LexicalUpdateJSON<SerializedLexicalNode>,
   ): this {
-    return $updateStateFromJSON(this, serializedNode[NODE_STATE_KEY]);
+    return $updateStateFromJSON(this, serializedNode);
   }
 
   /**

packages/lexical/src/LexicalNodeState.ts

@@ -13,7 +13,9 @@ import {
   type Klass,
   type LexicalNode,
   type LexicalNodeConfig,
+  type LexicalUpdateJSON,
   NODE_STATE_KEY,
+  type SerializedLexicalNode,
   type Spread,
   type StaticNodeConfigRecord,
 } from '.';
@@ -800,7 +802,9 @@ export function $getWritableNodeState<T extends LexicalNode>(
 export function $getSharedNodeState<T extends LexicalNode>(
   node: T,
 ): SharedNodeState {
-  return getRegisteredNodeOrThrow($getEditor(), node.getType()).sharedNodeState;
+  return node.__state
+    ? node.__state.sharedNodeState
+    : getRegisteredNodeOrThrow($getEditor(), node.getType()).sharedNodeState;
 }
 
 /**
@@ -815,11 +819,21 @@ export function $getSharedNodeState<T extends LexicalNode>(
  */
 export function $updateStateFromJSON<T extends LexicalNode>(
   node: T,
-  unknownState: undefined | UnknownStateRecord,
+  serialized: LexicalUpdateJSON<SerializedLexicalNode>,
 ): T {
   const writable = node.getWritable();
-  if (unknownState || writable.__state) {
-    $getWritableNodeState(node).updateFromJSON(unknownState);
+  const unknownState = serialized[NODE_STATE_KEY];
+  let parseState = unknownState;
+  for (const k of $getSharedNodeState(writable).flatKeys) {
+    if (k in serialized) {
+      if (parseState === undefined || parseState === unknownState) {
+        parseState = {...unknownState};
+      }
+      parseState[k] = serialized[k as keyof typeof serialized];
+    }
+  }
+  if (writable.__state || parseState) {
+    $getWritableNodeState(node).updateFromJSON(parseState);
   }
   return writable;
 }

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

@@ -236,6 +236,35 @@ describe('LexicalNode state', () => {
         });
       });
 
+      test('flat config serialization round-trip', () => {
+        const {editor} = testEnv;
+        editor.update(() => {
+          const flatJSON: LexicalExportJSON<StateNode> = {
+            [NODE_STATE_KEY]: {boolState: true},
+            numberState: 2,
+            type: 'state',
+            version: 1,
+          };
+          const nestedJSON: LexicalExportJSON<StateNode> = {
+            [NODE_STATE_KEY]: {boolState: true, numberState: 2},
+            type: 'state',
+            version: 1,
+          };
+          const bothJSON: LexicalExportJSON<StateNode> = {
+            [NODE_STATE_KEY]: {boolState: true, numberState: 3},
+            numberState: 2,
+            type: 'state',
+            version: 1,
+          };
+          for (const doc of [flatJSON, nestedJSON, bothJSON]) {
+            const node = StateNode.importJSON(doc);
+            expect(node.exportJSON()).toEqual(flatJSON);
+            expect($getState(node, boolState)).toBe(true);
+            expect($getState(node, numberState)).toBe(2);
+          }
+        });
+      });
+
       test('default value should not be exported', async () => {
         const {editor} = testEnv;
         editor.update(() => {