From e5fac45371f6df57cf5939e4278f009196bf06f4 Mon Sep 17 00:00:00 2001
From: peterkogo <peter@xyflow.com>
Date: Thu, 22 Aug 2024 16:07:37 +0200
Subject: [PATCH] fixed node extent flashing

---
 .../src/components/NodeWrapper/index.tsx      |  3 +-
 packages/react/src/store/index.ts             | 19 ++------
 packages/system/src/utils/store.ts            | 46 ++++++++++++++-----
 3 files changed, 39 insertions(+), 29 deletions(-)

diff --git a/packages/react/src/components/NodeWrapper/index.tsx b/packages/react/src/components/NodeWrapper/index.tsx
index 61e13166..5f528486 100644
--- a/packages/react/src/components/NodeWrapper/index.tsx
+++ b/packages/react/src/components/NodeWrapper/index.tsx
@@ -41,7 +41,7 @@ export function NodeWrapper<NodeType extends Node>({
   nodeClickDistance,
   onError,
 }: NodeWrapperProps<NodeType>) {
-  const { node, internals, isParent, origin } = useStore((s) => {
+  const { node, internals, isParent } = useStore((s) => {
     const node = s.nodeLookup.get(id)! as InternalNode<NodeType>;
     const isParent = s.parentLookup.has(id);
 
@@ -49,7 +49,6 @@ export function NodeWrapper<NodeType extends Node>({
       node,
       internals: node.internals,
       isParent,
-      origin: node.origin ?? s.nodeOrigin,
     };
   }, shallow);
 
diff --git a/packages/react/src/store/index.ts b/packages/react/src/store/index.ts
index 5990c2d7..2941e89f 100644
--- a/packages/react/src/store/index.ts
+++ b/packages/react/src/store/index.ts
@@ -95,6 +95,7 @@ const createStore = ({
           fitViewOnInitOptions,
           domNode,
           nodeOrigin,
+          nodeExtent,
           debug,
           fitViewSync,
         } = get();
@@ -104,7 +105,8 @@ const createStore = ({
           nodeLookup,
           parentLookup,
           domNode,
-          nodeOrigin
+          nodeOrigin,
+          nodeExtent
         );
 
         if (!updatedInternals) {
@@ -286,23 +288,10 @@ const createStore = ({
         triggerEdgeChanges(edgeChanges);
       },
       setNodeExtent: (nodeExtent) => {
-        const { nodeLookup } = get();
-
-        for (const [, node] of nodeLookup) {
-          const positionAbsolute = clampPosition(node.internals.positionAbsolute, nodeExtent, node.measured);
-
-          nodeLookup.set(node.id, {
-            ...node,
-            internals: {
-              ...node.internals,
-              positionAbsolute,
-            },
-          });
-        }
-
         set({
           nodeExtent,
         });
+        get().setNodes(get().nodes);
       },
       panBy: (delta): Promise<boolean> => {
         const { transform, width, height, panZoom, translateExtent } = get();
diff --git a/packages/system/src/utils/store.ts b/packages/system/src/utils/store.ts
index c0775af9..ebac335c 100644
--- a/packages/system/src/utils/store.ts
+++ b/packages/system/src/utils/store.ts
@@ -1,4 +1,4 @@
-import { infiniteExtent } from '..';
+import { Dimensions, infiniteExtent } from '..';
 import {
   NodeBase,
   CoordinateExtent,
@@ -153,6 +153,24 @@ function updateChildPosition<NodeType extends NodeBase>(
   }
 }
 
+function clampPositionToParent<NodeType extends NodeBase>(
+  childPosition: XYPosition,
+  childDimensions: Dimensions,
+  parent: InternalNodeBase<NodeType>
+) {
+  const parentDimensions = getNodeDimensions(parent);
+  const parentPosition = parent.internals.positionAbsolute;
+
+  return clampPosition(
+    childPosition,
+    [
+      [parentPosition.x, parentPosition.y],
+      [parentPosition.x + parentDimensions.width, parentPosition.y + parentDimensions.height],
+    ],
+    childDimensions
+  );
+}
+
 function calculateZ(node: NodeBase, selectedNodeZ: number) {
   return (isNumeric(node.zIndex) ? node.zIndex : 0) + (node.selected ? selectedNodeZ : 0);
 }
@@ -175,15 +193,7 @@ function calculateChildXYZ<NodeType extends NodeBase>(
   let absolutePosition = { x: parentPosition.x + position.x, y: parentPosition.y + position.y };
 
   if (childNode.extent === 'parent') {
-    const parentDimensions = getNodeDimensions(parentNode);
-    absolutePosition = clampPosition(
-      absolutePosition,
-      [
-        [parentPosition.x, parentPosition.y],
-        [parentPosition.x + parentDimensions.width, parentPosition.y + parentDimensions.height],
-      ],
-      childDimensions
-    );
+    absolutePosition = clampPositionToParent(absolutePosition, childDimensions, parentNode);
   }
 
   const childZ = calculateZ(childNode, selectedNodeZ);
@@ -287,7 +297,8 @@ export function updateNodeInternals<NodeType extends InternalNodeBase>(
   nodeLookup: NodeLookup<NodeType>,
   parentLookup: ParentLookup<NodeType>,
   domNode: HTMLElement | null,
-  nodeOrigin?: NodeOrigin
+  nodeOrigin?: NodeOrigin,
+  nodeExtent?: CoordinateExtent
 ): { changes: (NodeDimensionChange | NodePositionChange)[]; updatedInternals: boolean } {
   const viewportNode = domNode?.querySelector('.xyflow__viewport');
   let updatedInternals = false;
@@ -325,11 +336,22 @@ export function updateNodeInternals<NodeType extends InternalNodeBase>(
 
       if (doUpdate) {
         const nodeBounds = update.nodeElement.getBoundingClientRect();
+        let extent = isCoordinateExtent(node.extent) ? node.extent : nodeExtent;
+        let positionAbsolute = node.internals.positionAbsolute;
+        if (node.parentId && node.extent === 'parent') {
+          positionAbsolute = clampPositionToParent(
+            positionAbsolute,
+            getNodeDimensions(node),
+            nodeLookup.get(node.parentId)!
+          );
+        } else if (extent) {
+          positionAbsolute = clampPosition(positionAbsolute, extent, dimensions);
+        }
 
         node.measured = dimensions;
         node.internals = {
           ...node.internals,
-          positionAbsolute: getNodePositionWithOrigin(node, nodeOrigin),
+          positionAbsolute,
           handleBounds: {
             source: getHandleBounds('source', update.nodeElement, nodeBounds, zoom, node.id),
             target: getHandleBounds('target', update.nodeElement, nodeBounds, zoom, node.id),