From 2860feec112d52b39d90d4fc8f8ac1c3eecbdbcb Mon Sep 17 00:00:00 2001
From: moklick <info@moritzklack.com>
Date: Tue, 13 Feb 2024 16:12:31 +0100
Subject: [PATCH] fix(react): uncontrolled flows with Provider

---
 .../src/components/ReactFlowProvider/index.tsx |  1 -
 .../src/components/StoreUpdater/index.tsx      | 18 ++++++++++++++----
 .../react/src/container/ReactFlow/index.tsx    |  2 ++
 packages/react/src/store/index.ts              | 12 ++++++++++++
 packages/react/src/types/store.ts              |  1 +
 5 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/packages/react/src/components/ReactFlowProvider/index.tsx b/packages/react/src/components/ReactFlowProvider/index.tsx
index 829a304d..3d2577e7 100644
--- a/packages/react/src/components/ReactFlowProvider/index.tsx
+++ b/packages/react/src/components/ReactFlowProvider/index.tsx
@@ -26,7 +26,6 @@ export function ReactFlowProvider({
   fitView?: boolean;
 }) {
   const storeRef = useRef<UseBoundStoreWithEqualityFn<StoreApi<ReactFlowState>> | null>(null);
-
   if (!storeRef.current) {
     storeRef.current = createRFStore({
       nodes: initialNodes,
diff --git a/packages/react/src/components/StoreUpdater/index.tsx b/packages/react/src/components/StoreUpdater/index.tsx
index 98ad3bac..96080368 100644
--- a/packages/react/src/components/StoreUpdater/index.tsx
+++ b/packages/react/src/components/StoreUpdater/index.tsx
@@ -15,6 +15,8 @@ import { initNodeOrigin } from '../../container/ReactFlow';
 const reactFlowFieldsToTrack = [
   'nodes',
   'edges',
+  'defaultNodes',
+  'defaultEdges',
   'onConnect',
   'onConnectStart',
   'onConnectEnd',
@@ -79,16 +81,24 @@ const selector = (s: ReactFlowState) => ({
   setTranslateExtent: s.setTranslateExtent,
   setNodeExtent: s.setNodeExtent,
   reset: s.reset,
+  setDefaultNodesAndEdges: s.setDefaultNodesAndEdges,
 });
 
 export function StoreUpdater(props: StoreUpdaterProps) {
-  const { setNodes, setEdges, setMinZoom, setMaxZoom, setTranslateExtent, setNodeExtent, reset } = useStore(
-    selector,
-    shallow
-  );
+  const {
+    setNodes,
+    setEdges,
+    setMinZoom,
+    setMaxZoom,
+    setTranslateExtent,
+    setNodeExtent,
+    reset,
+    setDefaultNodesAndEdges,
+  } = useStore(selector, shallow);
   const store = useStoreApi();
 
   useEffect(() => {
+    setDefaultNodesAndEdges(props.defaultNodes, props.defaultEdges);
     return () => {
       reset();
     };
diff --git a/packages/react/src/container/ReactFlow/index.tsx b/packages/react/src/container/ReactFlow/index.tsx
index 9ac7b051..df9936a1 100644
--- a/packages/react/src/container/ReactFlow/index.tsx
+++ b/packages/react/src/container/ReactFlow/index.tsx
@@ -242,6 +242,8 @@ const ReactFlow = forwardRef<ReactFlowRefType, ReactFlowProps>(
           <StoreUpdater
             nodes={nodes}
             edges={edges}
+            defaultNodes={defaultNodes}
+            defaultEdges={defaultEdges}
             onConnect={onConnect}
             onConnectStart={onConnectStart}
             onConnectEnd={onConnectEnd}
diff --git a/packages/react/src/store/index.ts b/packages/react/src/store/index.ts
index 9e49f5d7..46447e00 100644
--- a/packages/react/src/store/index.ts
+++ b/packages/react/src/store/index.ts
@@ -63,6 +63,18 @@ const createRFStore = ({
 
         set({ edges });
       },
+      setDefaultNodesAndEdges: (nodes?: Node[], edges?: Edge[]) => {
+        if (nodes) {
+          const { setNodes } = get();
+          setNodes(nodes);
+          set({ hasDefaultNodes: true });
+        }
+        if (edges) {
+          const { setEdges } = get();
+          setEdges(edges);
+          set({ hasDefaultEdges: true });
+        }
+      },
       // 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.
diff --git a/packages/react/src/types/store.ts b/packages/react/src/types/store.ts
index 46f83128..f1c625ba 100644
--- a/packages/react/src/types/store.ts
+++ b/packages/react/src/types/store.ts
@@ -168,6 +168,7 @@ export type ReactFlowActions = {
   triggerEdgeChanges: (changes: EdgeChange[] | null) => void;
   panBy: PanBy;
   fitView: (nodes: Node[], options?: FitViewOptions) => boolean;
+  setDefaultNodesAndEdges: (nodes?: Node[], edges?: Edge[]) => void;
 };
 
 export type ReactFlowState = ReactFlowStore & ReactFlowActions;