[TreeView] Add a new lookup to access an item index without expansive…

… computation (mui#12729)

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

@@ -3,30 +3,29 @@ import { TreeViewPlugin } from '../../models';
 import {
   UseTreeViewItemsSignature,
   UseTreeViewItemsDefaultizedParameters,
-  TreeViewItemMetaMap,
   TreeViewItemIdAndChildren,
   UseTreeViewItemsState,
-  TreeViewItemMap,
 } from './useTreeViewItems.types';
 import { publishTreeViewEvent } from '../../utils/publishTreeViewEvent';
 import { TreeViewBaseItem } from '../../../models';
-import { TREE_VIEW_ROOT_PARENT_ID } from './useTreeViewItems.utils';
+import { buildSiblingIndexes, TREE_VIEW_ROOT_PARENT_ID } from './useTreeViewItems.utils';
 
 interface UpdateNodesStateParameters
   extends Pick<
     UseTreeViewItemsDefaultizedParameters<TreeViewBaseItem>,
     'items' | 'isItemDisabled' | 'getItemLabel' | 'getItemId'
   > {}
 
+type State = UseTreeViewItemsState<any>['items'];
 const updateItemsState = ({
   items,
   isItemDisabled,
   getItemLabel,
   getItemId,
 }: UpdateNodesStateParameters): UseTreeViewItemsState<any>['items'] => {
-  const itemMetaMap: TreeViewItemMetaMap = {};
-  const itemMap: TreeViewItemMap<any> = {};
-  const itemOrderedChildrenIds: { [parentId: string]: string[] } = {
+  const itemMetaMap: State['itemMetaMap'] = {};
+  const itemMap: State['itemMap'] = {};
+  const itemOrderedChildrenIds: State['itemOrderedChildrenIds'] = {
     [TREE_VIEW_ROOT_PARENT_ID]: [],
   };
 
@@ -95,11 +94,17 @@ const updateItemsState = ({
 
   const itemTree = items.map((item, itemIndex) => processItem(item, itemIndex, null));
 
+  const itemChildrenIndexes: State['itemChildrenIndexes'] = {};
+  Object.keys(itemOrderedChildrenIds).forEach((parentId) => {
+    itemChildrenIndexes[parentId] = buildSiblingIndexes(itemOrderedChildrenIds[parentId]);
+  });
+
   return {
     itemMetaMap,
     itemTree,
     itemMap,
     itemOrderedChildrenIds,
+    itemChildrenIndexes,
   };
 };
 
@@ -148,6 +153,14 @@ export const useTreeViewItems: TreeViewPlugin<UseTreeViewItemsSignature> = ({
     [instance],
   );
 
+  const getItemIndex = React.useCallback(
+    (itemId: string) => {
+      const parentId = instance.getItemMeta(itemId).parentId ?? TREE_VIEW_ROOT_PARENT_ID;
+      return state.items.itemChildrenIndexes[parentId][itemId];
+    },
+    [instance, state.items.itemChildrenIndexes],
+  );
+
   const getItemOrderedChildrenIds = React.useCallback(
     (itemId: string | null) =>
       state.items.itemOrderedChildrenIds[itemId ?? TREE_VIEW_ROOT_PARENT_ID] ?? [],
@@ -225,6 +238,7 @@ export const useTreeViewItems: TreeViewPlugin<UseTreeViewItemsSignature> = ({
       getItemMeta,
       getItem,
       getItemsToRender,
+      getItemIndex,
       getItemOrderedChildrenIds,
       getNavigableChildrenIds,
       isItemDisabled,

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

@@ -23,6 +23,7 @@ export interface UseTreeViewItemsInstance<R extends {}> extends UseTreeViewItems
   getItemOrderedChildrenIds: (parentId: string | null) => string[];
   getNavigableChildrenIds: (itemId: string | null) => string[];
   isItemDisabled: (itemId: string | null) => itemId is string;
+  getItemIndex: (itemId: string) => number;
   /**
    * Freeze any future update to the state based on the `items` prop.
    * This is useful when `useTreeViewJSXItems` is used to avoid having conflicting sources of truth.
@@ -92,6 +93,7 @@ export interface UseTreeViewItemsState<R extends {}> {
     itemMetaMap: TreeViewItemMetaMap;
     itemMap: TreeViewItemMap<R>;
     itemOrderedChildrenIds: { [parentItemId: string]: string[] };
+    itemChildrenIndexes: { [parentItemId: string]: { [itemId: string]: number } };
   };
 }
 

packages/x-tree-view/src/internals/plugins/useTreeViewItems/useTreeViewItems.utils.ts

@@ -1 +1,10 @@
 export const TREE_VIEW_ROOT_PARENT_ID = '__TREE_VIEW_ROOT_PARENT_ID__';
+
+export const buildSiblingIndexes = (siblings: string[]) => {
+  const siblingsIndexLookup: { [itemId: string]: number } = {};
+  siblings.forEach((childId, index) => {
+    siblingsIndexLookup[childId] = index;
+  });
+
+  return siblingsIndexLookup;
+};

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

@@ -10,7 +10,10 @@ import {
   TreeViewChildrenItemContext,
   TreeViewChildrenItemProvider,
 } from '../../TreeViewProvider/TreeViewChildrenItemProvider';
-import { TREE_VIEW_ROOT_PARENT_ID } from '../useTreeViewItems/useTreeViewItems.utils';
+import {
+  buildSiblingIndexes,
+  TREE_VIEW_ROOT_PARENT_ID,
+} from '../useTreeViewItems/useTreeViewItems.utils';
 import type { TreeItemProps } from '../../../TreeItem';
 import type { TreeItem2Props } from '../../../TreeItem2';
 import { UseTreeViewIdSignature } from '../useTreeViewId';
@@ -37,7 +40,7 @@ export const useTreeViewJSXItems: TreeViewPlugin<UseTreeViewJSXItemsSignature> =
         ...prevState,
         items: {
           ...prevState.items,
-          itemMetaMap: { ...prevState.items.itemMetaMap, [item.id]: { ...item, index: -1 } },
+          itemMetaMap: { ...prevState.items.itemMetaMap, [item.id]: item },
           // For `SimpleTreeView`, we don't have a proper `item` object, so we create a very basic one.
           itemMap: { ...prevState.items.itemMap, [item.id]: { id: item.id, label: item.label } },
         },
@@ -46,13 +49,19 @@ export const useTreeViewJSXItems: TreeViewPlugin<UseTreeViewJSXItemsSignature> =
   });
 
   const setJSXItemsOrderedChildrenIds = (parentId: string | null, orderedChildrenIds: string[]) => {
+    const parentIdWithDefault = parentId ?? TREE_VIEW_ROOT_PARENT_ID;
+
     setState((prevState) => ({
       ...prevState,
       items: {
         ...prevState.items,
         itemOrderedChildrenIds: {
           ...prevState.items.itemOrderedChildrenIds,
-          [parentId ?? TREE_VIEW_ROOT_PARENT_ID]: orderedChildrenIds,
+          [parentIdWithDefault]: orderedChildrenIds,
+        },
+        itemChildrenIndexes: {
+          ...prevState.items.itemChildrenIndexes,
+          [parentIdWithDefault]: buildSiblingIndexes(orderedChildrenIds),
         },
       },
     }));