[TreeView] Extract some logic outside of the useTreeView hook

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

@@ -6,8 +6,6 @@ import {
   TreeViewInstance,
   TreeViewPlugin,
   MergeSignaturesProperty,
-  TreeItemWrapper,
-  TreeRootWrapper,
   TreeViewPublicAPI,
   ConvertSignaturesIntoPlugins,
 } from '../models';
@@ -18,9 +16,9 @@ import {
   UseTreeViewRootSlotProps,
 } from './useTreeView.types';
 import { useTreeViewModels } from './useTreeViewModels';
-import { TreeViewContextValue, TreeViewItemPluginsRunner } from '../TreeViewProvider';
 import { TREE_VIEW_CORE_PLUGINS, TreeViewCorePluginSignatures } from '../corePlugins';
 import { extractPluginParamsFromProps } from './extractPluginParamsFromProps';
+import { useTreeViewBuildContext } from './useTreeViewBuildContext';
 
 export function useTreeViewApiInitialization<T>(
   inputApiRef: React.MutableRefObject<T | undefined> | undefined,
@@ -63,12 +61,17 @@ export const useTreeView = <
   const models = useTreeViewModels<TSignatures>(plugins, pluginParams);
   const instanceRef = React.useRef({} as TreeViewInstance<TSignatures>);
   const instance = instanceRef.current as TreeViewInstance<TSignatures>;
-
   const publicAPI = useTreeViewApiInitialization<TreeViewPublicAPI<TSignatures>>(apiRef);
-
   const innerRootRef: React.RefObject<HTMLUListElement> = React.useRef(null);
   const handleRootRef = useForkRef(innerRootRef, rootRef);
 
+  const contextValue = useTreeViewBuildContext({
+    plugins,
+    instance,
+    publicAPI,
+    rootRef: innerRootRef,
+  });
+
   const [state, setState] = React.useState(() => {
     const temp = {} as MergeSignaturesProperty<TSignaturesWithCorePluginSignatures, 'state'>;
     plugins.forEach((plugin) => {
@@ -80,69 +83,6 @@ export const useTreeView = <
     return temp;
   });
 
-  const itemWrappers = plugins
-    .map((plugin) => plugin.wrapItem)
-    .filter((wrapItem): wrapItem is TreeItemWrapper<any> => !!wrapItem);
-  const wrapItem: TreeItemWrapper<TSignatures> = ({ itemId, children }) => {
-    let finalChildren: React.ReactNode = children;
-    itemWrappers.forEach((itemWrapper) => {
-      finalChildren = itemWrapper({ itemId, children: finalChildren, instance });
-    });
-
-    return finalChildren;
-  };
-
-  const rootWrappers = plugins
-    .map((plugin) => plugin.wrapRoot)
-    .filter((wrapRoot): wrapRoot is TreeRootWrapper<any> => !!wrapRoot)
-    // The wrappers are reversed to ensure that the first wrapper is the outermost one.
-    .reverse();
-  const wrapRoot: TreeRootWrapper<TSignatures> = ({ children }) => {
-    let finalChildren: React.ReactNode = children;
-    rootWrappers.forEach((rootWrapper) => {
-      finalChildren = rootWrapper({ children: finalChildren, instance });
-    });
-
-    return finalChildren;
-  };
-
-  const runItemPlugins: TreeViewItemPluginsRunner = (itemPluginProps) => {
-    let finalRootRef: React.RefCallback<HTMLLIElement> | null = null;
-    let finalContentRef: React.RefCallback<HTMLElement> | null = null;
-
-    plugins.forEach((plugin) => {
-      if (!plugin.itemPlugin) {
-        return;
-      }
-
-      const itemPluginResponse = plugin.itemPlugin({
-        props: itemPluginProps,
-        rootRef: finalRootRef,
-        contentRef: finalContentRef,
-      });
-      if (itemPluginResponse?.rootRef) {
-        finalRootRef = itemPluginResponse.rootRef;
-      }
-      if (itemPluginResponse?.contentRef) {
-        finalContentRef = itemPluginResponse.contentRef;
-      }
-    });
-
-    return {
-      contentRef: finalContentRef,
-      rootRef: finalRootRef,
-    };
-  };
-
-  const contextValue = {
-    publicAPI,
-    wrapItem,
-    wrapRoot,
-    runItemPlugins,
-    instance: instance as TreeViewInstance<any>,
-    rootRef: innerRootRef,
-  } as TreeViewContextValue<TSignatures>;
-
   const rootPropsGetters: (<TOther extends EventHandlers = {}>(
     otherHandlers: TOther,
   ) => React.HTMLAttributes<HTMLUListElement>)[] = [];

packages/x-tree-view/src/internals/useTreeView/useTreeViewBuildContext.ts

@@ -0,0 +1,86 @@
+import * as React from 'react';
+import { TreeViewContextValue, TreeViewItemPluginsRunner } from '../TreeViewProvider';
+import {
+  ConvertSignaturesIntoPlugins,
+  TreeItemWrapper,
+  TreeRootWrapper,
+  TreeViewAnyPluginSignature,
+  TreeViewInstance,
+  TreeViewPublicAPI,
+} from '../models';
+import { TreeViewCorePluginSignatures } from '../corePlugins';
+
+export const useTreeViewBuildContext = <TSignatures extends readonly TreeViewAnyPluginSignature[]>({
+  plugins,
+  instance,
+  publicAPI,
+  rootRef,
+}: {
+  plugins: ConvertSignaturesIntoPlugins<readonly [...TreeViewCorePluginSignatures, ...TSignatures]>;
+  instance: TreeViewInstance<TSignatures>;
+  publicAPI: TreeViewPublicAPI<TSignatures>;
+  rootRef: React.RefObject<HTMLUListElement>;
+}): TreeViewContextValue<TSignatures> => {
+  const runItemPlugins: TreeViewItemPluginsRunner = (itemPluginProps) => {
+    let finalRootRef: React.RefCallback<HTMLLIElement> | null = null;
+    let finalContentRef: React.RefCallback<HTMLElement> | null = null;
+
+    plugins.forEach((plugin) => {
+      if (!plugin.itemPlugin) {
+        return;
+      }
+
+      const itemPluginResponse = plugin.itemPlugin({
+        props: itemPluginProps,
+        rootRef: finalRootRef,
+        contentRef: finalContentRef,
+      });
+      if (itemPluginResponse?.rootRef) {
+        finalRootRef = itemPluginResponse.rootRef;
+      }
+      if (itemPluginResponse?.contentRef) {
+        finalContentRef = itemPluginResponse.contentRef;
+      }
+    });
+
+    return {
+      contentRef: finalContentRef,
+      rootRef: finalRootRef,
+    };
+  };
+
+  const wrapItem: TreeItemWrapper<TSignatures> = ({ itemId, children }) => {
+    let finalChildren: React.ReactNode = children;
+    // The wrappers are reversed to ensure that the first wrapper is the outermost one.
+    for (let i = plugins.length - 1; i >= 0; i -= 1) {
+      const plugin = plugins[i];
+      if (plugin.wrapItem) {
+        finalChildren = plugin.wrapItem({ itemId, children: finalChildren, instance });
+      }
+    }
+
+    return finalChildren;
+  };
+
+  const wrapRoot: TreeRootWrapper<TSignatures> = ({ children }) => {
+    let finalChildren: React.ReactNode = children;
+    // The wrappers are reversed to ensure that the first wrapper is the outermost one.
+    for (let i = plugins.length - 1; i >= 0; i -= 1) {
+      const plugin = plugins[i];
+      if (plugin.wrapRoot) {
+        finalChildren = plugin.wrapRoot({ children: finalChildren, instance });
+      }
+    }
+
+    return finalChildren;
+  };
+
+  return {
+    runItemPlugins,
+    wrapItem,
+    wrapRoot,
+    instance,
+    rootRef,
+    publicAPI,
+  } as TreeViewContextValue<TSignatures>;
+};