From dc36ebf31ff35c4beafe9c959045de7fd62ae88f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Christopher=20M=C3=B6ller?= <chrtze@gmail.com>
Date: Fri, 22 Oct 2021 15:13:14 +0200
Subject: [PATCH] feat(nested-nodes): callback for parent node size change

---
 example/src/Basic/index.tsx          |  32 +++--
 src/container/NodeRenderer/index.tsx | 174 ++++++++++++++++++---------
 2 files changed, 139 insertions(+), 67 deletions(-)

diff --git a/example/src/Basic/index.tsx b/example/src/Basic/index.tsx
index 3b0cf32e..e8c650e1 100644
--- a/example/src/Basic/index.tsx
+++ b/example/src/Basic/index.tsx
@@ -1,4 +1,4 @@
-import { useState, MouseEvent, useCallback } from 'react';
+import { useState, useMemo, MouseEvent, useCallback } from 'react';
 
 import ReactFlow, {
   addEdge,
@@ -26,21 +26,23 @@ const initialNodes: Node[] = [
   {
     id: '4',
     data: { label: 'Node 4' },
-    position: { x: 400, y: 200 },
+    position: { x: 0, y: 0 },
     className: 'light',
     style: { backgroundColor: 'rgba(255, 0, 0, .2)' },
+    width: 0,
+    height: 0,
   },
   {
     id: '4a',
     data: { label: 'Node 4a' },
-    position: { x: 400, y: 400 },
+    position: { x: 70, y: 200 },
     className: 'light',
     parentNode: '4',
   },
   {
     id: '4b',
     data: { label: 'Node 4b' },
-    position: { x: 500, y: 500 },
+    position: { x: 0, y: 0 },
     className: 'light',
     style: { backgroundColor: 'rgba(255, 0, 0, .2)' },
     parentNode: '4',
@@ -48,14 +50,14 @@ const initialNodes: Node[] = [
   {
     id: '4b1',
     data: { label: 'Node 4b1' },
-    position: { x: 450, y: 450 },
+    position: { x: 150, y: 270 },
     className: 'light',
     parentNode: '4b',
   },
   {
     id: '4b2',
     data: { label: 'Node 4b2' },
-    position: { x: 550, y: 550 },
+    position: { x: 420, y: 370 },
     className: 'light',
     parentNode: '4b',
   },
@@ -64,6 +66,13 @@ const initialNodes: Node[] = [
 const initialEdges: Edge[] = [
   { id: 'e1-2', source: '1', target: '2', animated: true },
   { id: 'e1-3', source: '1', target: '3' },
+  { id: 'e2-4a', source: '2', target: '4a', animated: true },
+  { id: 'e3-4', source: '3', target: '4' },
+  { id: 'e3-4b', source: '3', target: '4b' },
+  { id: 'e3-4b2', source: '3', target: '4b2' },
+  { id: 'e4a-4b1', source: '4a', target: '4b1' },
+  { id: 'e4a-4b2', source: '4a', target: '4b2' },
+  { id: 'e4b1-4b2', source: '4b1', target: '4b2' },
 ];
 
 const BasicFlow = () => {
@@ -120,9 +129,18 @@ const BasicFlow = () => {
     setEdges((es) => applyEdgeChanges(changes, es));
   }, []);
 
+  const nodesWithLabel = useMemo(
+    () =>
+      nodes.map((n) => {
+        n.data = { ...n.data, label: `${n.width}x${n.height}` };
+        return n;
+      }),
+    [nodes]
+  );
+
   return (
     <ReactFlow
-      nodes={nodes}
+      nodes={nodesWithLabel}
       edges={edges}
       onLoad={onLoad}
       onNodesChange={onNodesChange}
diff --git a/src/container/NodeRenderer/index.tsx b/src/container/NodeRenderer/index.tsx
index 72bae44a..002b015a 100644
--- a/src/container/NodeRenderer/index.tsx
+++ b/src/container/NodeRenderer/index.tsx
@@ -1,4 +1,4 @@
-import React, { memo, useMemo, ComponentType, MouseEvent, Fragment } from 'react';
+import React, { memo, useMemo, ComponentType, MouseEvent, Fragment, useEffect } from 'react';
 import shallow from 'zustand/shallow';
 
 import { useStore } from '../../store';
@@ -44,6 +44,107 @@ interface NodesProps extends NodeRendererProps {
   parentId?: string;
 }
 
+interface NodeProps extends NodesProps {
+  nodes: Node[];
+  node: Node;
+  nodeType: string;
+}
+
+function Node({
+  nodes,
+  node,
+  nodeType,
+  isDraggable,
+  resizeObserver,
+  scale,
+  snapToGrid,
+  snapGrid,
+  nodesDraggable,
+  nodesConnectable,
+  elementsSelectable,
+  recursionDepth,
+  ...props
+}: NodeProps) {
+  const onNodesChange = useStore((s) => s.onNodesChange);
+  const NodeComponent = (props.nodeTypes[nodeType] || props.nodeTypes.default) as ComponentType<WrapNodeProps>;
+  const isNodeDraggable =
+    typeof isDraggable !== 'undefined'
+      ? isDraggable
+      : !!(node.draggable || (nodesDraggable && typeof node.draggable === 'undefined'));
+  const isSelectable = !!(node.selectable || (elementsSelectable && typeof node.selectable === 'undefined'));
+  const isConnectable = !!(node.connectable || (nodesConnectable && typeof node.connectable === 'undefined'));
+  const isInitialized =
+    node.width !== null &&
+    node.height !== null &&
+    typeof node.width !== 'undefined' &&
+    typeof node.height !== 'undefined';
+
+  const childNodes = useMemo(() => nodes.filter((n) => n.parentNode === node.id && !n.isHidden), [nodes, node.id]);
+  const childRect = useMemo(() => getRectOfNodes(childNodes), [childNodes]);
+  const isParentNode = !!childNodes.length;
+
+  node.style = useMemo(() => {
+    if (isParentNode) {
+      return {
+        ...node.style,
+        width: Math.floor(childRect.width) + 20,
+        height: Math.floor(childRect.height) + 20,
+        boxSizing: 'border-box',
+      };
+    }
+    return node.style;
+  }, [childRect.width, childRect.height, isParentNode, node.style]);
+
+  useEffect(() => {
+    if (onNodesChange && isParentNode) {
+      onNodesChange([
+        {
+          type: 'position',
+          id: node.id,
+          position: { x: Math.floor(childRect.x) - 10, y: Math.floor(childRect.y) - 10 },
+        },
+      ]);
+    }
+  }, [childRect.width, childRect.height, childRect.x, childRect.y, isParentNode]);
+
+  return (
+    <NodeComponent
+      id={node.id}
+      className={node.className}
+      style={node.style}
+      type={nodeType}
+      data={node.data}
+      sourcePosition={node.sourcePosition}
+      targetPosition={node.targetPosition}
+      isHidden={node.isHidden}
+      xPos={node.position.x}
+      yPos={node.position.y}
+      isDragging={node.isDragging}
+      isInitialized={isInitialized}
+      snapGrid={snapGrid}
+      snapToGrid={snapToGrid}
+      selectNodesOnDrag={props.selectNodesOnDrag}
+      onClick={props.onNodeClick}
+      onMouseEnter={props.onNodeMouseEnter}
+      onMouseMove={props.onNodeMouseMove}
+      onMouseLeave={props.onNodeMouseLeave}
+      onContextMenu={props.onNodeContextMenu}
+      onNodeDoubleClick={props.onNodeDoubleClick}
+      onNodeDragStart={props.onNodeDragStart}
+      onNodeDrag={props.onNodeDrag}
+      onNodeDragStop={props.onNodeDragStop}
+      scale={scale}
+      isSelected={!!node.isSelected}
+      isDraggable={isNodeDraggable}
+      isSelectable={isSelectable}
+      isConnectable={isConnectable}
+      resizeObserver={resizeObserver}
+      dragHandle={node.dragHandle}
+      zIndex={3 + recursionDepth}
+    />
+  );
+}
+
 function Nodes({
   nodes,
   parentId,
@@ -70,69 +171,22 @@ function Nodes({
       console.warn(`Node type "${nodeType}" not found. Using fallback type "default".`);
     }
 
-    const NodeComponent = (props.nodeTypes[nodeType] || props.nodeTypes.default) as ComponentType<WrapNodeProps>;
-    const isNodeDraggable =
-      typeof isDraggable !== 'undefined'
-        ? isDraggable
-        : !!(node.draggable || (nodesDraggable && typeof node.draggable === 'undefined'));
-    const isSelectable = !!(node.selectable || (elementsSelectable && typeof node.selectable === 'undefined'));
-    const isConnectable = !!(node.connectable || (nodesConnectable && typeof node.connectable === 'undefined'));
-    const isInitialized =
-      node.width !== null &&
-      node.height !== null &&
-      typeof node.width !== 'undefined' &&
-      typeof node.height !== 'undefined';
-
-    const childNodes = nodes.filter((n) => n.parentNode === node.id && !n.isHidden);
-
-    if (childNodes.length) {
-      const childRect = getRectOfNodes(childNodes);
-      node.position = node.isDragging
-        ? node.position
-        : { x: Math.round(childRect.x) - 10, y: Math.round(childRect.y) - 10 };
-      node.style = {
-        ...node.style,
-        width: Math.round(childRect.width) + 20,
-        height: Math.round(childRect.height) + 20,
-        boxSizing: 'border-box',
-      };
-    }
-
     return (
       <Fragment key={node.id}>
-        <NodeComponent
-          id={node.id}
-          className={node.className}
-          style={node.style}
-          type={nodeType}
-          data={node.data}
-          sourcePosition={node.sourcePosition}
-          targetPosition={node.targetPosition}
-          isHidden={node.isHidden}
-          xPos={node.position.x}
-          yPos={node.position.y}
-          isDragging={node.isDragging}
-          isInitialized={isInitialized}
-          snapGrid={snapGrid}
+        <Node
+          node={node}
+          nodes={nodes}
+          nodeType={nodeType}
+          parentId={node.id}
           snapToGrid={snapToGrid}
-          selectNodesOnDrag={props.selectNodesOnDrag}
-          onClick={props.onNodeClick}
-          onMouseEnter={props.onNodeMouseEnter}
-          onMouseMove={props.onNodeMouseMove}
-          onMouseLeave={props.onNodeMouseLeave}
-          onContextMenu={props.onNodeContextMenu}
-          onNodeDoubleClick={props.onNodeDoubleClick}
-          onNodeDragStart={props.onNodeDragStart}
-          onNodeDrag={props.onNodeDrag}
-          onNodeDragStop={props.onNodeDragStop}
+          snapGrid={snapGrid}
+          nodesDraggable={nodesDraggable}
+          nodesConnectable={nodesConnectable}
+          resizeObserver={resizeObserver}
+          elementsSelectable={elementsSelectable}
           scale={scale}
-          isSelected={!!node.isSelected}
-          isDraggable={isNodeDraggable}
-          isSelectable={isSelectable}
-          isConnectable={isConnectable}
-          resizeObserver={childNodes.length ? resizeObserver : null}
-          dragHandle={node.dragHandle}
-          zIndex={3 + recursionDepth}
+          recursionDepth={recursionDepth}
+          {...props}
         />
         <MemoizedNodes
           nodes={nodes}