const Tree = forwardRef<TreeInstance, TreeProps>((props, ref) => {
  const { checkable, treeData, checkedKeys, defaultExpandAll, multiple, singleSelected, selectable = true, selectedKeys: propsSelectedKeys, onCheck, onSelect, onRightClick, ...ret } = props
  
  // ... 省略部分內(nèi)容，只展示核心結(jié)構(gòu)

  // 遞歸渲染 tree 的列表 
  const renderTreeList = (list?: TreeNode[]) => {
    // checkTree 的說明見下面
    if(!checkTree) return null
    return list?.map(item => {
      const checkItem = checkTree![item.key]
      return (
        <div key={item.key} className={`node`}>
          <div className={`node-content`}>
            // checkItem.show 用來判斷展開
            <div>↓</div>
            // checkItem.checked 用來處理是否 check
            <Checkbox />
            <div>{item.title}</div>
          </div>
          <div className='children'>
            {renderTreeList(item.children)}
          </div>
        </div>
      )
    })
  }

  return (
    <div className={`${classPrefix} ${ret.className ?? ''}`} style={ret.style}>
      {renderTreeList(treeData)}
    </div>
  )
})

export type CheckTreeItem = {
  /** 父節(jié)點(diǎn)的 key 值 */
  parentKey?: string
  /** 子節(jié)點(diǎn)的 key 數(shù)組 */
  childKeys?: string[]
  /** 是否展開 */
  show: boolean
  /** 是否選中 */
  checked: boolean
  /** 是否有 checkbox */
  checkable?: boolean
  /** 禁用 checkbox */
  disableCheckbox?: boolean
  /** 禁止整個節(jié)點(diǎn)的選擇 */
  disabled?: boolean
}

export type CheckTree = Record<string, CheckTreeItem>
// ...
const [checkTree, setCheckTree] = useState<CheckTree>();

// ...
useEffect(() => {
  if(!treeData?.length) return
  const generateCheckTree = (list: TreeNode[], parentKey?: string) => {
    return list?.reduce((pre, cur) => {
      // checkedKeys 就是默認(rèn)傳入 check 項，用于默認(rèn)是否勾選
      const curChecked = Boolean(checkedKeys?.includes(cur.key));
      pre[cur.key] = {
        // 默認(rèn)是否展開該樹形結(jié)構(gòu)
        show: !!defaultExpandAll, 
        checked: curChecked, 
        parentKey,
      }
      // 一些屬性的默認(rèn)值
      if(cur.checkable) pre[cur.key].checkable = true
      if(cur.disableCheckbox) pre[cur.key].disableCheckbox = true
      if(cur.disabled) pre[cur.key].disabled = true
      // 有孩子節(jié)點(diǎn)就遞歸調(diào)用，生成數(shù)據(jù)
      if(cur.children?.length) {
        pre[cur.key].childKeys = cur.children.map(c => c.key)
        const treeChild = generateCheckTree(cur.children, cur.key)
        pre = {...pre, ...treeChild}
      }
      return pre
    }, {} as CheckTree)
  }
  const state = generateCheckTree(treeData)
  setCheckTree(state)
  // ...
}, [treeData])

// ...
{(checkable && item.checkable !== false) && (
  <CheckBox 
    checked={checkItem.checked} 
    disabled={item.disabled || item.disableCheckbox}
    // 先忽略，用來判斷當(dāng)前是否有孩子節(jié)點(diǎn)被選中了，true 則代表需要展示 checkbox 的半選樣式
    indeterminate={getIsSomeChildCheck(checkItem, checkTree)}
    onChange={() => {
      if(item.disabled || item.disableCheckbox) return
      onNodeCheck(item.key)
    }} 
  />   
)}

