refactor(react): cleanup provider

packages/react/src/components/ReactFlowProvider/index.tsx

@@ -1,22 +1,10 @@
-import { useRef, type ReactNode } from 'react';
-import { type StoreApi } from 'zustand';
-import { UseBoundStoreWithEqualityFn } from 'zustand/traditional';
+import { useState, type ReactNode } from 'react';
 
-import { Provider } from '../../contexts/RFStoreContext';
-import { createRFStore } from '../../store';
-import type { ReactFlowState, Node, Edge } from '../../types';
+import { Provider } from '../../contexts/StoreContext';
+import { createStore } from '../../store';
+import type { Node, Edge } from '../../types';
 
-export function ReactFlowProvider({
-  children,
-  initialNodes,
-  initialEdges,
-  defaultNodes,
-  defaultEdges,
-  initialWidth,
-  initialHeight,
-  fitView,
-}: {
-  children: ReactNode;
+export type ReactFlowProviderProps = {
   initialNodes?: Node[];
   initialEdges?: Edge[];
   defaultNodes?: Node[];
@@ -24,19 +12,30 @@ export function ReactFlowProvider({
   initialWidth?: number;
   initialHeight?: number;
   fitView?: boolean;
-}) {
-  const storeRef = useRef<UseBoundStoreWithEqualityFn<StoreApi<ReactFlowState>> | null>(null);
-  if (!storeRef.current) {
-    storeRef.current = createRFStore({
-      nodes: initialNodes,
-      edges: initialEdges,
+  children: ReactNode;
+};
+
+export function ReactFlowProvider({
+  initialNodes: nodes,
+  initialEdges: edges,
+  defaultNodes,
+  defaultEdges,
+  initialWidth: width,
+  initialHeight: height,
+  fitView,
+  children,
+}: ReactFlowProviderProps) {
+  const [store] = useState(() =>
+    createStore({
+      nodes,
+      edges,
       defaultNodes,
       defaultEdges,
-      width: initialWidth,
-      height: initialHeight,
+      width,
+      height,
       fitView,
-    });
-  }
+    })
+  );
 
-  return <Provider value={storeRef.current}>{children}</Provider>;
+  return <Provider value={store}>{children}</Provider>;
 }