[TreeView] Rework how items are being rendered in Rich Tree View components

packages/x-tree-view-pro/src/RichTreeViewPro/RichTreeViewPro.tsx

@@ -4,8 +4,7 @@ import PropTypes from 'prop-types';
 import composeClasses from '@mui/utils/composeClasses';
 import { useLicenseVerifier, Watermark } from '@mui/x-license';
 import useSlotProps from '@mui/utils/useSlotProps';
-import { TreeItem, TreeItemProps } from '@mui/x-tree-view/TreeItem';
-import { useTreeView, TreeViewProvider } from '@mui/x-tree-view/internals';
+import { useTreeView, TreeViewProvider, RichTreeViewItems } from '@mui/x-tree-view/internals';
 import { warnOnce } from '@mui/x-internals/warning';
 import { styled, createUseThemeProps } from '../internals/zero-styled';
 import { getRichTreeViewProUtilityClass } from './richTreeViewProClasses';
@@ -46,52 +45,6 @@ type RichTreeViewProComponent = (<R extends {}, Multiple extends boolean | undef
   props: RichTreeViewProProps<R, Multiple> & React.RefAttributes<HTMLUListElement>,
 ) => React.JSX.Element) & { propTypes?: any };
 
-function WrappedTreeItem<R extends {}>({
-  slots,
-  slotProps,
-  label,
-  id,
-  itemId,
-  children,
-}: Pick<RichTreeViewProProps<R, any>, 'slots' | 'slotProps'> &
-  Pick<TreeItemProps, 'id' | 'itemId' | 'children'> & { label: string }) {
-  const Item = slots?.item ?? TreeItem;
-  const itemProps = useSlotProps({
-    elementType: Item,
-    externalSlotProps: slotProps?.item,
-    additionalProps: { itemId, id, label },
-    ownerState: { itemId, label },
-  });
-
-  return <Item {...itemProps}>{children}</Item>;
-}
-
-WrappedTreeItem.propTypes = {
-  // ----------------------------- Warning --------------------------------
-  // | These PropTypes are generated from the TypeScript type definitions |
-  // | To update them edit the TypeScript types and run "pnpm proptypes"  |
-  // ----------------------------------------------------------------------
-  /**
-   * The content of the component.
-   */
-  children: PropTypes.node,
-  /**
-   * The id of the item.
-   */
-  itemId: PropTypes.string.isRequired,
-  label: PropTypes.string.isRequired,
-  /**
-   * The props used for each component slot.
-   * @default {}
-   */
-  slotProps: PropTypes.object,
-  /**
-   * Overridable component slots.
-   * @default {}
-   */
-  slots: PropTypes.object,
-} as any;
-
 const releaseInfo = getReleaseInfo();
 
 /**
@@ -143,32 +96,14 @@ const RichTreeViewPro = React.forwardRef(function RichTreeViewPro<
     ownerState: props as RichTreeViewProProps<any, any>,
   });
 
-  const itemsToRender = instance.getItemsToRender();
-
-  const renderItem = ({
-    label,
-    itemId,
-    id,
-    children,
-  }: ReturnType<typeof instance.getItemsToRender>[number]) => {
-    return (
-      <WrappedTreeItem
-        slots={slots}
-        slotProps={slotProps}
-        key={itemId}
-        label={label}
-        id={id}
-        itemId={itemId}
-      >
-        {children?.map(renderItem)}
-      </WrappedTreeItem>
-    );
-  };
-
   return (
     <TreeViewProvider value={contextValue}>
       <Root {...rootProps}>
-        {itemsToRender.map(renderItem)}
+        <RichTreeViewItems
+          slots={slots}
+          slotProps={slotProps}
+          itemsToRender={instance.getItemsToRender()}
+        />
         <Watermark packageName="x-tree-view-pro" releaseInfo={releaseInfo} />
       </Root>
     </TreeViewProvider>

packages/x-tree-view/src/RichTreeView/RichTreeView.tsx

@@ -10,7 +10,7 @@ import { styled, createUseThemeProps } from '../internals/zero-styled';
 import { useTreeView } from '../internals/useTreeView';
 import { TreeViewProvider } from '../internals/TreeViewProvider';
 import { RICH_TREE_VIEW_PLUGINS, RichTreeViewPluginSignatures } from './RichTreeView.plugins';
-import { TreeItem, TreeItemProps } from '../TreeItem';
+import { RichTreeViewItems } from '../internals/components/RichTreeViewItems';
 
 const useThemeProps = createUseThemeProps('MuiRichTreeView');
 
@@ -42,26 +42,6 @@ type RichTreeViewComponent = (<R extends {}, Multiple extends boolean | undefine
   props: RichTreeViewProps<R, Multiple> & React.RefAttributes<HTMLUListElement>,
 ) => React.JSX.Element) & { propTypes?: any };
 
-function WrappedTreeItem<R extends {}>({
-  slots,
-  slotProps,
-  label,
-  id,
-  itemId,
-  children,
-}: Pick<RichTreeViewProps<R, any>, 'slots' | 'slotProps'> &
-  Pick<TreeItemProps, 'id' | 'itemId' | 'children'> & { label: string }) {
-  const Item = slots?.item ?? TreeItem;
-  const itemProps = useSlotProps({
-    elementType: Item,
-    externalSlotProps: slotProps?.item,
-    additionalProps: { itemId, id, label },
-    ownerState: { itemId, label },
-  });
-
-  return <Item {...itemProps}>{children}</Item>;
-}
-
 /**
  *
  * Demos:
@@ -109,31 +89,15 @@ const RichTreeView = React.forwardRef(function RichTreeView<
     ownerState: props as RichTreeViewProps<any, any>,
   });
 
-  const itemsToRender = instance.getItemsToRender();
-
-  const renderItem = ({
-    label,
-    itemId,
-    id,
-    children,
-  }: ReturnType<typeof instance.getItemsToRender>[number]) => {
-    return (
-      <WrappedTreeItem
-        slots={slots}
-        slotProps={slotProps}
-        key={itemId}
-        label={label}
-        id={id}
-        itemId={itemId}
-      >
-        {children?.map(renderItem)}
-      </WrappedTreeItem>
-    );
-  };
-
   return (
     <TreeViewProvider value={contextValue}>
-      <Root {...rootProps}>{itemsToRender.map(renderItem)}</Root>
+      <Root {...rootProps}>
+        <RichTreeViewItems
+          slots={slots}
+          slotProps={slotProps}
+          itemsToRender={instance.getItemsToRender()}
+        />
+      </Root>
     </TreeViewProvider>
   );
 }) as RichTreeViewComponent;

packages/x-tree-view/src/internals/components/RichTreeViewItems.tsx

@@ -0,0 +1,91 @@
+import * as React from 'react';
+import useSlotProps from '@mui/utils/useSlotProps';
+import { SlotComponentProps } from '@mui/utils';
+
+import { TreeItem, TreeItemProps } from '../../TreeItem';
+import { TreeItem2Props } from '../../TreeItem2';
+import { TreeViewItemId } from '../../models';
+import { TreeViewItemToRenderProps } from '../plugins/useTreeViewItems';
+
+interface RichTreeViewItemsOwnerState {
+  itemId: TreeViewItemId;
+  label: string;
+}
+
+export interface RichTreeViewItemsSlots {
+  /**
+   * Custom component for the item.
+   * @default TreeItem.
+   */
+  item?: React.JSXElementConstructor<TreeItemProps> | React.JSXElementConstructor<TreeItem2Props>;
+}
+
+export interface RichTreeViewItemsSlotProps {
+  item?: SlotComponentProps<typeof TreeItem, {}, RichTreeViewItemsOwnerState>;
+}
+
+export interface RichTreeViewItemsProps {
+  itemsToRender: TreeViewItemToRenderProps[];
+  /**
+   * Overridable component slots.
+   * @default {}
+   */
+  slots?: RichTreeViewItemsSlots;
+  /**
+   * The props used for each component slot.
+   * @default {}
+   */
+  slotProps?: RichTreeViewItemsSlotProps;
+}
+
+function WrappedTreeItem({
+  slots,
+  slotProps,
+  label,
+  id,
+  itemId,
+  itemsToRender,
+}: Pick<RichTreeViewItemsProps, 'slots' | 'slotProps'> &
+  Pick<TreeItemProps, 'id' | 'itemId' | 'children'> & {
+    label: string;
+    isContentHidden?: boolean;
+    itemsToRender: TreeViewItemToRenderProps[] | undefined;
+  }) {
+  const Item = slots?.item ?? TreeItem;
+  const { ownerState, ...itemProps } = useSlotProps({
+    elementType: Item,
+    externalSlotProps: slotProps?.item,
+    additionalProps: { itemId, id, label },
+    ownerState: { itemId, label },
+  });
+
+  const children = React.useMemo(
+    () =>
+      itemsToRender ? (
+        <RichTreeViewItems itemsToRender={itemsToRender} slots={slots} slotProps={slotProps} />
+      ) : null,
+    [itemsToRender, slots, slotProps],
+  );
+
+  return <Item {...itemProps}>{children}</Item>;
+}
+
+export function RichTreeViewItems(props: RichTreeViewItemsProps) {
+  const { itemsToRender, slots, slotProps } = props;
+
+  return (
+    <React.Fragment>
+      {itemsToRender.map((item) => (
+        <WrappedTreeItem
+          slots={slots}
+          slotProps={slotProps}
+          key={item.itemId}
+          label={item.label}
+          id={item.id}
+          itemId={item.itemId}
+          itemsToRender={item.children}
+        />
+      ))}
+    </React.Fragment>
+  );
+}

