From 7e915ad39f9af475f81477b079932b0ddcce81ee Mon Sep 17 00:00:00 2001
From: moklick <info@moritzklack.com>
Date: Fri, 28 Jan 2022 14:48:14 +0100
Subject: [PATCH] refactor(nodes): re-init when sourcePos, targetPos or type
 changes

---
 src/components/Nodes/wrapNode.tsx    | 29 +++++++++++++++++++++++++++-
 src/container/NodeRenderer/index.tsx |  6 +++---
 src/types/nodes.ts                   |  4 ++--
 3 files changed, 33 insertions(+), 6 deletions(-)

diff --git a/src/components/Nodes/wrapNode.tsx b/src/components/Nodes/wrapNode.tsx
index 6c9e74c0..a90ce271 100644
--- a/src/components/Nodes/wrapNode.tsx
+++ b/src/components/Nodes/wrapNode.tsx
@@ -11,6 +11,7 @@ const selector = (s: ReactFlowState) => ({
   addSelectedNodes: s.addSelectedNodes,
   updateNodePosition: s.updateNodePosition,
   unselectNodesAndEdges: s.unselectNodesAndEdges,
+  updateNodeDimensions: s.updateNodeDimensions,
 });
 
 export default (NodeComponent: ComponentType<NodeProps>) => {
@@ -51,8 +52,14 @@ export default (NodeComponent: ComponentType<NodeProps>) => {
     noDragClassName,
   }: WrapNodeProps) => {
     const store = useStoreApi();
-    const { addSelectedNodes, unselectNodesAndEdges, updateNodePosition } = useStore(selector, shallow);
+    const { addSelectedNodes, unselectNodesAndEdges, updateNodePosition, updateNodeDimensions } = useStore(
+      selector,
+      shallow
+    );
     const nodeElement = useRef<HTMLDivElement>(null);
+    const prevSourcePosition = useRef(sourcePosition);
+    const prevTargetPosition = useRef(targetPosition);
+    const prevType = useRef(type);
     const hasPointerEvents = isSelectable || isDraggable || onClick || onMouseEnter || onMouseMove || onMouseLeave;
     const nodeStyle: CSSProperties = useMemo(
       () => ({
@@ -186,6 +193,26 @@ export default (NodeComponent: ComponentType<NodeProps>) => {
       }
     }, [hidden]);
 
+    useEffect(() => {
+      // when the user programmatically changes the source or handle position, we re-initialize the node
+      const typeChanged = prevType.current !== type;
+      const sourcePosChanged = prevSourcePosition.current !== sourcePosition;
+      const targetPosChanged = prevTargetPosition.current !== targetPosition;
+
+      if (nodeElement.current && (typeChanged || sourcePosChanged || targetPosChanged)) {
+        if (typeChanged) {
+          prevType.current = type;
+        }
+        if (sourcePosChanged) {
+          prevSourcePosition.current = sourcePosition;
+        }
+        if (targetPosChanged) {
+          prevTargetPosition.current = targetPosition;
+        }
+        updateNodeDimensions([{ id, nodeElement: nodeElement.current, forceUpdate: true }]);
+      }
+    }, [id, type, sourcePosition, targetPosition]);
+
     if (hidden) {
       return null;
     }
diff --git a/src/container/NodeRenderer/index.tsx b/src/container/NodeRenderer/index.tsx
index 6977ec4d..19078f1b 100644
--- a/src/container/NodeRenderer/index.tsx
+++ b/src/container/NodeRenderer/index.tsx
@@ -3,7 +3,7 @@ import shallow from 'zustand/shallow';
 
 import useVisibleNodes from '../../hooks/useVisibleNodes';
 import { useStore } from '../../store';
-import { Node, NodeTypesType, ReactFlowState, WrapNodeProps } from '../../types';
+import { Node, NodeTypesType, Position, ReactFlowState, WrapNodeProps } from '../../types';
 
 interface NodeRendererProps {
   nodeTypes: NodeTypesType;
@@ -87,8 +87,8 @@ const NodeRenderer = (props: NodeRendererProps) => {
             style={node.style}
             type={nodeType}
             data={node.data}
-            sourcePosition={node.sourcePosition}
-            targetPosition={node.targetPosition}
+            sourcePosition={node.sourcePosition || Position.Bottom}
+            targetPosition={node.targetPosition || Position.Top}
             hidden={node.hidden}
             xPos={node.positionAbsolute?.x ?? 0}
             yPos={node.positionAbsolute?.y ?? 0}
diff --git a/src/types/nodes.ts b/src/types/nodes.ts
index 28ecf914..b88a82d1 100644
--- a/src/types/nodes.ts
+++ b/src/types/nodes.ts
@@ -78,8 +78,8 @@ export interface WrapNodeProps<T = any> {
   onNodeDragStop?: NodeMouseHandler;
   style?: CSSProperties;
   className?: string;
-  sourcePosition?: Position;
-  targetPosition?: Position;
+  sourcePosition: Position;
+  targetPosition: Position;
   hidden?: boolean;
   snapToGrid?: boolean;
   snapGrid?: SnapGrid;