/** 點(diǎn)擊選中節(jié)點(diǎn) */
const onNodeCheck = (key: string) => {
  const checkItem = checkTree![key]
  const curChecked = !checkItem.checked
  checkItem.checked = curChecked;

  // 先忽略，用來判斷是否是單選的
  if(singleSelected) onSingleCheck(key, curChecked)
  else onCheckChildAndParent(key, curChecked)
  setCheckTree({...checkTree})

  // 先忽略，用來獲取當(dāng)前 check 的所有 key 值
  const keys = getCheckKeys(checkTree!)

  // check 觸發(fā)的組件回調(diào)
  onCheck?.(keys, {
    key, 
    // 這步判斷主要是單選時，選擇父節(jié)點(diǎn)時只會選中其子節(jié)點(diǎn)
    checked: keys.includes(key) ? curChecked : false,  
    parentKeys: getParentKeys(key, checkTree!),
    treeDataItem: getTreeDataItem(key, treeData),
  })
}

/** 處理父子節(jié)點(diǎn)的選中狀態(tài) */
const onCheckChildAndParent = (key: string, curChecked: boolean, cTree = checkTree!) => {
  const checkItem = cTree[key];

  // 全選/不選所有子節(jié)點(diǎn)
  (function checkAllChild(childKeys?: string[]) {
    childKeys?.forEach(childKey => {
      cTree[childKey].checked = curChecked
      checkAllChild(cTree[childKey].childKeys)
    })
  })(checkItem.childKeys);

  // 處理父節(jié)點(diǎn)的選中狀態(tài)
  (function checkAllParent(parentKey?: string) {
    if(!parentKey) return
    if(!curChecked) { // 取消所有父節(jié)點(diǎn)的選中
      cTree[parentKey].checked = false
      checkAllParent(cTree[parentKey].parentKey)
    } else { // 將所有子節(jié)點(diǎn)被全選的父節(jié)點(diǎn)也選中
      const isSiblingCheck = !!cTree[parentKey].childKeys?.every(childKey => cTree[childKey].checked)
      if(isSiblingCheck) { // 判斷兄弟節(jié)點(diǎn)是否也全被選中
        cTree[parentKey].checked = true
        checkAllParent(cTree[parentKey].parentKey)
      }
    }
  })(checkItem.parentKey);

  // 同層單選時，使兄弟節(jié)點(diǎn)取消選中
  if(singleSelected && curChecked) {
    const keys = cTree[key].parentKey ? cTree[cTree[key].parentKey!].childKeys : firstNodeKeys
    keys?.forEach(siblingKey => {
      if(siblingKey !== key) {
        cTree[siblingKey].checked = false
      }
    })
  }
}

<div 
  className={`node-children`} 
  // height: fit-content; 無法觸發(fā)過渡效果，需要準(zhǔn)確的值
  // 也可通過 maxHeight 設(shè)置一個很大的值來解決，但值過大又會使過度效果難看，所以這里需要獲取一個準(zhǔn)確的高度
  style={{maxHeight: checkItem.show ? `${getTreeChildHeight(item.children!)}px` : 0}}
>
  {renderTreeList(item.children)}
</div>

.node-children {
  padding-left: 24px;
  overflow-y: hidden;
  transition: max-height 0.3s ease;
}

/** 標(biāo)題的最小高度 */
const TITLE_MIN_HEIGHT = 24;
/** 每個標(biāo)題的下邊距 */
const TITLE_MB = 6;

// 等待樹形結(jié)構(gòu)渲染完畢，獲取 title 的高度
useEffect(() => {
  if(!checkTree || !isTreeRender.current) return
  const info: TitleNodeInfo = {};
  for(let key in checkTree) {
    // 每個標(biāo)題渲染的內(nèi)容，都要根據(jù) key 給一個唯一的類名。
    const titleNode = document.querySelector(`.node-title-${key}`)
    if(titleNode) {
      info[key] = {height: Math.max(titleNode.clientHeight, TITLE_MIN_HEIGHT) + TITLE_MB} 
    }
  }
  setTitleNodeInfo(info)
  isTreeRender.current = false
}, [checkTree])

const getTreeChildHeight = (list: TreeNode[]) => {
  return list?.reduce((pre, cur) => {
    pre += (titleNodeInfo[cur.key]?.height ?? (TITLE_MIN_HEIGHT + TITLE_MB))
    if(checkTree![cur.key].show && cur.children?.length) {
      pre += getTreeChildHeight(cur.children)
    }
    return pre
  }, 0) ?? 0
}