[tree view] Focus item sibling on unmount instead of the 1st item

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

@@ -24,7 +24,31 @@ describeTreeView<TreeViewAnyStore>(
         expect(view.getFocusedItemId()).to.equal('1');
       });
 
-      it('should move the focus when the focused item is removed', () => {
+      it('should focus the next sibling when a focused middle item is removed', () => {
+        const view = render({
+          items: [{ id: '1' }, { id: '2' }, { id: '3' }],
+        });
+
+        fireEvent.focus(view.getItemRoot('2'));
+        expect(view.getFocusedItemId()).to.equal('2');
+
+        view.setItems([{ id: '1' }, { id: '3' }]);
+        expect(view.getFocusedItemId()).to.equal('3');
+      });
+
+      it('should focus the previous sibling when the focused last item is removed', () => {
+        const view = render({
+          items: [{ id: '1' }, { id: '2' }, { id: '3' }],
+        });
+
+        fireEvent.focus(view.getItemRoot('3'));
+        expect(view.getFocusedItemId()).to.equal('3');
+
+        view.setItems([{ id: '1' }, { id: '2' }]);
+        expect(view.getFocusedItemId()).to.equal('2');
+      });
+
+      it('should focus the remaining sibling when the focused item is removed and only one sibling is left', () => {
         const view = render({
           items: [{ id: '1' }, { id: '2' }],
         });
@@ -35,6 +59,35 @@ describeTreeView<TreeViewAnyStore>(
         view.setItems([{ id: '1' }]);
         expect(view.getFocusedItemId()).to.equal('1');
       });
+
+      it('should focus the parent when the focused item is removed and has no siblings left', () => {
+        const view = render({
+          items: [{ id: '1' }, { id: '2', children: [{ id: '2.1' }] }],
+          defaultExpandedItems: ['2'],
+        });
+
+        fireEvent.focus(view.getItemRoot('2.1'));
+        expect(view.getFocusedItemId()).to.equal('2.1');
+
+        view.setItems([{ id: '1' }, { id: '2' }]);
+        expect(view.getFocusedItemId()).to.equal('2');
+      });
+
+      it('should focus the next nested sibling when a focused nested item is removed', () => {
+        const view = render({
+          items: [
+            { id: '1' },
+            { id: '2', children: [{ id: '2.1' }, { id: '2.2' }, { id: '2.3' }] },
+          ],
+          defaultExpandedItems: ['2'],
+        });
+
+        fireEvent.focus(view.getItemRoot('2.2'));
+        expect(view.getFocusedItemId()).to.equal('2.2');
+
+        view.setItems([{ id: '1' }, { id: '2', children: [{ id: '2.1' }, { id: '2.3' }] }]);
+        expect(view.getFocusedItemId()).to.equal('2.3');
+      });
     });
 
     describe('tabIndex HTML attribute', () => {

packages/x-tree-view/src/internals/plugins/focus/TreeViewFocusPlugin.ts

@@ -3,6 +3,11 @@ import { expansionSelectors } from '../expansion';
 import { focusSelectors } from './selectors';
 import { itemsSelectors } from '../items';
 import { MinimalTreeViewStore } from '../../MinimalTreeViewStore';
+import {
+  getFirstNavigableItem,
+  getNextNavigableItem,
+  getPreviousNavigableItem,
+} from '../../utils/tree';
 
 export class TreeViewFocusPlugin {
   private store: MinimalTreeViewStore<any, any>;
@@ -13,24 +18,36 @@ export class TreeViewFocusPlugin {
     this.store = store;
 
     // Whenever the items change, we need to ensure the focused item is still present.
-    this.store.registerStoreEffect(itemsSelectors.itemMetaLookup, () => {
-      const focusedItemId = focusSelectors.focusedItemId(store.state);
-      if (focusedItemId == null) {
+    // If the focused item was removed, focus the closest neighbor instead of the first item.
+    let previousState = store.state;
+    this.store.subscribe((newState) => {
+      // Only run when items actually changed.
+      if (newState.itemMetaLookup === previousState.itemMetaLookup) {
+        previousState = newState;
         return;
       }
 
-      const hasItemBeenRemoved = !itemsSelectors.itemMeta(store.state, focusedItemId);
-      if (!hasItemBeenRemoved) {
+      const focusedItemId = focusSelectors.focusedItemId(newState);
+      if (focusedItemId == null || itemsSelectors.itemMeta(newState, focusedItemId)) {
+        previousState = newState;
         return;
       }
 
-      const defaultFocusableItemId = focusSelectors.defaultFocusableItemId(store.state);
-      if (defaultFocusableItemId == null) {
+      const checkItemInNewTree = (itemId: TreeViewItemId | null) =>
+        itemId == null || !itemsSelectors.itemMeta(newState, itemId) ? null : itemId;
+
+      const itemToFocusId =
+        checkItemInNewTree(getNextNavigableItem(previousState, focusedItemId)) ??
+        checkItemInNewTree(getPreviousNavigableItem(previousState, focusedItemId)) ??
+        getFirstNavigableItem(newState);
+
+      if (itemToFocusId == null) {
         this.setFocusedItemId(null);
-        return;
+      } else {
+        this.applyItemFocus(null, itemToFocusId);
       }
 
-      this.applyItemFocus(null, defaultFocusableItemId);
+      previousState = newState;
     });
   }
 

packages/x-tree-view/src/internals/plugins/items/selectors.ts

@@ -87,7 +87,8 @@ export const itemsSelectors = {
    */
   canItemBeFocused: createSelector(
     (state: MinimalTreeViewState<any, any>, itemId: TreeViewItemId) =>
-      state.disabledItemsFocusable || !isItemDisabled(state.itemMetaLookup, itemId),
+      state.disabledItemsFocusable ||
+      (state.itemModelLookup[itemId] != null && !isItemDisabled(state.itemMetaLookup, itemId)),
   ),
   /**
    * Gets the identation between an item and its children.

packages/x-tree-view/src/internals/plugins/jsxItems/TreeViewJSXItemsPlugin.ts

@@ -1,24 +1,32 @@
 import { TreeViewItemId } from '../../../models';
 import { TreeViewItemMeta } from '../../models';
 import type { SimpleTreeViewStore } from '../../SimpleTreeViewStore';
-import { buildSiblingIndexes, TREE_VIEW_ROOT_PARENT_ID } from '../items';
+import { buildSiblingIndexes, itemsSelectors, TREE_VIEW_ROOT_PARENT_ID } from '../items';
 import { jsxItemsitemWrapper, useJSXItemsItemPlugin } from './itemPlugin';
 
 export class TreeViewJSXItemsPlugin {
   private store: SimpleTreeViewStore<any>;
 
+  /**
+   * Tracks which component instance owns each item id,
+   * so that duplicate ids from different components can be detected.
+   */
+  private itemOwners = new Map<string, symbol>();
+
   public constructor(store: SimpleTreeViewStore<any>) {
     this.store = store;
     store.itemPluginManager.register(useJSXItemsItemPlugin, jsxItemsitemWrapper);
   }
 
   /**
-   * Insert a new item in the state from a Tree Item component.
-   * @param {TreeViewItemMeta} item The meta-information of the item to insert.
-   * @returns {() => void} A function to remove the item from the state.
+   * Insert or update an item in the state from a Tree Item component.
+   * If the item already exists and belongs to the same owner (e.g. after a deps-change re-run of the layout effect),
+   * its meta is updated in place instead of removing and re-inserting.
    */
-  public insertJSXItem = (item: TreeViewItemMeta) => {
-    if (this.store.state.itemMetaLookup[item.id] != null) {
+  public upsertJSXItem = (item: TreeViewItemMeta, ownerToken: symbol) => {
+    const currentOwner = this.itemOwners.get(item.id);
+
+    if (currentOwner != null && currentOwner !== ownerToken) {
       throw new Error(
         [
           'MUI X: The Tree View component requires all items to have a unique `id` property.',
@@ -28,16 +36,41 @@ export class TreeViewJSXItemsPlugin {
       );
     }
 
-    this.store.update({
-      itemMetaLookup: { ...this.store.state.itemMetaLookup, [item.id]: item },
-      // For Simple Tree View, we don't have a proper `item` object, so we create a very basic one.
-      itemModelLookup: {
-        ...this.store.state.itemModelLookup,
-        [item.id]: { id: item.id, label: item.label ?? '' },
-      },
-    });
+    this.itemOwners.set(item.id, ownerToken);
+    const existingMeta = itemsSelectors.itemMeta(this.store.state, item.id);
+
+    if (existingMeta != null) {
+      // Update the existing item in place.
+      let hasChanges = false;
+      for (const key of Object.keys(item) as (keyof TreeViewItemMeta)[]) {
+        if (existingMeta[key] !== item[key]) {
+          hasChanges = true;
+          break;
+        }
+      }
+
+      if (hasChanges) {
+        this.store.update({
+          itemMetaLookup: {
+            ...this.store.state.itemMetaLookup,
+            [item.id]: { ...existingMeta, ...item },
+          },
+        });
+      }
+    } else {
+      this.store.update({
+        itemMetaLookup: { ...this.store.state.itemMetaLookup, [item.id]: item },
+        // For Simple Tree View, we don't have a proper `item` object, so we create a very basic one.
+        itemModelLookup: {
+          ...this.store.state.itemModelLookup,
+          [item.id]: { id: item.id, label: item.label ?? '' },
+        },
+      });
+    }
 
     return () => {
+      this.itemOwners.delete(item.id);
+
       const newItemMetaLookup = { ...this.store.state.itemMetaLookup };
       const newItemModelLookup = { ...this.store.state.itemModelLookup };
       delete newItemMetaLookup[item.id];

packages/x-tree-view/src/internals/plugins/jsxItems/itemPlugin.tsx

@@ -3,6 +3,7 @@ import * as React from 'react';
 import { useStore } from '@mui/x-internals/store';
 import { useMergedRefs } from '@base-ui/utils/useMergedRefs';
 import { useIsoLayoutEffect } from '@base-ui/utils/useIsoLayoutEffect';
+import { useRefWithInit } from '@base-ui/utils/useRefWithInit';
 import { TreeItemWrapper, TreeViewItemPlugin } from '../../models';
 import { useTreeViewContext } from '../../TreeViewProvider';
 import {
@@ -34,6 +35,8 @@ export const useJSXItemsItemPlugin: TreeViewItemPlugin = ({ props, rootRef, cont
   const pluginContentRef = React.useRef<HTMLDivElement>(null);
   const handleContentRef = useMergedRefs(pluginContentRef, contentRef);
   const idAttribute = useStore(store, idSelectors.treeItemIdAttribute, itemId, id);
+  const isMountedRef = React.useRef(true);
+  const ownerTokenRef = useRefWithInit(Symbol);
 
   // Prevent any flashing
   useIsoLayoutEffect(() => {
@@ -46,15 +49,32 @@ export const useJSXItemsItemPlugin: TreeViewItemPlugin = ({ props, rootRef, cont
   }, [store, registerChild, unregisterChild, idAttribute, itemId]);
 
   useIsoLayoutEffect(() => {
-    return store.jsxItems.insertJSXItem({
-      id: itemId,
-      idAttribute: id,
-      parentId,
-      expandable,
-      disabled,
-      selectable: !disableSelection,
-    });
-  }, [store, parentId, itemId, expandable, disabled, disableSelection, id]);
+    isMountedRef.current = true;
+    return () => {
+      isMountedRef.current = false;
+    };
+  }, []);
+
+  useIsoLayoutEffect(() => {
+    const remove = store.jsxItems.upsertJSXItem(
+      {
+        id: itemId,
+        idAttribute: id,
+        parentId,
+        expandable,
+        disabled,
+        selectable: !disableSelection,
+      },
+      ownerTokenRef.current,
+    );
+
+    return () => {
+      // Only remove the item if the component is unmounting, not if the dependencies are changing.
+      if (!isMountedRef.current) {
+        remove();
+      }
+    };
+  }, [store, parentId, itemId, expandable, disabled, disableSelection, id, ownerTokenRef]);
 
   React.useEffect(() => {
     if (label) {

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

@@ -1271,9 +1271,9 @@ describeTreeView<TreeViewAnyStore>(
         expect(view.getFocusedItemId()).to.equal('4');
 
         view.setItems([{ id: '1' }, { id: '2' }, { id: '3' }]);
-        expect(view.getFocusedItemId()).to.equal('1');
+        expect(view.getFocusedItemId()).to.equal('3');
 
-        fireEvent.keyDown(view.getItemRoot('1'), { key: '2' });
+        fireEvent.keyDown(view.getItemRoot('3'), { key: '2' });
         expect(view.getFocusedItemId()).to.equal('2');
 
         view.setItems([{ id: '1' }, { id: '2' }, { id: '3' }, { id: '4' }]);