refactor(applyChanges): use specific replace events instead of a rough reset

packages/react/src/utils/changes.ts

@@ -1,4 +1,5 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
+import { EdgeLookup, NodeLookup } from '@xyflow/system';
 import type { Node, Edge, EdgeChange, NodeChange, NodeSelectionChange, EdgeSelectionChange } from '../types';
 
 export function handleParentExpand(updatedElements: any[], updateItem: any) {
@@ -53,11 +54,6 @@ export function handleParentExpand(updatedElements: any[], updateItem: any) {
 // When you drag a node for example, React Flow will send a position change update.
 // This function then applies the changes and returns the updated elements.
 function applyChanges(changes: any[], elements: any[]): any[] {
-  // we need this hack to handle the setNodes and setEdges function of the useReactFlow hook for controlled flows
-  if (changes.some((c) => c.type === 'reset')) {
-    return changes.filter((c) => c.type === 'reset').map((c) => c.item);
-  }
-
   const updatedElements: any[] = [];
   // By storing a map of changes for each element, we can a quick lookup as we
   // iterate over the elements array!
@@ -67,7 +63,7 @@ function applyChanges(changes: any[], elements: any[]): any[] {
     if (change.type === 'add') {
       updatedElements.push(change.item);
       continue;
-    } else if (change.type === 'remove') {
+    } else if (change.type === 'remove' || change.type === 'replace') {
       // For a 'remove' change we can safely ignore any other changes queued for
       // the same element, it's going to be removed anyway!
       changesMap.set(change.id, [change]);
@@ -99,6 +95,11 @@ function applyChanges(changes: any[], elements: any[]): any[] {
       continue;
     }
 
+    if (changes[0].type === 'replace') {
+      updatedElements.push({ ...changes[0].item });
+      continue;
+    }
+
     // For other types of changes, we want to start with a shallow copy of the
     // object so React knows this element has changed. Sequential changes will
     /// each _mutate_ this object, so there's only ever one copy.
@@ -245,3 +246,59 @@ export function getSelectionChanges(
 
   return changes;
 }
+
+/**
+ * This function is used to find the changes between two sets of elements.
+ * It is used to determine which nodes or edges have been added, removed or replaced.
+ *
+ * @internal
+ * @param params.items = the next set of elements (nodes or edges)
+ * @param params.lookup = a lookup map of the current store elements
+ * @returns an array of changes
+ */
+export function getElementsDiffChanges({
+  items,
+  lookup,
+}: {
+  items: Node[] | undefined;
+  lookup: NodeLookup;
+}): NodeChange[];
+export function getElementsDiffChanges({
+  items,
+  lookup,
+}: {
+  items: Edge[] | undefined;
+  lookup: EdgeLookup;
+}): EdgeChange[];
+export function getElementsDiffChanges({
+  items = [],
+  lookup,
+}: {
+  items: any[] | undefined;
+  lookup: Map<string, any>;
+}): any[] {
+  const changes: any[] = [];
+  const itemsLookup = new Map<string, any>(items.map((item) => [item.id, item]));
+
+  for (const item of items) {
+    const storeItem = lookup.get(item.id);
+
+    if (storeItem !== undefined && storeItem !== item) {
+      changes.push({ id: item.id, item: item, type: 'replace' });
+    }
+
+    if (storeItem === undefined) {
+      changes.push({ item: item, type: 'add' });
+    }
+  }
+
+  for (const [id] of lookup) {
+    const nextNode = itemsLookup.get(id);
+
+    if (nextNode === undefined) {
+      changes.push({ id, type: 'remove' });
+    }
+  }
+
+  return changes;
+}