From 2d2d7ca1a6c8b4a64f5c8e42f6e1f83ff4d01a7e Mon Sep 17 00:00:00 2001
From: Liuzj <530604689@qq.com>
Date: Wed, 8 Feb 2023 14:17:28 +0800
Subject: [PATCH] fix(pro:tree): expandAll button state is wrong (#1429)

---
 packages/components/tree/src/Tree.tsx |  7 ++++
 packages/components/tree/src/types.ts |  3 ++
 packages/pro/tree/src/ProTree.tsx     | 55 +++++++++++++++++++++++----
 3 files changed, 58 insertions(+), 7 deletions(-)

diff --git a/packages/components/tree/src/Tree.tsx b/packages/components/tree/src/Tree.tsx
index 0126e0e2e..57dabb42c 100644
--- a/packages/components/tree/src/Tree.tsx
+++ b/packages/components/tree/src/Tree.tsx
@@ -146,6 +146,10 @@ export default defineComponent({
       return mergedNodeMap.value.get(key)?.rawNode
     }
 
+    const _getFlattedNodes = () => {
+      return flattedNodes.value
+    }
+
     expose({
       focus,
       blur,
@@ -153,6 +157,9 @@ export default defineComponent({
       collapseAll: expandableContext.collapseAll,
       scrollTo,
       getNode,
+
+      // private
+      _getFlattedNodes,
     })
 
     const handleScrolledChange = (startIndex: number, endIndex: number, visibleNodes: MergedNode[]) => {
diff --git a/packages/components/tree/src/types.ts b/packages/components/tree/src/types.ts
index 07238d201..9a94c6a81 100644
--- a/packages/components/tree/src/types.ts
+++ b/packages/components/tree/src/types.ts
@@ -127,6 +127,9 @@ export interface TreeBindings<K = VKey> {
    * @returns node
    */
   getNode: (key: K) => TreeNode<K> | undefined
+
+  //private
+  _getFlattedNodes: () => MergedNode[]
 }
 export type TreeComponent = DefineComponent<Omit<HTMLAttributes, keyof TreePublicProps> & TreePublicProps, TreeBindings>
 export type TreeInstance = InstanceType<DefineComponent<TreeProps, TreeBindings>>
diff --git a/packages/pro/tree/src/ProTree.tsx b/packages/pro/tree/src/ProTree.tsx
index ff85b81a8..7389bb7d6 100644
--- a/packages/pro/tree/src/ProTree.tsx
+++ b/packages/pro/tree/src/ProTree.tsx
@@ -5,11 +5,11 @@
  * found in the LICENSE file at https://github.com/IDuxFE/idux/blob/main/LICENSE
  */
 
-import { computed, defineComponent, normalizeClass, normalizeStyle, ref } from 'vue'
+import { computed, defineComponent, normalizeClass, normalizeStyle, onMounted, ref, watch } from 'vue'
 
 import { isNil } from 'lodash-es'
 
-import { type VKey, callEmit, useControlledProp } from '@idux/cdk/utils'
+import { type VKey, callEmit, useControlledProp, useState } from '@idux/cdk/utils'
 import { ɵHeader } from '@idux/components/_private/header'
 import { ɵInput } from '@idux/components/_private/input'
 import { IxButton } from '@idux/components/button'
@@ -32,12 +32,30 @@ export default defineComponent({
     const mergedClearIcon = computed(() => props.clearIcon ?? config.clearIcon)
     const mergedCollapseIcon = computed(() => props.collapseIcon ?? config.collapseIcon)
 
-    const expandAllBtnStatus = ref(false)
+    /**
+     * true: 显示展开全部  false：显示收起全部
+     */
+    const [expandAllBtnStatus, setExpandAllBtnStatus] = useState(false)
     const treeRef = ref<TreeInstance>()
 
     const [searchValue, setSearchValue] = useControlledProp(props, 'searchValue', '')
+    const [expandedKeys, setExpandedKeys] = useControlledProp(props, 'expandedKeys', [])
     const [collapsed, setCollapsed] = useControlledProp(props, 'collapsed')
 
+    onMounted(() => {
+      setExpandAllBtnStatusWithFlattedNodes()
+    })
+
+    watch(
+      expandedKeys,
+      () => {
+        setExpandAllBtnStatusWithFlattedNodes()
+      },
+      {
+        flush: 'post',
+      },
+    )
+
     const classes = computed(() => {
       const prefixCls = mergedPrefixCls.value
 
@@ -55,8 +73,8 @@ export default defineComponent({
 
     const handleExpandAll = () => {
       const currStatus = expandAllBtnStatus.value
-      expandAllBtnStatus.value = !currStatus
-      expandAllBtnStatus.value ? treeRef.value?.collapseAll() : treeRef.value?.expandAll()
+      setExpandAllBtnStatus(!currStatus)
+      currStatus ? treeRef.value?.expandAll() : treeRef.value?.collapseAll()
     }
 
     const handleCollapsed = () => {
@@ -75,6 +93,29 @@ export default defineComponent({
       callEmit(props.onClear, evt)
     }
 
+    /**
+     * 根据当前的树状态来动态改变【展开/收起】按钮的状态，例如搜索或者改变 expandedKeys；
+     * 若当前渲染的节点中存在展开的节点则按钮的操作为收起全部
+     */
+    const setExpandAllBtnStatusWithFlattedNodes = () => {
+      const nodes = treeRef.value?._getFlattedNodes() ?? []
+      const expandedKeysLength = expandedKeys.value.length
+
+      let nextStatus = !expandAllBtnStatus.value
+
+      // 若不存在expandedKeys，则需要将按钮设置为全部展开
+      if (!expandedKeysLength) {
+        nextStatus = true
+      } else {
+        // 若当前渲染的节点有一个存在于expandedKeys中时，则按钮需要设置收起全部
+        nextStatus = !nodes.some(node => {
+          return expandedKeys.value.includes(node.key)
+        })
+      }
+
+      setExpandAllBtnStatus(nextStatus)
+    }
+
     expose({
       collapseAll: () => treeRef.value?.collapseAll(),
       expandAll: () => treeRef.value?.expandAll(),
@@ -85,7 +126,7 @@ export default defineComponent({
       const prefixCls = mergedPrefixCls.value
       const treeProps = {
         checkedKeys: props.checkedKeys,
-        expandedKeys: props.expandedKeys,
+        expandedKeys: expandedKeys.value,
         selectedKeys: props.selectedKeys,
         loadedKeys: props.loadedKeys,
 
@@ -133,7 +174,7 @@ export default defineComponent({
         onScrolledBottom: props.onScrolledBottom,
         onScrolledChange: props.onScrolledChange,
         'onUpdate:checkedKeys': props['onUpdate:checkedKeys'],
-        'onUpdate:expandedKeys': props['onUpdate:expandedKeys'],
+        'onUpdate:expandedKeys': setExpandedKeys,
         'onUpdate:selectedKeys': props['onUpdate:selectedKeys'],
         'onUpdate:loadedKeys': props['onUpdate:loadedKeys'],
       }