packages/tree/src/hooks/useTree/useTree.ts - develop - bearury/ui-kit-ce

packages/tree/src/hooks/useTree/useTree.ts
324 строки8 KB
Pavel Uryadyshev
feat(tree-dnd): реализован новый компонент Tree с поддержкой drag-and-drop
09 июл 2025, 13:07
09 июл 2025, 13:07a06fe68
'use client'

import * as React from 'react'
import {
  TreeItem,
  UseTreeProps,
  UseTreeReturnProps,
  TreeNodeItem,
  TreeState,
  TreeNodePlain,
  TreeGetters,
} from '../types'
import * as Getters from '../getters'
import { useExpand } from '../useExpand'
import { useCheck } from '../useCheck'
import { useValue } from '@v-uik/hooks'
import { useSelect } from '../useSelect'

const makeFormatValueFromProp =
  <TItem = TreeItem>({
    getTreeItemChildren,
    getTreeItemKey,
    getTreeItemCheckDisabled,
    getTreeItemDisabled,
    getTreeItemSelectDisabled,
    getTreeItemLoadable,
  }: TreeGetters<TItem>) =>
  (dataSource: TItem[]): TreeState<TItem> => {
    const allowedKeysForExpand: Record<React.Key, boolean> = {}
    const disabledKeysForCheck: Record<React.Key, boolean> = {}
    const disabledKeysForSelect: Record<React.Key, boolean> = {}
    const disabledKeysForInitialSelect: Record<React.Key, boolean> = {}
    const treePlain: TreeNodePlain<TItem> = {}
    const nodeKeys: React.Key[] = []

    const processTreeItems = (dataSource: TItem[], depth = 0, parent?: TItem) =>
      dataSource.reduce((accum, item, index) => {
        const key = getTreeItemKey(item)

        if (!key) {
          return accum
        }

        const children = getTreeItemChildren(item)
        const disabled = getTreeItemDisabled(item)
        const selectDisabled = getTreeItemSelectDisabled(item)
        const parentKey = parent ? getTreeItemKey(parent) : undefined
        nodeKeys.push(key)

        if (getTreeItemCheckDisabled(item) || disabled) {
          disabledKeysForCheck[key] = true
        }

        if (selectDisabled) {
          disabledKeysForInitialSelect[key] = true
        }

        if (selectDisabled || disabled) {
          disabledKeysForSelect[key] = true
        }

        const node: TreeNodeItem<TItem> = {
          ...item,
        } as unknown as TreeNodeItem<TItem>

        const meta: TreeNodeItem<TItem>['$meta'] = {
          depth,
          id: String(key),
          order: nodeKeys.length,
          disabled,
          index,
        } as TreeNodeItem<TItem>['$meta']

        if (parentKey) {
          meta['parentKey'] = parentKey
        }

        if (children?.length || getTreeItemLoadable(item)) {
          allowedKeysForExpand[key] = true
        }

        if (children?.length) {
          meta['closestChildKey'] = getTreeItemKey(children[0])
          meta['boundaryChildrenKeys'] = children.map((item) =>
            String(getTreeItemKey(item))
          )
        }

        node['$meta'] = meta

        accum.push(node)
        treePlain[key] = node

        if (children?.length) {
          node['children'] = processTreeItems(
            children,
            depth + 1,
            node as unknown as TItem
          )
        }

        return accum
      }, [] as unknown as TreeNodeItem<TItem>[])

    const tree = processTreeItems(dataSource)

    return {
      tree,
      treePlain,
      allowedKeysForExpand,
      disabledKeysForCheck,
      disabledKeysForSelect,
      disabledKeysForInitialSelect,
      nodeKeys,
    }
  }

