[lexical-react] Bug Fix: TreeView re-renders when the editor prop changes

packages/lexical-react/src/LexicalTreeView.tsx

@@ -16,6 +16,8 @@ import {type EditorState, type LexicalEditor, mergeRegister} from 'lexical';
 import * as React from 'react';
 import {type JSX, useEffect, useState} from 'react';
 
+import useLayoutEffect from './shared/useLayoutEffect';
+
 /**
  * TreeView is a React component that provides a visual representation of
  * the Lexical editor's state and enables debugging features like time travel
@@ -65,7 +67,14 @@ export function TreeView({
 
   const commandsLog = useLexicalCommandsLog(editor);
 
-  useEffect(() => {
+  // useLayoutEffect, like useCanShowPlaceholder and ContentEditableElement, so
+  // the state can be re-derived before subscribing without the cascading
+  // render that setState inside a useEffect body causes.
+  useLayoutEffect(() => {
+    // The state was seeded on the first render only, so re-read it here: when
+    // the editor prop changes, neither listener has fired yet for the new
+    // editor and the view would keep rendering the previous editor's state.
+    setEditorCurrentState(editor.getEditorState());
     // Registers listeners to update the tree view when the editor state changes
     return mergeRegister(
       editor.registerUpdateListener(({editorState}) => {

packages/lexical-react/src/__tests__/unit/LexicalTreeView.test.tsx

@@ -0,0 +1,82 @@
+/**
+ * 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.
+ *
+ */
+
+import {TreeView} from '@lexical/react/LexicalTreeView';
+import {
+  $createParagraphNode,
+  $createTextNode,
+  $getRoot,
+  createEditor,
+  type LexicalEditor,
+} from 'lexical';
+import * as React from 'react';
+import {act} from 'react';
+import {createRoot, type Root} from 'react-dom/client';
+import {afterEach, beforeEach, describe, expect, it} from 'vitest';
+
+function makeEditor(text: string): LexicalEditor {
+  const editor = createEditor({
+    namespace: 'tree-view',
+    onError: error => {
+      throw error;
+    },
+  });
+  editor.setRootElement(document.createElement('div'));
+  editor.update(
+    () => {
+      $getRoot()
+        .clear()
+        .append($createParagraphNode().append($createTextNode(text)));
+    },
+    {discrete: true},
+  );
+  return editor;
+}
+
+describe('TreeView', () => {
+  let container: HTMLDivElement;
+  let reactRoot: Root;
+
+  beforeEach(() => {
+    container = document.createElement('div');
+    document.body.appendChild(container);
+    reactRoot = createRoot(container);
+  });
+
+  afterEach(async () => {
+    await act(async () => {
+      reactRoot.unmount();
+    });
+    container.remove();
+  });
+
+  /** The rendered tree is produced asynchronously, so let it settle. */
+  async function renderTreeFor(editor: LexicalEditor): Promise<string> {
+    await act(async () => {
+      reactRoot.render(<TreeView editor={editor} />);
+    });
+    await act(async () => {
+      await new Promise(resolve => setTimeout(resolve, 0));
+    });
+    const pre = container.querySelector('pre');
+    expect(pre).not.toBe(null);
+    return (pre as HTMLPreElement).textContent ?? '';
+  }
+
+  it('renders the tree of the editor it was given', async () => {
+    expect(await renderTreeFor(makeEditor('AAAAA'))).toContain('"AAAAA"');
+  });
+
+  it('re-renders when the editor prop changes', async () => {
+    expect(await renderTreeFor(makeEditor('AAAAA'))).toContain('"AAAAA"');
+
+    const tree = await renderTreeFor(makeEditor('BBBBB'));
+    expect(tree).toContain('"BBBBB"');
+    expect(tree).not.toContain('"AAAAA"');
+  });
+});