From 3ac253c2a5cd667931ccc8df45039b1a82a25963 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=BC=A0=E9=91=AB?= <zhang.xin@choiceform.com>
Date: Tue, 21 Mar 2023 09:23:00 +0800
Subject: [PATCH] Add preview ref api for Node

---
 .../src/components/drag-preview-container.tsx |  2 +-
 .../src/components/row-container.tsx          |  4 ++--
 .../react-arborist/src/components/tree.tsx    |  2 +-
 packages/react-arborist/src/dnd/drag-hook.ts  | 19 ++++++++++++-------
 .../react-arborist/src/interfaces/tree-api.ts |  3 +--
 .../react-arborist/src/types/renderers.ts     |  1 +
 6 files changed, 18 insertions(+), 13 deletions(-)

diff --git a/packages/react-arborist/src/components/drag-preview-container.tsx b/packages/react-arborist/src/components/drag-preview-container.tsx
index 1c174fe..56e9076 100644
--- a/packages/react-arborist/src/components/drag-preview-container.tsx
+++ b/packages/react-arborist/src/components/drag-preview-container.tsx
@@ -1,5 +1,5 @@
 import { useDragLayer } from "react-dnd";
-import { useDndContext, useTreeApi } from "../context";
+import { useTreeApi } from "../context";
 import { DefaultDragPreview } from "./default-drag-preview";
 
 export function DragPreviewContainer() {
diff --git a/packages/react-arborist/src/components/row-container.tsx b/packages/react-arborist/src/components/row-container.tsx
index 2ebcd43..baacfe0 100644
--- a/packages/react-arborist/src/components/row-container.tsx
+++ b/packages/react-arborist/src/components/row-container.tsx
@@ -35,7 +35,7 @@ export const RowContainer = React.memo(function RowContainer<T>({
   const node = useFreshNode<T>(index);
 
   const el = useRef<HTMLDivElement | null>(null);
-  const dragRef = useDragHook<T>(node);
+  const { dragRef, previewRef } = useDragHook<T>(node);
   const dropRef = useDropHook(el, node);
   const innerRef = useCallback(
     (n: any) => {
@@ -76,7 +76,7 @@ export const RowContainer = React.memo(function RowContainer<T>({
 
   return (
     <Row node={node} innerRef={innerRef} attrs={rowAttrs}>
-      <Node node={node} tree={tree} style={nodeStyle} dragHandle={dragRef} />
+      <Node node={node} tree={tree} style={nodeStyle}  dragHandle={dragRef} previewHandle={previewRef} />
     </Row>
   );
 });
diff --git a/packages/react-arborist/src/components/tree.tsx b/packages/react-arborist/src/components/tree.tsx
index 82601a7..7ab8240 100644
--- a/packages/react-arborist/src/components/tree.tsx
+++ b/packages/react-arborist/src/components/tree.tsx
@@ -18,7 +18,7 @@ function TreeComponent<T>(
       <OuterDrop>
         <TreeContainer />
       </OuterDrop>
-      <DragPreviewContainer />
+      { props.renderDragPreview && <DragPreviewContainer /> }
     </TreeProvider>
   );
 }
diff --git a/packages/react-arborist/src/dnd/drag-hook.ts b/packages/react-arborist/src/dnd/drag-hook.ts
index 8e8bb7c..5715f4f 100644
--- a/packages/react-arborist/src/dnd/drag-hook.ts
+++ b/packages/react-arborist/src/dnd/drag-hook.ts
@@ -1,4 +1,4 @@
-import { useEffect } from "react";
+import { useEffect, useMemo } from "react";
 import { ConnectDragSource, useDrag } from "react-dnd";
 import { getEmptyImage } from "react-dnd-html5-backend";
 import { useTreeApi } from "../context";
@@ -9,7 +9,7 @@ import { actions as dnd } from "../state/dnd-slice";
 import { safeRun } from "../utils";
 import { ROOT_ID } from "../data/create-root";
 
-export function useDragHook<T>(node: NodeApi<T>): ConnectDragSource {
+export function useDragHook<T>(node: NodeApi<T>) {
   const tree = useTreeApi();
   const ids = tree.selectedIds;
   const [_, ref, preview] = useDrag<DragItem, DropResult, void>(
@@ -42,10 +42,15 @@ export function useDragHook<T>(node: NodeApi<T>): ConnectDragSource {
     }),
     [ids, node]
   );
+  
+   useEffect(() => {
+    if (tree.renderDragPreview) {
+      preview(getEmptyImage());
+    }
+  }, []);
 
-  useEffect(() => {
-    preview(getEmptyImage());
-  }, [preview]);
-
-  return ref;
+  return {
+    dragRef: ref,
+    previewRef: preview,
+  };
 }
diff --git a/packages/react-arborist/src/interfaces/tree-api.ts b/packages/react-arborist/src/interfaces/tree-api.ts
index e9b6f29..39e51bc 100644
--- a/packages/react-arborist/src/interfaces/tree-api.ts
+++ b/packages/react-arborist/src/interfaces/tree-api.ts
@@ -15,7 +15,6 @@ import { createRoot, ROOT_ID } from "../data/create-root";
 import { actions as visibility } from "../state/open-slice";
 import { actions as selection } from "../state/selection-slice";
 import { actions as dnd } from "../state/dnd-slice";
-import { DefaultDragPreview } from "../components/default-drag-preview";
 import { DefaultContainer } from "../components/default-container";
 import { Cursor } from "../dnd/compute-drop";
 import { Store } from "redux";
@@ -628,7 +627,7 @@ export class TreeApi<T> {
   }
 
   get renderDragPreview() {
-    return this.props.renderDragPreview || DefaultDragPreview;
+    return this.props.renderDragPreview;
   }
 
   get renderCursor() {
diff --git a/packages/react-arborist/src/types/renderers.ts b/packages/react-arborist/src/types/renderers.ts
index 356ad57..ca7c291 100644
--- a/packages/react-arborist/src/types/renderers.ts
+++ b/packages/react-arborist/src/types/renderers.ts
@@ -9,6 +9,7 @@ export type NodeRendererProps<T> = {
   node: NodeApi<T>;
   tree: TreeApi<T>;
   dragHandle?: (el: HTMLDivElement | null) => void;
+  previewHandle?: (el: HTMLDivElement | null) => void; 
   preview?: boolean;
 };