From e02e3703c1ebf8a444cdf091b090f2f72777bb81 Mon Sep 17 00:00:00 2001
From: moklick <info@moritzklack.com>
Date: Wed, 28 Aug 2024 16:17:03 +0200
Subject: [PATCH] fix(setNodes): use correct index when using setNodes for
 inserting #4455

---
 .../src/examples/ReconnectEdge/CustomNode.tsx      | 14 ++++++++++++++
 packages/react/src/utils/changes.ts                | 13 ++++++++++---
 packages/system/src/types/changes.ts               |  2 ++
 3 files changed, 26 insertions(+), 3 deletions(-)
 create mode 100644 examples/react/src/examples/ReconnectEdge/CustomNode.tsx

diff --git a/examples/react/src/examples/ReconnectEdge/CustomNode.tsx b/examples/react/src/examples/ReconnectEdge/CustomNode.tsx
new file mode 100644
index 00000000..24e8b03c
--- /dev/null
+++ b/examples/react/src/examples/ReconnectEdge/CustomNode.tsx
@@ -0,0 +1,14 @@
+import { memo } from 'react';
+import { Handle, Position } from '@xyflow/react';
+
+const styles = { padding: '10px', background: 'green' };
+
+const CustomNode = ({ data }) => (
+  <>
+    <div style={styles}>{data.label}</div>
+    <Handle type="source" position={Position.Top} id="a" />
+    <Handle type="source" position={Position.Bottom} id="b" />
+  </>
+);
+
+export default memo(CustomNode);
diff --git a/packages/react/src/utils/changes.ts b/packages/react/src/utils/changes.ts
index df215340..f00fa7b3 100644
--- a/packages/react/src/utils/changes.ts
+++ b/packages/react/src/utils/changes.ts
@@ -19,10 +19,11 @@ function applyChanges(changes: any[], elements: any[]): any[] {
   // By storing a map of changes for each element, we can a quick lookup as we
   // iterate over the elements array!
   const changesMap = new Map<any, any[]>();
+  const newItems: any[] = [];
 
   for (const change of changes) {
     if (change.type === 'add') {
-      updatedElements.push(change.item);
+      newItems.push(change);
       continue;
     } else if (change.type === 'remove' || change.type === 'replace') {
       // For a 'remove' change we can safely ignore any other changes queued for
@@ -73,6 +74,12 @@ function applyChanges(changes: any[], elements: any[]): any[] {
     updatedElements.push(updatedElement);
   }
 
+  if (newItems.length) {
+    newItems.forEach((item) => {
+      updatedElements.splice(item.index, 0, { ...item.item });
+    });
+  }
+
   return updatedElements;
 }
 
@@ -237,7 +244,7 @@ export function getElementsDiffChanges({
   const changes: any[] = [];
   const itemsLookup = new Map<string, any>(items.map((item) => [item.id, item]));
 
-  for (const item of items) {
+  for (const [index, item] of items.entries()) {
     const lookupItem = lookup.get(item.id);
     const storeItem = lookupItem?.internals?.userNode ?? lookupItem;
 
@@ -246,7 +253,7 @@ export function getElementsDiffChanges({
     }
 
     if (storeItem === undefined) {
-      changes.push({ item: item, type: 'add' });
+      changes.push({ item: item, type: 'add', index });
     }
   }
 
diff --git a/packages/system/src/types/changes.ts b/packages/system/src/types/changes.ts
index 5db0b6de..9b232733 100644
--- a/packages/system/src/types/changes.ts
+++ b/packages/system/src/types/changes.ts
@@ -32,6 +32,7 @@ export type NodeRemoveChange = {
 export type NodeAddChange<NodeType extends NodeBase = NodeBase> = {
   item: NodeType;
   type: 'add';
+  index?: number;
 };
 
 export type NodeReplaceChange<NodeType extends NodeBase = NodeBase> = {
@@ -57,6 +58,7 @@ export type EdgeRemoveChange = NodeRemoveChange;
 export type EdgeAddChange<EdgeType extends EdgeBase = EdgeBase> = {
   item: EdgeType;
   type: 'add';
+  index?: number;
 };
 
 export type EdgeReplaceChange<EdgeType extends EdgeBase = EdgeBase> = {