refactor defaultNodes/defaultEdges handling (#3882)

* refactor(react): cleanup defaultnodes/defaultedges handling

* chore(examples): cleanup

* chore(examples): cleanup

* fix(react): selection for uncontrolled flows

packages/react/src/store/index.ts

@@ -10,8 +10,7 @@ import {
   updateConnectionLookup,
 } from '@xyflow/system';
 
-import { applyNodeChanges, createSelectionChange, getSelectionChanges } from '../utils/changes';
-import { updateNodesAndEdgesSelections } from './utils';
+import { applyEdgeChanges, applyNodeChanges, createSelectionChange, getSelectionChanges } from '../utils/changes';
 import getInitialState from './initialState';
 import type {
   ReactFlowState,
@@ -28,19 +27,23 @@ import type {
 const createRFStore = ({
   nodes,
   edges,
+  defaultNodes,
+  defaultEdges,
   width,
   height,
   fitView,
 }: {
   nodes?: Node[];
   edges?: Edge[];
+  defaultNodes?: Node[];
+  defaultEdges?: Edge[];
   width?: number;
   height?: number;
   fitView?: boolean;
 }) =>
   createWithEqualityFn<ReactFlowState>(
     (set, get) => ({
-      ...getInitialState({ nodes, edges, width, height, fitView }),
+      ...getInitialState({ nodes, edges, width, height, fitView, defaultNodes, defaultEdges }),
       setNodes: (nodes: Node[]) => {
         const { nodeLookup, nodeOrigin, elevateNodesOnSelect } = get();
         // setNodes() is called exclusively in response to user actions:
@@ -49,7 +52,6 @@ const createRFStore = ({
         //
         // When this happens, we take the note objects passed by the user and extend them with fields
         // relevant for internal React Flow operations.
-        // TODO: consider updating the types to reflect the distinction between user-provided nodes and internal nodes.
         const nodesWithInternalData = adoptUserProvidedNodes(nodes, nodeLookup, { nodeOrigin, elevateNodesOnSelect });
 
         set({ nodes: nodesWithInternalData });
@@ -61,38 +63,6 @@ const createRFStore = ({
 
         set({ edges });
       },
-      // when the user works with an uncontrolled flow,
-      // we set a flag `hasDefaultNodes` / `hasDefaultEdges`
-      setDefaultNodesAndEdges: (nodes?: Node[], edges?: Edge[]) => {
-        const hasDefaultNodes = typeof nodes !== 'undefined';
-        const hasDefaultEdges = typeof edges !== 'undefined';
-
-        const nextState: {
-          nodes?: Node[];
-          edges?: Edge[];
-          hasDefaultNodes: boolean;
-          hasDefaultEdges: boolean;
-        } = {
-          hasDefaultNodes,
-          hasDefaultEdges,
-        };
-
-        if (hasDefaultNodes) {
-          const { nodeLookup, nodeOrigin, elevateNodesOnSelect } = get();
-          nextState.nodes = adoptUserProvidedNodes(nodes, nodeLookup, {
-            nodeOrigin,
-            elevateNodesOnSelect,
-          });
-        }
-        if (hasDefaultEdges) {
-          const { connectionLookup, edgeLookup } = get();
-          updateConnectionLookup(connectionLookup, edgeLookup, edges);
-
-          nextState.edges = edges;
-        }
-
-        set(nextState);
-      },
       // Every node gets registerd at a ResizeObserver. Whenever a node
       // changes its dimensions, this function is called to measure the
       // new dimensions and update the nodes.
@@ -166,81 +136,67 @@ const createRFStore = ({
 
         get().triggerNodeChanges(changes);
       },
-
       triggerNodeChanges: (changes) => {
-        const { onNodesChange, nodeLookup, nodes, hasDefaultNodes, nodeOrigin, elevateNodesOnSelect } = get();
+        const { onNodesChange, setNodes, nodes, hasDefaultNodes } = get();
 
         if (changes?.length) {
           if (hasDefaultNodes) {
             const updatedNodes = applyNodeChanges(changes, nodes);
-            const nextNodes = adoptUserProvidedNodes(updatedNodes, nodeLookup, {
-              nodeOrigin,
-              elevateNodesOnSelect,
-            });
-            set({ nodes: nextNodes });
+            setNodes(updatedNodes);
           }
 
           onNodesChange?.(changes);
         }
       },
+      triggerEdgeChanges: (changes) => {
+        const { onEdgesChange, setEdges, edges, hasDefaultEdges } = get();
 
+        if (changes?.length) {
+          if (hasDefaultEdges) {
+            const updatedEdges = applyEdgeChanges(changes, edges);
+            setEdges(updatedEdges);
+          }
+
+          onEdgesChange?.(changes);
+        }
+      },
       addSelectedNodes: (selectedNodeIds) => {
-        const { multiSelectionActive, edges, nodes } = get();
-        let changedNodes: NodeSelectionChange[];
-        let changedEdges: EdgeSelectionChange[] | null = null;
+        const { multiSelectionActive, edges, nodes, triggerNodeChanges, triggerEdgeChanges } = get();
 
         if (multiSelectionActive) {
-          changedNodes = selectedNodeIds.map((nodeId) => createSelectionChange(nodeId, true)) as NodeSelectionChange[];
-        } else {
-          changedNodes = getSelectionChanges(nodes, new Set([...selectedNodeIds]), true);
-          changedEdges = getSelectionChanges(edges);
+          const nodeChanges = selectedNodeIds.map((nodeId) => createSelectionChange(nodeId, true));
+          triggerNodeChanges(nodeChanges as NodeSelectionChange[]);
+          return;
         }
 
-        updateNodesAndEdgesSelections({
-          changedNodes,
-          changedEdges,
-          get,
-          set,
-        });
+        triggerNodeChanges(getSelectionChanges(nodes, new Set([...selectedNodeIds]), true));
+        triggerEdgeChanges(getSelectionChanges(edges));
       },
       addSelectedEdges: (selectedEdgeIds) => {
-        const { multiSelectionActive, edges, nodes } = get();
-        let changedEdges: EdgeSelectionChange[];
-        let changedNodes: NodeSelectionChange[] | null = null;
+        const { multiSelectionActive, edges, nodes, triggerNodeChanges, triggerEdgeChanges } = get();
 
         if (multiSelectionActive) {
-          changedEdges = selectedEdgeIds.map((edgeId) => createSelectionChange(edgeId, true)) as EdgeSelectionChange[];
-        } else {
-          changedEdges = getSelectionChanges(edges, new Set([...selectedEdgeIds]));
-          changedNodes = getSelectionChanges(nodes, new Set(), true);
+          const changedEdges = selectedEdgeIds.map((edgeId) => createSelectionChange(edgeId, true));
+          triggerEdgeChanges(changedEdges as EdgeSelectionChange[]);
+          return;
         }
 
-        updateNodesAndEdgesSelections({
-          changedNodes,
-          changedEdges,
-          get,
-          set,
-        });
+        triggerEdgeChanges(getSelectionChanges(edges, new Set([...selectedEdgeIds])));
+        triggerNodeChanges(getSelectionChanges(nodes, new Set(), true));
       },
       unselectNodesAndEdges: ({ nodes, edges }: UnselectNodesAndEdgesParams = {}) => {
-        const { edges: storeEdges, nodes: storeNodes } = get();
+        const { edges: storeEdges, nodes: storeNodes, triggerNodeChanges, triggerEdgeChanges } = get();
         const nodesToUnselect = nodes ? nodes : storeNodes;
         const edgesToUnselect = edges ? edges : storeEdges;
 
-        const changedNodes = nodesToUnselect.map((n) => {
+        const nodeChanges = nodesToUnselect.map((n) => {
           n.selected = false;
           return createSelectionChange(n.id, false);
-        }) as NodeSelectionChange[];
-        const changedEdges = edgesToUnselect.map((edge) =>
-          createSelectionChange(edge.id, false)
-        ) as EdgeSelectionChange[];
-
-        updateNodesAndEdgesSelections({
-          changedNodes,
-          changedEdges,
-          get,
-          set,
         });
+        const edgeChanges = edgesToUnselect.map((edge) => createSelectionChange(edge.id, false));
+
+        triggerNodeChanges(nodeChanges as NodeSelectionChange[]);
+        triggerEdgeChanges(edgeChanges as EdgeSelectionChange[]);
       },
       setMinZoom: (minZoom) => {
         const { panZoom, maxZoom } = get();
@@ -260,21 +216,19 @@ const createRFStore = ({
         set({ translateExtent });
       },
       resetSelectedElements: () => {
-        const { edges, nodes } = get();
+        const { edges, nodes, triggerNodeChanges, triggerEdgeChanges } = get();
 
-        const nodesToUnselect = nodes
-          .filter((e) => e.selected)
-          .map((n) => createSelectionChange(n.id, false)) as NodeSelectionChange[];
-        const edgesToUnselect = edges
-          .filter((e) => e.selected)
-          .map((e) => createSelectionChange(e.id, false)) as EdgeSelectionChange[];
+        const nodeChanges = nodes.reduce<NodeSelectionChange[]>(
+          (res, node) => (node.selected ? [...res, createSelectionChange(node.id, false) as NodeSelectionChange] : res),
+          []
+        );
+        const edgeChanges = edges.reduce<EdgeSelectionChange[]>(
+          (res, edge) => (edge.selected ? [...res, createSelectionChange(edge.id, false) as EdgeSelectionChange] : res),
+          []
+        );
 
-        updateNodesAndEdgesSelections({
-          changedNodes: nodesToUnselect,
-          changedEdges: edgesToUnselect,
-          get,
-          set,
-        });
+        triggerNodeChanges(nodeChanges);
+        triggerEdgeChanges(edgeChanges);
       },
       setNodeExtent: (nodeExtent) => {
         const { nodes } = get();

packages/react/src/store/initialState.ts

@@ -12,14 +12,18 @@ import {
 import type { Edge, Node, ReactFlowStore } from '../types';
 
 const getInitialState = ({
-  nodes = [],
-  edges = [],
+  nodes,
+  edges,
+  defaultNodes,
+  defaultEdges,
   width,
   height,
   fitView,
 }: {
   nodes?: Node[];
   edges?: Edge[];
+  defaultNodes?: Node[];
+  defaultEdges?: Edge[];
   width?: number;
   height?: number;
   fitView?: boolean;
@@ -27,9 +31,11 @@ const getInitialState = ({
   const nodeLookup = new Map();
   const connectionLookup = new Map();
   const edgeLookup = new Map();
+  const storeEdges = defaultEdges ?? edges ?? [];
+  const storeNodes = defaultNodes ?? nodes ?? [];
 
-  updateConnectionLookup(connectionLookup, edgeLookup, edges);
-  const nextNodes = adoptUserProvidedNodes(nodes, nodeLookup, {
+  updateConnectionLookup(connectionLookup, edgeLookup, storeEdges);
+  const nextNodes = adoptUserProvidedNodes(storeNodes, nodeLookup, {
     nodeOrigin: [0, 0],
     elevateNodesOnSelect: false,
   });
@@ -51,13 +57,13 @@ const getInitialState = ({
     transform,
     nodes: nextNodes,
     nodeLookup,
-    edges,
+    edges: storeEdges,
     edgeLookup,
     connectionLookup,
     onNodesChange: null,
     onEdgesChange: null,
-    hasDefaultNodes: false,
-    hasDefaultEdges: false,
+    hasDefaultNodes: defaultNodes !== undefined,
+    hasDefaultEdges: defaultEdges !== undefined,
     panZoom: null,
     minZoom: 0.5,
     maxZoom: 2,

packages/react/src/store/utils.ts

@@ -1,44 +0,0 @@
-import type { StoreApi } from 'zustand';
-import type { Edge, EdgeSelectionChange, Node, NodeSelectionChange, ReactFlowState } from '../types';
-
-export function handleControlledSelectionChange<NodeOrEdge extends Node | Edge>(
-  changes: NodeSelectionChange[] | EdgeSelectionChange[],
-  items: NodeOrEdge[]
-): NodeOrEdge[] {
-  return items.map((item) => {
-    const change = changes.find((change) => change.id === item.id);
-
-    if (change) {
-      item.selected = change.selected;
-    }
-
-    return item;
-  });
-}
-
-type UpdateNodesAndEdgesParams = {
-  changedNodes: NodeSelectionChange[] | null;
-  changedEdges: EdgeSelectionChange[] | null;
-  get: StoreApi<ReactFlowState>['getState'];
-  set: StoreApi<ReactFlowState>['setState'];
-};
-
-export function updateNodesAndEdgesSelections({ changedNodes, changedEdges, get, set }: UpdateNodesAndEdgesParams) {
-  const { nodes, edges, onNodesChange, onEdgesChange, hasDefaultNodes, hasDefaultEdges } = get();
-
-  if (changedNodes?.length) {
-    if (hasDefaultNodes) {
-      set({ nodes: handleControlledSelectionChange(changedNodes, nodes) });
-    }
-
-    onNodesChange?.(changedNodes);
-  }
-
-  if (changedEdges?.length) {
-    if (hasDefaultEdges) {
-      set({ edges: handleControlledSelectionChange(changedEdges, edges) });
-    }
-
-    onEdgesChange?.(changedEdges);
-  }
-}