export const useTree = <TItem = TreeItem>({
  dataSource,
  getters: {
    getTreeItemChildren = Getters.getTreeItemChildren,
    getTreeItemKey = Getters.getTreeItemKey,
    getTreeItemLabel = Getters.getTreeItemLabel,
    getTreeItemCheckDisabled = Getters.getTreeItemCheckDisabled,
    getTreeItemDisabled = Getters.getTreeItemDisabled,
    getTreeItemSelectDisabled = Getters.getTreeItemSelectDisabled,
    getTreeItemLoadable = Getters.getTreeItemLoadable,
  } = Getters,
  onNodeExpand,
  defaultExpandedKeys,
  expandedKeys,
  checkedKeys,
  defaultCheckedKeys,
  onNodeCheck,
  defaultSelectedKeys,
  onNodeSelect,
  selectedKeys,
  multipleSelect,
  selectMode,
  checkMode,
  onLoadData,
}: UseTreeProps<TItem>): UseTreeReturnProps<TItem> => {
  const getters = React.useMemo(
    () => ({
      getTreeItemChildren,
      getTreeItemKey,
      getTreeItemLabel,
      getTreeItemCheckDisabled,
      getTreeItemDisabled,
      getTreeItemSelectDisabled,
      getTreeItemLoadable,
    }),
    [
      getTreeItemChildren,
      getTreeItemKey,
      getTreeItemLabel,
      getTreeItemCheckDisabled,
      getTreeItemDisabled,
      getTreeItemSelectDisabled,
      getTreeItemLoadable,
    ]
  )

  const [
    {
      tree,
      allowedKeysForExpand,
      treePlain,
      disabledKeysForCheck,
      nodeKeys,
      disabledKeysForSelect,
      disabledKeysForInitialSelect,
    },
  ] = useValue(dataSource, {
    formatValueFromProp: makeFormatValueFromProp<TItem>(getters),
  })

  const { getIsExpanded, toggleExpand, getIsLoading } = useExpand({
    allowedKeysForExpand,
    treePlain,
    onNodeExpand,
    defaultExpandedKeys,
    expandedKeys,
    onLoadData,
    getters,
  })

  const { getIsChecked, getIsIndeterminate, toggleCheck } = useCheck({
    disabledKeysForCheck,
    treePlain,
    checkedKeys,
    defaultCheckedKeys,
    onNodeCheck,
    getTreeItemCheckDisabled,
    checkMode,
  })

  const {
    getIsSelected,
    toggleSelect,
    toggleSelectRange,
    getFirstSelectedKey,
    getLastSelectedKey,
  } = useSelect<TItem>({
    defaultSelectedKeys,
    onNodeSelect,
    selectedKeys,
    treePlain,
    disabledKeysForSelect,
    multipleSelect,
    selectMode,
    nodeKeys,
    getIsExpanded,
    disabledKeysForInitialSelect,
  })

  const getTreeNodeData = React.useCallback(
    (key: React.Key) => treePlain[key],
    [treePlain]
  )

  const getTreeNodeVisibleChildren = React.useCallback(
    (key: React.Key) => {
      const targetNode = treePlain[key]

      if (
        !targetNode ||
        !getIsExpanded(key) ||
        !targetNode.$meta.boundaryChildrenKeys?.length
      ) {
        return []
      }

      const childrenKeys = [...targetNode.$meta.boundaryChildrenKeys]
      const visibleChildren = [] as TreeNodeItem<TItem>[]

      while (childrenKeys.length > 0) {
        const childrenKey = childrenKeys.shift()

        if (!childrenKey) {
          continue
        }

        visibleChildren.push(treePlain[childrenKey])

        if (
          getIsExpanded(childrenKey) &&
          treePlain[childrenKey].$meta.boundaryChildrenKeys
        ) {
          childrenKeys.push(
            ...(treePlain[childrenKey].$meta.boundaryChildrenKeys as string[])
          )
        }
      }

      return visibleChildren
    },
    [treePlain, getIsExpanded]
  )

  const getParentRootNode = React.useCallback(
    (key: React.Key) => {
      let targetNode = treePlain[key]

      if (!targetNode) {
        return undefined
      }

      if (!targetNode.$meta.parentKey) {
        return targetNode
      }

      const parentKeys = [targetNode.$meta.parentKey] as React.Key[]

      while (parentKeys.length) {
        const parentKey = parentKeys.shift()

        if (!parentKey) {
          return targetNode
        }

        const parentNode = treePlain[parentKey]

        if (parentNode.$meta.parentKey) {
          parentKeys.push(parentNode.$meta.parentKey)
        }

        targetNode = parentNode
      }

      return targetNode
    },
    [treePlain]
  )

  const getFlattenedTree = React.useCallback(
    () =>
      Object.values(treePlain).filter(
        ({ $meta: { id, parentKey } }) =>
          !parentKey || getIsExpanded(id) || getIsExpanded(parentKey)
      ),
    [treePlain, getIsExpanded]
  )

  return {
    tree,
    getIsExpanded,
    toggleExpand,
    getIsChecked,
    getIsIndeterminate,
    toggleCheck,
    nodeKeys,
    getIsSelected,
    toggleSelect,
    toggleSelectRange,
    getTreeNodeData,
    getFirstSelectedKey,
    getLastSelectedKey,
    getIsLoading,
    getTreeNodeVisibleChildren,
    getParentRootNode,
    getFlattenedTree,
  }
}