packages/x-tree-view/src/internals/index.ts

@@ -1,6 +1,8 @@
 export { useTreeView } from './useTreeView';
 export { TreeViewProvider, useTreeViewContext } from './TreeViewProvider';
 
+export { RichTreeViewItems } from './components/RichTreeViewItems';
+
 export { unstable_resetCleanupTracking } from './hooks/useInstanceEventHandler';
 
 export type {

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

@@ -4,5 +4,6 @@ export type {
   UseTreeViewItemsParameters,
   UseTreeViewItemsDefaultizedParameters,
   UseTreeViewItemsState,
+  TreeViewItemToRenderProps,
 } from './useTreeViewItems.types';
 export { buildSiblingIndexes, TREE_VIEW_ROOT_PARENT_ID } from './useTreeViewItems.utils';

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

@@ -2,11 +2,11 @@ import * as React from 'react';
 import { TreeViewItemMeta, DefaultizedProps, TreeViewPluginSignature } from '../../models';
 import { TreeViewBaseItem, TreeViewItemId } from '../../../models';
 
-interface TreeViewItemProps {
+export interface TreeViewItemToRenderProps {
   label: string;
   itemId: string;
   id: string | undefined;
-  children?: TreeViewItemProps[];
+  children?: TreeViewItemToRenderProps[];
 }
 
 export interface UseTreeViewItemsPublicAPI<R extends {}> {
@@ -33,7 +33,7 @@ export interface UseTreeViewItemsPublicAPI<R extends {}> {
   getItemOrderedChildrenIds: (itemId: TreeViewItemId | null) => TreeViewItemId[];
   /**
    * Get all the items in the same format as provided by `props.items`.
-   * @returns {TreeViewItemProps[]} The items in the tree.
+   * @returns {TreeViewItemToRenderProps[]} The items in the tree.
    */
   getItemTree: () => TreeViewBaseItem[];
 }
@@ -49,10 +49,10 @@ export interface UseTreeViewItemsInstance<R extends {}> extends UseTreeViewItems
   /**
    * Get the item that should be rendered.
    * This method is only used on Rich Tree View components.
-   * Check the `TreeViewItemProps` type for more information.
-   * @returns {TreeViewItemProps[]} The items to render.
+   * Check the `TreeViewItemToRenderProps` type for more information.
+   * @returns {TreeViewItemToRenderProps[]} The items to render.
    */
-  getItemsToRender: () => TreeViewItemProps[];
+  getItemsToRender: () => TreeViewItemToRenderProps[];
   /**
    * Check if a given item is disabled.
    * An item is disabled if it was marked as disabled or if one of its ancestors is disabled.