labelmain-demo/components/tree-select/VirtualizedTree.tsx

// VirtualizedTree.tsx
import {
  RenderTreeNodePayload,
  TreeNodeData,
  UseTreeReturnType,
} from "@mantine/core"
import React, { useMemo } from "react"
import { List, RowComponentProps } from "react-window"
import "./mantine-scrollbar.css"

export interface VirtualizedTreeProps {
  /** Mantine tree controller (from useTree()) — optional, if not provided this component won't control state */
  tree: UseTreeReturnType

  /** Mantine tree data */
  data: TreeNodeData[]

  /** renderNode keeps same signature as Mantine: receives RenderTreeNodePayload */
  renderNode: (payload: RenderTreeNodePayload) => React.ReactNode

  /** pixels to indent per level (matches Mantine's levelOffset) */
  levelOffset?: number

  /** react-window container height */
  height?: string | number

  /** react-window item (row) height */
  itemSize?: number

  /** overscan rows */
  overscanCount?: number
}

/** 小工具：把 TreeNodeData 扁平化为 [node, level] 列表，依据 tree.expandedState */
function flattenByExpanded(
  nodes: TreeNodeData[],
  expandedState: Record<string, boolean>,
  level = 0
): { node: TreeNodeData; level: number }[] {
  const res: { node: TreeNodeData; level: number }[] = []

  for (const node of nodes) {
    res.push({ node, level })
    const value = node.value
    const isExpanded = !!expandedState?.[String(value)]
    if (isExpanded && node.children && node.children.length > 0) {
      res.push(...flattenByExpanded(node.children, expandedState, level + 1))
    }
  }

  return res
}

/** 行渲染所需的闭包数据，通过 react-window 的 rowProps 传入 */
interface RowData {
  flattened: { node: TreeNodeData; level: number }[]
  tree: UseTreeReturnType
  renderNode: (payload: RenderTreeNodePayload) => React.ReactNode
}

/**
 * 单行渲染组件 —— 提到模块作用域，避免每次父组件渲染都生成新的组件实例
 * （否则 react-window 会在每次渲染时 remount 所有行，销毁行内状态）
 */
function Row({
  index,
  style,
  flattened,
  tree,
  renderNode,
}: RowComponentProps<RowData>) {
  const { node, level } = flattened[index]

  // 构造 Mantine 的 RenderTreeNodePayload
  const nodeValue = String(node.value)
  const hasChildren = Array.isArray(node.children) && node.children.length > 0
  const expanded = !!(tree && (tree as any).expandedState?.[nodeValue])
  const selected = !!(
    tree && (tree as any).selectedState?.includes?.(nodeValue)
  )
  const hovered = tree ? (tree as any).hoveredNode === nodeValue : false

  // elementProps: 与 Mantine Tree 内部的 props shape 对齐，外部的 renderNode 会 spread 这个对象
  const elementProps = {
    className: "",
    style: {
      display: "flex",
      alignItems: "center",
      width: "100%",
    } as React.CSSProperties,
    onClick: (e: React.MouseEvent) => {
      // 保持与 Mantine Tree 一致的基础行为：选中、展开切换由 tree 控制
      // 点击行为在业务端通常由 renderNode 决定（例如中间区域用于 toggleExpanded）
      // 这里我们调用 tree.toggleExpanded 以保证基础交互一致（不会阻止外部覆盖）
      if (tree && typeof tree.toggleExpanded === "function") {
        tree.toggleExpanded(nodeValue)
      }
      false && console.log(e)
    },
    "data-selected": selected as boolean | undefined,
    "data-value": nodeValue,
    "data-hovered": hovered as boolean | undefined,
  }

  const payload: RenderTreeNodePayload = {
    level,
    expanded,
    hasChildren,
    selected,
    node,
    tree,
    elementProps,
  }

  // paddingLeft 模拟 Mantine 的 levelOffset 行为（renderNode 可使用 elementProps.style 覆盖）
  const paddedStyle: React.CSSProperties = {
    ...style,
    boxSizing: "border-box",
    overflow: "hidden",
  }

  // 将 style 作为容器传入，Mantine 的 renderNode 也会接受 elementProps.style
  return <div style={paddedStyle}>{renderNode(payload)}</div>
}

/**
 * VirtualizedTree - Mantine v8 compatible virtualized Tree
 *  - Accepts tree (useTree() result) and data: TreeNodeData[]
 *  - Passes RenderTreeNodePayload to renderNode exactly like Mantine
 *  - Uses react-window FixedSizeList for virtualization
 */
export function VirtualizedTree({
  tree,
  data,
  renderNode,
  // levelOffset = 24,
  height = 400,
  itemSize = 36,
  overscanCount = 5,
}: VirtualizedTreeProps) {
  // 扁平化数据：根据 tree.expandedState 决定哪些子节点可见
  const flattened = useMemo(
    () =>
      flattenByExpanded(data, (tree && (tree as any).expandedState) || {}, 0),
    [data, tree]
  )

  return (
    <List
      rowCount={flattened.length}
      style={{ height, width: "100%" }}
      className="custom-scrollbar"
      rowHeight={itemSize}
      rowComponent={Row}
      rowProps={{ flattened, tree, renderNode }}
      overscanCount={overscanCount}></List>
  )
}