chore(zustand): upgrade to v4

packages/core/src/components/ConnectionLine/index.tsx

@@ -1,10 +1,10 @@
 import React, { CSSProperties, useCallback } from 'react';
 import shallow from 'zustand/shallow';
 
-import { useStore } from '../../store';
+import { useStore } from '../../hooks/useStore';
 import { getBezierPath } from '../Edges/BezierEdge';
 import { getSmoothStepPath } from '../Edges/SmoothStepEdge';
-import { ConnectionLineType, ConnectionLineComponent, HandleType, Position } from '../../types';
+import { ConnectionLineType, ConnectionLineComponent, HandleType, Position, ReactFlowStore } from '../../types';
 import { getSimpleBezierPath } from '../Edges/SimpleBezierEdge';
 import { internalsSymbol } from '../../utils';
 
@@ -34,7 +34,7 @@ export default ({
 }: ConnectionLineProps) => {
   const { fromNode, handleId, toX, toY } = useStore(
     useCallback(
-      (s) => ({
+      (s: ReactFlowStore) => ({
         fromNode: s.nodeInternals.get(connectionNodeId),
         handleId: s.connectionHandleId,
         toX: (s.connectionPosition.x - s.transform[0]) / s.transform[2],

packages/core/src/components/Edges/wrapEdge.tsx

@@ -1,7 +1,7 @@
 import React, { memo, ComponentType, useState, useMemo, KeyboardEvent, useRef } from 'react';
 import cc from 'classcat';
 
-import { useStoreApi } from '../../store';
+import { useStoreApi } from '../../hooks/useStore';
 import { ARIA_EDGE_DESC_KEY } from '../A11yDescriptions';
 import { handleMouseDown } from '../Handle/handler';
 import { EdgeAnchor } from './EdgeAnchor';

packages/core/src/components/Handle/index.tsx

@@ -2,7 +2,7 @@ import React, { memo, useContext, HTMLAttributes, forwardRef } from 'react';
 import cc from 'classcat';
 import shallow from 'zustand/shallow';
 
-import { useStore, useStoreApi } from '../../store';
+import { useStore, useStoreApi } from '../../hooks/useStore';
 import NodeIdContext from '../../contexts/NodeIdContext';
 import { HandleProps, Connection, ReactFlowState, Position } from '../../types';
 import { checkElementBelowIsValid, handleMouseDown } from './handler';

packages/core/src/components/Nodes/wrapNode.tsx

@@ -1,7 +1,7 @@
 import React, { useEffect, useRef, memo, ComponentType, MouseEvent, KeyboardEvent } from 'react';
 import cc from 'classcat';
 
-import { useStoreApi } from '../../store';
+import { useStoreApi } from '../../hooks/useStore';
 import { Provider } from '../../contexts/NodeIdContext';
 import { ARIA_NODE_DESC_KEY } from '../A11yDescriptions';
 import useDrag from '../../hooks/useDrag';

packages/core/src/components/NodesSelection/index.tsx

@@ -7,7 +7,7 @@ import React, { memo, useRef, MouseEvent, KeyboardEvent, useEffect } from 'react
 import cc from 'classcat';
 import shallow from 'zustand/shallow';
 
-import { useStore, useStoreApi } from '../../store';
+import { useStore, useStoreApi } from '../../hooks/useStore';
 import { Node, ReactFlowState } from '../../types';
 import { getRectOfNodes } from '../../utils/graph';
 import useDrag from '../../hooks/useDrag';

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

@@ -1,10 +1,19 @@
-import React, { FC, PropsWithChildren } from 'react';
+import React, { FC, PropsWithChildren, useRef } from 'react';
+import { StoreApi } from 'zustand';
 
-import { Provider, createStore } from '../../store';
+import { Provider } from '../../contexts/RFStoreContext';
+import { createRFStore } from '../../store';
+import { ReactFlowState } from '../../types';
 
-const ReactFlowProvider: FC<PropsWithChildren<{}>> = ({ children }) => (
-  <Provider createStore={createStore}>{children}</Provider>
-);
+const ReactFlowProvider: FC<PropsWithChildren<{}>> = ({ children }) => {
+  const storeRef = useRef<StoreApi<ReactFlowState> | null>(null);
+
+  if (!storeRef.current) {
+    storeRef.current = createRFStore();
+  }
+
+  return <Provider value={storeRef.current}>{children}</Provider>;
+};
 
 ReactFlowProvider.displayName = 'ReactFlowProvider';
 

packages/core/src/components/SelectionListener/index.tsx

@@ -2,7 +2,7 @@ import { memo, useEffect } from 'react';
 import shallow from 'zustand/shallow';
 
 import { ReactFlowState, OnSelectionChangeFunc, Node, Edge } from '../../types';
-import { useStore } from '../../store';
+import { useStore } from '../../hooks/useStore';
 
 interface SelectionListenerProps {
   onSelectionChange: OnSelectionChangeFunc;

packages/core/src/components/StoreUpdater/index.tsx

@@ -1,8 +1,8 @@
 import { useEffect } from 'react';
-import { SetState } from 'zustand';
+import { StoreApi } from 'zustand';
 import shallow from 'zustand/shallow';
 
-import { useStore, useStoreApi } from '../../store';
+import { useStore, useStoreApi } from '../../hooks/useStore';
 import { Node, Edge, ReactFlowState, CoordinateExtent, ReactFlowProps } from '../../types';
 
 type StoreUpdaterProps = Pick<
@@ -61,7 +61,7 @@ function useStoreUpdater<T>(value: T | undefined, setStoreState: (param: T) => v
   }, [value]);
 }
 
-function useDirectStoreUpdater(key: keyof ReactFlowState, value: any, setState: SetState<ReactFlowState>) {
+function useDirectStoreUpdater(key: keyof ReactFlowState, value: any, setState: StoreApi<ReactFlowState>['setState']) {
   useEffect(() => {
     if (typeof value !== 'undefined') {
       // @ts-ignore

packages/core/src/components/UserSelection/index.tsx

@@ -5,7 +5,7 @@
 import React, { memo, useState, useRef } from 'react';
 import shallow from 'zustand/shallow';
 
-import { useStore, useStoreApi } from '../../store';
+import { useStore, useStoreApi } from '../../hooks/useStore';
 import { getSelectionChanges } from '../../utils/changes';
 import { XYPosition, ReactFlowState, NodeChange, EdgeChange, Rect } from '../../types';
 import { getConnectedEdges, getNodesInside } from '../../utils/graph';

packages/core/src/container/EdgeRenderer/MarkerDefinitions.tsx

@@ -1,6 +1,6 @@
 import React, { memo, useCallback } from 'react';
 
-import { useStore } from '../../store';
+import { useStore } from '../../hooks/useStore';
 import { EdgeMarker, ReactFlowState } from '../../types';
 import { getMarkerId } from '../../utils/graph';
 import { useMarkerSymbol } from './MarkerSymbols';

packages/core/src/container/EdgeRenderer/index.tsx

@@ -2,7 +2,7 @@ import React, { memo, CSSProperties } from 'react';
 import shallow from 'zustand/shallow';
 import cc from 'classcat';
 
-import { useStore } from '../../store';
+import { useStore } from '../../hooks/useStore';
 import ConnectionLine from '../../components/ConnectionLine/index';
 import MarkerDefinitions from './MarkerDefinitions';
 import { getEdgePositions, getHandle, getNodeData } from './utils';

packages/core/src/container/FlowRenderer/index.tsx

@@ -1,6 +1,6 @@
 import React, { memo, ReactNode, WheelEvent, MouseEvent } from 'react';
 
-import { useStore, useStoreApi } from '../../store';
+import { useStore, useStoreApi } from '../../hooks/useStore';
 import useGlobalKeyHandler from '../../hooks/useGlobalKeyHandler';
 import useKeyPress from '../../hooks/useKeyPress';
 import { GraphViewProps } from '../GraphView';

packages/core/src/container/NodeRenderer/index.tsx

@@ -2,10 +2,10 @@ import React, { memo, useMemo, ComponentType, useEffect, useRef } from 'react';
 import shallow from 'zustand/shallow';
 
 import useVisibleNodes from '../../hooks/useVisibleNodes';
-import { useStore } from '../../store';
+import { useStore } from '../../hooks/useStore';
+import { internalsSymbol } from '../../utils';
 import { containerStyle } from '../../styles';
 import { NodeMouseHandler, NodeTypesWrapped, Position, ReactFlowState, WrapNodeProps } from '../../types';
-import { internalsSymbol } from '../../utils';
 
 interface NodeRendererProps {
   nodeTypes: NodeTypesWrapped;

packages/core/src/container/ReactFlow/Wrapper.tsx

@@ -1,6 +1,7 @@
 import React, { FC, PropsWithChildren } from 'react';
 
-import { Provider, createStore, useStoreApi } from '../../store';
+import { useStoreApi } from '../../hooks/useStore';
+import ReactFlowProvider from '../../components/ReactFlowProvider';
 
 const Wrapper: FC<PropsWithChildren<{}>> = ({ children }) => {
   let isWrapped = true;
@@ -17,7 +18,7 @@ const Wrapper: FC<PropsWithChildren<{}>> = ({ children }) => {
     return <>{children}</>;
   }
 
-  return <Provider createStore={createStore}>{children}</Provider>;
+  return <ReactFlowProvider>{children}</ReactFlowProvider>;
 };
 
 Wrapper.displayName = 'ReactFlowWrapper';

packages/core/src/container/Viewport/index.tsx

@@ -1,6 +1,6 @@
 import React, { ReactNode } from 'react';
 
-import { useStore } from '../../store';
+import { useStore } from '../../hooks/useStore';
 import { ReactFlowState } from '../../types';
 
 const selector = (s: ReactFlowState) => `translate(${s.transform[0]}px,${s.transform[1]}px) scale(${s.transform[2]})`;

packages/core/src/container/ZoomPane/index.tsx

@@ -6,10 +6,10 @@ import shallow from 'zustand/shallow';
 import { clamp } from '../../utils';
 import useKeyPress from '../../hooks/useKeyPress';
 import useResizeHandler from '../../hooks/useResizeHandler';
-import { useStore, useStoreApi } from '../../store';
-import { Viewport, PanOnScrollMode, ReactFlowState } from '../../types';
+import { useStore, useStoreApi } from '../../hooks/useStore';
 import { FlowRendererProps } from '../FlowRenderer';
 import { containerStyle } from '../../styles';
+import { Viewport, PanOnScrollMode, ReactFlowState } from '../../types';
 
 type ZoomPaneProps = Omit<
   FlowRendererProps,

packages/core/src/contexts/RFStoreContext.ts

@@ -0,0 +1,8 @@
+import { createContext } from 'react';
+
+import { createRFStore } from '../store';
+
+const StoreContext = createContext<ReturnType<typeof createRFStore> | null>(null);
+
+export const Provider = StoreContext.Provider;
+export default StoreContext;

packages/core/src/hooks/useDrag/index.ts

@@ -2,7 +2,7 @@ import { RefObject, useEffect, useRef, MouseEvent, useState, useCallback } from
 import { D3DragEvent, drag, SubjectPosition } from 'd3-drag';
 import { select } from 'd3-selection';
 
-import { useStoreApi } from '../../store';
+import { useStoreApi } from '../../hooks/useStore';
 import { pointToRendererPoint } from '../../utils/graph';
 import { NodeDragItem, Node, SelectionDragHandler } from '../../types';
 import { getDragItems, getEventHandlerParams, hasSelector, calcNextPosition } from './utils';

packages/core/src/hooks/useEdges.ts

@@ -1,4 +1,4 @@
-import { useStore } from '../store';
+import { useStore } from '../hooks/useStore';
 import { Edge, ReactFlowState } from '../types';
 
 const edgesSelector = (state: ReactFlowState) => state.edges;

packages/core/src/hooks/useGlobalKeyHandler.ts

@@ -1,6 +1,6 @@
 import { useEffect } from 'react';
 
-import { useStoreApi } from '../store';
+import { useStoreApi } from '../hooks/useStore';
 import useKeyPress from './useKeyPress';
 import { getConnectedEdges } from '../utils/graph';
 import { KeyCode, NodeChange, Node } from '../types';

packages/core/src/hooks/useNodes.ts

@@ -1,4 +1,4 @@
-import { useStore } from '../store';
+import { useStore } from '../hooks/useStore';
 import { Node, ReactFlowState } from '../types';
 
 const nodesSelector = (state: ReactFlowState) => Array.from(state.nodeInternals.values());

packages/core/src/hooks/useReactFlow.ts

@@ -1,7 +1,7 @@
 import { useCallback } from 'react';
 
 import useViewportHelper from './useViewportHelper';
-import { useStoreApi } from '../store';
+import { useStoreApi } from '../hooks/useStore';
 import {
   ReactFlowInstance,
   Instance,

packages/core/src/hooks/useResizeHandler.ts

@@ -1,6 +1,6 @@
 import { useEffect, MutableRefObject } from 'react';
 
-import { useStoreApi } from '../store';
+import { useStoreApi } from '../hooks/useStore';
 import { getDimensions } from '../utils';
 
 function useResizeHandler(rendererNode: MutableRefObject<HTMLDivElement | null>): void {

packages/core/src/hooks/useStore.ts

@@ -0,0 +1,42 @@
+import { useContext, useMemo } from 'react';
+import { StoreApi, useStore as useZustandStore } from 'zustand';
+
+import StoreContext from '../contexts/RFStoreContext';
+import { ReactFlowState } from '../types';
+
+const errorMessage = 'Seems like you have not used zustand provider as an ancestor.';
+
+type ExtractState = StoreApi<ReactFlowState> extends { getState: () => infer T } ? T : never;
+
+function useStore<StateSlice = ExtractState>(
+  selector: (state: ExtractState) => StateSlice,
+  equalityFn?: (a: StateSlice, b: StateSlice) => boolean
+) {
+  const store = useContext(StoreContext);
+
+  if (store === null) {
+    throw new Error(errorMessage);
+  }
+
+  return useZustandStore(store, selector, equalityFn);
+}
+
+const useStoreApi = () => {
+  const store = useContext(StoreContext);
+
+  if (store === null) {
+    throw new Error(errorMessage);
+  }
+
+  return useMemo(
+    () => ({
+      getState: store.getState,
+      setState: store.setState,
+      subscribe: store.subscribe,
+      destroy: store.destroy,
+    }),
+    [store]
+  );
+};
+
+export { useStore, useStoreApi };

packages/core/src/hooks/useUpdateNodeInternals.ts

@@ -1,6 +1,6 @@
 import { useCallback } from 'react';
 
-import { useStore, useStoreApi } from '../store';
+import { useStore, useStoreApi } from '../hooks/useStore';
 import { UpdateNodeInternals, ReactFlowState } from '../types';
 
 const selector = (state: ReactFlowState) => state.updateNodeDimensions;

packages/core/src/hooks/useUpdateNodePositions.ts

@@ -1,6 +1,6 @@
 import { useCallback } from 'react';
 
-import { useStoreApi } from '../store';
+import { useStoreApi } from '../hooks/useStore';
 import { calcNextPosition } from './useDrag/utils';
 
 import { XYPosition } from '../types';

packages/core/src/hooks/useViewport.ts

@@ -1,6 +1,6 @@
 import shallow from 'zustand/shallow';
 
-import { useStore } from '../store';
+import { useStore } from '../hooks/useStore';
 import { Viewport, ReactFlowState } from '../types';
 
 const viewportSelector = (state: ReactFlowState) => ({

packages/core/src/hooks/useViewportHelper.ts

@@ -2,7 +2,7 @@ import { useMemo } from 'react';
 import { zoomIdentity } from 'd3-zoom';
 import shallow from 'zustand/shallow';
 
-import { useStoreApi, useStore } from '../store';
+import { useStoreApi, useStore } from '../hooks/useStore';
 import { pointToRendererPoint, getTransformForBounds, getD3Transition } from '../utils/graph';
 import { FitViewOptions, Viewport, ViewportHelperFunctions, ReactFlowState, Rect, XYPosition } from '../types';
 import { fitView as fitViewStore } from '../store/utils';

packages/core/src/hooks/useVisibleEdges.ts

@@ -1,6 +1,6 @@
 import { useCallback } from 'react';
 
-import { useStore } from '../store';
+import { useStore } from '../hooks/useStore';
 import { isEdgeVisible } from '../container/EdgeRenderer/utils';
 import { ReactFlowState, NodeInternals, Edge } from '../types';
 import { internalsSymbol, isNumeric } from '../utils';

packages/core/src/hooks/useVisibleNodes.ts

@@ -1,6 +1,6 @@
 import { useCallback } from 'react';
 
-import { useStore } from '../store';
+import { useStore } from '../hooks/useStore';
 import { getNodesInside } from '../utils/graph';
 
 import { ReactFlowState } from '../types';

packages/core/src/index.ts

@@ -39,6 +39,6 @@ export { default as useEdges } from './hooks/useEdges';
 export { default as useViewport } from './hooks/useViewport';
 export { default as useKeyPress } from './hooks/useKeyPress';
 export * from './hooks/useNodesEdgesState';
-export { useStore, useStoreApi } from './store';
+export { useStore, useStoreApi } from './hooks/useStore';
 
 export * from './types';

packages/core/src/store/index.ts

@@ -1,5 +1,4 @@
-import create from 'zustand';
-import createContext from 'zustand/context';
+import { createStore } from 'zustand';
 
 import { clampPosition, getDimensions, internalsSymbol } from '../utils';
 import { applyNodeChanges } from '../utils/changes';
@@ -21,10 +20,8 @@ import { createSelectionChange, getSelectionChanges } from '../utils/changes';
 import { createNodeInternals, fitView, updateNodesAndEdgesSelections } from './utils';
 import initialState from './initialState';
 
-const { Provider, useStore, useStoreApi } = createContext<ReactFlowState>();
-
-const createStore = () =>
-  create<ReactFlowState>((set, get) => ({
+const createRFStore = () =>
+  createStore<ReactFlowState>((set, get) => ({
     ...initialState,
     setNodes: (nodes: Node[]) => {
       set({ nodeInternals: createNodeInternals(nodes, get().nodeInternals) });
@@ -236,4 +233,4 @@ const createStore = () =>
     reset: () => set({ ...initialState }),
   }));
 
-export { Provider, useStore, createStore, useStoreApi };
+export { createRFStore };