From cb698975e4d32e12f8a2b72030e304ae6deab062 Mon Sep 17 00:00:00 2001
From: moklick <info@moritzklack.com>
Date: Tue, 16 Aug 2022 12:58:33 +0200
Subject: [PATCH] refactor(useUpdateNodeInternals): make it easier to use #2008

---
 .../src/UseUpdateNodeInternals/CustomNode.tsx | 33 +++++++++--------
 example/src/UseUpdateNodeInternals/index.tsx  | 37 ++-----------------
 src/hooks/useUpdateNodeInternals.ts           |  9 +----
 3 files changed, 23 insertions(+), 56 deletions(-)

diff --git a/example/src/UseUpdateNodeInternals/CustomNode.tsx b/example/src/UseUpdateNodeInternals/CustomNode.tsx
index da95f259e..6f0bd296b 100644
--- a/example/src/UseUpdateNodeInternals/CustomNode.tsx
+++ b/example/src/UseUpdateNodeInternals/CustomNode.tsx
@@ -1,30 +1,33 @@
-import React, { memo, FC, useMemo, CSSProperties } from 'react';
-import { Handle, Position, NodeProps } from 'react-flow-renderer';
+import React, { memo, FC, useMemo, CSSProperties, useState } from 'react';
+import { Handle, Position, NodeProps, useUpdateNodeInternals } from 'react-flow-renderer';
 
 const nodeStyles: CSSProperties = { padding: 10, border: '1px solid #ddd' };
 
-const CustomNode: FC<NodeProps> = ({ data }) => {
+const CustomNode: FC<NodeProps> = ({ id }) => {
+  const [handleCount, setHandleCount] = useState(1);
+  const updateNodeInternals = useUpdateNodeInternals();
+
   const handles = useMemo(
     () =>
-      Array.from({ length: data.handleCount }, (x, i) => {
+      Array.from({ length: handleCount }, (x, i) => {
         const handleId = `handle-${i}`;
-        return (
-          <Handle
-            key={handleId}
-            type="source"
-            position={Position.Right}
-            id={handleId}
-            style={{ top: 10 * i + data.handlePosition * 10 }}
-          />
-        );
+        return <Handle key={handleId} type="source" position={Position.Right} id={handleId} style={{ top: 10 * i }} />;
       }),
-    [data.handleCount, data.handlePosition]
+    [handleCount]
   );
 
   return (
     <div style={nodeStyles}>
       <Handle type="target" position={Position.Left} />
-      <div>output handle count: {data.handleCount}</div>
+      <div>output handle count: {handleCount}</div>
+      <button
+        onClick={() => {
+          setHandleCount((c) => c + 1);
+          updateNodeInternals(id);
+        }}
+      >
+        add handle
+      </button>
       {handles}
     </div>
   );
diff --git a/example/src/UseUpdateNodeInternals/index.tsx b/example/src/UseUpdateNodeInternals/index.tsx
index 2c90ce2d0..539681c9e 100644
--- a/example/src/UseUpdateNodeInternals/index.tsx
+++ b/example/src/UseUpdateNodeInternals/index.tsx
@@ -1,4 +1,4 @@
-import { useCallback, CSSProperties, MouseEvent } from 'react';
+import { useCallback, MouseEvent } from 'react';
 
 import ReactFlow, {
   NodeTypes,
@@ -8,26 +8,21 @@ import ReactFlow, {
   Node,
   Connection,
   Edge,
-  useUpdateNodeInternals,
   Position,
   useNodesState,
   useEdgesState,
 } from 'react-flow-renderer';
 import CustomNode from './CustomNode';
 
-const initialHandleCount = 1;
-
 const initialNodes: Node[] = [
   {
     id: '1',
     type: 'custom',
-    data: { label: 'Node 1', handleCount: initialHandleCount, handlePosition: 0 },
+    data: { label: 'Node 1' },
     position: { x: 250, y: 5 },
   },
 ];
 
-const buttonWrapperStyles: CSSProperties = { position: 'absolute', right: 10, top: 10, zIndex: 10 };
-
 const nodeTypes: NodeTypes = {
   custom: CustomNode,
 };
@@ -40,7 +35,6 @@ const UpdateNodeInternalsFlow = () => {
   const [edges, setEdges, onEdgesChange] = useEdgesState([]);
   const onConnect = (params: Edge | Connection) => setEdges((els) => addEdge(params, els));
 
-  const updateNodeInternals = useUpdateNodeInternals();
   const { project } = useReactFlow();
 
   const onPaneClick = useCallback(
@@ -56,25 +50,6 @@ const UpdateNodeInternalsFlow = () => {
       ),
     [project, setNodes]
   );
-
-  const toggleHandleCount = useCallback(() => {
-    setNodes((nds) =>
-      nds.map((node) => {
-        return { ...node, data: { ...node.data, handleCount: node.data?.handleCount === 1 ? 2 : 1 } };
-      })
-    );
-  }, [setNodes]);
-
-  const toggleHandlePosition = useCallback(() => {
-    setNodes((nds) =>
-      nds.map((node) => {
-        return { ...node, data: { ...node.data, handlePosition: node.data?.handlePosition === 0 ? 1 : 0 } };
-      })
-    );
-  }, [setNodes]);
-
-  const updateNode = useCallback(() => updateNodeInternals('1'), [updateNodeInternals]);
-
   return (
     <ReactFlow
       nodes={nodes}
@@ -84,13 +59,7 @@ const UpdateNodeInternalsFlow = () => {
       nodeTypes={nodeTypes}
       onConnect={onConnect}
       onPaneClick={onPaneClick}
-    >
-      <div style={buttonWrapperStyles}>
-        <button onClick={toggleHandleCount}>toggle handle count</button>
-        <button onClick={toggleHandlePosition}>toggle handle position</button>
-        <button onClick={updateNode}>update node internals</button>
-      </div>
-    </ReactFlow>
+    />
   );
 };
 
diff --git a/src/hooks/useUpdateNodeInternals.ts b/src/hooks/useUpdateNodeInternals.ts
index 71861036a..951103928 100644
--- a/src/hooks/useUpdateNodeInternals.ts
+++ b/src/hooks/useUpdateNodeInternals.ts
@@ -10,15 +10,10 @@ function useUpdateNodeInternals(): UpdateNodeInternals {
   const updateNodeDimensions = useStore(selector);
 
   return useCallback<UpdateNodeInternals>((id: string) => {
-    const { domNode } = store.getState();
-    if (!domNode) {
-      return;
-    }
-
-    const nodeElement = domNode.querySelector(`.react-flow__node[data-id="${id}"]`) as HTMLDivElement;
+    const nodeElement = store.getState().domNode?.querySelector(`.react-flow__node[data-id="${id}"]`) as HTMLDivElement;
 
     if (nodeElement) {
-      updateNodeDimensions([{ id, nodeElement, forceUpdate: true }]);
+      requestAnimationFrame(() => updateNodeDimensions([{ id, nodeElement, forceUpdate: true }]));
     }
   }, []);
 }