From e1da14cec6c320c21f43b92ba6fa318299906a14 Mon Sep 17 00:00:00 2001
From: zhouyun1 <zhouyun1@xiaomi.com>
Date: Mon, 25 Sep 2023 16:10:05 +0800
Subject: [PATCH] =?UTF-8?q?feat(tree):=20=E5=A2=9E=E5=8A=A0=20scrollTo=20?=
 =?UTF-8?q?=E6=96=B9=E6=B3=95?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .changeset/green-teachers-travel.md           |  5 ++
 .changeset/sixty-fishes-tan.md                |  5 ++
 packages/ui/tree/src/Tree.tsx                 | 17 +++++-
 packages/ui/tree/src/types.ts                 | 17 ++++++
 packages/ui/tree/stories/index.stories.tsx    |  1 +
 .../ui/tree/stories/scroll-to.stories.tsx     | 59 +++++++++++++++++++
 6 files changed, 102 insertions(+), 2 deletions(-)
 create mode 100644 .changeset/green-teachers-travel.md
 create mode 100644 .changeset/sixty-fishes-tan.md
 create mode 100644 packages/ui/tree/stories/scroll-to.stories.tsx

diff --git a/.changeset/green-teachers-travel.md b/.changeset/green-teachers-travel.md
new file mode 100644
index 000000000..390bfd561
--- /dev/null
+++ b/.changeset/green-teachers-travel.md
@@ -0,0 +1,5 @@
+---
+"@hi-ui/tree": patch
+---
+
+feat: 增加 TreeHelper,支持 scrollTo 方法
diff --git a/.changeset/sixty-fishes-tan.md b/.changeset/sixty-fishes-tan.md
new file mode 100644
index 000000000..34bcac21f
--- /dev/null
+++ b/.changeset/sixty-fishes-tan.md
@@ -0,0 +1,5 @@
+---
+"@hi-ui/hiui": patch
+---
+
+Tree feat: 增加 TreeHelper,支持 scrollTo 方法
diff --git a/packages/ui/tree/src/Tree.tsx b/packages/ui/tree/src/Tree.tsx
index 43ef482aa..f3cc97652 100644
--- a/packages/ui/tree/src/Tree.tsx
+++ b/packages/ui/tree/src/Tree.tsx
@@ -1,4 +1,4 @@
-import React, { forwardRef, useMemo, useRef } from 'react'
+import React, { forwardRef, useMemo, useRef, useImperativeHandle } from 'react'
 import { HiBaseSizeEnum } from '@hi-ui/core'
 import { cx, getPrefixCls } from '@hi-ui/classname'
 import { __DEV__ } from '@hi-ui/env'
@@ -20,9 +20,10 @@ import {
   TreeNodeRequiredProps,
   TreeNodeEventData,
   FlattedTreeNodeData,
+  TreeHelper,
 } from './types'
 import { TreeProvider } from './context'
-import VirtualList from '@hi-ui/virtual-list'
+import VirtualList, { ListRef } from '@hi-ui/virtual-list'
 import { MotionTreeNode } from './MotionTreeNode'
 import { TreeNode } from './TreeNode'
 import { useLatestCallback } from '@hi-ui/use-latest'
@@ -86,6 +87,7 @@ export const Tree = forwardRef<HTMLUListElement | null, TreeProps>(
       checkedMode = 'ALL',
       expandOnSelect,
       size = 'lg',
+      innerRef,
       ...rest
     },
     ref
@@ -217,6 +219,12 @@ export const Tree = forwardRef<HTMLUListElement | null, TreeProps>(
       return Math.ceil(height / itemHeight) + 1
     }, [virtual, height, itemHeight])
 
+    const listRef = useRef<ListRef>(null)
+
+    useImperativeHandle(innerRef, () => ({
+      scrollTo: listRef.current?.scrollTo,
+    }))
+
     const cls = cx(prefixCls, className, `${prefixCls}--size-${size}`)
 
     return (
@@ -231,6 +239,7 @@ export const Tree = forwardRef<HTMLUListElement | null, TreeProps>(
           {...rest}
         >
           <VirtualList
+            ref={listRef}
             itemKey="id"
             fullHeight={false}
             height={height}
@@ -456,6 +465,10 @@ export interface TreeProps {
    * 设置大小
    */
   size?: HiBaseSizeEnum
+  /**
+   * 提供辅助方法的内部引用
+   */
+  innerRef?: React.Ref<TreeHelper>
 }
 
 if (__DEV__) {
diff --git a/packages/ui/tree/src/types.ts b/packages/ui/tree/src/types.ts
index c9207da9c..3ae53cb29 100644
--- a/packages/ui/tree/src/types.ts
+++ b/packages/ui/tree/src/types.ts
@@ -156,3 +156,20 @@ export type TreeEditActions = {
    */
   closeMenu: () => void
 }
+
+export type ScrollAlign = 'top' | 'bottom' | 'auto'
+export type ScrollConfig =
+  | {
+      index: number
+      align?: ScrollAlign
+      offset?: number
+    }
+  | {
+      key: React.Key
+      align?: ScrollAlign
+      offset?: number
+    }
+
+export interface TreeHelper {
+  scrollTo?: (arg: number | ScrollConfig) => void
+}
diff --git a/packages/ui/tree/stories/index.stories.tsx b/packages/ui/tree/stories/index.stories.tsx
index 18e219745..d9781f46e 100644
--- a/packages/ui/tree/stories/index.stories.tsx
+++ b/packages/ui/tree/stories/index.stories.tsx
@@ -15,6 +15,7 @@ export * from './custom-title.stories'
 export * from './custom-icon.stories'
 export * from './expand-on-click.stories'
 export * from './size.stories'
+export * from './scroll-to.stories'
 
 export default {
   title: 'Data Display/Tree',
diff --git a/packages/ui/tree/stories/scroll-to.stories.tsx b/packages/ui/tree/stories/scroll-to.stories.tsx
new file mode 100644
index 000000000..d614bdd7f
--- /dev/null
+++ b/packages/ui/tree/stories/scroll-to.stories.tsx
@@ -0,0 +1,59 @@
+import React from 'react'
+import Tree, { TreeHelper } from '../src'
+import Button from '@hi-ui/button'
+
+/**
+ * @title 设置滚动位置
+ * @desc 仅在设置 height 属性即虚拟列表下有效
+ */
+export const ScrollTo = () => {
+  const treeRef = React.useRef<TreeHelper>(null)
+
+  const [treeData] = React.useState(() => {
+    function dig(path = '0', level) {
+      const list: { title: string; id: string; children: any[] }[] = []
+
+      for (let i = 0; i < 5; i += 1) {
+        const id = `${path}-${i}`
+        const treeNode = {
+          title: id,
+          id,
+          children: [] as any[],
+        }
+
+        if (level > 0) {
+          treeNode.children = dig(id, level - 1)
+        }
+
+        list.push(treeNode)
+      }
+      return list
+    }
+
+    const treeData = dig('0', 4)
+    return treeData
+  })
+
+  return (
+    <>
+      <h1>ScrollTo for Tree</h1>
+      <div className="tree-scroll-to__wrap">
+        <Button
+          onClick={() => {
+            // key 为节点 id
+            treeRef.current?.scrollTo?.({ key: '0-2-2-0-0', align: 'top' })
+          }}
+        >
+          scroll to key: 0-2-2-0-0
+        </Button>
+        <Tree
+          innerRef={treeRef}
+          height={300}
+          defaultExpandAll
+          data={treeData}
+          expandOnSelect
+        ></Tree>
+      </div>
+    </>
+  )
+}