[TreeView] Do not use outdated version of the state to compute new label first char in RichTreeView

packages/x-tree-view/src/internals/plugins/useTreeViewFocus/useTreeViewFocus.ts

@@ -107,11 +107,13 @@ export const useTreeViewFocus: TreeViewPlugin<UseTreeViewFocusSignature> = ({
     }
 
     const node = instance.getNode(state.focusedNodeId);
-    const itemElement = document.getElementById(
-      instance.getTreeItemId(state.focusedNodeId, node.idAttribute),
-    );
-    if (itemElement) {
-      itemElement.blur();
+    if (node) {
+      const itemElement = document.getElementById(
+        instance.getTreeItemId(state.focusedNodeId, node.idAttribute),
+      );
+      if (itemElement) {
+        itemElement.blur();
+      }
     }
 
     setFocusedItemId(null);

packages/x-tree-view/src/internals/plugins/useTreeViewJSXNodes/useTreeViewJSXNodes.tsx

@@ -18,6 +18,8 @@ export const useTreeViewJSXNodes: TreeViewPlugin<UseTreeViewJSXNodesSignature> =
   instance,
   setState,
 }) => {
+  instance.preventItemUpdate();
+
   const insertJSXNode = useEventCallback((node: TreeViewNode) => {
     setState((prevState) => {
       if (prevState.nodes.nodeMap[node.id] != null) {

packages/x-tree-view/src/internals/plugins/useTreeViewKeyboardNavigation/useTreeViewKeyboardNavigation.test.tsx

@@ -0,0 +1,53 @@
+import * as React from 'react';
+import { expect } from 'chai';
+import { act, createRenderer, fireEvent } from '@mui-internal/test-utils';
+import { RichTreeView } from '@mui/x-tree-view/RichTreeView';
+
+describe('useTreeViewKeyboardNavigation', () => {
+  const { render } = createRenderer();
+
+  it('should work after adding / removing items', () => {
+    const { getByRole, setProps } = render(
+      <RichTreeView
+        items={[
+          { id: 'one', label: 'one' },
+          { id: 'two', label: 'two' },
+          { id: 'three', label: 'three' },
+          { id: 'four', label: 'four' },
+        ]}
+      />,
+    );
+
+    act(() => {
+      getByRole('treeitem', { name: 'one' }).focus();
+    });
+
+    fireEvent.keyDown(getByRole('treeitem', { name: 'one' }), { key: 'f' });
+    expect(getByRole('treeitem', { name: 'four' })).toHaveFocus();
+
+    setProps({
+      items: [
+        { id: 'one', label: 'one' },
+        { id: 'two', label: 'two' },
+        { id: 'three', label: 'three' },
+      ],
+    });
+    expect(getByRole('treeitem', { name: 'one' })).toHaveFocus();
+
+    fireEvent.keyDown(getByRole('treeitem', { name: 'one' }), { key: 't' });
+    expect(getByRole('treeitem', { name: 'two' })).toHaveFocus();
+
+    setProps({
+      items: [
+        { id: 'one', label: 'one' },
+        { id: 'two', label: 'two' },
+        { id: 'three', label: 'three' },
+        { id: 'four', label: 'four' },
+      ],
+    });
+    expect(getByRole('treeitem', { name: 'two' })).toHaveFocus();
+
+    fireEvent.keyDown(getByRole('treeitem', { name: 'two' }), { key: 'f' });
+    expect(getByRole('treeitem', { name: 'four' })).toHaveFocus();
+  });
+});

packages/x-tree-view/src/internals/plugins/useTreeViewKeyboardNavigation/useTreeViewKeyboardNavigation.ts

@@ -1,7 +1,7 @@
 import * as React from 'react';
 import { useTheme } from '@mui/material/styles';
 import useEventCallback from '@mui/utils/useEventCallback';
-import { TreeViewPlugin } from '../../models';
+import { TreeViewNode, TreeViewPlugin } from '../../models';
 import {
   getFirstNode,
   getLastNode,
@@ -13,7 +13,6 @@ import {
   TreeViewFirstCharMap,
   UseTreeViewKeyboardNavigationSignature,
 } from './useTreeViewKeyboardNavigation.types';
-import { TreeViewBaseItem } from '../../../models';
 import { MuiCancellableEvent } from '../../models/MuiCancellableEvent';
 
 function isPrintableCharacter(string: string) {
@@ -31,36 +30,31 @@ function findNextFirstChar(firstChars: string[], startIndex: number, char: strin
 
 export const useTreeViewKeyboardNavigation: TreeViewPlugin<
   UseTreeViewKeyboardNavigationSignature
-> = ({ instance, params }) => {
+> = ({ instance, params, state }) => {
   const theme = useTheme();
   const isRTL = theme.direction === 'rtl';
   const firstCharMap = React.useRef<TreeViewFirstCharMap>({});
-  const hasFirstCharMapBeenUpdatedImperatively = React.useRef(false);
 
   const updateFirstCharMap = useEventCallback(
     (callback: (firstCharMap: TreeViewFirstCharMap) => TreeViewFirstCharMap) => {
-      hasFirstCharMapBeenUpdatedImperatively.current = true;
       firstCharMap.current = callback(firstCharMap.current);
     },
   );
 
   React.useEffect(() => {
-    if (hasFirstCharMapBeenUpdatedImperatively.current) {
+    if (instance.isItemUpdatePrevented()) {
       return;
     }
 
     const newFirstCharMap: { [itemId: string]: string } = {};
 
-    const processItem = (item: TreeViewBaseItem) => {
-      const getItemId = params.getItemId;
-      const itemId = getItemId ? getItemId(item) : (item as { id: string }).id;
-      newFirstCharMap[itemId] = instance.getNode(itemId).label!.substring(0, 1).toLowerCase();
-      item.children?.forEach(processItem);
+    const processItem = (node: TreeViewNode) => {
+      newFirstCharMap[node.id] = node.label!.substring(0, 1).toLowerCase();
     };
 
-    params.items.forEach(processItem);
+    Object.values(state.nodes.nodeMap).forEach(processItem);
     firstCharMap.current = newFirstCharMap;
-  }, [params.items, params.getItemId, instance]);
+  }, [state.nodes.nodeMap, params.getItemId, instance]);
 
   const getFirstMatchingItem = (itemId: string, firstChar: string) => {
     let start: number;

packages/x-tree-view/src/internals/plugins/useTreeViewNodes/useTreeViewNodes.ts

@@ -155,7 +155,18 @@ export const useTreeViewNodes: TreeViewPlugin<UseTreeViewNodesSignature> = ({
     return childrenIds;
   };
 
+  const isItemUpdatedPreventedRef = React.useRef(false);
+  const preventItemUpdate = React.useCallback(() => {
+    isItemUpdatedPreventedRef.current = true;
+  }, []);
+
+  const isItemUpdatePrevented = React.useCallback(() => isItemUpdatedPreventedRef.current, []);
+
   React.useEffect(() => {
+    if (instance.isItemUpdatePrevented()) {
+      return;
+    }
+
     setState((prevState) => {
       const newState = updateNodesState({
         items: params.items,
@@ -205,6 +216,8 @@ export const useTreeViewNodes: TreeViewPlugin<UseTreeViewNodesSignature> = ({
     getChildrenIds,
     getNavigableChildrenIds,
     isNodeDisabled,
+    preventItemUpdate,
+    isItemUpdatePrevented,
   });
 
   populatePublicAPI<UseTreeViewNodesSignature>(publicAPI, {

packages/x-tree-view/src/internals/plugins/useTreeViewNodes/useTreeViewNodes.types.ts

@@ -15,6 +15,17 @@ export interface UseTreeViewNodesInstance<R extends {}> {
   getChildrenIds: (itemId: string | null) => string[];
   getNavigableChildrenIds: (itemId: string | null) => string[];
   isNodeDisabled: (itemId: string | null) => itemId is string;
+  /**
+   * Freeze any future update to the state based on the `items` prop.
+   * This is useful when `useTreeViewJSXNodes` is used to avoid having conflicting sources of truth.
+   */
+  preventItemUpdate: () => void;
+  /**
+   * Check if the updates to the state based on the `items` prop are prevented.
+   * This is useful when `useTreeViewJSXNodes` is used to avoid having conflicting sources of truth.
+   * @returns {boolean} `true` if the updates to the state based on the `items` prop are prevented.
+   */
+  isItemUpdatePrevented: () => boolean;
 }
 
 export interface UseTreeViewNodesPublicAPI<R extends {}>