diff --git a/cypress/integration/flow/graph-utils.spec.js b/cypress/integration/flow/graph-utils.spec.js
index e582231b..e6df1455 100644
--- a/cypress/integration/flow/graph-utils.spec.js
+++ b/cypress/integration/flow/graph-utils.spec.js
@@ -1,4 +1,4 @@
-import { isNode, isEdge, getOutgoers, getIncomers, removeElements, addEdge } from '../../../dist/ReactFlow.js';
+import { isNode, isEdge, getOutgoers, getIncomers, addEdge } from '../../../dist/ReactFlow.js';
 
 const nodes = [
   { id: '1', type: 'input', data: { label: 'Node 1' }, position: { x: 250, y: 5 } },
@@ -69,66 +69,4 @@ describe('Graph Utils Testing', () => {
       }
     });
   });
-
-  describe('tests removeElements function', () => {
-    it('removes a node', () => {
-      const nextElements = removeElements([nodes[0]], elements);
-
-      const nextNodes = nextElements.filter((e) => isNode(e));
-      const nextEdges = nextElements.filter((e) => isEdge(e));
-
-      expect(nextNodes.length).to.be.equal(nodes.length - 1);
-      expect(nextEdges.length).to.be.equal(edges.length - 2);
-    });
-
-    it('removes multiple nodes', () => {
-      const elementsToRemove = [nodes[0], nodes[1]];
-      const nextElements = removeElements(elementsToRemove, elements);
-      const nextNodes = nextElements.filter((e) => isNode(e));
-      const nextEdges = nextElements.filter((e) => isEdge(e));
-
-      expect(nextNodes.length).to.be.equal(nodes.length - 2);
-      expect(nextEdges.length).to.be.equal(0);
-    });
-
-    it('removes no node', () => {
-      const nextElementsNoRemove = removeElements([], elements);
-      expect(nextElementsNoRemove.length).to.be.equal(elements.length);
-    });
-
-    it('tries to removes node that does not exist', () => {
-      const nextElementsNoRemove = removeElements([{ id: 'id-that-does-not-exist' }], elements);
-      expect(nextElementsNoRemove.length).to.be.equal(elements.length);
-    });
-
-    it('removes an edge', () => {
-      const nextElements = removeElements([edges[0]], elements);
-
-      const nextNodes = nextElements.filter((e) => isNode(e));
-      const nextEdges = nextElements.filter((e) => isEdge(e));
-
-      expect(nextNodes.length).to.be.equal(nodes.length);
-      expect(nextEdges.length).to.be.equal(edges.length - 1);
-    });
-
-    it('removes multiple edges', () => {
-      const nextElements = removeElements([edges[0], edges[1]], elements);
-
-      const nextNodes = nextElements.filter((e) => isNode(e));
-      const nextEdges = nextElements.filter((e) => isEdge(e));
-
-      expect(nextNodes.length).to.be.equal(nodes.length);
-      expect(nextEdges.length).to.be.equal(edges.length - 2);
-    });
-
-    it('removes node and edge', () => {
-      const nextElements = removeElements([nodes[0], edges[0]], elements);
-
-      const nextNodes = nextElements.filter((e) => isNode(e));
-      const nextEdges = nextElements.filter((e) => isEdge(e));
-
-      expect(nextNodes.length).to.be.equal(nodes.length - 1);
-      expect(nextEdges.length).to.be.equal(edges.length - 2);
-    });
-  });
 });
diff --git a/example/src/Basic/DebugNode.tsx b/example/src/Basic/DebugNode.tsx
new file mode 100644
index 00000000..f7cc4b36
--- /dev/null
+++ b/example/src/Basic/DebugNode.tsx
@@ -0,0 +1,27 @@
+import { memo, FC, CSSProperties } from 'react';
+
+import { Handle, NodeProps, Position } from 'react-flow-renderer';
+
+const infoStyle: CSSProperties = { fontSize: 11 };
+const idStyle: CSSProperties = {
+  fontSize: 10,
+  color: '#888899',
+  position: 'absolute',
+  top: 2,
+  left: 2,
+};
+
+const ColorSelectorNode: FC<NodeProps> = ({ zIndex, xPos, yPos, id }) => {
+  return (
+    <>
+      <Handle type="target" position={Position.Top} />
+      <div style={idStyle}>{id}</div>
+      <div style={infoStyle}>
+        x:{Math.round(xPos || 0)} y:{Math.round(yPos || 0)} z:{zIndex}
+      </div>
+      <Handle type="source" position={Position.Bottom} />
+    </>
+  );
+};
+
+export default memo(ColorSelectorNode);
diff --git a/example/src/Basic/index.tsx b/example/src/Basic/index.tsx
index ae003545..2335443f 100644
--- a/example/src/Basic/index.tsx
+++ b/example/src/Basic/index.tsx
@@ -1,49 +1,126 @@
-import React, { useState, MouseEvent } from 'react';
+import { useState, MouseEvent, useCallback } from 'react';
 
 import ReactFlow, {
-  removeElements,
   addEdge,
-  isNode,
   Background,
-  Elements,
-  BackgroundVariant,
-  FlowElement,
+  applyNodeChanges,
+  applyEdgeChanges,
+  MiniMap,
+  Controls,
   Node,
   Edge,
-  Connection,
+  NodeChange,
+  EdgeChange,
   OnLoadParams,
+  Connection,
 } from 'react-flow-renderer';
+import DebugNode from './DebugNode';
 
 const onNodeDragStop = (_: MouseEvent, node: Node) => console.log('drag stop', node);
-const onElementClick = (_: MouseEvent, element: FlowElement) => console.log('click', element);
+const onNodeClick = (_: MouseEvent, node: Node) => console.log('click', node);
+const onEdgeClick = (_: MouseEvent, edge: Edge) => console.log('click', edge);
 
-const initialElements: Elements = [
+const initialNodes: Node[] = [
   { id: '1', type: 'input', data: { label: 'Node 1' }, position: { x: 250, y: 5 }, className: 'light' },
+  {
+    id: '4',
+    data: { label: 'Node 4' },
+    position: { x: 100, y: 200 },
+    className: 'light',
+    style: { backgroundColor: 'rgba(255,50, 50, 0.5)', width: 500, height: 300 },
+  },
+  {
+    id: '4a',
+    data: { label: 'Node 4a' },
+    position: { x: 15, y: 15 },
+    className: 'light',
+    parentNode: '4',
+    extent: 'parent',
+  },
+  {
+    id: '4b',
+    data: { label: 'Node 4b' },
+    position: { x: 150, y: 50 },
+    className: 'light',
+    style: { backgroundColor: 'rgba(50, 50, 255, 0.5)', height: 200, width: 300 },
+    parentNode: '4',
+  },
+  {
+    id: '4b1',
+    data: { label: 'Node 4b1' },
+    position: { x: 20, y: 20 },
+    className: 'light',
+    parentNode: '4b',
+  },
+  {
+    id: '4b2',
+    data: { label: 'Node 4b2' },
+    position: { x: 100, y: 100 },
+    className: 'light',
+    parentNode: '4b',
+  },
+  {
+    id: '5',
+    data: { label: 'Node 5' },
+    position: { x: 650, y: 250 },
+    className: 'light',
+    style: { backgroundColor: 'rgba(20 ,200, 255, 1.5)', width: 400, height: 150 },
+    zIndex: 1000,
+  },
+  {
+    id: '5a',
+    data: { label: 'Node 5a' },
+    position: { x: 25, y: 50 },
+    className: 'light',
+    parentNode: '5',
+  },
+  {
+    id: '5b',
+    data: { label: 'Node 5b' },
+    position: { x: 225, y: 50 },
+    className: 'light',
+    parentNode: '5',
+  },
   { id: '2', data: { label: 'Node 2' }, position: { x: 100, y: 100 }, className: 'light' },
   { id: '3', data: { label: 'Node 3' }, position: { x: 400, y: 100 }, className: 'light' },
-  { id: '4', data: { label: 'Node 4' }, position: { x: 400, y: 200 }, className: 'light' },
+];
+
+const initialEdges: Edge[] = [
   { id: 'e1-2', source: '1', target: '2', animated: true },
   { id: 'e1-3', source: '1', target: '3' },
+  { id: 'e3-4', source: '3', target: '4', zIndex: 100 },
+  { id: 'e3-4b', source: '3', target: '4b' },
+  { id: 'e4a-4b1', source: '4a', target: '4b1' },
+  { id: 'e4a-4b2', source: '4a', target: '4b2', zIndex: 100 },
+  { id: 'e4b1-4b2', source: '4b1', target: '4b2' },
+  { id: '3-5', source: '3', target: '5' },
 ];
 
+const nodeTypes = {
+  default: DebugNode,
+};
+
 const BasicFlow = () => {
   const [rfInstance, setRfInstance] = useState<OnLoadParams | null>(null);
-  const [elements, setElements] = useState<Elements>(initialElements);
-  const onElementsRemove = (elementsToRemove: Elements) => setElements((els) => removeElements(elementsToRemove, els));
-  const onConnect = (params: Edge | Connection) => setElements((els) => addEdge(params, els));
-  const onLoad = (reactFlowInstance: OnLoadParams) => setRfInstance(reactFlowInstance);
+  const [nodes, setNodes] = useState<Node[]>(initialNodes);
+  const [edges, setEdges] = useState<Edge[]>(initialEdges);
+
+  const onConnect = useCallback((params: Edge | Connection) => {
+    setEdges((eds) => {
+      return addEdge(params, eds);
+    });
+  }, []);
+  const onLoad = useCallback((reactFlowInstance: OnLoadParams) => setRfInstance(reactFlowInstance), []);
 
   const updatePos = () => {
-    setElements((elms) => {
-      return elms.map((el) => {
-        if (isNode(el)) {
-          el.position = {
-            x: Math.random() * 400,
-            y: Math.random() * 400,
-          };
-        }
+    setNodes((nds) => {
+      return nds.map((n) => {
+        n.position = {
+          x: Math.random() * 400,
+          y: Math.random() * 400,
+        };
 
-        return el;
+        return n;
       });
     });
   };
@@ -52,31 +129,53 @@ const BasicFlow = () => {
   const resetTransform = () => rfInstance?.setTransform({ x: 0, y: 0, zoom: 1 });
 
   const toggleClassnames = () => {
-    setElements((elms) => {
-      return elms.map((el) => {
-        if (isNode(el)) {
-          el.className = el.className === 'light' ? 'dark' : 'light';
-        }
-
-        return el;
+    setNodes((nds) => {
+      return nds.map((n) => {
+        n.className = n.className === 'light' ? 'dark' : 'light';
+        return n;
       });
     });
   };
 
+  const toggleChildNodes = () => {
+    setNodes((nds) => {
+      return nds.map((n) => {
+        n.hidden = !!n.parentNode && !n.hidden;
+        return n;
+      });
+    });
+  };
+
+  const onNodesChange = useCallback((changes: NodeChange[]) => {
+    console.log('node change', changes);
+    setNodes((ns) => applyNodeChanges(changes, ns));
+  }, []);
+
+  const onEdgesChange = useCallback((changes: EdgeChange[]) => {
+    setEdges((es) => applyEdgeChanges(changes, es));
+  }, []);
+
   return (
     <ReactFlow
-      elements={elements}
+      nodes={nodes}
+      edges={edges}
       onLoad={onLoad}
-      onElementClick={onElementClick}
-      onElementsRemove={onElementsRemove}
+      onNodesChange={onNodesChange}
+      onEdgesChange={onEdgesChange}
+      onNodeClick={onNodeClick}
+      onEdgeClick={onEdgeClick}
       onConnect={onConnect}
       onNodeDragStop={onNodeDragStop}
       className="react-flow-basic-example"
       defaultZoom={1.5}
       minZoom={0.2}
       maxZoom={4}
+      onlyRenderVisibleElements={false}
+      nodeTypes={nodeTypes}
     >
-      <Background variant={BackgroundVariant.Lines} />
+      <MiniMap />
+      <Controls />
+      <Background />
 
       <div style={{ position: 'absolute', right: 10, top: 10, zIndex: 4 }}>
         <button onClick={resetTransform} style={{ marginRight: 5 }}>
@@ -88,6 +187,9 @@ const BasicFlow = () => {
         <button onClick={toggleClassnames} style={{ marginRight: 5 }}>
           toggle classnames
         </button>
+        <button style={{ marginRight: 5 }} onClick={toggleChildNodes}>
+          toggleChildNodes
+        </button>
         <button onClick={logToObject}>toObject</button>
       </div>
     </ReactFlow>
diff --git a/example/src/CustomNode/index.tsx b/example/src/CustomNode/index.tsx
index 804aabea..9f444c27 100644
--- a/example/src/CustomNode/index.tsx
+++ b/example/src/CustomNode/index.tsx
@@ -1,27 +1,30 @@
-import React, { useState, useEffect, MouseEvent } from 'react';
+import { useState, useEffect, MouseEvent, useCallback } from 'react';
 import { ChangeEvent } from 'react';
 
 import ReactFlow, {
-  isEdge,
-  removeElements,
   addEdge,
   MiniMap,
   Controls,
   Node,
-  FlowElement,
   OnLoadParams,
-  Elements,
   Position,
   SnapGrid,
   Connection,
   Edge,
+  NodeChange,
+  applyNodeChanges,
+  applyEdgeChanges,
+  EdgeChange,
 } from 'react-flow-renderer';
 
 import ColorSelectorNode from './ColorSelectorNode';
 
-const onLoad = (reactFlowInstance: OnLoadParams) => console.log('flow loaded:', reactFlowInstance);
+const onLoad = (reactFlowInstance: OnLoadParams) => {
+  console.log('flow loaded:', reactFlowInstance);
+  reactFlowInstance.fitView();
+};
 const onNodeDragStop = (_: MouseEvent, node: Node) => console.log('drag stop', node);
-const onElementClick = (_: MouseEvent, element: FlowElement) => console.log('click', element);
+const onNodeClick = (_: MouseEvent, node: Node) => console.log('click', node);
 
 const initBgColor = '#1A192B';
 
@@ -32,15 +35,16 @@ const nodeTypes = {
 };
 
 const CustomNodeFlow = () => {
-  const [elements, setElements] = useState<Elements>([]);
+  const [nodes, setNodes] = useState<Node[]>([]);
+  const [edges, setEdges] = useState<Edge[]>([]);
   const [bgColor, setBgColor] = useState<string>(initBgColor);
 
   useEffect(() => {
     const onChange = (event: ChangeEvent<HTMLInputElement>) => {
-      setElements((els) =>
-        els.map((e) => {
-          if (isEdge(e) || e.id !== '2') {
-            return e;
+      setNodes((nds) =>
+        nds.map((node) => {
+          if (node.id !== '2') {
+            return node;
           }
 
           const color = event.target.value;
@@ -48,9 +52,9 @@ const CustomNodeFlow = () => {
           setBgColor(color);
 
           return {
-            ...e,
+            ...node,
             data: {
-              ...e.data,
+              ...node.data,
               color,
             },
           };
@@ -58,7 +62,7 @@ const CustomNodeFlow = () => {
       );
     };
 
-    setElements([
+    setNodes([
       {
         id: '1',
         type: 'input',
@@ -87,22 +91,33 @@ const CustomNodeFlow = () => {
         position: { x: 550, y: 100 },
         targetPosition: Position.Left,
       },
+    ]);
 
+    setEdges([
       { id: 'e1-2', source: '1', target: '2', animated: true, style: { stroke: '#fff' } },
       { id: 'e2a-3', source: '2', sourceHandle: 'a', target: '3', animated: true, style: { stroke: '#fff' } },
       { id: 'e2b-4', source: '2', sourceHandle: 'b', target: '4', animated: true, style: { stroke: '#fff' } },
     ]);
   }, []);
 
-  const onElementsRemove = (elementsToRemove: Elements) => setElements((els) => removeElements(elementsToRemove, els));
   const onConnect = (params: Connection | Edge) =>
-    setElements((els) => addEdge({ ...params, animated: true, style: { stroke: '#fff' } }, els));
+    setEdges((eds) => addEdge({ ...params, animated: true, style: { stroke: '#fff' } }, eds));
+
+  const onNodesChange = useCallback((changes: NodeChange[]) => {
+    setNodes((ns) => applyNodeChanges(changes, ns));
+  }, []);
+
+  const onEdgesChange = useCallback((changes: EdgeChange[]) => {
+    setEdges((es) => applyEdgeChanges(changes, es));
+  }, []);
 
   return (
     <ReactFlow
-      elements={elements}
-      onElementClick={onElementClick}
-      onElementsRemove={onElementsRemove}
+      nodes={nodes}
+      edges={edges}
+      onNodesChange={onNodesChange}
+      onEdgesChange={onEdgesChange}
+      onNodeClick={onNodeClick}
       onConnect={onConnect}
       onNodeDragStop={onNodeDragStop}
       style={{ background: bgColor }}
diff --git a/example/src/FloatingEdges/FloatingConnectionLine.tsx b/example/src/FloatingEdges/FloatingConnectionLine.tsx
index 987af43d..ef202589 100644
--- a/example/src/FloatingEdges/FloatingConnectionLine.tsx
+++ b/example/src/FloatingEdges/FloatingConnectionLine.tsx
@@ -16,7 +16,9 @@ const FloatingConnectionLine: FC<ConnectionLineComponentProps> = ({
 
   const targetNode = {
     id: 'connection-target',
-    __rf: { width: 1, height: 1, position: { x: targetX, y: targetY } },
+    width: 1,
+    height: 1,
+    position: { x: targetX, y: targetY },
   } as Node;
 
   const { sx, sy } = getEdgeParams(sourceNode, targetNode);
diff --git a/example/src/FloatingEdges/FloatingEdge.tsx b/example/src/FloatingEdges/FloatingEdge.tsx
index b03241d1..1812ee2b 100644
--- a/example/src/FloatingEdges/FloatingEdge.tsx
+++ b/example/src/FloatingEdges/FloatingEdge.tsx
@@ -1,13 +1,15 @@
 import { FC, useMemo, CSSProperties } from 'react';
-import { EdgeProps, useStoreState, getBezierPath } from 'react-flow-renderer';
+import { EdgeProps, useStore, getBezierPath, ReactFlowState } from 'react-flow-renderer';
 
 import { getEdgeParams } from './utils';
 
-const FloatingEdge: FC<EdgeProps> = ({ id, source, target, style }) => {
-  const nodes = useStoreState((state) => state.nodes);
+const nodeSelector = (s: ReactFlowState) => s.nodeInternals;
 
-  const sourceNode = useMemo(() => nodes.find((n) => n.id === source), [source, nodes]);
-  const targetNode = useMemo(() => nodes.find((n) => n.id === target), [target, nodes]);
+const FloatingEdge: FC<EdgeProps> = ({ id, source, target, style }) => {
+  const nodeInternals = useStore(nodeSelector);
+
+  const sourceNode = useMemo(() => nodeInternals.get(source), [source, nodeInternals]);
+  const targetNode = useMemo(() => nodeInternals.get(target), [target, nodeInternals]);
 
   if (!sourceNode || !targetNode) {
     return null;
diff --git a/example/src/FloatingEdges/index.tsx b/example/src/FloatingEdges/index.tsx
index f1902515..8ec3ebc2 100644
--- a/example/src/FloatingEdges/index.tsx
+++ b/example/src/FloatingEdges/index.tsx
@@ -1,14 +1,17 @@
-import React, { useState } from 'react';
+import { useState, useCallback } from 'react';
 
 import ReactFlow, {
-  removeElements,
   addEdge,
   Background,
   OnLoadParams,
   EdgeTypesType,
-  Elements,
+  Node,
   Connection,
   Edge,
+  applyNodeChanges,
+  applyEdgeChanges,
+  NodeChange,
+  EdgeChange,
 } from 'react-flow-renderer';
 
 import './style.css';
@@ -19,24 +22,35 @@ import { createElements } from './utils';
 
 const onLoad = (reactFlowInstance: OnLoadParams) => reactFlowInstance.fitView();
 
-const initialElements: Elements = createElements();
+const { nodes: initialNodes, edges: initialEdges } = createElements();
 
 const edgeTypes: EdgeTypesType = {
   floating: FloatingEdge,
 };
 
-const NodeAsHandleFlow = () => {
-  const [elements, setElements] = useState<Elements>(initialElements);
+const FloatingEdges = () => {
+  const [nodes, setNodes] = useState<Node[]>(initialNodes);
+  const [edges, setEdges] = useState<Edge[]>(initialEdges);
 
-  const onElementsRemove = (elementsToRemove: Elements) => setElements((els) => removeElements(elementsToRemove, els));
+  const onConnect = useCallback((params: Edge | Connection) => {
+    setEdges((eds) => addEdge(params, eds));
+  }, []);
 
-  const onConnect = (params: Connection | Edge) => setElements((els) => addEdge({ ...params, type: 'floating' }, els));
+  const onNodesChange = useCallback((changes: NodeChange[]) => {
+    setNodes((ns) => applyNodeChanges(changes, ns));
+  }, []);
+
+  const onEdgesChange = useCallback((changes: EdgeChange[]) => {
+    setEdges((es) => applyEdgeChanges(changes, es));
+  }, []);
 
   return (
     <div className="floatingedges">
       <ReactFlow
-        elements={elements}
-        onElementsRemove={onElementsRemove}
+        nodes={nodes}
+        edges={edges}
+        onNodesChange={onNodesChange}
+        onEdgesChange={onEdgesChange}
         onConnect={onConnect}
         onLoad={onLoad}
         edgeTypes={edgeTypes}
@@ -48,4 +62,4 @@ const NodeAsHandleFlow = () => {
   );
 };
 
-export default NodeAsHandleFlow;
+export default FloatingEdges;
diff --git a/example/src/FloatingEdges/utils.ts b/example/src/FloatingEdges/utils.ts
index c8962b72..eb8797d2 100644
--- a/example/src/FloatingEdges/utils.ts
+++ b/example/src/FloatingEdges/utils.ts
@@ -1,18 +1,19 @@
-import { Position, ArrowHeadType, Node, XYPosition } from 'react-flow-renderer';
+import { Position, ArrowHeadType, XYPosition, Node, Edge } from 'react-flow-renderer';
 
 // this helper function returns the intersection point
 // of the line between the center of the intersectionNode and the target node
 function getNodeIntersection(intersectionNode: Node, targetNode: Node): XYPosition {
   // https://math.stackexchange.com/questions/1724792/an-algorithm-for-finding-the-intersection-point-between-a-center-of-vision-and-a
+
   const {
     width: intersectionNodeWidth,
     height: intersectionNodeHeight,
     position: intersectionNodePosition,
-  } = intersectionNode.__rf;
-  const targetPosition = targetNode.__rf.position;
+  } = intersectionNode;
+  const targetPosition = targetNode.position;
 
-  const w = intersectionNodeWidth / 2;
-  const h = intersectionNodeHeight / 2;
+  const w = (intersectionNodeWidth ?? 0) / 2;
+  const h = (intersectionNodeHeight ?? 0) / 2;
 
   const x2 = intersectionNodePosition.x + w;
   const y2 = intersectionNodePosition.y + h;
@@ -32,7 +33,7 @@ function getNodeIntersection(intersectionNode: Node, targetNode: Node): XYPositi
 
 // returns the position (top,right,bottom or right) passed node compared to the intersection point
 function getEdgePosition(node: Node, intersectionPoint: XYPosition) {
-  const n = { ...node.__rf.position, ...node.__rf };
+  const n = { ...node.position, ...node };
   const nx = Math.round(n.x);
   const ny = Math.round(n.y);
   const px = Math.round(intersectionPoint.x);
@@ -41,13 +42,13 @@ function getEdgePosition(node: Node, intersectionPoint: XYPosition) {
   if (px <= nx + 1) {
     return Position.Left;
   }
-  if (px >= nx + n.width - 1) {
+  if (px >= nx + (n.width ?? 0) - 1) {
     return Position.Right;
   }
   if (py <= ny + 1) {
     return Position.Top;
   }
-  if (py >= n.y + n.height - 1) {
+  if (py >= n.y + (n.height ?? 0) - 1) {
     return Position.Bottom;
   }
 
@@ -72,11 +73,18 @@ export function getEdgeParams(source: Node, target: Node) {
   };
 }
 
-export function createElements() {
-  const elements = [];
+type NodesAndEdges = {
+  nodes: Node[];
+  edges: Edge[];
+};
+
+export function createElements(): NodesAndEdges {
+  const nodes: Node[] = [];
+  const edges: Edge[] = [];
+
   const center = { x: window.innerWidth / 2, y: window.innerHeight / 2 };
 
-  elements.push({ id: 'target', data: { label: 'Target' }, position: center });
+  nodes.push({ id: 'target', data: { label: 'Target' }, position: center });
 
   for (let i = 0; i < 8; i++) {
     const degrees = i * (360 / 8);
@@ -84,16 +92,15 @@ export function createElements() {
     const x = 250 * Math.cos(radians) + center.x;
     const y = 250 * Math.sin(radians) + center.y;
 
-    elements.push({ id: `${i}`, data: { label: 'Source' }, position: { x, y } });
+    nodes.push({ id: `${i}`, data: { label: 'Source' }, position: { x, y } });
 
-    elements.push({
+    edges.push({
       id: `edge-${i}`,
       target: 'target',
       source: `${i}`,
       type: 'floating',
-      arrowHeadType: ArrowHeadType.Arrow,
     });
   }
 
-  return elements;
+  return { nodes, edges };
 }
diff --git a/example/src/Hidden/index.tsx b/example/src/Hidden/index.tsx
index 9869b2bc..4ac8eaa4 100644
--- a/example/src/Hidden/index.tsx
+++ b/example/src/Hidden/index.tsx
@@ -1,34 +1,43 @@
-import React, { useState } from 'react';
+import { useState, useCallback } from 'react';
 
 import { useEffect } from 'react';
-import ReactFlow, { addEdge, MiniMap, Controls, Connection, Edge, Elements } from 'react-flow-renderer';
+import ReactFlow, { addEdge, MiniMap, Controls, Connection, Edge, Node } from 'react-flow-renderer';
 
-const initialElements: Elements = [
-  { id: '1', type: 'input', data: { label: 'Node 1' }, position: { x: 250, y: 5 } },
-  { id: '2', data: { label: 'Node 2' }, position: { x: 100, y: 100 } },
-  { id: '3', data: { label: 'Node 3' }, position: { x: 400, y: 100 } },
-  { id: '4', data: { label: 'Node 4' }, position: { x: 400, y: 200 } },
+const initialNodes: Node[] = [
+  { id: '1', type: 'input', hidden: true, data: { label: 'Node 1' }, position: { x: 250, y: 5 } },
+  { id: '2', hidden: true, data: { label: 'Node 2' }, position: { x: 100, y: 100 } },
+  { id: '3', hidden: true, data: { label: 'Node 3' }, position: { x: 400, y: 100 } },
+  { id: '4', hidden: true, data: { label: 'Node 4' }, position: { x: 400, y: 200 } },
+];
+
+const initialEdges: Edge[] = [
   { id: 'e1-2', source: '1', target: '2' },
   { id: 'e1-3', source: '1', target: '3' },
   { id: 'e3-4', source: '3', target: '4' },
 ];
 
+const setHidden = (hidden: boolean) => (els: any[]) =>
+  els.map((e: any) => {
+    e.hidden = hidden;
+    return e;
+  });
+
 const HiddenFlow = () => {
-  const [elements, setElements] = useState<Elements>(initialElements);
-  const [isHidden, setIsHidden] = useState<boolean>(false);
-  const onConnect = (params: Connection | Edge) => setElements((els) => addEdge(params, els));
+  const [nodes, setNodes] = useState<Node[]>(initialNodes);
+  const [edges, setEdges] = useState<Edge[]>(initialEdges);
+  const [isHidden, setIsHidden] = useState<boolean>(true);
+
+  const onConnect = useCallback((params: Edge | Connection) => {
+    setEdges((eds) => addEdge(params, eds));
+  }, []);
 
   useEffect(() => {
-    setElements((els) =>
-      els.map((e) => {
-        e.isHidden = isHidden;
-        return e;
-      })
-    );
+    setNodes(setHidden(isHidden));
+    setNodes(setHidden(isHidden));
   }, [isHidden]);
 
   return (
-    <ReactFlow elements={elements} onConnect={onConnect}>
+    <ReactFlow nodes={nodes} edges={edges} onConnect={onConnect}>
       <MiniMap />
       <Controls />
 
diff --git a/example/src/Layouting/index.tsx b/example/src/Layouting/index.tsx
index 34ee23e6..56484640 100644
--- a/example/src/Layouting/index.tsx
+++ b/example/src/Layouting/index.tsx
@@ -1,74 +1,87 @@
-import React, { useState } from 'react';
+import { useState, useCallback } from 'react';
 import ReactFlow, {
   ReactFlowProvider,
   addEdge,
-  removeElements,
+  applyNodeChanges,
+  applyEdgeChanges,
   Controls,
-  isNode,
-  Elements,
+  NodeChange,
+  EdgeChange,
+  Node,
   Connection,
   Edge,
-  NodeExtent,
+  CoordinateExtent,
   Position,
 } from 'react-flow-renderer';
 import dagre from 'dagre';
 
-import initialElements from './initial-elements';
+import initialNodesAndEdges from './initial-elements';
 
 import './layouting.css';
 
 const dagreGraph = new dagre.graphlib.Graph();
 dagreGraph.setDefaultEdgeLabel(() => ({}));
 
-const nodeExtent: NodeExtent = [
+const nodeExtent: CoordinateExtent = [
   [0, 0],
   [1000, 1000],
 ];
 
 const LayoutFlow = () => {
-  const [elements, setElements] = useState<Elements>(initialElements);
-  const onConnect = (params: Connection | Edge) => setElements((els) => addEdge(params, els));
-  const onElementsRemove = (elementsToRemove: Elements) => setElements((els) => removeElements(elementsToRemove, els));
+  const [nodes, setNodes] = useState<Node[]>(initialNodesAndEdges.nodes);
+  const [edges, setEdges] = useState<Edge[]>(initialNodesAndEdges.edges);
+
+  const onConnect = useCallback((params: Edge | Connection) => {
+    setEdges((eds) => {
+      return addEdge(params, eds);
+    });
+  }, []);
 
   const onLayout = (direction: string) => {
     const isHorizontal = direction === 'LR';
     dagreGraph.setGraph({ rankdir: direction });
 
-    elements.forEach((el) => {
-      if (isNode(el)) {
-        dagreGraph.setNode(el.id, { width: 150, height: 50 });
-      } else {
-        dagreGraph.setEdge(el.source, el.target);
-      }
+    nodes.forEach((node) => {
+      dagreGraph.setNode(node.id, { width: 150, height: 50 });
+    });
+
+    edges.forEach((edge) => {
+      dagreGraph.setEdge(edge.source, edge.target);
     });
 
     dagre.layout(dagreGraph);
 
-    const layoutedElements = elements.map((el) => {
-      if (isNode(el)) {
-        const nodeWithPosition = dagreGraph.node(el.id);
-        el.targetPosition = isHorizontal ? Position.Left : Position.Top;
-        el.sourcePosition = isHorizontal ? Position.Right : Position.Bottom;
-        // we need to pass a slightly different position in order to notify react flow about the change
-        // @TODO how can we change the position handling so that we dont need this hack?
-        el.position = { x: nodeWithPosition.x + Math.random() / 1000, y: nodeWithPosition.y };
-      }
+    const layoutedNodes = nodes.map((node) => {
+      const nodeWithPosition = dagreGraph.node(node.id);
+      node.targetPosition = isHorizontal ? Position.Left : Position.Top;
+      node.sourcePosition = isHorizontal ? Position.Right : Position.Bottom;
+      // we need to pass a slightly different position in order to notify react flow about the change
+      // @TODO how can we change the position handling so that we dont need this hack?
+      node.position = { x: nodeWithPosition.x + Math.random() / 1000, y: nodeWithPosition.y };
 
-      return el;
+      return node;
     });
 
-    setElements(layoutedElements);
+    setNodes(layoutedNodes);
   };
 
+  const onNodesChange = useCallback((changes: NodeChange[]) => setNodes((ns) => applyNodeChanges(changes, ns)), []);
+
+  const onEdgesChange = useCallback((changes: EdgeChange[]) => {
+    setEdges((es) => applyEdgeChanges(changes, es));
+  }, []);
+
   return (
     <div className="layoutflow">
       <ReactFlowProvider>
         <ReactFlow
-          elements={elements}
+          nodes={nodes}
+          edges={edges}
           onConnect={onConnect}
-          onElementsRemove={onElementsRemove}
           nodeExtent={nodeExtent}
           onLoad={() => onLayout('TB')}
+          onNodesChange={onNodesChange}
+          onEdgesChange={onEdgesChange}
         >
           <Controls />
         </ReactFlow>
diff --git a/example/src/Layouting/initial-elements.ts b/example/src/Layouting/initial-elements.ts
index 404b1b8b..ce763320 100644
--- a/example/src/Layouting/initial-elements.ts
+++ b/example/src/Layouting/initial-elements.ts
@@ -1,8 +1,8 @@
-import { Elements, XYPosition } from 'react-flow-renderer';
+import { Node, Edge, XYPosition } from 'react-flow-renderer';
 
 const position: XYPosition = { x: 0, y: 0 };
 
-const elements: Elements = [
+const nodes: Node[] = [
   {
     id: '1',
     type: 'input',
@@ -56,6 +56,9 @@ const elements: Elements = [
     position,
   },
   { id: '7', type: 'output', data: { label: 'output' }, position: { x: 400, y: 450 } },
+];
+
+const edges: Edge[] = [
   { id: 'e12', source: '1', target: '2', type: 'smoothstep' },
   { id: 'e13', source: '1', target: '3', type: 'smoothstep' },
   { id: 'e22a', source: '2', target: '2a', type: 'smoothstep' },
@@ -68,4 +71,6 @@ const elements: Elements = [
   { id: 'e57', source: '5', target: '7', type: 'smoothstep' },
 ];
 
-export default elements;
+const nodesAndEdges = { nodes, edges };
+
+export default nodesAndEdges;
diff --git a/example/src/NestedNodes/index.tsx b/example/src/NestedNodes/index.tsx
new file mode 100644
index 00000000..ed445a64
--- /dev/null
+++ b/example/src/NestedNodes/index.tsx
@@ -0,0 +1,186 @@
+import { useState, MouseEvent, useCallback } from 'react';
+
+import ReactFlow, {
+  addEdge,
+  Background,
+  applyNodeChanges,
+  applyEdgeChanges,
+  MiniMap,
+  Controls,
+  Node,
+  Edge,
+  NodeChange,
+  EdgeChange,
+  OnLoadParams,
+  Connection,
+} from 'react-flow-renderer';
+
+const onNodeDragStop = (_: MouseEvent, node: Node) => console.log('drag stop', node);
+const onNodeClick = (_: MouseEvent, node: Node) => console.log('click', node);
+const onEdgeClick = (_: MouseEvent, edge: Edge) => console.log('click', edge);
+
+const initialNodes: Node[] = [
+  { id: '1', type: 'input', data: { label: 'Node 1' }, position: { x: 250, y: 5 }, className: 'light' },
+  {
+    id: '2',
+    data: { label: 'Node 2' },
+    position: { x: 100, y: 100 },
+    className: 'light',
+    style: { backgroundColor: 'rgba(255, 0, 0, 0.8)', width: 200, height: 200 },
+  },
+  {
+    id: '2a',
+    data: { label: 'Node 2a' },
+    position: { x: 10, y: 50 },
+    parentNode: '2',
+  },
+  { id: '3', data: { label: 'Node 3' }, position: { x: 320, y: 100 }, className: 'light' },
+  {
+    id: '4',
+    data: { label: 'Node 4' },
+    position: { x: 320, y: 200 },
+    className: 'light',
+    style: { backgroundColor: 'rgba(255, 0, 0, 0.7)', width: 300, height: 300 },
+  },
+  {
+    id: '4a',
+    data: { label: 'Node 4a' },
+    position: { x: 15, y: 65 },
+    className: 'light',
+    parentNode: '4',
+    extent: 'parent',
+  },
+  {
+    id: '4b',
+    data: { label: 'Node 4b' },
+    position: { x: 15, y: 120 },
+    className: 'light',
+    style: { backgroundColor: 'rgba(255, 0, 255, 0.7)', height: 150, width: 270 },
+    parentNode: '4',
+  },
+  {
+    id: '4b1',
+    data: { label: 'Node 4b1' },
+    position: { x: 20, y: 40 },
+    className: 'light',
+    parentNode: '4b',
+  },
+  {
+    id: '4b2',
+    data: { label: 'Node 4b2' },
+    position: { x: 100, y: 100 },
+    className: 'light',
+    parentNode: '4b',
+  },
+];
+
+const initialEdges: Edge[] = [
+  { id: 'e1-2', source: '1', target: '2', animated: true },
+  { id: 'e1-3', source: '1', target: '3' },
+  { id: 'e2a-4a', source: '2a', target: '4a' },
+  { id: 'e3-4', source: '3', target: '4' },
+  { id: 'e3-4b', source: '3', target: '4b' },
+  { id: 'e4a-4b1', source: '4a', target: '4b1' },
+  { id: 'e4a-4b2', source: '4a', target: '4b2' },
+  { id: 'e4b1-4b2', source: '4b1', target: '4b2' },
+];
+
+function shuffle() {
+  return Math.random() - Math.random();
+}
+
+const NestedFlow = () => {
+  const [rfInstance, setRfInstance] = useState<OnLoadParams | null>(null);
+  const [nodes, setNodes] = useState<Node[]>(initialNodes);
+  const [edges, setEdges] = useState<Edge[]>(initialEdges);
+
+  const onConnect = useCallback((params: Edge | Connection) => {
+    setEdges((eds) => {
+      return addEdge(params, eds);
+    });
+  }, []);
+  const onLoad = useCallback((reactFlowInstance: OnLoadParams) => setRfInstance(reactFlowInstance), []);
+
+  const updatePos = () => {
+    setNodes((nds) => {
+      return nds.map((n) => {
+        n.position = {
+          x: Math.random() * 400,
+          y: Math.random() * 400,
+        };
+
+        return n;
+      });
+    });
+  };
+
+  const logToObject = () => console.log(rfInstance?.toObject());
+  const resetTransform = () => rfInstance?.setTransform({ x: 0, y: 0, zoom: 1 });
+
+  const toggleClassnames = () => {
+    setNodes((nds) => {
+      return nds.map((n) => {
+        n.className = n.className === 'light' ? 'dark' : 'light';
+        return n;
+      });
+    });
+  };
+
+  const toggleChildNodes = () => {
+    setNodes((nds) => {
+      return nds.map((n) => {
+        n.hidden = !!n.parentNode && !n.hidden;
+        return n;
+      });
+    });
+  };
+
+  const onNodesChange = useCallback((changes: NodeChange[]) => {
+    setNodes((ns) => applyNodeChanges(changes, ns));
+  }, []);
+
+  const onEdgesChange = useCallback((changes: EdgeChange[]) => {
+    setEdges((es) => applyEdgeChanges(changes, es));
+  }, []);
+
+  return (
+    <ReactFlow
+      nodes={nodes}
+      edges={edges}
+      onLoad={onLoad}
+      onNodesChange={onNodesChange}
+      onEdgesChange={onEdgesChange}
+      onNodeClick={onNodeClick}
+      onEdgeClick={onEdgeClick}
+      onConnect={onConnect}
+      onNodeDragStop={onNodeDragStop}
+      className="react-flow-basic-example"
+      defaultZoom={1.5}
+      minZoom={0.2}
+      maxZoom={4}
+      onlyRenderVisibleElements={false}
+    >
+      <MiniMap />
+      <Controls />
+      <Background />
+
+      <div style={{ position: 'absolute', right: 10, top: 10, zIndex: 4 }}>
+        <button onClick={resetTransform} style={{ marginRight: 5 }}>
+          reset transform
+        </button>
+        <button onClick={updatePos} style={{ marginRight: 5 }}>
+          change pos
+        </button>
+        <button onClick={toggleClassnames} style={{ marginRight: 5 }}>
+          toggle classnames
+        </button>
+        <button style={{ marginRight: 5 }} onClick={toggleChildNodes}>
+          toggleChildNodes
+        </button>
+        <button onClick={logToObject}>toObject</button>
+      </div>
+    </ReactFlow>
+  );
+};
+
+export default NestedFlow;
diff --git a/example/src/Stress/index.tsx b/example/src/Stress/index.tsx
index b734c6b6..5b97298c 100644
--- a/example/src/Stress/index.tsx
+++ b/example/src/Stress/index.tsx
@@ -1,58 +1,63 @@
-import React, { useState, CSSProperties } from 'react';
+import { useState, CSSProperties, useCallback } from 'react';
 import ReactFlow, {
-  removeElements,
-  addEdge,
   MiniMap,
-  isNode,
   Controls,
   Background,
   OnLoadParams,
-  Elements,
-  Connection,
   Edge,
+  Node,
+  NodeChange,
+  applyNodeChanges,
+  Connection,
+  addEdge,
 } from 'react-flow-renderer';
 
-import { getElements } from './utils';
+import { getNodesAndEdges } from './utils';
 
 const buttonWrapperStyles: CSSProperties = { position: 'absolute', right: 10, top: 10, zIndex: 4 };
 
 const onLoad = (reactFlowInstance: OnLoadParams) => {
   reactFlowInstance.fitView();
-  console.log(reactFlowInstance.getElements());
+  console.log(reactFlowInstance.getNodes());
 };
 
-const initialElements: Elements = getElements(30, 30);
+const { nodes: initialNodes, edges: initialEdges } = getNodesAndEdges(30, 30);
 
 const StressFlow = () => {
-  const [elements, setElements] = useState<Elements>(initialElements);
-  const onElementsRemove = (elementsToRemove: Elements) => setElements((els) => removeElements(elementsToRemove, els));
-  const onConnect = (params: Connection | Edge) => setElements((els) => addEdge(params, els));
-
+  const [nodes, setNodes] = useState<Node[]>(initialNodes);
+  const [edges, setEdges] = useState<Edge[]>(initialEdges);
+  const onConnect = useCallback((params: Edge | Connection) => {
+    setEdges((eds) => addEdge(params, eds));
+  }, []);
   const updatePos = () => {
-    setElements((elms) => {
-      return elms.map((el) => {
-        if (isNode(el)) {
-          return {
-            ...el,
-            position: {
-              x: Math.random() * window.innerWidth,
-              y: Math.random() * window.innerHeight,
-            },
-          };
-        }
-
-        return el;
+    setNodes((nds) => {
+      return nds.map((n) => {
+        return {
+          ...n,
+          position: {
+            x: Math.random() * window.innerWidth,
+            y: Math.random() * window.innerHeight,
+          },
+        };
       });
     });
   };
 
   const updateElements = () => {
     const grid = Math.ceil(Math.random() * 10);
-    setElements(getElements(grid, grid));
+    const initialElements = getNodesAndEdges(grid, grid);
+    setNodes(initialElements.nodes);
+    setEdges(initialElements.edges);
   };
 
+  const onNodesChange = useCallback((changes: NodeChange[]) => {
+    console.log('node change', changes);
+
+    setNodes((ns) => applyNodeChanges(changes, ns));
+  }, []);
+
   return (
-    <ReactFlow elements={elements} onLoad={onLoad} onElementsRemove={onElementsRemove} onConnect={onConnect}>
+    <ReactFlow nodes={nodes} edges={edges} onLoad={onLoad} onConnect={onConnect} onNodesChange={onNodesChange}>
       <MiniMap />
       <Controls />
       <Background />
diff --git a/example/src/Stress/utils.ts b/example/src/Stress/utils.ts
index a37eb8de..c5c6e307 100644
--- a/example/src/Stress/utils.ts
+++ b/example/src/Stress/utils.ts
@@ -1,7 +1,13 @@
-import { Elements } from 'react-flow-renderer';
+import { Node, Edge } from 'react-flow-renderer';
 
-export function getElements(xElements: number = 10, yElements: number = 10): Elements {
-  const initialElements = [];
+type ElementsCollection = {
+  nodes: Node[];
+  edges: Edge[];
+};
+
+export function getNodesAndEdges(xElements: number = 10, yElements: number = 10): ElementsCollection {
+  const initialNodes = [];
+  const initialEdges: Edge[] = [];
   let nodeId = 1;
   let recentNodeId = null;
 
@@ -15,10 +21,10 @@ export function getElements(xElements: number = 10, yElements: number = 10): Ele
         data,
         position,
       };
-      initialElements.push(node);
+      initialNodes.push(node);
 
       if (recentNodeId && nodeId <= xElements * yElements) {
-        initialElements.push({ id: `${x}-${y}`, source: recentNodeId.toString(), target: nodeId.toString() });
+        initialEdges.push({ id: `${x}-${y}`, source: recentNodeId.toString(), target: nodeId.toString() });
       }
 
       recentNodeId = nodeId;
@@ -26,5 +32,8 @@ export function getElements(xElements: number = 10, yElements: number = 10): Ele
     }
   }
 
-  return initialElements;
+  return {
+    nodes: initialNodes,
+    edges: initialEdges,
+  };
 }
diff --git a/example/src/UpdateNode/index.tsx b/example/src/UpdateNode/index.tsx
index a7cce8f2..e38babad 100644
--- a/example/src/UpdateNode/index.tsx
+++ b/example/src/UpdateNode/index.tsx
@@ -1,64 +1,71 @@
-import React, { useEffect, useState } from 'react';
-import ReactFlow, { Elements } from 'react-flow-renderer';
+import { useEffect, useState, useCallback } from 'react';
+import ReactFlow, { Node, Edge, applyNodeChanges, NodeChange } from 'react-flow-renderer';
 
 import './updatenode.css';
 
-const initialElements: Elements = [
+const initialNodes: Node[] = [
   { id: '1', data: { label: '-' }, position: { x: 100, y: 100 } },
   { id: '2', data: { label: 'Node 2' }, position: { x: 100, y: 200 } },
-  { id: 'e1-2', source: '1', target: '2' },
 ];
 
+const initialEdges: Edge[] = [{ id: 'e1-2', source: '1', target: '2' }];
+
 const UpdateNode = () => {
-  const [elements, setElements] = useState<Elements>(initialElements);
+  const [nodes, setNodes] = useState<Node[]>(initialNodes);
+  const [edges, setEdges] = useState<Edge[]>(initialEdges);
+
   const [nodeName, setNodeName] = useState<string>('Node 1');
   const [nodeBg, setNodeBg] = useState<string>('#eee');
   const [nodeHidden, setNodeHidden] = useState<boolean>(false);
 
   useEffect(() => {
-    setElements((els) =>
-      els.map((el) => {
-        if (el.id === '1') {
+    setNodes((nds) =>
+      nds.map((n) => {
+        if (n.id === '1') {
           // it's important that you create a new object here in order to notify react flow about the change
-          el.data = {
-            ...el.data,
+          n.data = {
+            ...n.data,
             label: nodeName,
           };
         }
 
-        return el;
+        return n;
       })
     );
-  }, [nodeName, setElements]);
+  }, [nodeName]);
 
   useEffect(() => {
-    setElements((els) =>
-      els.map((el) => {
-        if (el.id === '1') {
+    setNodes((nds) =>
+      nds.map((n) => {
+        if (n.id === '1') {
           // it's important that you create a new object here in order to notify react flow about the change
-          el.style = { ...el.style, backgroundColor: nodeBg };
+          n.style = { ...n.style, backgroundColor: nodeBg };
         }
 
-        return el;
+        return n;
       })
     );
-  }, [nodeBg, setElements]);
+  }, [nodeBg]);
 
   useEffect(() => {
-    setElements((els) =>
-      els.map((el) => {
-        if (el.id === '1' || el.id === 'e1-2') {
+    setNodes((nds) =>
+      nds.map((n) => {
+        if (n.id === '1' || n.id === 'e1-2') {
           // when you update a simple type you can just update the value
-          el.isHidden = nodeHidden;
+          n.hidden = nodeHidden;
         }
 
-        return el;
+        return n;
       })
     );
-  }, [nodeHidden, setElements]);
+  }, [nodeHidden]);
+
+  const onNodesChange = useCallback((changes: NodeChange[]) => {
+    setNodes((ns) => applyNodeChanges(changes, ns));
+  }, []);
 
   return (
-    <ReactFlow elements={elements} defaultZoom={1.5} minZoom={0.2} maxZoom={4}>
+    <ReactFlow nodes={nodes} edges={edges} defaultZoom={1.5} minZoom={0.2} maxZoom={4} onNodesChange={onNodesChange}>
       <div className="updatenode__controls">
         <label>label:</label>
         <input value={nodeName} onChange={(evt) => setNodeName(evt.target.value)} />
diff --git a/example/src/index.tsx b/example/src/index.tsx
index 8371c0af..45aa0a16 100644
--- a/example/src/index.tsx
+++ b/example/src/index.tsx
@@ -1,140 +1,50 @@
-import React, { ChangeEvent } from 'react';
+import { ChangeEvent } from 'react';
 import ReactDOM from 'react-dom';
 import { BrowserRouter as Router, Route, Switch, withRouter } from 'react-router-dom';
 
-import Overview from './Overview';
 import Basic from './Basic';
-import CustomNode from './CustomNode';
-import Stress from './Stress';
-import Interaction from './Interaction';
-import Empty from './Empty';
-import Edges from './Edges';
-import Validation from './Validation';
-import Provider from './Provider';
-import Hidden from './Hidden';
-import EdgeTypes from './EdgeTypes';
-import CustomConnectionLine from './CustomConnectionLine';
-import FloatingEdges from './FloatingEdges';
-import NodeTypeChange from './NodeTypeChange';
-import NodeTypesObjectChange from './NodeTypesObjectChange';
-import UpdatableEdge from './UpdatableEdge';
 import UpdateNode from './UpdateNode';
-import SaveRestore from './SaveRestore';
-import DragNDrop from './DragNDrop';
-import Layout from './Layouting';
-import SwitchFlows from './Switch';
-import UseZoomPanHelper from './UseZoomPanHelper';
-import UseUpdateNodeInternals from './UseUpdateNodeInternals';
-import Undirectional from './Undirectional';
-import MultiFlows from './MultiFlows';
-import DragHandle from './DragHandle';
+import Stress from './Stress';
+import CustomNode from './CustomNode';
+import FloatingEdges from './FloatingEdges';
+import Layouting from './Layouting';
+import NestedNodes from './NestedNodes';
+import Hidden from './Hidden';
 
 import './index.css';
 
 const routes = [
   {
     path: '/',
-    component: Overview,
-  },
-  {
-    path: '/edges',
-    component: Edges,
-  },
-  {
-    path: '/custom-node',
-    component: CustomNode,
-  },
-  {
-    path: '/validation',
-    component: Validation,
-  },
-  {
-    path: '/provider',
-    component: Provider,
-  },
-  {
-    path: '/stress',
-    component: Stress,
-  },
-  {
-    path: '/interaction',
-    component: Interaction,
-  },
-  {
-    path: '/basic',
     component: Basic,
   },
-  {
-    path: '/empty',
-    component: Empty,
-  },
-  {
-    path: '/hidden',
-    component: Hidden,
-  },
-  {
-    path: '/edge-types',
-    component: EdgeTypes,
-  },
-  {
-    path: '/custom-connectionline',
-    component: CustomConnectionLine,
-  },
-  {
-    path: '/floating-edges',
-    component: FloatingEdges,
-  },
-  {
-    path: '/nodetype-change',
-    component: NodeTypeChange,
-  },
-  {
-    path: '/nodetypesobject-change',
-    component: NodeTypesObjectChange,
-  },
-  {
-    path: '/updatable-edge',
-    component: UpdatableEdge,
-  },
   {
     path: '/update-node',
     component: UpdateNode,
   },
   {
-    path: '/save-restore',
-    component: SaveRestore,
+    path: '/stress',
+    component: Stress,
   },
   {
-    path: '/drag-and-drop',
-    component: DragNDrop,
+    path: '/custom-node',
+    component: CustomNode,
+  },
+  {
+    path: '/floating-edges',
+    component: FloatingEdges,
   },
   {
     path: '/layouting',
-    component: Layout,
+    component: Layouting,
   },
   {
-    path: '/switch',
-    component: SwitchFlows,
+    path: '/nested-nodes',
+    component: NestedNodes,
   },
   {
-    path: '/usezoompanhelper',
-    component: UseZoomPanHelper,
-  },
-  {
-    path: '/useupdatenodeinternals',
-    component: UseUpdateNodeInternals,
-  },
-  {
-    path: '/undirectional',
-    component: Undirectional,
-  },
-  {
-    path: '/multiflows',
-    component: MultiFlows,
-  },
-  {
-    path: '/draghandle',
-    component: DragHandle,
+    path: '/hidden',
+    component: Hidden,
   },
 ];
 
diff --git a/example/src_oldapi/Basic/index.tsx b/example/src_oldapi/Basic/index.tsx
new file mode 100644
index 00000000..19523f9b
--- /dev/null
+++ b/example/src_oldapi/Basic/index.tsx
@@ -0,0 +1,98 @@
+import React, { useState, MouseEvent } from 'react';
+
+import ReactFlow, {
+  removeElements,
+  addEdge,
+  isNode,
+  Background,
+  Elements,
+  BackgroundVariant,
+  FlowElement,
+  Node,
+  Edge,
+  Connection,
+  OnLoadParams,
+} from 'react-flow-renderer';
+
+const onNodeDragStop = (_: MouseEvent, node: Node) => console.log('drag stop', node);
+const onElementClick = (_: MouseEvent, element: FlowElement) => console.log('click', element);
+
+const initialNodes: Node[] = [
+  { id: '1', type: 'input', data: { label: 'Node 1' }, position: { x: 250, y: 5 }, className: 'light' },
+  { id: '2', data: { label: 'Node 2' }, position: { x: 100, y: 100 }, className: 'light' },
+  { id: '3', data: { label: 'Node 3' }, position: { x: 400, y: 100 }, className: 'light' },
+  { id: '4', data: { label: 'Node 4' }, position: { x: 400, y: 200 }, className: 'light' },
+];
+
+const initialEdges: Edge[] = [
+  { id: 'e1-2', source: '1', target: '2', animated: true },
+  { id: 'e1-3', source: '1', target: '3' },
+];
+
+const BasicFlow = () => {
+  const [rfInstance, setRfInstance] = useState<OnLoadParams | null>(null);
+  const [nodes, setNodes] = useState<Node[]>(initialNodes);
+  const [edges, setEdges] = useState<Edge[]>(initialEdges);
+  // const onElementsRemove = (elementsToRemove: Elements) => setNodes((els) => removeElements(elementsToRemove, els));
+  // const onConnect = (params: Edge | Connection) => setNodes((els) => addEdge(params, els));
+  const onLoad = (reactFlowInstance: OnLoadParams) => setRfInstance(reactFlowInstance);
+
+  const updatePos = () => {
+    setNodes((nds) => {
+      return nds.map((n) => {
+        n.position = {
+          x: Math.random() * 400,
+          y: Math.random() * 400,
+        };
+
+        return n;
+      });
+    });
+  };
+
+  const logToObject = () => console.log(rfInstance?.toObject());
+  const resetTransform = () => rfInstance?.setTransform({ x: 0, y: 0, zoom: 1 });
+
+  const toggleClassnames = () => {
+    setNodes((nds) => {
+      return nds.map((n) => {
+        n.className = n.className === 'light' ? 'dark' : 'light';
+
+        return n;
+      });
+    });
+  };
+
+  return (
+    <ReactFlow
+      nodes={nodes}
+      edges={edges}
+      onLoad={onLoad}
+      onElementClick={onElementClick}
+      // onElementsRemove={onElementsRemove}
+      // onConnect={onConnect}
+      onNodeDragStop={onNodeDragStop}
+      className="react-flow-basic-example"
+      defaultZoom={1.5}
+      minZoom={0.2}
+      maxZoom={4}
+    >
+      <Background variant={BackgroundVariant.Lines} />
+
+      <div style={{ position: 'absolute', right: 10, top: 10, zIndex: 4 }}>
+        <button onClick={resetTransform} style={{ marginRight: 5 }}>
+          reset transform
+        </button>
+        <button onClick={updatePos} style={{ marginRight: 5 }}>
+          change pos
+        </button>
+        <button onClick={toggleClassnames} style={{ marginRight: 5 }}>
+          toggle classnames
+        </button>
+        <button onClick={logToObject}>toObject</button>
+      </div>
+    </ReactFlow>
+  );
+};
+
+export default BasicFlow;
diff --git a/example/src/CustomConnectionLine/ConnectionLine.tsx b/example/src_oldapi/CustomConnectionLine/ConnectionLine.tsx
similarity index 100%
rename from example/src/CustomConnectionLine/ConnectionLine.tsx
rename to example/src_oldapi/CustomConnectionLine/ConnectionLine.tsx
diff --git a/example/src/CustomConnectionLine/index.tsx b/example/src_oldapi/CustomConnectionLine/index.tsx
similarity index 100%
rename from example/src/CustomConnectionLine/index.tsx
rename to example/src_oldapi/CustomConnectionLine/index.tsx
diff --git a/example/src_oldapi/CustomNode/ColorSelectorNode.tsx b/example/src_oldapi/CustomNode/ColorSelectorNode.tsx
new file mode 100644
index 00000000..7ad20340
--- /dev/null
+++ b/example/src_oldapi/CustomNode/ColorSelectorNode.tsx
@@ -0,0 +1,25 @@
+import React, { memo, FC, CSSProperties } from 'react';
+
+import { Handle, Position, NodeProps, Connection, Edge } from 'react-flow-renderer';
+
+const targetHandleStyle: CSSProperties = { background: '#555' };
+const sourceHandleStyleA: CSSProperties = { ...targetHandleStyle, top: 10 };
+const sourceHandleStyleB: CSSProperties = { ...targetHandleStyle, bottom: 10, top: 'auto' };
+
+const onConnect = (params: Connection | Edge) => console.log('handle onConnect', params);
+
+const ColorSelectorNode: FC<NodeProps> = ({ data, isConnectable }) => {
+  return (
+    <>
+      <Handle type="target" position={Position.Left} style={targetHandleStyle} onConnect={onConnect} />
+      <div>
+        Custom Color Picker Node: <strong>{data.color}</strong>
+      </div>
+      <input className="nodrag" type="color" onChange={data.onChange} defaultValue={data.color} />
+      <Handle type="source" position={Position.Right} id="a" style={sourceHandleStyleA} isConnectable={isConnectable} />
+      <Handle type="source" position={Position.Right} id="b" style={sourceHandleStyleB} isConnectable={isConnectable} />
+    </>
+  );
+};
+
+export default memo(ColorSelectorNode);
diff --git a/example/src_oldapi/CustomNode/index.tsx b/example/src_oldapi/CustomNode/index.tsx
new file mode 100644
index 00000000..804aabea
--- /dev/null
+++ b/example/src_oldapi/CustomNode/index.tsx
@@ -0,0 +1,135 @@
+import React, { useState, useEffect, MouseEvent } from 'react';
+import { ChangeEvent } from 'react';
+
+import ReactFlow, {
+  isEdge,
+  removeElements,
+  addEdge,
+  MiniMap,
+  Controls,
+  Node,
+  FlowElement,
+  OnLoadParams,
+  Elements,
+  Position,
+  SnapGrid,
+  Connection,
+  Edge,
+} from 'react-flow-renderer';
+
+import ColorSelectorNode from './ColorSelectorNode';
+
+const onLoad = (reactFlowInstance: OnLoadParams) => console.log('flow loaded:', reactFlowInstance);
+const onNodeDragStop = (_: MouseEvent, node: Node) => console.log('drag stop', node);
+const onElementClick = (_: MouseEvent, element: FlowElement) => console.log('click', element);
+
+const initBgColor = '#1A192B';
+
+const connectionLineStyle = { stroke: '#fff' };
+const snapGrid: SnapGrid = [16, 16];
+const nodeTypes = {
+  selectorNode: ColorSelectorNode,
+};
+
+const CustomNodeFlow = () => {
+  const [elements, setElements] = useState<Elements>([]);
+  const [bgColor, setBgColor] = useState<string>(initBgColor);
+
+  useEffect(() => {
+    const onChange = (event: ChangeEvent<HTMLInputElement>) => {
+      setElements((els) =>
+        els.map((e) => {
+          if (isEdge(e) || e.id !== '2') {
+            return e;
+          }
+
+          const color = event.target.value;
+
+          setBgColor(color);
+
+          return {
+            ...e,
+            data: {
+              ...e.data,
+              color,
+            },
+          };
+        })
+      );
+    };
+
+    setElements([
+      {
+        id: '1',
+        type: 'input',
+        data: { label: 'An input node' },
+        position: { x: 0, y: 50 },
+        sourcePosition: Position.Right,
+      },
+      {
+        id: '2',
+        type: 'selectorNode',
+        data: { onChange: onChange, color: initBgColor },
+        style: { border: '1px solid #777', padding: 10 },
+        position: { x: 250, y: 50 },
+      },
+      {
+        id: '3',
+        type: 'output',
+        data: { label: 'Output A' },
+        position: { x: 550, y: 25 },
+        targetPosition: Position.Left,
+      },
+      {
+        id: '4',
+        type: 'output',
+        data: { label: 'Output B' },
+        position: { x: 550, y: 100 },
+        targetPosition: Position.Left,
+      },
+
+      { id: 'e1-2', source: '1', target: '2', animated: true, style: { stroke: '#fff' } },
+      { id: 'e2a-3', source: '2', sourceHandle: 'a', target: '3', animated: true, style: { stroke: '#fff' } },
+      { id: 'e2b-4', source: '2', sourceHandle: 'b', target: '4', animated: true, style: { stroke: '#fff' } },
+    ]);
+  }, []);
+
+  const onElementsRemove = (elementsToRemove: Elements) => setElements((els) => removeElements(elementsToRemove, els));
+  const onConnect = (params: Connection | Edge) =>
+    setElements((els) => addEdge({ ...params, animated: true, style: { stroke: '#fff' } }, els));
+
+  return (
+    <ReactFlow
+      elements={elements}
+      onElementClick={onElementClick}
+      onElementsRemove={onElementsRemove}
+      onConnect={onConnect}
+      onNodeDragStop={onNodeDragStop}
+      style={{ background: bgColor }}
+      onLoad={onLoad}
+      nodeTypes={nodeTypes}
+      connectionLineStyle={connectionLineStyle}
+      snapToGrid={true}
+      snapGrid={snapGrid}
+      defaultZoom={1.5}
+    >
+      <MiniMap
+        nodeStrokeColor={(n: Node): string => {
+          if (n.type === 'input') return '#0041d0';
+          if (n.type === 'selectorNode') return bgColor;
+          if (n.type === 'output') return '#ff0072';
+
+          return '#eee';
+        }}
+        nodeColor={(n: Node): string => {
+          if (n.type === 'selectorNode') return bgColor;
+
+          return '#fff';
+        }}
+      />
+      <Controls />
+    </ReactFlow>
+  );
+};
+
+export default CustomNodeFlow;
diff --git a/example/src/DragHandle/DragHandleNode.tsx b/example/src_oldapi/DragHandle/DragHandleNode.tsx
similarity index 100%
rename from example/src/DragHandle/DragHandleNode.tsx
rename to example/src_oldapi/DragHandle/DragHandleNode.tsx
diff --git a/example/src/DragHandle/index.tsx b/example/src_oldapi/DragHandle/index.tsx
similarity index 100%
rename from example/src/DragHandle/index.tsx
rename to example/src_oldapi/DragHandle/index.tsx
diff --git a/example/src/DragNDrop/Sidebar.tsx b/example/src_oldapi/DragNDrop/Sidebar.tsx
similarity index 100%
rename from example/src/DragNDrop/Sidebar.tsx
rename to example/src_oldapi/DragNDrop/Sidebar.tsx
diff --git a/example/src/DragNDrop/dnd.css b/example/src_oldapi/DragNDrop/dnd.css
similarity index 100%
rename from example/src/DragNDrop/dnd.css
rename to example/src_oldapi/DragNDrop/dnd.css
diff --git a/example/src/DragNDrop/index.tsx b/example/src_oldapi/DragNDrop/index.tsx
similarity index 100%
rename from example/src/DragNDrop/index.tsx
rename to example/src_oldapi/DragNDrop/index.tsx
diff --git a/example/src/EdgeTypes/index.tsx b/example/src_oldapi/EdgeTypes/index.tsx
similarity index 100%
rename from example/src/EdgeTypes/index.tsx
rename to example/src_oldapi/EdgeTypes/index.tsx
diff --git a/example/src/EdgeTypes/utils.ts b/example/src_oldapi/EdgeTypes/utils.ts
similarity index 100%
rename from example/src/EdgeTypes/utils.ts
rename to example/src_oldapi/EdgeTypes/utils.ts
diff --git a/example/src/Edges/CustomEdge.tsx b/example/src_oldapi/Edges/CustomEdge.tsx
similarity index 100%
rename from example/src/Edges/CustomEdge.tsx
rename to example/src_oldapi/Edges/CustomEdge.tsx
diff --git a/example/src/Edges/CustomEdge2.tsx b/example/src_oldapi/Edges/CustomEdge2.tsx
similarity index 100%
rename from example/src/Edges/CustomEdge2.tsx
rename to example/src_oldapi/Edges/CustomEdge2.tsx
diff --git a/example/src/Edges/index.tsx b/example/src_oldapi/Edges/index.tsx
similarity index 100%
rename from example/src/Edges/index.tsx
rename to example/src_oldapi/Edges/index.tsx
diff --git a/example/src/Empty/index.tsx b/example/src_oldapi/Empty/index.tsx
similarity index 100%
rename from example/src/Empty/index.tsx
rename to example/src_oldapi/Empty/index.tsx
diff --git a/example/src/Interaction/index.tsx b/example/src_oldapi/Interaction/index.tsx
similarity index 100%
rename from example/src/Interaction/index.tsx
rename to example/src_oldapi/Interaction/index.tsx
diff --git a/example/src/MultiFlows/index.tsx b/example/src_oldapi/MultiFlows/index.tsx
similarity index 100%
rename from example/src/MultiFlows/index.tsx
rename to example/src_oldapi/MultiFlows/index.tsx
diff --git a/example/src/MultiFlows/multiflows.css b/example/src_oldapi/MultiFlows/multiflows.css
similarity index 100%
rename from example/src/MultiFlows/multiflows.css
rename to example/src_oldapi/MultiFlows/multiflows.css
diff --git a/example/src/NodeTypeChange/index.tsx b/example/src_oldapi/NodeTypeChange/index.tsx
similarity index 100%
rename from example/src/NodeTypeChange/index.tsx
rename to example/src_oldapi/NodeTypeChange/index.tsx
diff --git a/example/src/NodeTypesObjectChange/index.tsx b/example/src_oldapi/NodeTypesObjectChange/index.tsx
similarity index 100%
rename from example/src/NodeTypesObjectChange/index.tsx
rename to example/src_oldapi/NodeTypesObjectChange/index.tsx
diff --git a/example/src/Overview/index.tsx b/example/src_oldapi/Overview/index.tsx
similarity index 100%
rename from example/src/Overview/index.tsx
rename to example/src_oldapi/Overview/index.tsx
diff --git a/example/src/Provider/Sidebar.tsx b/example/src_oldapi/Provider/Sidebar.tsx
similarity index 100%
rename from example/src/Provider/Sidebar.tsx
rename to example/src_oldapi/Provider/Sidebar.tsx
diff --git a/example/src/Provider/index.tsx b/example/src_oldapi/Provider/index.tsx
similarity index 100%
rename from example/src/Provider/index.tsx
rename to example/src_oldapi/Provider/index.tsx
diff --git a/example/src/Provider/provider.css b/example/src_oldapi/Provider/provider.css
similarity index 100%
rename from example/src/Provider/provider.css
rename to example/src_oldapi/Provider/provider.css
diff --git a/example/src/SaveRestore/Controls.tsx b/example/src_oldapi/SaveRestore/Controls.tsx
similarity index 100%
rename from example/src/SaveRestore/Controls.tsx
rename to example/src_oldapi/SaveRestore/Controls.tsx
diff --git a/example/src/SaveRestore/index.tsx b/example/src_oldapi/SaveRestore/index.tsx
similarity index 100%
rename from example/src/SaveRestore/index.tsx
rename to example/src_oldapi/SaveRestore/index.tsx
diff --git a/example/src/SaveRestore/save.css b/example/src_oldapi/SaveRestore/save.css
similarity index 100%
rename from example/src/SaveRestore/save.css
rename to example/src_oldapi/SaveRestore/save.css
diff --git a/example/src_oldapi/Stress/index.tsx b/example/src_oldapi/Stress/index.tsx
new file mode 100644
index 00000000..b734c6b6
--- /dev/null
+++ b/example/src_oldapi/Stress/index.tsx
@@ -0,0 +1,70 @@
+import React, { useState, CSSProperties } from 'react';
+import ReactFlow, {
+  removeElements,
+  addEdge,
+  MiniMap,
+  isNode,
+  Controls,
+  Background,
+  OnLoadParams,
+  Elements,
+  Connection,
+  Edge,
+} from 'react-flow-renderer';
+
+import { getElements } from './utils';
+
+const buttonWrapperStyles: CSSProperties = { position: 'absolute', right: 10, top: 10, zIndex: 4 };
+
+const onLoad = (reactFlowInstance: OnLoadParams) => {
+  reactFlowInstance.fitView();
+  console.log(reactFlowInstance.getElements());
+};
+
+const initialElements: Elements = getElements(30, 30);
+
+const StressFlow = () => {
+  const [elements, setElements] = useState<Elements>(initialElements);
+  const onElementsRemove = (elementsToRemove: Elements) => setElements((els) => removeElements(elementsToRemove, els));
+  const onConnect = (params: Connection | Edge) => setElements((els) => addEdge(params, els));
+
+  const updatePos = () => {
+    setElements((elms) => {
+      return elms.map((el) => {
+        if (isNode(el)) {
+          return {
+            ...el,
+            position: {
+              x: Math.random() * window.innerWidth,
+              y: Math.random() * window.innerHeight,
+            },
+          };
+        }
+
+        return el;
+      });
+    });
+  };
+
+  const updateElements = () => {
+    const grid = Math.ceil(Math.random() * 10);
+    setElements(getElements(grid, grid));
+  };
+
+  return (
+    <ReactFlow elements={elements} onLoad={onLoad} onElementsRemove={onElementsRemove} onConnect={onConnect}>
+      <MiniMap />
+      <Controls />
+      <Background />
+
+      <div style={buttonWrapperStyles}>
+        <button onClick={updatePos} style={{ marginRight: 5 }}>
+          change pos
+        </button>
+        <button onClick={updateElements}>update elements</button>
+      </div>
+    </ReactFlow>
+  );
+};
+
+export default StressFlow;
diff --git a/example/src_oldapi/Stress/utils.ts b/example/src_oldapi/Stress/utils.ts
new file mode 100644
index 00000000..a37eb8de
--- /dev/null
+++ b/example/src_oldapi/Stress/utils.ts
@@ -0,0 +1,30 @@
+import { Elements } from 'react-flow-renderer';
+
+export function getElements(xElements: number = 10, yElements: number = 10): Elements {
+  const initialElements = [];
+  let nodeId = 1;
+  let recentNodeId = null;
+
+  for (let y = 0; y < yElements; y++) {
+    for (let x = 0; x < xElements; x++) {
+      const position = { x: x * 100, y: y * 50 };
+      const data = { label: `Node ${nodeId}` };
+      const node = {
+        id: nodeId.toString(),
+        style: { width: 50, fontSize: 11 },
+        data,
+        position,
+      };
+      initialElements.push(node);
+
+      if (recentNodeId && nodeId <= xElements * yElements) {
+        initialElements.push({ id: `${x}-${y}`, source: recentNodeId.toString(), target: nodeId.toString() });
+      }
+
+      recentNodeId = nodeId;
+      nodeId++;
+    }
+  }
+
+  return initialElements;
+}
diff --git a/example/src/Switch/index.tsx b/example/src_oldapi/Switch/index.tsx
similarity index 100%
rename from example/src/Switch/index.tsx
rename to example/src_oldapi/Switch/index.tsx
diff --git a/example/src/Undirectional/CustomNode.tsx b/example/src_oldapi/Undirectional/CustomNode.tsx
similarity index 100%
rename from example/src/Undirectional/CustomNode.tsx
rename to example/src_oldapi/Undirectional/CustomNode.tsx
diff --git a/example/src/Undirectional/index.tsx b/example/src_oldapi/Undirectional/index.tsx
similarity index 100%
rename from example/src/Undirectional/index.tsx
rename to example/src_oldapi/Undirectional/index.tsx
diff --git a/example/src/UpdatableEdge/index.tsx b/example/src_oldapi/UpdatableEdge/index.tsx
similarity index 100%
rename from example/src/UpdatableEdge/index.tsx
rename to example/src_oldapi/UpdatableEdge/index.tsx
diff --git a/example/src_oldapi/UpdateNode/index.tsx b/example/src_oldapi/UpdateNode/index.tsx
new file mode 100644
index 00000000..a7cce8f2
--- /dev/null
+++ b/example/src_oldapi/UpdateNode/index.tsx
@@ -0,0 +1,78 @@
+import React, { useEffect, useState } from 'react';
+import ReactFlow, { Elements } from 'react-flow-renderer';
+
+import './updatenode.css';
+
+const initialElements: Elements = [
+  { id: '1', data: { label: '-' }, position: { x: 100, y: 100 } },
+  { id: '2', data: { label: 'Node 2' }, position: { x: 100, y: 200 } },
+  { id: 'e1-2', source: '1', target: '2' },
+];
+
+const UpdateNode = () => {
+  const [elements, setElements] = useState<Elements>(initialElements);
+  const [nodeName, setNodeName] = useState<string>('Node 1');
+  const [nodeBg, setNodeBg] = useState<string>('#eee');
+  const [nodeHidden, setNodeHidden] = useState<boolean>(false);
+
+  useEffect(() => {
+    setElements((els) =>
+      els.map((el) => {
+        if (el.id === '1') {
+          // it's important that you create a new object here in order to notify react flow about the change
+          el.data = {
+            ...el.data,
+            label: nodeName,
+          };
+        }
+
+        return el;
+      })
+    );
+  }, [nodeName, setElements]);
+
+  useEffect(() => {
+    setElements((els) =>
+      els.map((el) => {
+        if (el.id === '1') {
+          // it's important that you create a new object here in order to notify react flow about the change
+          el.style = { ...el.style, backgroundColor: nodeBg };
+        }
+
+        return el;
+      })
+    );
+  }, [nodeBg, setElements]);
+
+  useEffect(() => {
+    setElements((els) =>
+      els.map((el) => {
+        if (el.id === '1' || el.id === 'e1-2') {
+          // when you update a simple type you can just update the value
+          el.isHidden = nodeHidden;
+        }
+
+        return el;
+      })
+    );
+  }, [nodeHidden, setElements]);
+
+  return (
+    <ReactFlow elements={elements} defaultZoom={1.5} minZoom={0.2} maxZoom={4}>
+      <div className="updatenode__controls">
+        <label>label:</label>
+        <input value={nodeName} onChange={(evt) => setNodeName(evt.target.value)} />
+
+        <label className="updatenode__bglabel">background:</label>
+        <input value={nodeBg} onChange={(evt) => setNodeBg(evt.target.value)} />
+
+        <div className="updatenode__checkboxwrapper">
+          <label>hidden:</label>
+          <input type="checkbox" checked={nodeHidden} onChange={(evt) => setNodeHidden(evt.target.checked)} />
+        </div>
+      </div>
+    </ReactFlow>
+  );
+};
+
+export default UpdateNode;
diff --git a/example/src_oldapi/UpdateNode/updatenode.css b/example/src_oldapi/UpdateNode/updatenode.css
new file mode 100644
index 00000000..95438054
--- /dev/null
+++ b/example/src_oldapi/UpdateNode/updatenode.css
@@ -0,0 +1,21 @@
+.updatenode__controls {
+  position: absolute;
+  right: 10px;
+  top: 10px;
+  z-index: 4;
+  font-size: 12px;
+}
+
+.updatenode__controls label {
+  display: block;
+}
+
+.updatenode__bglabel {
+  margin-top: 10px;
+}
+
+.updatenode__checkboxwrapper {
+  margin-top: 10px;
+  display: flex;
+  align-items: center;
+}
diff --git a/example/src/UseUpdateNodeInternals/CustomNode.tsx b/example/src_oldapi/UseUpdateNodeInternals/CustomNode.tsx
similarity index 100%
rename from example/src/UseUpdateNodeInternals/CustomNode.tsx
rename to example/src_oldapi/UseUpdateNodeInternals/CustomNode.tsx
diff --git a/example/src/UseUpdateNodeInternals/index.tsx b/example/src_oldapi/UseUpdateNodeInternals/index.tsx
similarity index 100%
rename from example/src/UseUpdateNodeInternals/index.tsx
rename to example/src_oldapi/UseUpdateNodeInternals/index.tsx
diff --git a/example/src/UseZoomPanHelper/index.tsx b/example/src_oldapi/UseZoomPanHelper/index.tsx
similarity index 100%
rename from example/src/UseZoomPanHelper/index.tsx
rename to example/src_oldapi/UseZoomPanHelper/index.tsx
diff --git a/example/src/Validation/index.tsx b/example/src_oldapi/Validation/index.tsx
similarity index 100%
rename from example/src/Validation/index.tsx
rename to example/src_oldapi/Validation/index.tsx
diff --git a/example/src/Validation/validation.css b/example/src_oldapi/Validation/validation.css
similarity index 100%
rename from example/src/Validation/validation.css
rename to example/src_oldapi/Validation/validation.css
diff --git a/example/src_oldapi/index.css b/example/src_oldapi/index.css
new file mode 100644
index 00000000..3f6f9960
--- /dev/null
+++ b/example/src_oldapi/index.css
@@ -0,0 +1,82 @@
+body {
+  font-family: sans-serif;
+  color: #111;
+}
+
+html,
+body,
+#root {
+  margin: 0;
+  height: 100%;
+}
+
+#root {
+  display: flex;
+  flex-direction: column;
+}
+
+header {
+  padding: 10px;
+  border-bottom: 1px solid #eee;
+  display: flex;
+  font-weight: 700;
+  align-items: center;
+}
+
+.logo {
+  text-decoration: none;
+  display: block;
+  line-height: 1;
+}
+
+header a,
+header a:focus,
+header a:active,
+header a:visited {
+  color: #111;
+}
+
+header a:hover {
+  color: #333;
+}
+
+header select {
+  margin-left: 1em;
+}
+
+.overview-example__add {
+  display: none;
+}
+
+.react-flow__node a {
+  font-weight: 700;
+  color: #111;
+}
+
+.react-flow__node.dark-node {
+  background: #0041d0;
+  color: #f8f8f8;
+}
+
+.react-flow__node.dark {
+  background: #557;
+  color: #f8f8f8;
+}
+
+.react-flow__node-selectorNode {
+  font-size: 12px;
+  background: #f0f2f3;
+  border: 1px solid 555;
+  border-radius: 5px;
+  text-align: center;
+}
+
+.react-flow__node-selectorNode .react-flow__handle {
+  border-color: #f0f2f3;
+}
+
+@media screen and (min-width: 768px) {
+  .overview-example__add {
+    display: block;
+  }
+}
diff --git a/example/src_oldapi/index.tsx b/example/src_oldapi/index.tsx
new file mode 100644
index 00000000..f81278ba
--- /dev/null
+++ b/example/src_oldapi/index.tsx
@@ -0,0 +1,165 @@
+import React, { ChangeEvent } from 'react';
+import ReactDOM from 'react-dom';
+import { BrowserRouter as Router, Route, Switch, withRouter } from 'react-router-dom';
+
+import Overview from './Overview';
+import Basic from './Basic';
+import CustomNode from './CustomNode';
+import Stress from './Stress';
+import Interaction from './Interaction';
+import Empty from './Empty';
+import Edges from './Edges';
+import Validation from './Validation';
+import Provider from './Provider';
+import Hidden from './Hidden';
+import EdgeTypes from './EdgeTypes';
+import CustomConnectionLine from './CustomConnectionLine';
+import NodeTypeChange from './NodeTypeChange';
+import NodeTypesObjectChange from './NodeTypesObjectChange';
+import UpdatableEdge from './UpdatableEdge';
+import UpdateNode from './UpdateNode';
+import SaveRestore from './SaveRestore';
+import DragNDrop from './DragNDrop';
+import Layout from './Layouting';
+import SwitchFlows from './Switch';
+import UseZoomPanHelper from './UseZoomPanHelper';
+import UseUpdateNodeInternals from './UseUpdateNodeInternals';
+import Undirectional from './Undirectional';
+import MultiFlows from './MultiFlows';
+import DragHandle from './DragHandle';
+
+import './index.css';
+
+const routes = [
+  {
+    path: '/',
+    component: Overview,
+  },
+  {
+    path: '/edges',
+    component: Edges,
+  },
+  {
+    path: '/custom-node',
+    component: CustomNode,
+  },
+  {
+    path: '/validation',
+    component: Validation,
+  },
+  {
+    path: '/provider',
+    component: Provider,
+  },
+  {
+    path: '/stress',
+    component: Stress,
+  },
+  {
+    path: '/interaction',
+    component: Interaction,
+  },
+  {
+    path: '/basic',
+    component: Basic,
+  },
+  {
+    path: '/empty',
+    component: Empty,
+  },
+  {
+    path: '/hidden',
+    component: Hidden,
+  },
+  {
+    path: '/edge-types',
+    component: EdgeTypes,
+  },
+  {
+    path: '/custom-connectionline',
+    component: CustomConnectionLine,
+  },
+  {
+    path: '/nodetype-change',
+    component: NodeTypeChange,
+  },
+  {
+    path: '/nodetypesobject-change',
+    component: NodeTypesObjectChange,
+  },
+  {
+    path: '/updatable-edge',
+    component: UpdatableEdge,
+  },
+  {
+    path: '/update-node',
+    component: UpdateNode,
+  },
+  {
+    path: '/save-restore',
+    component: SaveRestore,
+  },
+  {
+    path: '/drag-and-drop',
+    component: DragNDrop,
+  },
+  {
+    path: '/layouting',
+    component: Layout,
+  },
+  {
+    path: '/switch',
+    component: SwitchFlows,
+  },
+  {
+    path: '/usezoompanhelper',
+    component: UseZoomPanHelper,
+  },
+  {
+    path: '/useupdatenodeinternals',
+    component: UseUpdateNodeInternals,
+  },
+  {
+    path: '/undirectional',
+    component: Undirectional,
+  },
+  {
+    path: '/multiflows',
+    component: MultiFlows,
+  },
+  {
+    path: '/draghandle',
+    component: DragHandle,
+  },
+];
+
+const Header = withRouter(({ history, location }) => {
+  const onChange = (event: ChangeEvent<HTMLSelectElement>) => history.push(event.target.value);
+
+  return (
+    <header>
+      <a className="logo" href="https://github.com/wbkd/react-flow">
+        React Flow Dev
+      </a>
+      <select defaultValue={location.pathname} onChange={onChange}>
+        {routes.map((route) => (
+          <option value={route.path} key={route.path}>
+            {route.path === '/' ? 'overview' : route.path.substr(1, route.path.length)}
+          </option>
+        ))}
+      </select>
+    </header>
+  );
+});
+
+ReactDOM.render(
+  <Router forceRefresh={true}>
+    <Header />
+    <Switch>
+      {routes.map((route) => (
+        <Route exact path={route.path} render={() => <route.component />} key={route.path} />
+      ))}
+    </Switch>
+  </Router>,
+  document.getElementById('root')
+);
diff --git a/example/src_oldapi/react-app-env.d.ts b/example/src_oldapi/react-app-env.d.ts
new file mode 100644
index 00000000..6431bc5f
--- /dev/null
+++ b/example/src_oldapi/react-app-env.d.ts
@@ -0,0 +1 @@
+/// <reference types="react-scripts" />
diff --git a/example/tsconfig.json b/example/tsconfig.json
index a273b0cf..edc77511 100644
--- a/example/tsconfig.json
+++ b/example/tsconfig.json
@@ -1,11 +1,7 @@
 {
   "compilerOptions": {
     "target": "es5",
-    "lib": [
-      "dom",
-      "dom.iterable",
-      "esnext"
-    ],
+    "lib": ["dom", "dom.iterable", "esnext"],
     "allowJs": true,
     "skipLibCheck": true,
     "esModuleInterop": true,
@@ -20,7 +16,6 @@
     "noEmit": true,
     "jsx": "react-jsx"
   },
-  "include": [
-    "src"
-  ]
+  "include": ["src"],
+  "exclude": ["src_oldapi"]
 }
diff --git a/package-lock.json b/package-lock.json
index 0b755877..17711e21 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -1,12 +1,12 @@
 {
   "name": "react-flow-renderer",
-  "version": "10.0.0-next.2",
+  "version": "10.0.0-next.16",
   "lockfileVersion": 2,
   "requires": true,
   "packages": {
     "": {
       "name": "react-flow-renderer",
-      "version": "10.0.0-next.2",
+      "version": "10.0.0-next.16",
       "license": "MIT",
       "dependencies": {
         "@babel/runtime": "^7.15.4",
@@ -15,10 +15,11 @@
         "classcat": "^5.0.3",
         "d3-selection": "^3.0.0",
         "d3-zoom": "^3.0.0",
-        "fast-deep-equal": "^3.1.3",
         "react-draggable": "^4.4.4",
         "react-redux": "^7.2.5",
-        "redux": "^4.1.1"
+        "redux": "^4.1.1",
+        "redux-thunk": "^2.3.0",
+        "zustand": "^3.5.13"
       },
       "devDependencies": {
         "@babel/core": "^7.15.8",
@@ -32,7 +33,8 @@
         "@svgr/rollup": "^5.5.0",
         "@types/redux": "^3.6.31",
         "@types/resize-observer-browser": "^0.1.6",
-        "autoprefixer": "^10.3.7",
+        "@welldone-software/why-did-you-render": "^6.2.1",
+        "autoprefixer": "^10.3.6",
         "babel-preset-react-app": "^10.0.0",
         "cypress": "^8.5.0",
         "postcss": "^8.3.9",
@@ -2836,6 +2838,18 @@
         "@types/node": "*"
       }
     },
+    "node_modules/@welldone-software/why-did-you-render": {
+      "version": "6.2.1",
+      "resolved": "https://registry.npmjs.org/@welldone-software/why-did-you-render/-/why-did-you-render-6.2.1.tgz",
+      "integrity": "sha512-eIVKeK6ueS3tuzCqMVTaaNrPYvb9cA8NHiNgLA7Op8SD4TiT31zqNjxmhzLEK+y3sBxcwr6YhsiQGX9EThrvaw==",
+      "dev": true,
+      "dependencies": {
+        "lodash": "^4"
+      },
+      "peerDependencies": {
+        "react": "^16 || ^17"
+      }
+    },
     "node_modules/aggregate-error": {
       "version": "3.1.0",
       "resolved": "https://registry.npmjs.org/aggregate-error/-/aggregate-error-3.1.0.tgz",
@@ -5255,7 +5269,8 @@
     "node_modules/fast-deep-equal": {
       "version": "3.1.3",
       "resolved": "https://registry.npmjs.org/fast-deep-equal/-/fast-deep-equal-3.1.3.tgz",
-      "integrity": "sha512-f3qQ9oQy9j2AhBe/H9VC91wLmKBCCU/gDOnKNAYG5hswO7BLKj09Hc5HYNz9cGI++xlpDCIgDaitVs03ATR84Q=="
+      "integrity": "sha512-f3qQ9oQy9j2AhBe/H9VC91wLmKBCCU/gDOnKNAYG5hswO7BLKj09Hc5HYNz9cGI++xlpDCIgDaitVs03ATR84Q==",
+      "dev": true
     },
     "node_modules/fast-glob": {
       "version": "3.2.7",
@@ -9255,6 +9270,11 @@
         "@babel/runtime": "^7.9.2"
       }
     },
+    "node_modules/redux-thunk": {
+      "version": "2.3.0",
+      "resolved": "https://registry.npmjs.org/redux-thunk/-/redux-thunk-2.3.0.tgz",
+      "integrity": "sha512-km6dclyFnmcvxhAcrQV2AkZmPQjzPDjgVlQtR0EQjxZPyJ0BnMf3in1ryuR8A2qU0HldVRfxYXbFSKlI3N7Slw=="
+    },
     "node_modules/regenerate": {
       "version": "1.4.2",
       "resolved": "https://registry.npmjs.org/regenerate/-/regenerate-1.4.2.tgz",
@@ -11229,6 +11249,22 @@
         "buffer-crc32": "~0.2.3",
         "fd-slicer": "~1.1.0"
       }
+    },
+    "node_modules/zustand": {
+      "version": "3.5.13",
+      "resolved": "https://registry.npmjs.org/zustand/-/zustand-3.5.13.tgz",
+      "integrity": "sha512-orO/XcYwSWffsrPVTdCtuKM/zkUaOIyKDasOk/lecsD3R0euELsj+cB65uKZ1KyinrK2STHIuUhRoLpH8QprQg==",
+      "engines": {
+        "node": ">=12.7.0"
+      },
+      "peerDependencies": {
+        "react": ">=16.8"
+      },
+      "peerDependenciesMeta": {
+        "react": {
+          "optional": true
+        }
+      }
     }
   },
   "dependencies": {
@@ -13290,6 +13326,15 @@
         "@types/node": "*"
       }
     },
+    "@welldone-software/why-did-you-render": {
+      "version": "6.2.1",
+      "resolved": "https://registry.npmjs.org/@welldone-software/why-did-you-render/-/why-did-you-render-6.2.1.tgz",
+      "integrity": "sha512-eIVKeK6ueS3tuzCqMVTaaNrPYvb9cA8NHiNgLA7Op8SD4TiT31zqNjxmhzLEK+y3sBxcwr6YhsiQGX9EThrvaw==",
+      "dev": true,
+      "requires": {
+        "lodash": "^4"
+      }
+    },
     "aggregate-error": {
       "version": "3.1.0",
       "resolved": "https://registry.npmjs.org/aggregate-error/-/aggregate-error-3.1.0.tgz",
@@ -15147,7 +15192,8 @@
     "fast-deep-equal": {
       "version": "3.1.3",
       "resolved": "https://registry.npmjs.org/fast-deep-equal/-/fast-deep-equal-3.1.3.tgz",
-      "integrity": "sha512-f3qQ9oQy9j2AhBe/H9VC91wLmKBCCU/gDOnKNAYG5hswO7BLKj09Hc5HYNz9cGI++xlpDCIgDaitVs03ATR84Q=="
+      "integrity": "sha512-f3qQ9oQy9j2AhBe/H9VC91wLmKBCCU/gDOnKNAYG5hswO7BLKj09Hc5HYNz9cGI++xlpDCIgDaitVs03ATR84Q==",
+      "dev": true
     },
     "fast-glob": {
       "version": "3.2.7",
@@ -18028,6 +18074,11 @@
         "@babel/runtime": "^7.9.2"
       }
     },
+    "redux-thunk": {
+      "version": "2.3.0",
+      "resolved": "https://registry.npmjs.org/redux-thunk/-/redux-thunk-2.3.0.tgz",
+      "integrity": "sha512-km6dclyFnmcvxhAcrQV2AkZmPQjzPDjgVlQtR0EQjxZPyJ0BnMf3in1ryuR8A2qU0HldVRfxYXbFSKlI3N7Slw=="
+    },
     "regenerate": {
       "version": "1.4.2",
       "resolved": "https://registry.npmjs.org/regenerate/-/regenerate-1.4.2.tgz",
@@ -19538,6 +19589,11 @@
         "buffer-crc32": "~0.2.3",
         "fd-slicer": "~1.1.0"
       }
+    },
+    "zustand": {
+      "version": "3.5.13",
+      "resolved": "https://registry.npmjs.org/zustand/-/zustand-3.5.13.tgz",
+      "integrity": "sha512-orO/XcYwSWffsrPVTdCtuKM/zkUaOIyKDasOk/lecsD3R0euELsj+cB65uKZ1KyinrK2STHIuUhRoLpH8QprQg=="
     }
   }
 }
diff --git a/package.json b/package.json
index 0a50f64e..5f33733e 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "react-flow-renderer",
-  "version": "10.0.0-next.2",
+  "version": "10.0.0-next.16",
   "engines": {
     "node": ">=12"
   },
@@ -37,10 +37,11 @@
     "classcat": "^5.0.3",
     "d3-selection": "^3.0.0",
     "d3-zoom": "^3.0.0",
-    "fast-deep-equal": "^3.1.3",
     "react-draggable": "^4.4.4",
     "react-redux": "^7.2.5",
-    "redux": "^4.1.1"
+    "redux": "^4.1.1",
+    "redux-thunk": "^2.3.0",
+    "zustand": "^3.5.13"
   },
   "devDependencies": {
     "@babel/core": "^7.15.8",
@@ -54,7 +55,8 @@
     "@svgr/rollup": "^5.5.0",
     "@types/redux": "^3.6.31",
     "@types/resize-observer-browser": "^0.1.6",
-    "autoprefixer": "^10.3.7",
+    "@welldone-software/why-did-you-render": "^6.2.1",
+    "autoprefixer": "^10.3.6",
     "babel-preset-react-app": "^10.0.0",
     "cypress": "^8.5.0",
     "postcss": "^8.3.9",
diff --git a/src/additional-components/Background/index.tsx b/src/additional-components/Background/index.tsx
index 0e99a835..8f73da12 100644
--- a/src/additional-components/Background/index.tsx
+++ b/src/additional-components/Background/index.tsx
@@ -1,8 +1,8 @@
 import React, { memo, useMemo, FC, HTMLAttributes } from 'react';
 import cc from 'classcat';
 
-import { useStoreState } from '../../store/hooks';
-import { BackgroundVariant } from '../../types';
+import { useStore } from '../../store';
+import { BackgroundVariant, ReactFlowState } from '../../types';
 import { createGridLinesPath, createGridDotsPath } from './utils';
 
 export interface BackgroundProps extends HTMLAttributes<SVGElement> {
@@ -17,6 +17,8 @@ const defaultColors = {
   [BackgroundVariant.Lines]: '#eee',
 };
 
+const transformSelector = (s: ReactFlowState) => s.transform;
+
 const Background: FC<BackgroundProps> = ({
   variant = BackgroundVariant.Dots,
   gap = 15,
@@ -25,11 +27,11 @@ const Background: FC<BackgroundProps> = ({
   style,
   className,
 }) => {
-  const [x, y, scale] = useStoreState((s) => s.transform);
+  const [x, y, scale] = useStore(transformSelector);
   // when there are multiple flows on a page we need to make sure that every background gets its own pattern.
   const patternId = useMemo(() => `pattern-${Math.floor(Math.random() * 100000)}`, []);
 
-  const bgClasses = cc(['react-flow__background', className]);
+  const bgClasses = cc(['react-flow__background', 'react-flow__container', className]);
   const scaledGap = gap * scale;
   const xOffset = x % scaledGap;
   const yOffset = y % scaledGap;
diff --git a/src/additional-components/Controls/index.tsx b/src/additional-components/Controls/index.tsx
index 5e4ece7f..850ee781 100644
--- a/src/additional-components/Controls/index.tsx
+++ b/src/additional-components/Controls/index.tsx
@@ -1,7 +1,7 @@
 import React, { memo, useCallback, HTMLAttributes, FC, useEffect, useState } from 'react';
 import cc from 'classcat';
 
-import { useStoreState, useStoreActions } from '../../store/hooks';
+import { useStore } from '../../store';
 
 import PlusIcon from '../../../assets/icons/plus.svg';
 import MinusIcon from '../../../assets/icons/minus.svg';
@@ -10,7 +10,7 @@ import LockIcon from '../../../assets/icons/lock.svg';
 import UnlockIcon from '../../../assets/icons/unlock.svg';
 
 import useZoomPanHelper from '../../hooks/useZoomPanHelper';
-import { FitViewParams } from '../../types';
+import { FitViewParams, ReactFlowState } from '../../types';
 
 export interface ControlProps extends HTMLAttributes<HTMLDivElement> {
   showZoom?: boolean;
@@ -31,6 +31,9 @@ export const ControlButton: FC<ControlButtonProps> = ({ children, className, ...
   </button>
 );
 
+const setInteractiveSelector = (s: ReactFlowState) => s.setInteractive;
+const isInteractiveSelector = (s: ReactFlowState) => s.nodesDraggable && s.nodesConnectable && s.elementsSelectable;
+
 const Controls: FC<ControlProps> = ({
   style,
   showZoom = true,
@@ -45,10 +48,10 @@ const Controls: FC<ControlProps> = ({
   children,
 }) => {
   const [isVisible, setIsVisible] = useState<boolean>(false);
-  const setInteractive = useStoreActions((actions) => actions.setInteractive);
+  const setInteractive = useStore(setInteractiveSelector);
+  const isInteractive = useStore(isInteractiveSelector);
   const { zoomIn, zoomOut, fitView } = useZoomPanHelper();
 
-  const isInteractive = useStoreState((s) => s.nodesDraggable && s.nodesConnectable && s.elementsSelectable);
   const mapClasses = cc(['react-flow__controls', className]);
 
   const onZoomInHandler = useCallback(() => {
diff --git a/src/additional-components/MiniMap/MiniMapNode.tsx b/src/additional-components/MiniMap/MiniMapNode.tsx
index d8a48e9e..e879211b 100644
--- a/src/additional-components/MiniMap/MiniMapNode.tsx
+++ b/src/additional-components/MiniMap/MiniMapNode.tsx
@@ -26,7 +26,7 @@ const MiniMapNode = ({
   strokeWidth,
   className,
   borderRadius,
-  shapeRendering
+  shapeRendering,
 }: MiniMapNodeProps) => {
   const { background, backgroundColor } = style || {};
   const fill = (color || background || backgroundColor) as string;
diff --git a/src/additional-components/MiniMap/index.tsx b/src/additional-components/MiniMap/index.tsx
index dbb86f6c..174caf4f 100644
--- a/src/additional-components/MiniMap/index.tsx
+++ b/src/additional-components/MiniMap/index.tsx
@@ -1,9 +1,11 @@
 import React, { memo, HTMLAttributes } from 'react';
 import cc from 'classcat';
+import shallow from 'zustand/shallow';
 
-import { useStoreState } from '../../store/hooks';
-import { getRectOfNodes, getBoundsofRects } from '../../utils/graph';
-import { Node, Rect } from '../../types';
+import { useStore } from '../../store';
+import { getRectOfNodes } from '../../utils/graph';
+import { getBoundsofRects } from '../../utils';
+import { Node, ReactFlowState, Rect } from '../../types';
 import MiniMapNode from './MiniMapNode';
 
 type StringFunc = (node: Node) => string;
@@ -22,6 +24,13 @@ declare const window: any;
 const defaultWidth = 200;
 const defaultHeight = 150;
 
+const selector = (s: ReactFlowState) => ({
+  width: s.width,
+  height: s.height,
+  transform: s.transform,
+  nodeInternals: s.nodeInternals,
+});
+
 const MiniMap = ({
   style,
   className,
@@ -32,20 +41,20 @@ const MiniMap = ({
   nodeStrokeWidth = 2,
   maskColor = 'rgb(240, 242, 243, 0.7)',
 }: MiniMapProps) => {
-  const containerWidth = useStoreState((s) => s.width);
-  const containerHeight = useStoreState((s) => s.height);
-  const [tX, tY, tScale] = useStoreState((s) => s.transform);
-  const nodes = useStoreState((s) => s.nodes);
+  const { width: containerWidth, height: containerHeight, transform, nodeInternals } = useStore(selector, shallow);
+  const [tX, tY, tScale] = transform;
 
   const mapClasses = cc(['react-flow__minimap', className]);
   const elementWidth = (style?.width || defaultWidth)! as number;
   const elementHeight = (style?.height || defaultHeight)! as number;
   const nodeColorFunc = (nodeColor instanceof Function ? nodeColor : () => nodeColor) as StringFunc;
-  const nodeStrokeColorFunc = (nodeStrokeColor instanceof Function
-    ? nodeStrokeColor
-    : () => nodeStrokeColor) as StringFunc;
+  const nodeStrokeColorFunc = (
+    nodeStrokeColor instanceof Function ? nodeStrokeColor : () => nodeStrokeColor
+  ) as StringFunc;
   const nodeClassNameFunc = (nodeClassName instanceof Function ? nodeClassName : () => nodeClassName) as StringFunc;
-  const hasNodes = nodes && nodes.length;
+  const hasNodes = nodeInternals && nodeInternals.size > 0;
+  // @TODO: work with nodeInternals instead of converting it to an array
+  const nodes = Array.from(nodeInternals).map(([_, node]) => node);
   const bb = getRectOfNodes(nodes);
   const viewBB: Rect = {
     x: -tX / tScale,
@@ -64,7 +73,7 @@ const MiniMap = ({
   const y = boundingRect.y - (viewHeight - boundingRect.height) / 2 - offset;
   const width = viewWidth + offset * 2;
   const height = viewHeight + offset * 2;
-  const shapeRendering = (typeof window === "undefined" || !!window.chrome) ?  "crispEdges" : "geometricPrecision";
+  const shapeRendering = typeof window === 'undefined' || !!window.chrome ? 'crispEdges' : 'geometricPrecision';
 
   return (
     <svg
@@ -74,24 +83,28 @@ const MiniMap = ({
       style={style}
       className={mapClasses}
     >
-      {nodes
-        .filter((node) => !node.isHidden)
-        .map((node) => (
-          <MiniMapNode
-            key={node.id}
-            x={node.__rf.position.x}
-            y={node.__rf.position.y}
-            width={node.__rf.width}
-            height={node.__rf.height}
-            style={node.style}
-            className={nodeClassNameFunc(node)}
-            color={nodeColorFunc(node)}
-            borderRadius={nodeBorderRadius}
-            strokeColor={nodeStrokeColorFunc(node)}
-            strokeWidth={nodeStrokeWidth}
-            shapeRendering={shapeRendering}
-          />
-        ))}
+      {Array.from(nodeInternals)
+        .filter(([_, node]) => !node.hidden && node.width && node.height)
+        .map(([_, node]) => {
+          const positionAbsolute = nodeInternals.get(node.id)?.positionAbsolute;
+
+          return (
+            <MiniMapNode
+              key={node.id}
+              x={positionAbsolute?.x || 0}
+              y={positionAbsolute?.y || 0}
+              width={node.width!}
+              height={node.height!}
+              style={node.style}
+              className={nodeClassNameFunc(node)}
+              color={nodeColorFunc(node)}
+              borderRadius={nodeBorderRadius}
+              strokeColor={nodeStrokeColorFunc(node)}
+              strokeWidth={nodeStrokeWidth}
+              shapeRendering={shapeRendering}
+            />
+          );
+        })}
       <path
         className="react-flow__minimap-mask"
         d={`M${x - offset},${y - offset}h${width + offset * 2}v${height + offset * 2}h${-width - offset * 2}z
diff --git a/src/additional-components/ReactFlowProvider/index.tsx b/src/additional-components/ReactFlowProvider/index.tsx
index d1abc923..a7044a3f 100644
--- a/src/additional-components/ReactFlowProvider/index.tsx
+++ b/src/additional-components/ReactFlowProvider/index.tsx
@@ -1,16 +1,8 @@
-import React, { FC, useMemo } from 'react';
-import { Provider } from 'react-redux';
+import React, { FC } from 'react';
 
-import { initialState } from '../../store';
-import configureStore from '../../store/configure-store';
+import { Provider, createStore } from '../../store';
 
-const ReactFlowProvider: FC = ({ children }) => {
-  const store = useMemo(() => {
-    return configureStore(initialState);
-  }, []);
-
-  return <Provider store={store}>{children}</Provider>;
-};
+const ReactFlowProvider: FC = ({ children }) => <Provider createStore={createStore}>{children}</Provider>;
 
 ReactFlowProvider.displayName = 'ReactFlowProvider';
 
diff --git a/src/components/ConnectionLine/index.tsx b/src/components/ConnectionLine/index.tsx
index e3c71f9e..390132d1 100644
--- a/src/components/ConnectionLine/index.tsx
+++ b/src/components/ConnectionLine/index.tsx
@@ -1,32 +1,34 @@
-import React, { useEffect, useState, CSSProperties } from 'react';
+import React, { useRef, CSSProperties } from 'react';
+import shallow from 'zustand/shallow';
 
+import { useStore } from '../../store';
 import { getBezierPath } from '../Edges/BezierEdge';
 import { getSmoothStepPath } from '../Edges/SmoothStepEdge';
 import {
-  ElementId,
-  Node,
-  Transform,
+  NodeInternalsItem,
   HandleElement,
-  Position,
   ConnectionLineType,
   ConnectionLineComponent,
   HandleType,
+  Node,
+  ReactFlowState,
+  Position,
 } from '../../types';
 
 interface ConnectionLineProps {
-  connectionNodeId: ElementId;
-  connectionHandleId: ElementId | null;
+  connectionNodeId: string;
+  connectionHandleId: string | null;
   connectionHandleType: HandleType;
   connectionPositionX: number;
   connectionPositionY: number;
   connectionLineType: ConnectionLineType;
-  nodes: Node[];
-  transform: Transform;
   isConnectable: boolean;
   connectionLineStyle?: CSSProperties;
   CustomConnectionLineComponent?: ConnectionLineComponent;
 }
 
+const selector = (s: ReactFlowState) => ({ nodeInternals: s.nodeInternals, transform: s.transform });
+
 export default ({
   connectionNodeId,
   connectionHandleId,
@@ -35,31 +37,32 @@ export default ({
   connectionPositionX,
   connectionPositionY,
   connectionLineType = ConnectionLineType.Bezier,
-  nodes = [],
-  transform,
   isConnectable,
   CustomConnectionLineComponent,
 }: ConnectionLineProps) => {
-  const [sourceNode, setSourceNode] = useState<Node | null>(null);
   const nodeId = connectionNodeId;
   const handleId = connectionHandleId;
 
-  useEffect(() => {
-    const nextSourceNode = nodes.find((n) => n.id === nodeId) || null;
-    setSourceNode(nextSourceNode);
-  }, []);
+  const { nodeInternals, transform } = useStore(selector, shallow);
+  const sourceNodeInternals = useRef<NodeInternalsItem | undefined>(nodeInternals.get(nodeId));
+  const sourceNode = useRef<Node | undefined>(nodeInternals.get(nodeId));
 
-  if (!sourceNode || !isConnectable) {
+  if (
+    !sourceNode.current ||
+    !sourceNodeInternals.current ||
+    !isConnectable ||
+    !sourceNodeInternals.current.handleBounds?.[connectionHandleType]
+  ) {
     return null;
   }
 
   const sourceHandle = handleId
-    ? sourceNode.__rf.handleBounds[connectionHandleType].find((d: HandleElement) => d.id === handleId)
-    : sourceNode.__rf.handleBounds[connectionHandleType][0];
-  const sourceHandleX = sourceHandle ? sourceHandle.x + sourceHandle.width / 2 : sourceNode.__rf.width / 2;
-  const sourceHandleY = sourceHandle ? sourceHandle.y + sourceHandle.height / 2 : sourceNode.__rf.height;
-  const sourceX = sourceNode.__rf.position.x + sourceHandleX;
-  const sourceY = sourceNode.__rf.position.y + sourceHandleY;
+    ? sourceNodeInternals.current.handleBounds[connectionHandleType]!.find((d: HandleElement) => d.id === handleId)
+    : sourceNodeInternals.current.handleBounds[connectionHandleType]![0];
+  const sourceHandleX = sourceHandle ? sourceHandle.x + sourceHandle.width / 2 : sourceNodeInternals.current.width! / 2;
+  const sourceHandleY = sourceHandle ? sourceHandle.y + sourceHandle.height / 2 : sourceNodeInternals.current.height!;
+  const sourceX = sourceNodeInternals.current.positionAbsolute!.x + sourceHandleX;
+  const sourceY = sourceNodeInternals.current.positionAbsolute!.y + sourceHandleY;
 
   const targetX = (connectionPositionX - transform[0]) / transform[2];
   const targetY = (connectionPositionY - transform[1]) / transform[2];
@@ -79,7 +82,7 @@ export default ({
           targetPosition={targetPosition}
           connectionLineType={connectionLineType}
           connectionLineStyle={connectionLineStyle}
-          sourceNode={sourceNode}
+          sourceNode={sourceNode.current as Node}
           sourceHandle={sourceHandle}
         />
       </g>
diff --git a/src/components/Edges/BezierEdge.tsx b/src/components/Edges/BezierEdge.tsx
index 2313767f..7edb5a42 100644
--- a/src/components/Edges/BezierEdge.tsx
+++ b/src/components/Edges/BezierEdge.tsx
@@ -1,7 +1,6 @@
 import React, { memo } from 'react';
 
 import EdgeText from './EdgeText';
-
 import { getCenter } from './utils';
 import { EdgeProps, Position } from '../../types';
 
diff --git a/src/components/Edges/wrapEdge.tsx b/src/components/Edges/wrapEdge.tsx
index 4e191cf7..a89eb28e 100644
--- a/src/components/Edges/wrapEdge.tsx
+++ b/src/components/Edges/wrapEdge.tsx
@@ -1,12 +1,21 @@
 import React, { memo, ComponentType, useCallback, useState, useMemo } from 'react';
 import cc from 'classcat';
+import shallow from 'zustand/shallow';
 
-import { useStoreActions, useStoreState } from '../../store/hooks';
-import { Edge, EdgeProps, WrapEdgeProps } from '../../types';
+import { useStore, useStoreApi } from '../../store';
+import { Edge, EdgeProps, WrapEdgeProps, ReactFlowState } from '../../types';
 import { onMouseDown } from '../../components/Handle/handler';
 import { EdgeAnchor } from './EdgeAnchor';
 import { getMarkerId } from '../../utils/graph';
 
+const selector = (s: ReactFlowState) => ({
+  addSelectedElements: s.addSelectedElements,
+  setConnectionNodeId: s.setConnectionNodeId,
+  unsetNodesSelection: s.unsetNodesSelection,
+  setPosition: s.setConnectionPosition,
+  connectionMode: s.connectionMode,
+});
+
 export default (EdgeComponent: ComponentType<EdgeProps>) => {
   const EdgeWrapper = ({
     id,
@@ -33,7 +42,7 @@ export default (EdgeComponent: ComponentType<EdgeProps>) => {
     sourcePosition,
     targetPosition,
     elementsSelectable,
-    isHidden,
+    hidden,
     sourceHandleId,
     targetHandleId,
     handleEdgeUpdate,
@@ -48,11 +57,11 @@ export default (EdgeComponent: ComponentType<EdgeProps>) => {
     markerEnd,
     markerStart,
   }: WrapEdgeProps): JSX.Element | null => {
-    const addSelectedElements = useStoreActions((actions) => actions.addSelectedElements);
-    const setConnectionNodeId = useStoreActions((actions) => actions.setConnectionNodeId);
-    const unsetNodesSelection = useStoreActions((actions) => actions.unsetNodesSelection);
-    const setPosition = useStoreActions((actions) => actions.setConnectionPosition);
-    const connectionMode = useStoreState((state) => state.connectionMode);
+    const store = useStoreApi();
+    const { addSelectedElements, setConnectionNodeId, unsetNodesSelection, setPosition, connectionMode } = useStore(
+      selector,
+      shallow
+    );
 
     const [updating, setUpdating] = useState<boolean>(false);
 
@@ -91,7 +100,7 @@ export default (EdgeComponent: ComponentType<EdgeProps>) => {
       (event: React.MouseEvent<SVGGElement, MouseEvent>): void => {
         if (elementsSelectable) {
           unsetNodesSelection();
-          addSelectedElements(edgeElement);
+          addSelectedElements([edgeElement]);
         }
 
         onClick?.(event, edgeElement);
@@ -158,7 +167,8 @@ export default (EdgeComponent: ComponentType<EdgeProps>) => {
           isValidConnection,
           connectionMode,
           isSourceHandle ? 'target' : 'source',
-          _onEdgeUpdate
+          _onEdgeUpdate,
+          store.getState
         );
       },
       [
@@ -194,7 +204,7 @@ export default (EdgeComponent: ComponentType<EdgeProps>) => {
     const markerStartUrl = useMemo(() => `url(#${getMarkerId(markerStart)})`, [markerStart]);
     const markerEndUrl = useMemo(() => `url(#${getMarkerId(markerEnd)})`, [markerEnd]);
 
-    if (isHidden) {
+    if (hidden) {
       return null;
     }
 
diff --git a/src/components/ElementUpdater/index.tsx b/src/components/ElementUpdater/index.tsx
deleted file mode 100644
index 5bbe70a6..00000000
--- a/src/components/ElementUpdater/index.tsx
+++ /dev/null
@@ -1,20 +0,0 @@
-import { useEffect } from 'react';
-
-import { useStoreActions } from '../../store/hooks';
-import { Elements } from '../../types';
-
-interface ElementUpdaterProps {
-  elements: Elements;
-}
-
-const ElementUpdater = ({ elements }: ElementUpdaterProps) => {
-  const setElements = useStoreActions((actions) => actions.setElements);
-
-  useEffect(() => {
-    setElements(elements);
-  }, [elements]);
-
-  return null;
-};
-
-export default ElementUpdater;
diff --git a/src/components/Handle/handler.ts b/src/components/Handle/handler.ts
index 03fd37df..2ef44055 100644
--- a/src/components/Handle/handler.ts
+++ b/src/components/Handle/handler.ts
@@ -3,16 +3,15 @@ import { MouseEvent as ReactMouseEvent } from 'react';
 import { getHostForElement } from '../../utils';
 
 import {
-  ElementId,
-  XYPosition,
-  OnConnectFunc,
-  OnConnectStartFunc,
-  OnConnectStopFunc,
-  OnConnectEndFunc,
+  OnConnect,
+  OnConnectStart,
+  OnConnectStop,
+  OnConnectEnd,
   ConnectionMode,
   SetConnectionId,
   Connection,
   HandleType,
+  XYPosition,
 } from '../../types';
 
 type ValidConnectionFunc = (connection: Connection) => boolean;
@@ -32,8 +31,8 @@ function checkElementBelowIsValid(
   event: MouseEvent,
   connectionMode: ConnectionMode,
   isTarget: boolean,
-  nodeId: ElementId,
-  handleId: ElementId | null,
+  nodeId: string,
+  handleId: string | null,
   isValidConnection: ValidConnectionFunc,
   doc: Document | ShadowRoot
 ) {
@@ -91,19 +90,19 @@ function resetRecentHandle(hoveredHandle: Element): void {
 
 export function onMouseDown(
   event: ReactMouseEvent,
-  handleId: ElementId | null,
-  nodeId: ElementId,
+  handleId: string | null,
+  nodeId: string,
   setConnectionNodeId: SetSourceIdFunc,
   setPosition: SetPosition,
-  onConnect: OnConnectFunc,
+  onConnect: OnConnect,
   isTarget: boolean,
   isValidConnection: ValidConnectionFunc,
   connectionMode: ConnectionMode,
   elementEdgeUpdaterType?: HandleType,
   onEdgeUpdateEnd?: (evt: MouseEvent) => void,
-  onConnectStart?: OnConnectStartFunc,
-  onConnectStop?: OnConnectStopFunc,
-  onConnectEnd?: OnConnectEndFunc
+  onConnectStart?: OnConnectStart,
+  onConnectStop?: OnConnectStop,
+  onConnectEnd?: OnConnectEnd
 ): void {
   const reactFlowNode = (event.target as Element).closest('.react-flow');
   // when react-flow is used inside a shadow root we can't use document
diff --git a/src/components/Handle/index.tsx b/src/components/Handle/index.tsx
index 9a38c550..e3b0111c 100644
--- a/src/components/Handle/index.tsx
+++ b/src/components/Handle/index.tsx
@@ -1,9 +1,10 @@
 import React, { memo, useContext, useCallback, HTMLAttributes, forwardRef } from 'react';
 import cc from 'classcat';
+import shallow from 'zustand/shallow';
 
-import { useStoreActions, useStoreState } from '../../store/hooks';
+import { useStore } from '../../store';
 import NodeIdContext from '../../contexts/NodeIdContext';
-import { HandleProps, Connection, ElementId, Position } from '../../types';
+import { HandleProps, Connection, ReactFlowState, Position } from '../../types';
 
 import { onMouseDown, SetSourceIdFunc, SetPosition } from './handler';
 
@@ -11,6 +12,16 @@ const alwaysValid = () => true;
 
 export type HandleComponentProps = HandleProps & Omit<HTMLAttributes<HTMLDivElement>, 'id'>;
 
+const selector = (s: ReactFlowState) => ({
+  setPosition: s.setConnectionPosition,
+  setConnectionNodeId: s.setConnectionNodeId,
+  onConnectAction: s.onConnect,
+  onConnectStart: s.onConnectStart,
+  onConnectStop: s.onConnectStop,
+  onConnectEnd: s.onConnectEnd,
+  connectionMode: s.connectionMode,
+});
+
 const Handle = forwardRef<HTMLDivElement, HandleComponentProps>(
   (
     {
@@ -26,14 +37,17 @@ const Handle = forwardRef<HTMLDivElement, HandleComponentProps>(
     },
     ref
   ) => {
-    const nodeId = useContext(NodeIdContext) as ElementId;
-    const setPosition = useStoreActions((actions) => actions.setConnectionPosition);
-    const setConnectionNodeId = useStoreActions((actions) => actions.setConnectionNodeId);
-    const onConnectAction = useStoreState((state) => state.onConnect);
-    const onConnectStart = useStoreState((state) => state.onConnectStart);
-    const onConnectStop = useStoreState((state) => state.onConnectStop);
-    const onConnectEnd = useStoreState((state) => state.onConnectEnd);
-    const connectionMode = useStoreState((state) => state.connectionMode);
+    const nodeId = useContext(NodeIdContext) as string;
+    const {
+      setPosition,
+      setConnectionNodeId,
+      onConnectAction,
+      onConnectStart,
+      onConnectStop,
+      onConnectEnd,
+      connectionMode,
+    } = useStore(selector, shallow);
+
     const handleId = id || null;
     const isTarget = type === 'target';
 
diff --git a/src/components/Nodes/DefaultNode.tsx b/src/components/Nodes/DefaultNode.tsx
index 11d9af2f..7d9dc5d5 100644
--- a/src/components/Nodes/DefaultNode.tsx
+++ b/src/components/Nodes/DefaultNode.tsx
@@ -8,13 +8,15 @@ const DefaultNode = ({
   isConnectable,
   targetPosition = Position.Top,
   sourcePosition = Position.Bottom,
-}: NodeProps) => (
-  <>
-    <Handle type="target" position={targetPosition} isConnectable={isConnectable} />
-    {data.label}
-    <Handle type="source" position={sourcePosition} isConnectable={isConnectable} />
-  </>
-);
+}: NodeProps) => {
+  return (
+    <>
+      <Handle type="target" position={targetPosition} isConnectable={isConnectable} />
+      {data.label}
+      <Handle type="source" position={sourcePosition} isConnectable={isConnectable} />
+    </>
+  );
+};
 
 DefaultNode.displayName = 'DefaultNode';
 
diff --git a/src/components/Nodes/utils.ts b/src/components/Nodes/utils.ts
index e917e853..1ed82013 100644
--- a/src/components/Nodes/utils.ts
+++ b/src/components/Nodes/utils.ts
@@ -13,7 +13,7 @@ export const getHandleBounds = (nodeElement: HTMLDivElement, scale: number) => {
 export const getHandleBoundsByHandleType = (
   selector: string,
   nodeElement: HTMLDivElement,
-  parentBounds: ClientRect | DOMRect,
+  parentBounds: DOMRect,
   k: number
 ): HandleElement[] | null => {
   const handles = nodeElement.querySelectorAll(selector);
@@ -24,20 +24,18 @@ export const getHandleBoundsByHandleType = (
 
   const handlesArray = Array.from(handles) as HTMLDivElement[];
 
-  return handlesArray.map(
-    (handle): HandleElement => {
-      const bounds = handle.getBoundingClientRect();
-      const dimensions = getDimensions(handle);
-      const handleId = handle.getAttribute('data-handleid');
-      const handlePosition = (handle.getAttribute('data-handlepos') as unknown) as Position;
+  return handlesArray.map((handle): HandleElement => {
+    const bounds = handle.getBoundingClientRect();
+    const dimensions = getDimensions(handle);
+    const handleId = handle.getAttribute('data-handleid');
+    const handlePosition = handle.getAttribute('data-handlepos') as unknown as Position;
 
-      return {
-        id: handleId,
-        position: handlePosition,
-        x: (bounds.left - parentBounds.left) / k,
-        y: (bounds.top - parentBounds.top) / k,
-        ...dimensions,
-      };
-    }
-  );
+    return {
+      id: handleId,
+      position: handlePosition,
+      x: (bounds.left - parentBounds.left) / k,
+      y: (bounds.top - parentBounds.top) / k,
+      ...dimensions,
+    };
+  });
 };
diff --git a/src/components/Nodes/wrapNode.tsx b/src/components/Nodes/wrapNode.tsx
index a6902416..20348106 100644
--- a/src/components/Nodes/wrapNode.tsx
+++ b/src/components/Nodes/wrapNode.tsx
@@ -1,22 +1,21 @@
-import React, {
-  useEffect,
-  useLayoutEffect,
-  useRef,
-  memo,
-  ComponentType,
-  CSSProperties,
-  useMemo,
-  MouseEvent,
-  useCallback,
-} from 'react';
+import React, { useEffect, useRef, memo, ComponentType, CSSProperties, useMemo, MouseEvent, useCallback } from 'react';
 import { DraggableCore, DraggableData, DraggableEvent } from 'react-draggable';
 import cc from 'classcat';
+import shallow from 'zustand/shallow';
 
-import { useStoreActions } from '../../store/hooks';
+import { useStore } from '../../store';
 import { Provider } from '../../contexts/NodeIdContext';
-import { NodeComponentProps, WrapNodeProps } from '../../types';
+import { NodeProps, WrapNodeProps, ReactFlowState } from '../../types';
 
-export default (NodeComponent: ComponentType<NodeComponentProps>) => {
+const selector = (s: ReactFlowState) => ({
+  addSelectedElements: s.addSelectedElements,
+  unsetNodesSelection: s.unsetNodesSelection,
+  updateNodePosition: s.updateNodePosition,
+  updateNodeDimensions: s.updateNodeDimensions,
+  unselectNodesAndEdges: s.unselectNodesAndEdges,
+});
+
+export default (NodeComponent: ComponentType<NodeProps>) => {
   const NodeWrapper = ({
     id,
     type,
@@ -42,19 +41,25 @@ export default (NodeComponent: ComponentType<NodeComponentProps>) => {
     selectNodesOnDrag,
     sourcePosition,
     targetPosition,
-    isHidden,
+    hidden,
     isInitialized,
     snapToGrid,
     snapGrid,
-    isDragging,
+    dragging,
     resizeObserver,
     dragHandle,
+    zIndex,
+    isParent,
+    noPanClassName,
+    noDragClassName,
   }: WrapNodeProps) => {
-    const updateNodeDimensions = useStoreActions((actions) => actions.updateNodeDimensions);
-    const addSelectedElements = useStoreActions((actions) => actions.addSelectedElements);
-    const updateNodePosDiff = useStoreActions((actions) => actions.updateNodePosDiff);
-    const unsetNodesSelection = useStoreActions((actions) => actions.unsetNodesSelection);
-
+    const {
+      addSelectedElements,
+      unselectNodesAndEdges,
+      unsetNodesSelection,
+      updateNodePosition,
+      //  updateNodeDimensions,
+    } = useStore(selector, shallow);
     const nodeElement = useRef<HTMLDivElement>(null);
 
     const node = useMemo(() => ({ id, type, position: { x: xPos, y: yPos }, data }), [id, type, xPos, yPos, data]);
@@ -62,7 +67,7 @@ export default (NodeComponent: ComponentType<NodeComponentProps>) => {
 
     const nodeStyle: CSSProperties = useMemo(
       () => ({
-        zIndex: selected ? 10 : 3,
+        zIndex,
         transform: `translate(${xPos}px,${yPos}px)`,
         pointerEvents:
           isSelectable || isDraggable || onClick || onMouseEnter || onMouseMove || onMouseLeave ? 'all' : 'none',
@@ -71,7 +76,6 @@ export default (NodeComponent: ComponentType<NodeComponentProps>) => {
         ...style,
       }),
       [
-        selected,
         xPos,
         yPos,
         isSelectable,
@@ -82,31 +86,34 @@ export default (NodeComponent: ComponentType<NodeComponentProps>) => {
         onMouseEnter,
         onMouseMove,
         onMouseLeave,
+        isParent,
+        zIndex,
       ]
     );
+
     const onMouseEnterHandler = useMemo(() => {
-      if (!onMouseEnter || isDragging) {
+      if (!onMouseEnter || dragging) {
         return;
       }
 
       return (event: MouseEvent) => onMouseEnter(event, node);
-    }, [onMouseEnter, isDragging, node]);
+    }, [onMouseEnter, dragging, node]);
 
     const onMouseMoveHandler = useMemo(() => {
-      if (!onMouseMove || isDragging) {
+      if (!onMouseMove || dragging) {
         return;
       }
 
       return (event: MouseEvent) => onMouseMove(event, node);
-    }, [onMouseMove, isDragging, node]);
+    }, [onMouseMove, dragging, node]);
 
     const onMouseLeaveHandler = useMemo(() => {
-      if (!onMouseLeave || isDragging) {
+      if (!onMouseLeave || dragging) {
         return;
       }
 
       return (event: MouseEvent) => onMouseLeave(event, node);
-    }, [onMouseLeave, isDragging, node]);
+    }, [onMouseLeave, dragging, node]);
 
     const onContextMenuHandler = useMemo(() => {
       if (!onContextMenu) {
@@ -123,7 +130,7 @@ export default (NodeComponent: ComponentType<NodeComponentProps>) => {
             unsetNodesSelection();
 
             if (!selected) {
-              addSelectedElements(node);
+              addSelectedElements([node]);
             }
           }
 
@@ -141,11 +148,11 @@ export default (NodeComponent: ComponentType<NodeComponentProps>) => {
           unsetNodesSelection();
 
           if (!selected) {
-            addSelectedElements(node);
+            addSelectedElements([node]);
           }
         } else if (!selectNodesOnDrag && !selected && isSelectable) {
+          unselectNodesAndEdges();
           unsetNodesSelection();
-          addSelectedElements([]);
         }
       },
       [node, selected, selectNodesOnDrag, isSelectable, onNodeDragStart]
@@ -153,20 +160,14 @@ export default (NodeComponent: ComponentType<NodeComponentProps>) => {
 
     const onDrag = useCallback(
       (event: DraggableEvent, draggableData: DraggableData) => {
+        node.position.x += draggableData.deltaX;
+        node.position.y += draggableData.deltaY;
+
         if (onNodeDrag) {
-          node.position.x += draggableData.deltaX;
-          node.position.y += draggableData.deltaY;
           onNodeDrag(event as MouseEvent, node);
         }
 
-        updateNodePosDiff({
-          id,
-          diff: {
-            x: draggableData.deltaX,
-            y: draggableData.deltaY,
-          },
-          isDragging: true,
-        });
+        updateNodePosition({ id, dragging: true, diff: { x: draggableData.deltaX, y: draggableData.deltaY } });
       },
       [id, node, onNodeDrag]
     );
@@ -175,9 +176,9 @@ export default (NodeComponent: ComponentType<NodeComponentProps>) => {
       (event: DraggableEvent) => {
         // onDragStop also gets called when user just clicks on a node.
         // Because of that we set dragging to true inside the onDrag handler and handle the click here
-        if (!isDragging) {
+        if (!dragging) {
           if (isSelectable && !selectNodesOnDrag && !selected) {
-            addSelectedElements(node);
+            addSelectedElements([node]);
           }
 
           onClick?.(event as MouseEvent, node);
@@ -185,14 +186,14 @@ export default (NodeComponent: ComponentType<NodeComponentProps>) => {
           return;
         }
 
-        updateNodePosDiff({
+        updateNodePosition({
           id: node.id,
-          isDragging: false,
+          dragging: false,
         });
 
         onNodeDragStop?.(event as MouseEvent, node);
       },
-      [node, isSelectable, selectNodesOnDrag, onClick, onNodeDragStop, isDragging, selected]
+      [node, isSelectable, selectNodesOnDrag, onClick, onNodeDragStop, dragging, selected]
     );
 
     const onNodeDoubleClickHandler = useCallback(
@@ -202,12 +203,6 @@ export default (NodeComponent: ComponentType<NodeComponentProps>) => {
       [node, onNodeDoubleClick]
     );
 
-    useLayoutEffect(() => {
-      if (nodeElement.current && !isHidden) {
-        updateNodeDimensions([{ id, nodeElement: nodeElement.current, forceUpdate: true }]);
-      }
-    }, [id, isHidden, sourcePosition, targetPosition]);
-
     useEffect(() => {
       if (nodeElement.current) {
         const currNode = nodeElement.current;
@@ -217,17 +212,19 @@ export default (NodeComponent: ComponentType<NodeComponentProps>) => {
       }
     }, []);
 
-    if (isHidden) {
+    if (hidden) {
       return null;
     }
 
     const nodeClasses = cc([
       'react-flow__node',
       `react-flow__node-${type}`,
+      noPanClassName,
       className,
       {
         selected,
         selectable: isSelectable,
+        parent: isParent,
       },
     ]);
 
@@ -238,7 +235,7 @@ export default (NodeComponent: ComponentType<NodeComponentProps>) => {
         onStop={onDragStop}
         scale={scale}
         disabled={!isDraggable}
-        cancel=".nodrag"
+        cancel={`.${noDragClassName}`}
         nodeRef={nodeElement}
         grid={grid}
         enableUserSelectHack={false}
@@ -267,8 +264,9 @@ export default (NodeComponent: ComponentType<NodeComponentProps>) => {
               isConnectable={isConnectable}
               sourcePosition={sourcePosition}
               targetPosition={targetPosition}
-              isDragging={isDragging}
+              dragging={dragging}
               dragHandle={dragHandle}
+              zIndex={zIndex}
             />
           </Provider>
         </div>
diff --git a/src/components/NodesSelection/index.tsx b/src/components/NodesSelection/index.tsx
index efb780e7..b75c5ea6 100644
--- a/src/components/NodesSelection/index.tsx
+++ b/src/components/NodesSelection/index.tsx
@@ -5,53 +5,45 @@
 
 import React, { useMemo, useCallback, useRef, MouseEvent } from 'react';
 import ReactDraggable, { DraggableData } from 'react-draggable';
+import cc from 'classcat';
 
-import { useStoreState, useStoreActions } from '../../store/hooks';
-import { isNode } from '../../utils/graph';
-import { Node } from '../../types';
+import { useStore } from '../../store';
+import { Node, ReactFlowState } from '../../types';
 
 export interface NodesSelectionProps {
   onSelectionDragStart?: (event: MouseEvent, nodes: Node[]) => void;
   onSelectionDrag?: (event: MouseEvent, nodes: Node[]) => void;
   onSelectionDragStop?: (event: MouseEvent, nodes: Node[]) => void;
   onSelectionContextMenu?: (event: MouseEvent, nodes: Node[]) => void;
+  noPanClassName?: string;
 }
+// @TODO: work with nodeInternals instead of converting it to an array
+const selector = (s: ReactFlowState) => ({
+  transform: s.transform,
+  selectedNodesBbox: s.selectedNodesBbox,
+  selectionActive: s.selectionActive,
+  selectedNodes: Array.from(s.nodeInternals)
+    .filter(([_, n]) => n.selected)
+    .map(([_, n]) => n),
+  snapToGrid: s.snapToGrid,
+  snapGrid: s.snapGrid,
+  updateNodePosition: s.updateNodePosition,
+});
 
 export default ({
   onSelectionDragStart,
   onSelectionDrag,
   onSelectionDragStop,
   onSelectionContextMenu,
+  noPanClassName,
 }: NodesSelectionProps) => {
-  const [tX, tY, tScale] = useStoreState((state) => state.transform);
-  const selectedNodesBbox = useStoreState((state) => state.selectedNodesBbox);
-  const selectionActive = useStoreState((state) => state.selectionActive);
-  const selectedElements = useStoreState((state) => state.selectedElements);
-  const snapToGrid = useStoreState((state) => state.snapToGrid);
-  const snapGrid = useStoreState((state) => state.snapGrid);
-  const nodes = useStoreState((state) => state.nodes);
-
-  const updateNodePosDiff = useStoreActions((actions) => actions.updateNodePosDiff);
-
+  const { transform, selectedNodesBbox, selectionActive, selectedNodes, snapToGrid, snapGrid, updateNodePosition } =
+    useStore(selector);
+  const [tX, tY, tScale] = transform;
   const nodeRef = useRef(null);
 
   const grid = useMemo(() => (snapToGrid ? snapGrid : [1, 1])! as [number, number], [snapToGrid, snapGrid]);
 
-  const selectedNodes = useMemo(
-    () =>
-      selectedElements
-        ? selectedElements.filter(isNode).map((selectedNode) => {
-            const matchingNode = nodes.find((node) => node.id === selectedNode.id);
-
-            return {
-              ...matchingNode,
-              position: matchingNode?.__rf.position,
-            } as Node;
-          })
-        : [],
-    [selectedElements, nodes]
-  );
-
   const style = useMemo(
     () => ({
       transform: `translate(${tX}px,${tY}px) scale(${tScale})`,
@@ -82,21 +74,21 @@ export default ({
         onSelectionDrag(event, selectedNodes);
       }
 
-      updateNodePosDiff({
+      updateNodePosition({
         diff: {
           x: data.deltaX,
           y: data.deltaY,
         },
-        isDragging: true,
+        dragging: true,
       });
     },
-    [onSelectionDrag, selectedNodes, updateNodePosDiff]
+    [onSelectionDrag, selectedNodes, updateNodePosition]
   );
 
   const onStop = useCallback(
     (event: MouseEvent) => {
-      updateNodePosDiff({
-        isDragging: false,
+      updateNodePosition({
+        dragging: false,
       });
 
       onSelectionDragStop?.(event, selectedNodes);
@@ -106,21 +98,17 @@ export default ({
 
   const onContextMenu = useCallback(
     (event: MouseEvent) => {
-      const selectedNodes = selectedElements
-        ? selectedElements.filter(isNode).map((selectedNode) => nodes.find((node) => node.id === selectedNode.id)!)
-        : [];
-
       onSelectionContextMenu?.(event, selectedNodes);
     },
-    [onSelectionContextMenu]
+    [onSelectionContextMenu, selectedNodes]
   );
 
-  if (!selectedElements || selectionActive) {
+  if (!selectedNodes || selectionActive) {
     return null;
   }
 
   return (
-    <div className="react-flow__nodesselection" style={style}>
+    <div className={cc(['react-flow__nodesselection', 'react-flow__container', noPanClassName])} style={style}>
       <ReactDraggable
         scale={tScale}
         grid={grid}
diff --git a/src/components/SelectionListener/index.tsx b/src/components/SelectionListener/index.tsx
index 9f912794..c8b59795 100644
--- a/src/components/SelectionListener/index.tsx
+++ b/src/components/SelectionListener/index.tsx
@@ -1,20 +1,29 @@
 import { useEffect } from 'react';
+import shallow from 'zustand/shallow';
 
-import { Elements } from '../../types';
-import { useStoreState } from '../../store/hooks';
+import { ReactFlowState, OnSelectionChangeFunc } from '../../types';
+import { useStore } from '../../store';
 
 interface SelectionListenerProps {
-  onSelectionChange: (elements: Elements | null) => void;
+  onSelectionChange: OnSelectionChangeFunc;
 }
 
-// This is a helper component for calling the onSelectionChange listener
+// @TODO: work with nodeInternals instead of converting it to an array
+const selectedElementsSelector = (s: ReactFlowState) => ({
+  selectedNodes: Array.from(s.nodeInternals)
+    .filter(([_, n]) => n.selected)
+    .map(([_, node]) => node),
+  selectedEdges: s.edges.filter((e) => e.selected),
+});
+
+// This is just a helper component for calling the onSelectionChange listener.
 
 export default ({ onSelectionChange }: SelectionListenerProps) => {
-  const selectedElements = useStoreState((s) => s.selectedElements);
+  const { selectedNodes, selectedEdges } = useStore(selectedElementsSelector, shallow);
 
   useEffect(() => {
-    onSelectionChange(selectedElements);
-  }, [selectedElements]);
+    onSelectionChange({ nodes: selectedNodes, edges: selectedEdges });
+  }, [selectedNodes, selectedEdges]);
 
   return null;
 };
diff --git a/src/components/StoreUpdater/index.tsx b/src/components/StoreUpdater/index.tsx
new file mode 100644
index 00000000..894e3335
--- /dev/null
+++ b/src/components/StoreUpdater/index.tsx
@@ -0,0 +1,210 @@
+import { useEffect } from 'react';
+import shallow from 'zustand/shallow';
+
+import { useStore } from '../../store';
+import {
+  Node,
+  Edge,
+  ReactFlowState,
+  OnConnect,
+  OnConnectStart,
+  OnConnectStop,
+  OnConnectEnd,
+  CoordinateExtent,
+  OnNodesChange,
+  OnEdgesChange,
+  ConnectionMode,
+  SnapGrid,
+} from '../../types';
+
+interface StoreUpdaterProps {
+  nodes: Node[];
+  edges: Edge[];
+  onConnect?: OnConnect;
+  onConnectStart?: OnConnectStart;
+  onConnectStop?: OnConnectStop;
+  onConnectEnd?: OnConnectEnd;
+  nodesDraggable?: boolean;
+  nodesConnectable?: boolean;
+  minZoom?: number;
+  maxZoom?: number;
+  nodeExtent?: CoordinateExtent;
+  onNodesChange?: OnNodesChange;
+  onEdgesChange?: OnEdgesChange;
+  elementsSelectable?: boolean;
+  connectionMode?: ConnectionMode;
+  snapToGrid?: boolean;
+  snapGrid?: SnapGrid;
+  translateExtent?: CoordinateExtent;
+}
+
+const selector = (s: ReactFlowState) => ({
+  setNodes: s.setNodes,
+  setEdges: s.setEdges,
+  setOnConnect: s.setOnConnect,
+  setOnConnectStart: s.setOnConnectStart,
+  setOnConnectStop: s.setOnConnectStop,
+  setOnConnectEnd: s.setOnConnectEnd,
+  setSnapGrid: s.setSnapGrid,
+  setSnapToGrid: s.setSnapToGrid,
+  setNodesDraggable: s.setNodesDraggable,
+  setNodesConnectable: s.setNodesConnectable,
+  setElementsSelectable: s.setElementsSelectable,
+  setMinZoom: s.setMinZoom,
+  setMaxZoom: s.setMaxZoom,
+  setTranslateExtent: s.setTranslateExtent,
+  setNodeExtent: s.setNodeExtent,
+  setConnectionMode: s.setConnectionMode,
+  setOnNodesChange: s.setOnNodesChange,
+  setOnEdgesChange: s.setOnEdgesChange,
+});
+
+const StoreUpdater = ({
+  nodes,
+  edges,
+  onConnect,
+  onConnectStart,
+  onConnectStop,
+  onConnectEnd,
+  nodesDraggable,
+  nodesConnectable,
+  minZoom,
+  maxZoom,
+  nodeExtent,
+  onNodesChange,
+  onEdgesChange,
+  elementsSelectable,
+  connectionMode,
+  snapGrid,
+  snapToGrid,
+  translateExtent,
+}: StoreUpdaterProps) => {
+  const {
+    setNodes,
+    setEdges,
+    setOnConnect,
+    setOnConnectStart,
+    setOnConnectStop,
+    setOnConnectEnd,
+    setNodesDraggable,
+    setSnapGrid,
+    setSnapToGrid,
+    setNodesConnectable,
+    setElementsSelectable,
+    setMinZoom,
+    setMaxZoom,
+    setTranslateExtent,
+    setNodeExtent,
+    setOnNodesChange,
+    setOnEdgesChange,
+    setConnectionMode,
+  } = useStore(selector, shallow);
+
+  useEffect(() => {
+    setNodes(nodes);
+  }, [nodes]);
+
+  useEffect(() => {
+    setEdges(edges, nodes);
+  }, [edges, nodes]);
+
+  useEffect(() => {
+    if (onConnect) {
+      setOnConnect(onConnect);
+    }
+  }, [onConnect]);
+
+  useEffect(() => {
+    if (onConnectStart) {
+      setOnConnectStart(onConnectStart);
+    }
+  }, [onConnectStart]);
+
+  useEffect(() => {
+    if (onConnectStop) {
+      setOnConnectStop(onConnectStop);
+    }
+  }, [onConnectStop]);
+
+  useEffect(() => {
+    if (onConnectEnd) {
+      setOnConnectEnd(onConnectEnd);
+    }
+  }, [onConnectEnd]);
+
+  useEffect(() => {
+    if (typeof snapToGrid !== 'undefined') {
+      setSnapToGrid(snapToGrid);
+    }
+  }, [snapToGrid]);
+
+  useEffect(() => {
+    if (typeof snapGrid !== 'undefined') {
+      setSnapGrid(snapGrid);
+    }
+  }, [snapGrid]);
+
+  useEffect(() => {
+    if (typeof nodesDraggable !== 'undefined') {
+      setNodesDraggable(nodesDraggable);
+    }
+  }, [nodesDraggable]);
+
+  useEffect(() => {
+    if (typeof nodesConnectable !== 'undefined') {
+      setNodesConnectable(nodesConnectable);
+    }
+  }, [nodesConnectable]);
+
+  useEffect(() => {
+    if (typeof elementsSelectable !== 'undefined') {
+      setElementsSelectable(elementsSelectable);
+    }
+  }, [elementsSelectable]);
+
+  useEffect(() => {
+    if (typeof minZoom !== 'undefined') {
+      setMinZoom(minZoom);
+    }
+  }, [minZoom]);
+
+  useEffect(() => {
+    if (typeof maxZoom !== 'undefined') {
+      setMaxZoom(maxZoom);
+    }
+  }, [maxZoom]);
+
+  useEffect(() => {
+    if (typeof translateExtent !== 'undefined') {
+      setTranslateExtent(translateExtent);
+    }
+  }, [translateExtent]);
+
+  useEffect(() => {
+    if (typeof nodeExtent !== 'undefined') {
+      setNodeExtent(nodeExtent);
+    }
+  }, [nodeExtent]);
+
+  useEffect(() => {
+    if (typeof connectionMode !== 'undefined') {
+      setConnectionMode(connectionMode);
+    }
+  }, [connectionMode]);
+
+  useEffect(() => {
+    if (typeof onNodesChange !== 'undefined') {
+      setOnNodesChange(onNodesChange);
+    }
+  }, [onNodesChange]);
+
+  useEffect(() => {
+    if (typeof onEdgesChange !== 'undefined') {
+      setOnEdgesChange(onEdgesChange);
+    }
+  }, [onEdgesChange]);
+
+  return null;
+};
+
+export default StoreUpdater;
diff --git a/src/components/UserSelection/index.tsx b/src/components/UserSelection/index.tsx
index 14a8a13a..5ffbef7e 100644
--- a/src/components/UserSelection/index.tsx
+++ b/src/components/UserSelection/index.tsx
@@ -4,8 +4,8 @@
 
 import React, { memo } from 'react';
 
-import { useStoreActions, useStoreState } from '../../store/hooks';
-import { XYPosition } from '../../types';
+import { useStore } from '../../store';
+import { XYPosition, ReactFlowState } from '../../types';
 
 type UserSelectionProps = {
   selectionKeyPressed: boolean;
@@ -25,8 +25,10 @@ function getMousePosition(event: React.MouseEvent): XYPosition | void {
   };
 }
 
+const userSelectionRectSelector = (state: ReactFlowState) => state.userSelectionRect;
+
 const SelectionRect = () => {
-  const userSelectionRect = useStoreState((state) => state.userSelectionRect);
+  const userSelectionRect = useStore(userSelectionRectSelector);
 
   if (!userSelectionRect.draw) {
     return null;
@@ -34,7 +36,7 @@ const SelectionRect = () => {
 
   return (
     <div
-      className="react-flow__selection"
+      className="react-flow__selection react-flow__container"
       style={{
         width: userSelectionRect.width,
         height: userSelectionRect.height,
@@ -44,14 +46,25 @@ const SelectionRect = () => {
   );
 };
 
+const selector = (s: ReactFlowState) => ({
+  selectionActive: s.selectionActive,
+  elementsSelectable: s.elementsSelectable,
+  setUserSelection: s.setUserSelection,
+  updateUserSelection: s.updateUserSelection,
+  unsetUserSelection: s.unsetUserSelection,
+  unsetNodesSelection: s.unsetNodesSelection,
+});
+
 export default memo(({ selectionKeyPressed }: UserSelectionProps) => {
-  const selectionActive = useStoreState((state) => state.selectionActive);
-  const elementsSelectable = useStoreState((state) => state.elementsSelectable);
+  const {
+    selectionActive,
+    elementsSelectable,
+    setUserSelection,
+    updateUserSelection,
+    unsetUserSelection,
+    unsetNodesSelection,
+  } = useStore(selector);
 
-  const setUserSelection = useStoreActions((actions) => actions.setUserSelection);
-  const updateUserSelection = useStoreActions((actions) => actions.updateUserSelection);
-  const unsetUserSelection = useStoreActions((actions) => actions.unsetUserSelection);
-  const unsetNodesSelection = useStoreActions((actions) => actions.unsetNodesSelection);
   const renderUserSelectionPane = selectionActive || selectionKeyPressed;
 
   if (!elementsSelectable || !renderUserSelectionPane) {
@@ -89,7 +102,7 @@ export default memo(({ selectionKeyPressed }: UserSelectionProps) => {
 
   return (
     <div
-      className="react-flow__selectionpane"
+      className="react-flow__selectionpane react-flow__container"
       onMouseDown={onMouseDown}
       onMouseMove={onMouseMove}
       onMouseUp={onMouseUp}
diff --git a/src/container/EdgeRenderer/MarkerDefinitions.tsx b/src/container/EdgeRenderer/MarkerDefinitions.tsx
index 222435f4..2cdf6203 100644
--- a/src/container/EdgeRenderer/MarkerDefinitions.tsx
+++ b/src/container/EdgeRenderer/MarkerDefinitions.tsx
@@ -1,6 +1,7 @@
 import React, { useMemo } from 'react';
-import { useStoreState } from '../../store/hooks';
-import { EdgeMarker, ArrowHeadType } from '../../types';
+
+import { useStore } from '../../store';
+import { EdgeMarker, ArrowHeadType, ReactFlowState } from '../../types';
 import { getMarkerId } from '../../utils/graph';
 interface MarkerProps extends EdgeMarker {
   id: string;
@@ -71,8 +72,10 @@ const Marker = ({
   );
 };
 
+const edgesSelector = (s: ReactFlowState) => s.edges;
+
 const MarkerDefinitions = ({ defaultColor }: MarkerDefinitionsProps) => {
-  const edges = useStoreState((state) => state.edges);
+  const edges = useStore(edgesSelector);
   const markers = useMemo(() => {
     const ids: string[] = [];
 
diff --git a/src/container/EdgeRenderer/index.tsx b/src/container/EdgeRenderer/index.tsx
index eaea06ad..ff6794da 100644
--- a/src/container/EdgeRenderer/index.tsx
+++ b/src/container/EdgeRenderer/index.tsx
@@ -1,30 +1,30 @@
 import React, { memo, CSSProperties, useCallback } from 'react';
+import shallow from 'zustand/shallow';
+import cc from 'classcat';
 
-import { useStoreState } from '../../store/hooks';
+import { useStore } from '../../store';
 import ConnectionLine from '../../components/ConnectionLine/index';
-import { isEdge } from '../../utils/graph';
 import MarkerDefinitions from './MarkerDefinitions';
-import { getEdgePositions, getHandle, isEdgeVisible, getSourceTargetNodes } from './utils';
+import { getEdgePositions, getHandle } from './utils';
 import {
   Position,
   Edge,
-  Node,
-  Elements,
   Connection,
   ConnectionLineType,
   ConnectionLineComponent,
   ConnectionMode,
-  Transform,
   OnEdgeUpdateFunc,
+  ReactFlowState,
+  NodeHandleBounds,
 } from '../../types';
+import useVisibleEdges from '../../hooks/useVisibleEdges';
 
 interface EdgeRendererProps {
   edgeTypes: any;
   connectionLineType: ConnectionLineType;
   connectionLineStyle?: CSSProperties;
   connectionLineComponent?: ConnectionLineComponent;
-  connectionMode?: ConnectionMode;
-  onElementClick?: (event: React.MouseEvent, element: Node | Edge) => void;
+  onEdgeClick?: (event: React.MouseEvent, node: Edge) => void;
   onEdgeDoubleClick?: (event: React.MouseEvent, edge: Edge) => void;
   defaultMarkerColor: string;
   onlyRenderVisibleElements: boolean;
@@ -36,215 +36,278 @@ interface EdgeRendererProps {
   onEdgeUpdateStart?: (event: React.MouseEvent, edge: Edge) => void;
   onEdgeUpdateEnd?: (event: MouseEvent, edge: Edge) => void;
   edgeUpdaterRadius?: number;
+  noPanClassName?: string;
 }
 
 interface EdgeWrapperProps {
   edge: Edge;
-  props: EdgeRendererProps;
-  nodes: Node[];
-  selectedElements: Elements | null;
+  edgeTypes: any;
+  markerEndId?: string;
+  onEdgeClick?: (event: React.MouseEvent, node: Edge) => void;
+  onEdgeContextMenu?: (event: React.MouseEvent, edge: Edge) => void;
+  onEdgeMouseEnter?: (event: React.MouseEvent, edge: Edge) => void;
+  onEdgeMouseMove?: (event: React.MouseEvent, edge: Edge) => void;
+  onEdgeMouseLeave?: (event: React.MouseEvent, edge: Edge) => void;
+  edgeUpdaterRadius?: number;
+  onEdgeDoubleClick?: (event: React.MouseEvent, edge: Edge) => void;
+  onEdgeUpdateStart?: (event: React.MouseEvent, edge: Edge) => void;
+  onEdgeUpdateEnd?: (event: MouseEvent, edge: Edge) => void;
+  onEdgeUpdate?: OnEdgeUpdateFunc;
   elementsSelectable: boolean;
-  transform: Transform;
-  width: number;
-  height: number;
-  onlyRenderVisibleElements: boolean;
   connectionMode?: ConnectionMode;
+  sourceNodeWidth?: number | null;
+  sourceNodeHeight?: number | null;
+  sourceNodeX?: number;
+  sourceNodeY?: number;
+  sourceNodeHandleBounds?: NodeHandleBounds;
+  targetNodeWidth?: number | null;
+  targetNodeHeight?: number | null;
+  targetNodeX?: number;
+  targetNodeY?: number;
+  targetNodeHandleBounds?: NodeHandleBounds;
+  noPanClassName?: string;
 }
 
-const Edge = ({
-  edge,
-  props,
-  nodes,
-  selectedElements,
-  elementsSelectable,
-  transform,
-  width,
-  height,
-  onlyRenderVisibleElements,
-  connectionMode,
-}: EdgeWrapperProps) => {
-  const sourceHandleId = edge.sourceHandle || null;
-  const targetHandleId = edge.targetHandle || null;
-  const { sourceNode, targetNode } = getSourceTargetNodes(edge, nodes);
+const Edge = memo(
+  ({
+    edge,
+    edgeTypes,
+    markerEndId,
+    onEdgeClick,
+    onEdgeContextMenu,
+    onEdgeMouseEnter,
+    onEdgeMouseMove,
+    onEdgeMouseLeave,
+    edgeUpdaterRadius,
+    onEdgeDoubleClick,
+    onEdgeUpdateStart,
+    onEdgeUpdateEnd,
+    onEdgeUpdate,
+    connectionMode,
+    elementsSelectable,
+    sourceNodeWidth,
+    sourceNodeHeight,
+    sourceNodeX,
+    sourceNodeY,
+    sourceNodeHandleBounds,
+    targetNodeWidth,
+    targetNodeHeight,
+    targetNodeX,
+    targetNodeY,
+    targetNodeHandleBounds,
+    noPanClassName,
+  }: EdgeWrapperProps) => {
+    const sourceHandleId = edge.sourceHandle || null;
+    const targetHandleId = edge.targetHandle || null;
 
-  const onConnectEdge = useCallback(
-    (connection: Connection) => {
-      props.onEdgeUpdate?.(edge, connection);
-    },
-    [edge, props.onEdgeUpdate]
-  );
+    const onConnectEdge = useCallback(
+      (connection: Connection) => {
+        onEdgeUpdate?.(edge, connection);
+      },
+      [edge, onEdgeUpdate]
+    );
 
-  if (!sourceNode) {
-    console.warn(`couldn't create edge for source id: ${edge.source}; edge id: ${edge.id}`);
-    return null;
+    // source and target node need to be initialized
+    if (!sourceNodeHandleBounds || !targetNodeHandleBounds) {
+      return null;
+    }
+
+    if (
+      !sourceNodeWidth ||
+      !sourceNodeHeight ||
+      typeof sourceNodeX === 'undefined' ||
+      typeof sourceNodeY === 'undefined'
+    ) {
+      console.warn(`couldn't create edge for source id: ${edge.source}; edge id: ${edge.id}`);
+      return null;
+    }
+
+    if (
+      !targetNodeWidth ||
+      !targetNodeHeight ||
+      typeof targetNodeX === 'undefined' ||
+      typeof targetNodeY === 'undefined'
+    ) {
+      console.warn(`couldn't create edge for target id: ${edge.target}; edge id: ${edge.id}`);
+      return null;
+    }
+
+    const edgeType = edge.type || 'default';
+    const EdgeComponent = edgeTypes[edgeType] || edgeTypes.default;
+    // when connection type is loose we can define all handles as sources
+    const targetNodeHandles =
+      connectionMode === ConnectionMode.Strict
+        ? targetNodeHandleBounds.target
+        : targetNodeHandleBounds.target || targetNodeHandleBounds.source;
+    const sourceHandle = getHandle(sourceNodeHandleBounds.source!, sourceHandleId);
+    const targetHandle = getHandle(targetNodeHandles!, targetHandleId);
+    const sourcePosition = sourceHandle ? sourceHandle.position : Position.Bottom;
+    const targetPosition = targetHandle ? targetHandle.position : Position.Top;
+
+    if (!sourceHandle) {
+      console.warn(`couldn't create edge for source handle id: ${sourceHandleId}; edge id: ${edge.id}`);
+      return null;
+    }
+
+    if (!targetHandle) {
+      console.warn(`couldn't create edge for target handle id: ${targetHandleId}; edge id: ${edge.id}`);
+      return null;
+    }
+
+    const { sourceX, sourceY, targetX, targetY } = getEdgePositions(
+      { x: sourceNodeX, y: sourceNodeY, width: sourceNodeWidth, height: sourceNodeHeight },
+      sourceHandle,
+      sourcePosition,
+      { x: targetNodeX, y: targetNodeY, width: targetNodeWidth, height: targetNodeHeight },
+      targetHandle,
+      targetPosition
+    );
+
+    return (
+      <EdgeComponent
+        key={edge.id}
+        id={edge.id}
+        className={cc([edge.className, noPanClassName])}
+        type={edgeType}
+        data={edge.data}
+        onClick={onEdgeClick}
+        selected={!!edge.selected}
+        animated={edge.animated}
+        label={edge.label}
+        labelStyle={edge.labelStyle}
+        labelShowBg={edge.labelShowBg}
+        labelBgStyle={edge.labelBgStyle}
+        labelBgPadding={edge.labelBgPadding}
+        labelBgBorderRadius={edge.labelBgBorderRadius}
+        style={edge.style}
+        source={edge.source}
+        target={edge.target}
+        sourceHandleId={sourceHandleId}
+        targetHandleId={targetHandleId}
+        markerEnd={edge.markerEnd}
+        markerStart={edge.markerStart}
+        sourceX={sourceX}
+        sourceY={sourceY}
+        targetX={targetX}
+        targetY={targetY}
+        sourcePosition={sourcePosition}
+        targetPosition={targetPosition}
+        elementsSelectable={elementsSelectable}
+        markerEndId={markerEndId}
+        hidden={edge.hidden}
+        onConnectEdge={onConnectEdge}
+        handleEdgeUpdate={typeof onEdgeUpdate !== 'undefined'}
+        onContextMenu={onEdgeContextMenu}
+        onMouseEnter={onEdgeMouseEnter}
+        onMouseMove={onEdgeMouseMove}
+        onMouseLeave={onEdgeMouseLeave}
+        edgeUpdaterRadius={edgeUpdaterRadius}
+        onEdgeDoubleClick={onEdgeDoubleClick}
+        onEdgeUpdateStart={onEdgeUpdateStart}
+        onEdgeUpdateEnd={onEdgeUpdateEnd}
+      />
+    );
   }
+);
 
-  if (!targetNode) {
-    console.warn(`couldn't create edge for target id: ${edge.target}; edge id: ${edge.id}`);
-    return null;
-  }
-
-  // source and target node need to be initialized
-  if (!sourceNode.__rf.width || !targetNode.__rf.width) {
-    return null;
-  }
-
-  const edgeType = edge.type || 'default';
-  const EdgeComponent = props.edgeTypes[edgeType] || props.edgeTypes.default;
-  const targetNodeBounds = targetNode.__rf.handleBounds;
-  // when connection type is loose we can define all handles as sources
-  const targetNodeHandles =
-    connectionMode === ConnectionMode.Strict
-      ? targetNodeBounds.target
-      : targetNodeBounds.target || targetNodeBounds.source;
-  const sourceHandle = getHandle(sourceNode.__rf.handleBounds.source, sourceHandleId);
-  const targetHandle = getHandle(targetNodeHandles, targetHandleId);
-  const sourcePosition = sourceHandle ? sourceHandle.position : Position.Bottom;
-  const targetPosition = targetHandle ? targetHandle.position : Position.Top;
-
-  if (!sourceHandle) {
-    console.warn(`couldn't create edge for source handle id: ${sourceHandleId}; edge id: ${edge.id}`);
-    return null;
-  }
-
-  if (!targetHandle) {
-    console.warn(`couldn't create edge for target handle id: ${targetHandleId}; edge id: ${edge.id}`);
-    return null;
-  }
-
-  const { sourceX, sourceY, targetX, targetY } = getEdgePositions(
-    sourceNode,
-    sourceHandle,
-    sourcePosition,
-    targetNode,
-    targetHandle,
-    targetPosition
-  );
-
-  const isVisible = onlyRenderVisibleElements
-    ? isEdgeVisible({
-        sourcePos: { x: sourceX, y: sourceY },
-        targetPos: { x: targetX, y: targetY },
-        width,
-        height,
-        transform,
-      })
-    : true;
-
-  if (!isVisible) {
-    return null;
-  }
-
-  const isSelected = selectedElements?.some((elm) => isEdge(elm) && elm.id === edge.id) || false;
-
-  return (
-    <EdgeComponent
-      key={edge.id}
-      id={edge.id}
-      className={edge.className}
-      type={edge.type}
-      data={edge.data}
-      onClick={props.onElementClick}
-      selected={isSelected}
-      animated={edge.animated}
-      label={edge.label}
-      labelStyle={edge.labelStyle}
-      labelShowBg={edge.labelShowBg}
-      labelBgStyle={edge.labelBgStyle}
-      labelBgPadding={edge.labelBgPadding}
-      labelBgBorderRadius={edge.labelBgBorderRadius}
-      style={edge.style}
-      markerEnd={edge.markerEnd}
-      markerStart={edge.markerStart}
-      source={edge.source}
-      target={edge.target}
-      sourceHandleId={sourceHandleId}
-      targetHandleId={targetHandleId}
-      sourceX={sourceX}
-      sourceY={sourceY}
-      targetX={targetX}
-      targetY={targetY}
-      sourcePosition={sourcePosition}
-      targetPosition={targetPosition}
-      elementsSelectable={elementsSelectable}
-      isHidden={edge.isHidden}
-      onConnectEdge={onConnectEdge}
-      handleEdgeUpdate={typeof props.onEdgeUpdate !== 'undefined'}
-      onContextMenu={props.onEdgeContextMenu}
-      onMouseEnter={props.onEdgeMouseEnter}
-      onMouseMove={props.onEdgeMouseMove}
-      onMouseLeave={props.onEdgeMouseLeave}
-      edgeUpdaterRadius={props.edgeUpdaterRadius}
-      onEdgeDoubleClick={props.onEdgeDoubleClick}
-      onEdgeUpdateStart={props.onEdgeUpdateStart}
-      onEdgeUpdateEnd={props.onEdgeUpdateEnd}
-    />
-  );
-};
+const selector = (s: ReactFlowState) => ({
+  connectionNodeId: s.connectionNodeId,
+  connectionHandleId: s.connectionHandleId,
+  connectionHandleType: s.connectionHandleType,
+  connectionPosition: s.connectionPosition,
+  nodesConnectable: s.nodesConnectable,
+  elementsSelectable: s.elementsSelectable,
+  width: s.width,
+  height: s.height,
+  connectionMode: s.connectionMode,
+  nodeInternals: s.nodeInternals,
+});
 
 const EdgeRenderer = (props: EdgeRendererProps) => {
-  const transform = useStoreState((state) => state.transform);
-  const nodes = useStoreState((state) => state.nodes);
-  const edges = useStoreState((state) => state.edges);
-  const connectionNodeId = useStoreState((state) => state.connectionNodeId);
-  const connectionHandleId = useStoreState((state) => state.connectionHandleId);
-  const connectionHandleType = useStoreState((state) => state.connectionHandleType);
-  const connectionPosition = useStoreState((state) => state.connectionPosition);
-  const selectedElements = useStoreState((state) => state.selectedElements);
-  const nodesConnectable = useStoreState((state) => state.nodesConnectable);
-  const elementsSelectable = useStoreState((state) => state.elementsSelectable);
-  const width = useStoreState((state) => state.width);
-  const height = useStoreState((state) => state.height);
+  const {
+    connectionNodeId,
+    connectionHandleId,
+    connectionHandleType,
+    connectionPosition,
+    nodesConnectable,
+    elementsSelectable,
+    width,
+    height,
+    connectionMode,
+    nodeInternals,
+  } = useStore(selector, shallow);
+  const edgeTree = useVisibleEdges(props.onlyRenderVisibleElements, nodeInternals);
 
   if (!width) {
     return null;
   }
 
-  const {
-    connectionLineType,
-    defaultMarkerColor,
-    connectionLineStyle,
-    connectionLineComponent,
-    onlyRenderVisibleElements,
-  } = props;
-  const transformStyle = `translate(${transform[0]},${transform[1]}) scale(${transform[2]})`;
+  const { connectionLineType, defaultMarkerColor, connectionLineStyle, connectionLineComponent } = props;
   const renderConnectionLine = connectionNodeId && connectionHandleType;
 
   return (
-    <svg width={width} height={height} className="react-flow__edges">
-      <MarkerDefinitions defaultColor={defaultMarkerColor} />
-      <g transform={transformStyle}>
-        {edges.map((edge: Edge) => (
-          <Edge
-            key={edge.id}
-            edge={edge}
-            props={props}
-            nodes={nodes}
-            selectedElements={selectedElements}
-            elementsSelectable={elementsSelectable}
-            transform={transform}
-            width={width}
-            height={height}
-            onlyRenderVisibleElements={onlyRenderVisibleElements}
-          />
-        ))}
-        {renderConnectionLine && (
-          <ConnectionLine
-            nodes={nodes}
-            connectionNodeId={connectionNodeId!}
-            connectionHandleId={connectionHandleId}
-            connectionHandleType={connectionHandleType!}
-            connectionPositionX={connectionPosition.x}
-            connectionPositionY={connectionPosition.y}
-            transform={transform}
-            connectionLineStyle={connectionLineStyle}
-            connectionLineType={connectionLineType}
-            isConnectable={nodesConnectable}
-            CustomConnectionLineComponent={connectionLineComponent}
-          />
-        )}
-      </g>
-    </svg>
+    <>
+      {edgeTree.map(({ level, edges, isMaxLevel }) => (
+        <svg
+          key={level}
+          style={{ zIndex: level }}
+          width={width}
+          height={height}
+          className="react-flow__edges react-flow__container"
+        >
+          {isMaxLevel && <MarkerDefinitions defaultColor={defaultMarkerColor} />}
+          <g>
+            {edges.map((edge: Edge) => {
+              const sourceNode = nodeInternals.get(edge.source);
+              const targetNode = nodeInternals.get(edge.target);
+
+              return (
+                <Edge
+                  key={edge.id}
+                  edge={edge}
+                  sourceNodeWidth={sourceNode?.width}
+                  sourceNodeHeight={sourceNode?.height}
+                  sourceNodeX={sourceNode?.positionAbsolute?.x}
+                  sourceNodeY={sourceNode?.positionAbsolute?.y}
+                  sourceNodeHandleBounds={sourceNode?.handleBounds}
+                  targetNodeWidth={targetNode?.width}
+                  targetNodeHeight={targetNode?.height}
+                  targetNodeX={targetNode?.positionAbsolute?.x}
+                  targetNodeY={targetNode?.positionAbsolute?.y}
+                  targetNodeHandleBounds={targetNode?.handleBounds}
+                  elementsSelectable={elementsSelectable}
+                  onEdgeContextMenu={props.onEdgeContextMenu}
+                  onEdgeMouseEnter={props.onEdgeMouseEnter}
+                  onEdgeMouseMove={props.onEdgeMouseMove}
+                  onEdgeMouseLeave={props.onEdgeMouseLeave}
+                  edgeUpdaterRadius={props.edgeUpdaterRadius}
+                  onEdgeDoubleClick={props.onEdgeDoubleClick}
+                  onEdgeUpdateStart={props.onEdgeUpdateStart}
+                  onEdgeUpdateEnd={props.onEdgeUpdateEnd}
+                  onEdgeUpdate={props.onEdgeUpdate}
+                  edgeTypes={props.edgeTypes}
+                  connectionMode={connectionMode}
+                  noPanClassName={props.noPanClassName}
+                />
+              );
+            })}
+            {renderConnectionLine && isMaxLevel && (
+              <ConnectionLine
+                connectionNodeId={connectionNodeId!}
+                connectionHandleId={connectionHandleId}
+                connectionHandleType={connectionHandleType!}
+                connectionPositionX={connectionPosition.x}
+                connectionPositionY={connectionPosition.y}
+                connectionLineStyle={connectionLineStyle}
+                connectionLineType={connectionLineType}
+                isConnectable={nodesConnectable}
+                CustomConnectionLineComponent={connectionLineComponent}
+              />
+            )}
+          </g>
+        </svg>
+      ))}
+    </>
   );
 };
 
diff --git a/src/container/EdgeRenderer/utils.ts b/src/container/EdgeRenderer/utils.ts
index 0344aff8..29a29425 100644
--- a/src/container/EdgeRenderer/utils.ts
+++ b/src/container/EdgeRenderer/utils.ts
@@ -2,19 +2,9 @@ import { ComponentType } from 'react';
 
 import { BezierEdge, StepEdge, SmoothStepEdge, StraightEdge } from '../../components/Edges';
 import wrapEdge from '../../components/Edges/wrapEdge';
-import { rectToBox } from '../../utils/graph';
+import { rectToBox } from '../../utils';
 
-import {
-  EdgeTypesType,
-  EdgeProps,
-  Position,
-  Node,
-  XYPosition,
-  ElementId,
-  HandleElement,
-  Transform,
-  Edge,
-} from '../../types';
+import { EdgeTypesType, EdgeProps, HandleElement, Position, XYPosition, Transform, Rect } from '../../types';
 
 export function createEdgeTypes(edgeTypes: EdgeTypesType): EdgeTypesType {
   const standardTypes: EdgeTypesType = {
@@ -39,11 +29,11 @@ export function createEdgeTypes(edgeTypes: EdgeTypesType): EdgeTypesType {
   };
 }
 
-export function getHandlePosition(position: Position, node: Node, handle: any | null = null): XYPosition {
-  const x = (handle?.x || 0) + node.__rf.position.x;
-  const y = (handle?.y || 0) + node.__rf.position.y;
-  const width = handle?.width || node.__rf.width;
-  const height = handle?.height || node.__rf.height;
+export function getHandlePosition(position: Position, nodeRect: Rect, handle: any | null = null): XYPosition {
+  const x = (handle?.x || 0) + nodeRect.x;
+  const y = (handle?.y || 0) + nodeRect.y;
+  const width = handle?.width || nodeRect.width;
+  const height = handle?.height || nodeRect.height;
 
   switch (position) {
     case Position.Top:
@@ -69,7 +59,7 @@ export function getHandlePosition(position: Position, node: Node, handle: any |
   }
 }
 
-export function getHandle(bounds: HandleElement[], handleId: ElementId | null): HandleElement | null {
+export function getHandle(bounds: HandleElement[], handleId: string | null): HandleElement | null {
   if (!bounds) {
     return null;
   }
@@ -94,15 +84,15 @@ interface EdgePositions {
 }
 
 export const getEdgePositions = (
-  sourceNode: Node,
+  sourceNodeRect: Rect,
   sourceHandle: HandleElement | unknown,
   sourcePosition: Position,
-  targetNode: Node,
+  targetNodeRect: Rect,
   targetHandle: HandleElement | unknown,
   targetPosition: Position
 ): EdgePositions => {
-  const sourceHandlePos = getHandlePosition(sourcePosition, sourceNode, sourceHandle);
-  const targetHandlePos = getHandlePosition(targetPosition, targetNode, targetHandle);
+  const sourceHandlePos = getHandlePosition(sourcePosition, sourceNodeRect, sourceHandle);
+  const targetHandlePos = getHandlePosition(targetPosition, targetNodeRect, targetHandle);
 
   return {
     sourceX: sourceHandlePos.x,
@@ -115,17 +105,31 @@ export const getEdgePositions = (
 interface IsEdgeVisibleParams {
   sourcePos: XYPosition;
   targetPos: XYPosition;
+  sourceWidth: number;
+  sourceHeight: number;
+  targetWidth: number;
+  targetHeight: number;
   width: number;
   height: number;
   transform: Transform;
 }
 
-export function isEdgeVisible({ sourcePos, targetPos, width, height, transform }: IsEdgeVisibleParams): boolean {
+export function isEdgeVisible({
+  sourcePos,
+  targetPos,
+  sourceWidth,
+  sourceHeight,
+  targetWidth,
+  targetHeight,
+  width,
+  height,
+  transform,
+}: IsEdgeVisibleParams): boolean {
   const edgeBox = {
     x: Math.min(sourcePos.x, targetPos.x),
     y: Math.min(sourcePos.y, targetPos.y),
-    x2: Math.max(sourcePos.x, targetPos.x),
-    y2: Math.max(sourcePos.y, targetPos.y),
+    x2: Math.max(sourcePos.x + sourceWidth, targetPos.x + targetWidth),
+    y2: Math.max(sourcePos.y + sourceHeight, targetPos.y + targetHeight),
   };
 
   if (edgeBox.x === edgeBox.x2) {
@@ -149,23 +153,3 @@ export function isEdgeVisible({ sourcePos, targetPos, width, height, transform }
 
   return overlappingArea > 0;
 }
-
-type SourceTargetNode = {
-  sourceNode: Node | null;
-  targetNode: Node | null;
-};
-
-export const getSourceTargetNodes = (edge: Edge, nodes: Node[]): SourceTargetNode => {
-  return nodes.reduce(
-    (res, node) => {
-      if (node.id === edge.source) {
-        res.sourceNode = node;
-      }
-      if (node.id === edge.target) {
-        res.targetNode = node;
-      }
-      return res;
-    },
-    { sourceNode: null, targetNode: null } as SourceTargetNode
-  );
-};
diff --git a/src/container/FlowRenderer/index.tsx b/src/container/FlowRenderer/index.tsx
index 37f983ba..458e9c59 100644
--- a/src/container/FlowRenderer/index.tsx
+++ b/src/container/FlowRenderer/index.tsx
@@ -1,18 +1,18 @@
 import React, { useCallback, memo, ReactNode, WheelEvent, MouseEvent } from 'react';
-import { useStoreActions, useStoreState } from '../../store/hooks';
+import shallow from 'zustand/shallow';
 
+import { useStore } from '../../store';
 import useGlobalKeyHandler from '../../hooks/useGlobalKeyHandler';
 import useKeyPress from '../../hooks/useKeyPress';
-
 import { GraphViewProps } from '../GraphView';
 import ZoomPane from '../ZoomPane';
 import UserSelection from '../../components/UserSelection';
 import NodesSelection from '../../components/NodesSelection';
+import { ReactFlowState } from '../../types';
 
 interface FlowRendererProps
   extends Omit<
     GraphViewProps,
-    | 'elements'
     | 'snapToGrid'
     | 'nodeTypes'
     | 'edgeTypes'
@@ -26,12 +26,17 @@ interface FlowRendererProps
   children: ReactNode;
 }
 
+const selector = (s: ReactFlowState) => ({
+  unsetNodesSelection: s.unsetNodesSelection,
+  resetSelectedElements: s.resetSelectedElements,
+  nodesSelectionActive: s.nodesSelectionActive,
+});
+
 const FlowRenderer = ({
   children,
   onPaneClick,
   onPaneContextMenu,
   onPaneScroll,
-  onElementsRemove,
   deleteKeyCode,
   onMove,
   onMoveStart,
@@ -49,20 +54,19 @@ const FlowRenderer = ({
   paneMoveable,
   defaultPosition,
   defaultZoom,
-  translateExtent,
   preventScrolling,
   onSelectionDragStart,
   onSelectionDrag,
   onSelectionDragStop,
   onSelectionContextMenu,
+  noWheelClassName,
+  noPanClassName,
 }: FlowRendererProps) => {
-  const unsetNodesSelection = useStoreActions((actions) => actions.unsetNodesSelection);
-  const resetSelectedElements = useStoreActions((actions) => actions.resetSelectedElements);
-  const nodesSelectionActive = useStoreState((state) => state.nodesSelectionActive);
+  const { unsetNodesSelection, resetSelectedElements, nodesSelectionActive } = useStore(selector, shallow);
 
   const selectionKeyPressed = useKeyPress(selectionKeyCode);
 
-  useGlobalKeyHandler({ onElementsRemove, deleteKeyCode, multiSelectionKeyCode });
+  useGlobalKeyHandler({ deleteKeyCode, multiSelectionKeyCode });
 
   const onClick = useCallback(
     (event: MouseEvent) => {
@@ -103,9 +107,10 @@ const FlowRenderer = ({
       paneMoveable={paneMoveable}
       defaultPosition={defaultPosition}
       defaultZoom={defaultZoom}
-      translateExtent={translateExtent}
       zoomActivationKeyCode={zoomActivationKeyCode}
       preventScrolling={preventScrolling}
+      noWheelClassName={noWheelClassName}
+      noPanClassName={noPanClassName}
     >
       {children}
       <UserSelection selectionKeyPressed={selectionKeyPressed} />
@@ -115,9 +120,15 @@ const FlowRenderer = ({
           onSelectionDrag={onSelectionDrag}
           onSelectionDragStop={onSelectionDragStop}
           onSelectionContextMenu={onSelectionContextMenu}
+          noPanClassName={noPanClassName}
         />
       )}
-      <div className="react-flow__pane" onClick={onClick} onContextMenu={onContextMenu} onWheel={onWheel} />
+      <div
+        className="react-flow__pane react-flow__container"
+        onClick={onClick}
+        onContextMenu={onContextMenu}
+        onWheel={onWheel}
+      />
     </ZoomPane>
   );
 };
diff --git a/src/container/GraphView/index.tsx b/src/container/GraphView/index.tsx
index 51663eb5..5561833d 100644
--- a/src/container/GraphView/index.tsx
+++ b/src/container/GraphView/index.tsx
@@ -1,30 +1,28 @@
-import React, { useEffect, useRef, memo } from 'react';
+import React, { memo } from 'react';
 
-import { useStoreActions, useStore } from '../../store/hooks';
 import FlowRenderer from '../FlowRenderer';
 import NodeRenderer from '../NodeRenderer';
 import EdgeRenderer from '../EdgeRenderer';
-import { onLoadProject, onLoadGetElements, onLoadToObject } from '../../utils/graph';
-import useZoomPanHelper from '../../hooks/useZoomPanHelper';
-
+import Viewport from '../Viewport';
 import { ReactFlowProps } from '../ReactFlow';
-
+import useOnLoadHandler from '../../hooks/useOnLoadHandler';
 import { NodeTypesType, EdgeTypesType, ConnectionLineType, KeyCode } from '../../types';
 
-export interface GraphViewProps extends Omit<ReactFlowProps, 'onSelectionChange' | 'elements'> {
+export interface GraphViewProps extends Omit<ReactFlowProps, 'onSelectionChange' | 'nodes' | 'edges'> {
   nodeTypes: NodeTypesType;
   edgeTypes: EdgeTypesType;
   selectionKeyCode: KeyCode;
   deleteKeyCode: KeyCode;
   multiSelectionKeyCode: KeyCode;
   connectionLineType: ConnectionLineType;
-  snapToGrid: boolean;
-  snapGrid: [number, number];
   onlyRenderVisibleElements: boolean;
   defaultZoom: number;
   defaultPosition: [number, number];
   defaultMarkerColor: string;
   selectNodesOnDrag: boolean;
+  noDragClassName: string;
+  noWheelClassName: string;
+  noPanClassName: string;
 }
 
 const GraphView = ({
@@ -34,7 +32,8 @@ const GraphView = ({
   onMoveStart,
   onMoveEnd,
   onLoad,
-  onElementClick,
+  onNodeClick,
+  onEdgeClick,
   onNodeDoubleClick,
   onEdgeDoubleClick,
   onNodeMouseEnter,
@@ -48,33 +47,19 @@ const GraphView = ({
   onSelectionDrag,
   onSelectionDragStop,
   onSelectionContextMenu,
-  connectionMode,
   connectionLineType,
   connectionLineStyle,
   connectionLineComponent,
   selectionKeyCode,
   multiSelectionKeyCode,
   zoomActivationKeyCode,
-  onElementsRemove,
   deleteKeyCode,
-  onConnect,
-  onConnectStart,
-  onConnectStop,
-  onConnectEnd,
-  snapToGrid,
-  snapGrid,
   onlyRenderVisibleElements,
-  nodesDraggable,
-  nodesConnectable,
   elementsSelectable,
   selectNodesOnDrag,
-  minZoom,
-  maxZoom,
   defaultZoom,
   defaultPosition,
-  translateExtent,
   preventScrolling,
-  nodeExtent,
   defaultMarkerColor,
   zoomOnScroll,
   zoomOnPinch,
@@ -94,134 +79,17 @@ const GraphView = ({
   edgeUpdaterRadius,
   onEdgeUpdateStart,
   onEdgeUpdateEnd,
+  noDragClassName,
+  noWheelClassName,
+  noPanClassName,
 }: GraphViewProps) => {
-  const isInitialized = useRef<boolean>(false);
-  const setOnConnect = useStoreActions((actions) => actions.setOnConnect);
-  const setOnConnectStart = useStoreActions((actions) => actions.setOnConnectStart);
-  const setOnConnectStop = useStoreActions((actions) => actions.setOnConnectStop);
-  const setOnConnectEnd = useStoreActions((actions) => actions.setOnConnectEnd);
-  const setSnapGrid = useStoreActions((actions) => actions.setSnapGrid);
-  const setSnapToGrid = useStoreActions((actions) => actions.setSnapToGrid);
-  const setNodesDraggable = useStoreActions((actions) => actions.setNodesDraggable);
-  const setNodesConnectable = useStoreActions((actions) => actions.setNodesConnectable);
-  const setElementsSelectable = useStoreActions((actions) => actions.setElementsSelectable);
-  const setMinZoom = useStoreActions((actions) => actions.setMinZoom);
-  const setMaxZoom = useStoreActions((actions) => actions.setMaxZoom);
-  const setTranslateExtent = useStoreActions((actions) => actions.setTranslateExtent);
-  const setNodeExtent = useStoreActions((actions) => actions.setNodeExtent);
-  const setConnectionMode = useStoreActions((actions) => actions.setConnectionMode);
-  const currentStore = useStore();
-  const { zoomIn, zoomOut, zoomTo, transform, fitView, initialized } = useZoomPanHelper();
-
-  useEffect(() => {
-    if (!isInitialized.current && initialized) {
-      if (onLoad) {
-        onLoad({
-          fitView: (params = { padding: 0.1 }) => fitView(params),
-          zoomIn,
-          zoomOut,
-          zoomTo,
-          setTransform: transform,
-          project: onLoadProject(currentStore),
-          getElements: onLoadGetElements(currentStore),
-          toObject: onLoadToObject(currentStore),
-        });
-      }
-
-      isInitialized.current = true;
-    }
-  }, [onLoad, zoomIn, zoomOut, zoomTo, transform, fitView, initialized]);
-
-  useEffect(() => {
-    if (onConnect) {
-      setOnConnect(onConnect);
-    }
-  }, [onConnect]);
-
-  useEffect(() => {
-    if (onConnectStart) {
-      setOnConnectStart(onConnectStart);
-    }
-  }, [onConnectStart]);
-
-  useEffect(() => {
-    if (onConnectStop) {
-      setOnConnectStop(onConnectStop);
-    }
-  }, [onConnectStop]);
-
-  useEffect(() => {
-    if (onConnectEnd) {
-      setOnConnectEnd(onConnectEnd);
-    }
-  }, [onConnectEnd]);
-
-  useEffect(() => {
-    if (typeof snapToGrid !== 'undefined') {
-      setSnapToGrid(snapToGrid);
-    }
-  }, [snapToGrid]);
-
-  useEffect(() => {
-    if (typeof snapGrid !== 'undefined') {
-      setSnapGrid(snapGrid);
-    }
-  }, [snapGrid]);
-
-  useEffect(() => {
-    if (typeof nodesDraggable !== 'undefined') {
-      setNodesDraggable(nodesDraggable);
-    }
-  }, [nodesDraggable]);
-
-  useEffect(() => {
-    if (typeof nodesConnectable !== 'undefined') {
-      setNodesConnectable(nodesConnectable);
-    }
-  }, [nodesConnectable]);
-
-  useEffect(() => {
-    if (typeof elementsSelectable !== 'undefined') {
-      setElementsSelectable(elementsSelectable);
-    }
-  }, [elementsSelectable]);
-
-  useEffect(() => {
-    if (typeof minZoom !== 'undefined') {
-      setMinZoom(minZoom);
-    }
-  }, [minZoom]);
-
-  useEffect(() => {
-    if (typeof maxZoom !== 'undefined') {
-      setMaxZoom(maxZoom);
-    }
-  }, [maxZoom]);
-
-  useEffect(() => {
-    if (typeof translateExtent !== 'undefined') {
-      setTranslateExtent(translateExtent);
-    }
-  }, [translateExtent]);
-
-  useEffect(() => {
-    if (typeof nodeExtent !== 'undefined') {
-      setNodeExtent(nodeExtent);
-    }
-  }, [nodeExtent]);
-
-  useEffect(() => {
-    if (typeof connectionMode !== 'undefined') {
-      setConnectionMode(connectionMode);
-    }
-  }, [connectionMode]);
+  useOnLoadHandler(onLoad);
 
   return (
     <FlowRenderer
       onPaneClick={onPaneClick}
       onPaneContextMenu={onPaneContextMenu}
       onPaneScroll={onPaneScroll}
-      onElementsRemove={onElementsRemove}
       deleteKeyCode={deleteKeyCode}
       selectionKeyCode={selectionKeyCode}
       multiSelectionKeyCode={multiSelectionKeyCode}
@@ -239,48 +107,52 @@ const GraphView = ({
       paneMoveable={paneMoveable}
       defaultPosition={defaultPosition}
       defaultZoom={defaultZoom}
-      translateExtent={translateExtent}
       onSelectionDragStart={onSelectionDragStart}
       onSelectionDrag={onSelectionDrag}
       onSelectionDragStop={onSelectionDragStop}
       onSelectionContextMenu={onSelectionContextMenu}
       preventScrolling={preventScrolling}
+      noDragClassName={noDragClassName}
+      noWheelClassName={noWheelClassName}
+      noPanClassName={noPanClassName}
     >
-      <NodeRenderer
-        nodeTypes={nodeTypes}
-        onElementClick={onElementClick}
-        onNodeDoubleClick={onNodeDoubleClick}
-        onNodeMouseEnter={onNodeMouseEnter}
-        onNodeMouseMove={onNodeMouseMove}
-        onNodeMouseLeave={onNodeMouseLeave}
-        onNodeContextMenu={onNodeContextMenu}
-        onNodeDragStop={onNodeDragStop}
-        onNodeDrag={onNodeDrag}
-        onNodeDragStart={onNodeDragStart}
-        selectNodesOnDrag={selectNodesOnDrag}
-        snapToGrid={snapToGrid}
-        snapGrid={snapGrid}
-        onlyRenderVisibleElements={onlyRenderVisibleElements}
-      />
-      <EdgeRenderer
-        edgeTypes={edgeTypes}
-        onElementClick={onElementClick}
-        onEdgeDoubleClick={onEdgeDoubleClick}
-        connectionLineType={connectionLineType}
-        connectionLineStyle={connectionLineStyle}
-        connectionLineComponent={connectionLineComponent}
-        connectionMode={connectionMode}
-        defaultMarkerColor={defaultMarkerColor}
-        onEdgeUpdate={onEdgeUpdate}
-        onlyRenderVisibleElements={onlyRenderVisibleElements}
-        onEdgeContextMenu={onEdgeContextMenu}
-        onEdgeMouseEnter={onEdgeMouseEnter}
-        onEdgeMouseMove={onEdgeMouseMove}
-        onEdgeMouseLeave={onEdgeMouseLeave}
-        onEdgeUpdateStart={onEdgeUpdateStart}
-        onEdgeUpdateEnd={onEdgeUpdateEnd}
-        edgeUpdaterRadius={edgeUpdaterRadius}
-      />
+      <Viewport>
+        <EdgeRenderer
+          edgeTypes={edgeTypes}
+          onEdgeClick={onEdgeClick}
+          onEdgeDoubleClick={onEdgeDoubleClick}
+          connectionLineType={connectionLineType}
+          connectionLineStyle={connectionLineStyle}
+          connectionLineComponent={connectionLineComponent}
+          onEdgeUpdate={onEdgeUpdate}
+          onlyRenderVisibleElements={onlyRenderVisibleElements}
+          onEdgeContextMenu={onEdgeContextMenu}
+          onEdgeMouseEnter={onEdgeMouseEnter}
+          onEdgeMouseMove={onEdgeMouseMove}
+          onEdgeMouseLeave={onEdgeMouseLeave}
+          onEdgeUpdateStart={onEdgeUpdateStart}
+          onEdgeUpdateEnd={onEdgeUpdateEnd}
+          edgeUpdaterRadius={edgeUpdaterRadius}
+          defaultMarkerColor={defaultMarkerColor}
+          noPanClassName={noPanClassName}
+        />
+        <NodeRenderer
+          nodeTypes={nodeTypes}
+          onNodeClick={onNodeClick}
+          onNodeDoubleClick={onNodeDoubleClick}
+          onNodeMouseEnter={onNodeMouseEnter}
+          onNodeMouseMove={onNodeMouseMove}
+          onNodeMouseLeave={onNodeMouseLeave}
+          onNodeContextMenu={onNodeContextMenu}
+          onNodeDragStop={onNodeDragStop}
+          onNodeDrag={onNodeDrag}
+          onNodeDragStart={onNodeDragStart}
+          selectNodesOnDrag={selectNodesOnDrag}
+          onlyRenderVisibleElements={onlyRenderVisibleElements}
+          noPanClassName={noPanClassName}
+          noDragClassName={noDragClassName}
+        />
+      </Viewport>
     </FlowRenderer>
   );
 };
diff --git a/src/container/NodeRenderer/index.tsx b/src/container/NodeRenderer/index.tsx
index 8b10d69a..a60fdb88 100644
--- a/src/container/NodeRenderer/index.tsx
+++ b/src/container/NodeRenderer/index.tsx
@@ -1,12 +1,13 @@
 import React, { memo, useMemo, ComponentType, MouseEvent } from 'react';
+import shallow from 'zustand/shallow';
+
+import { useStore } from '../../store';
+import { Node, NodeTypesType, ReactFlowState, WrapNodeProps } from '../../types';
 
-import { getNodesInside } from '../../utils/graph';
-import { useStoreState, useStoreActions } from '../../store/hooks';
-import { Node, NodeTypesType, WrapNodeProps, Edge } from '../../types';
 interface NodeRendererProps {
   nodeTypes: NodeTypesType;
   selectNodesOnDrag: boolean;
-  onElementClick?: (event: MouseEvent, element: Node | Edge) => void;
+  onNodeClick?: (event: MouseEvent, element: Node) => void;
   onNodeDoubleClick?: (event: MouseEvent, element: Node) => void;
   onNodeMouseEnter?: (event: MouseEvent, node: Node) => void;
   onNodeMouseMove?: (event: MouseEvent, node: Node) => void;
@@ -15,32 +16,33 @@ interface NodeRendererProps {
   onNodeDragStart?: (event: MouseEvent, node: Node) => void;
   onNodeDrag?: (event: MouseEvent, node: Node) => void;
   onNodeDragStop?: (event: MouseEvent, node: Node) => void;
-  snapToGrid: boolean;
-  snapGrid: [number, number];
   onlyRenderVisibleElements: boolean;
+  noPanClassName: string;
+  noDragClassName: string;
 }
 
+const selector = (s: ReactFlowState) => ({
+  scale: s.transform[2],
+  nodesDraggable: s.nodesDraggable,
+  nodesConnectable: s.nodesConnectable,
+  elementsSelectable: s.elementsSelectable,
+  updateNodeDimensions: s.updateNodeDimensions,
+  snapGrid: s.snapGrid,
+  snapToGrid: s.snapToGrid,
+  nodeInternals: s.nodeInternals,
+});
+
 const NodeRenderer = (props: NodeRendererProps) => {
-  const transform = useStoreState((state) => state.transform);
-  const selectedElements = useStoreState((state) => state.selectedElements);
-  const nodesDraggable = useStoreState((state) => state.nodesDraggable);
-  const nodesConnectable = useStoreState((state) => state.nodesConnectable);
-  const elementsSelectable = useStoreState((state) => state.elementsSelectable);
-  const width = useStoreState((state) => state.width);
-  const height = useStoreState((state) => state.height);
-  const nodes = useStoreState((state) => state.nodes);
-  const updateNodeDimensions = useStoreActions((actions) => actions.updateNodeDimensions);
-
-  const visibleNodes = props.onlyRenderVisibleElements
-    ? getNodesInside(nodes, { x: 0, y: 0, width, height }, transform, true)
-    : nodes;
-
-  const transformStyle = useMemo(
-    () => ({
-      transform: `translate(${transform[0]}px,${transform[1]}px) scale(${transform[2]})`,
-    }),
-    [transform[0], transform[1], transform[2]]
-  );
+  const {
+    scale,
+    nodesDraggable,
+    nodesConnectable,
+    elementsSelectable,
+    updateNodeDimensions,
+    snapGrid,
+    snapToGrid,
+    nodeInternals,
+  } = useStore(selector, shallow);
 
   const resizeObserver = useMemo(() => {
     if (typeof ResizeObserver === 'undefined') {
@@ -58,18 +60,24 @@ const NodeRenderer = (props: NodeRendererProps) => {
   }, []);
 
   return (
-    <div className="react-flow__nodes" style={transformStyle}>
-      {visibleNodes.map((node) => {
+    <div className="react-flow__nodes react-flow__container">
+      {Array.from(nodeInternals).map(([_, node]) => {
         const nodeType = node.type || 'default';
-        const NodeComponent = (props.nodeTypes[nodeType] || props.nodeTypes.default) as ComponentType<WrapNodeProps>;
+        const internals = nodeInternals.get(node.id);
 
         if (!props.nodeTypes[nodeType]) {
           console.warn(`Node type "${nodeType}" not found. Using fallback type "default".`);
         }
 
+        const NodeComponent = (props.nodeTypes[nodeType] || props.nodeTypes.default) as ComponentType<WrapNodeProps>;
         const isDraggable = !!(node.draggable || (nodesDraggable && typeof node.draggable === 'undefined'));
         const isSelectable = !!(node.selectable || (elementsSelectable && typeof node.selectable === 'undefined'));
         const isConnectable = !!(node.connectable || (nodesConnectable && typeof node.connectable === 'undefined'));
+        const isInitialized =
+          node.width !== null &&
+          node.height !== null &&
+          typeof node.width !== 'undefined' &&
+          typeof node.height !== 'undefined';
 
         return (
           <NodeComponent
@@ -81,15 +89,15 @@ const NodeRenderer = (props: NodeRendererProps) => {
             data={node.data}
             sourcePosition={node.sourcePosition}
             targetPosition={node.targetPosition}
-            isHidden={node.isHidden}
-            xPos={node.__rf.position.x}
-            yPos={node.__rf.position.y}
-            isDragging={node.__rf.isDragging}
-            isInitialized={node.__rf.width !== null && node.__rf.height !== null}
-            snapGrid={props.snapGrid}
-            snapToGrid={props.snapToGrid}
+            hidden={node.hidden}
+            xPos={internals?.positionAbsolute?.x || 0}
+            yPos={internals?.positionAbsolute?.y || 0}
+            dragging={!!node.dragging}
+            isInitialized={isInitialized}
+            snapGrid={snapGrid}
+            snapToGrid={snapToGrid}
             selectNodesOnDrag={props.selectNodesOnDrag}
-            onClick={props.onElementClick}
+            onClick={props.onNodeClick}
             onMouseEnter={props.onNodeMouseEnter}
             onMouseMove={props.onNodeMouseMove}
             onMouseLeave={props.onNodeMouseLeave}
@@ -98,13 +106,17 @@ const NodeRenderer = (props: NodeRendererProps) => {
             onNodeDragStart={props.onNodeDragStart}
             onNodeDrag={props.onNodeDrag}
             onNodeDragStop={props.onNodeDragStop}
-            scale={transform[2]}
-            selected={selectedElements?.some(({ id }) => id === node.id) || false}
+            scale={scale}
+            selected={!!node.selected}
             isDraggable={isDraggable}
             isSelectable={isSelectable}
             isConnectable={isConnectable}
             resizeObserver={resizeObserver}
             dragHandle={node.dragHandle}
+            zIndex={internals?.z || 0}
+            isParent={!!internals?.isParent}
+            noDragClassName={props.noDragClassName}
+            noPanClassName={props.noPanClassName}
           />
         );
       })}
diff --git a/src/container/NodeRenderer/utils.ts b/src/container/NodeRenderer/utils.ts
index 9af887a9..ede9dc72 100644
--- a/src/container/NodeRenderer/utils.ts
+++ b/src/container/NodeRenderer/utils.ts
@@ -4,28 +4,20 @@ import DefaultNode from '../../components/Nodes/DefaultNode';
 import InputNode from '../../components/Nodes/InputNode';
 import OutputNode from '../../components/Nodes/OutputNode';
 import wrapNode from '../../components/Nodes/wrapNode';
-import { NodeTypesType, NodeComponentProps } from '../../types';
+import { NodeTypesType, NodeProps } from '../../types';
 
 export function createNodeTypes(nodeTypes: NodeTypesType): NodeTypesType {
   const standardTypes: NodeTypesType = {
-    input: wrapNode((nodeTypes.input || InputNode) as ComponentType<
-      NodeComponentProps
-    >),
-    default: wrapNode((nodeTypes.default || DefaultNode) as ComponentType<
-      NodeComponentProps
-    >),
-    output: wrapNode((nodeTypes.output || OutputNode) as ComponentType<
-      NodeComponentProps
-    >),
+    input: wrapNode((nodeTypes.input || InputNode) as ComponentType<NodeProps>),
+    default: wrapNode((nodeTypes.default || DefaultNode) as ComponentType<NodeProps>),
+    output: wrapNode((nodeTypes.output || OutputNode) as ComponentType<NodeProps>),
   };
 
   const wrappedTypes = {} as NodeTypesType;
   const specialTypes: NodeTypesType = Object.keys(nodeTypes)
-    .filter(k => !['input', 'default', 'output'].includes(k))
+    .filter((k) => !['input', 'default', 'output'].includes(k))
     .reduce((res, key) => {
-      res[key] = wrapNode((nodeTypes[key] || DefaultNode) as ComponentType<
-        NodeComponentProps
-      >);
+      res[key] = wrapNode((nodeTypes[key] || DefaultNode) as ComponentType<NodeProps>);
 
       return res;
     }, wrappedTypes);
diff --git a/src/container/ReactFlow/Wrapper.tsx b/src/container/ReactFlow/Wrapper.tsx
index 4a41a47c..43aac450 100644
--- a/src/container/ReactFlow/Wrapper.tsx
+++ b/src/container/ReactFlow/Wrapper.tsx
@@ -1,21 +1,23 @@
-import React, { FC, useContext, useMemo } from 'react';
-import { Provider, ReactReduxContext } from 'react-redux';
+import React, { FC } from 'react';
 
-import store from '../../store';
+import { Provider, createStore, useStoreApi } from '../../store';
 
 const Wrapper: FC = ({ children }) => {
-  const contextValue = useContext(ReactReduxContext);
-  const isWrappedWithReactFlowProvider = useMemo(() => contextValue?.store?.getState()?.reactFlowVersion, [
-    contextValue,
-  ]);
+  let isWrapped = true;
 
-  if (isWrappedWithReactFlowProvider) {
+  try {
+    useStoreApi();
+  } catch (e) {
+    isWrapped = false;
+  }
+
+  if (isWrapped) {
     // we need to wrap it with a fragment because it's not allowed for children to be a ReactNode
     // https://github.com/DefinitelyTyped/DefinitelyTyped/issues/18051
     return <>{children}</>;
   }
 
-  return <Provider store={store}>{children}</Provider>;
+  return <Provider createStore={createStore}>{children}</Provider>;
 };
 
 Wrapper.displayName = 'ReactFlowWrapper';
diff --git a/src/container/ReactFlow/index.tsx b/src/container/ReactFlow/index.tsx
index 9b29e4c8..87c3bbfb 100644
--- a/src/container/ReactFlow/index.tsx
+++ b/src/container/ReactFlow/index.tsx
@@ -5,11 +5,12 @@ import React, {
   MouseEvent as ReactMouseEvent,
   WheelEvent,
   forwardRef,
+  FunctionComponent,
 } from 'react';
 import cc from 'classcat';
 
 import GraphView from '../GraphView';
-import ElementUpdater from '../../components/ElementUpdater';
+import StoreUpdater from '../../components/StoreUpdater';
 import DefaultNode from '../../components/Nodes/DefaultNode';
 import InputNode from '../../components/Nodes/InputNode';
 import OutputNode from '../../components/Nodes/OutputNode';
@@ -19,25 +20,26 @@ import { BezierEdge, StepEdge, SmoothStepEdge, StraightEdge } from '../../compon
 import { createEdgeTypes } from '../EdgeRenderer/utils';
 import Wrapper from './Wrapper';
 import {
-  Elements,
+  OnSelectionChangeFunc,
   NodeTypesType,
   EdgeTypesType,
-  OnLoadFunc,
+  OnLoad,
   Node,
   Edge,
-  Connection,
   ConnectionMode,
   ConnectionLineType,
   ConnectionLineComponent,
   FlowTransform,
-  OnConnectStartFunc,
-  OnConnectStopFunc,
-  OnConnectEndFunc,
-  TranslateExtent,
+  OnConnectStart,
+  OnConnectStop,
+  OnConnectEnd,
+  OnConnect,
+  CoordinateExtent,
   KeyCode,
   PanOnScrollMode,
   OnEdgeUpdateFunc,
-  NodeExtent,
+  NodeChange,
+  EdgeChange,
 } from '../../types';
 
 import '../../style.css';
@@ -57,9 +59,12 @@ const defaultEdgeTypes = {
 };
 
 export interface ReactFlowProps extends Omit<HTMLAttributes<HTMLDivElement>, 'onLoad'> {
-  elements: Elements;
-  onElementClick?: (event: ReactMouseEvent, element: Node | Edge) => void;
-  onElementsRemove?: (elements: Elements) => void;
+  nodes: Node[];
+  edges: Edge[];
+  onNodesChange?: (nodeChanges: NodeChange[]) => void;
+  onEdgesChange?: (edgeChanges: EdgeChange[]) => void;
+  onNodeClick?: (event: React.MouseEvent, node: Node) => void;
+  onEdgeClick?: (event: React.MouseEvent, node: Edge) => void;
   onNodeDoubleClick?: (event: ReactMouseEvent, node: Node) => void;
   onNodeMouseEnter?: (event: ReactMouseEvent, node: Node) => void;
   onNodeMouseMove?: (event: ReactMouseEvent, node: Node) => void;
@@ -68,15 +73,15 @@ export interface ReactFlowProps extends Omit<HTMLAttributes<HTMLDivElement>, 'on
   onNodeDragStart?: (event: ReactMouseEvent, node: Node) => void;
   onNodeDrag?: (event: ReactMouseEvent, node: Node) => void;
   onNodeDragStop?: (event: ReactMouseEvent, node: Node) => void;
-  onConnect?: (connection: Edge | Connection) => void;
-  onConnectStart?: OnConnectStartFunc;
-  onConnectStop?: OnConnectStopFunc;
-  onConnectEnd?: OnConnectEndFunc;
-  onLoad?: OnLoadFunc;
+  onConnect?: OnConnect;
+  onConnectStart?: OnConnectStart;
+  onConnectStop?: OnConnectStop;
+  onConnectEnd?: OnConnectEnd;
+  onLoad?: OnLoad;
   onMove?: (flowTransform?: FlowTransform) => void;
   onMoveStart?: (flowTransform?: FlowTransform) => void;
   onMoveEnd?: (flowTransform?: FlowTransform) => void;
-  onSelectionChange?: (elements: Elements | null) => void;
+  onSelectionChange?: OnSelectionChangeFunc;
   onSelectionDragStart?: (event: ReactMouseEvent, nodes: Node[]) => void;
   onSelectionDrag?: (event: ReactMouseEvent, nodes: Node[]) => void;
   onSelectionDragStop?: (event: ReactMouseEvent, nodes: Node[]) => void;
@@ -106,9 +111,9 @@ export interface ReactFlowProps extends Omit<HTMLAttributes<HTMLDivElement>, 'on
   maxZoom?: number;
   defaultZoom?: number;
   defaultPosition?: [number, number];
-  translateExtent?: TranslateExtent;
+  translateExtent?: CoordinateExtent;
   preventScrolling?: boolean;
-  nodeExtent?: NodeExtent;
+  nodeExtent?: CoordinateExtent;
   defaultMarkerColor?: string;
   zoomOnScroll?: boolean;
   zoomOnPinch?: boolean;
@@ -125,25 +130,30 @@ export interface ReactFlowProps extends Omit<HTMLAttributes<HTMLDivElement>, 'on
   onEdgeUpdateStart?: (event: ReactMouseEvent, edge: Edge) => void;
   onEdgeUpdateEnd?: (event: MouseEvent, edge: Edge) => void;
   edgeUpdaterRadius?: number;
-  nodeTypesId?: string;
-  edgeTypesId?: string;
+  noDragClassName?: string;
+  noWheelClassName?: string;
+  noPanClassName?: string;
 }
 
 export type ReactFlowRefType = HTMLDivElement;
 
-const ReactFlow = forwardRef<ReactFlowRefType, ReactFlowProps>(
+const initSnapGrid: [number, number] = [15, 15];
+const initDefaultPosition: [number, number] = [0, 0];
+
+const ReactFlow: FunctionComponent<ReactFlowProps> = forwardRef<ReactFlowRefType, ReactFlowProps>(
   (
     {
-      elements = [],
+      nodes = [],
+      edges = [],
       className,
       nodeTypes = defaultNodeTypes,
       edgeTypes = defaultEdgeTypes,
-      onElementClick,
+      onNodeClick,
+      onEdgeClick,
       onLoad,
       onMove,
       onMoveStart,
       onMoveEnd,
-      onElementsRemove,
       onConnect,
       onConnectStart,
       onConnectStop,
@@ -170,7 +180,7 @@ const ReactFlow = forwardRef<ReactFlowRefType, ReactFlowProps>(
       multiSelectionKeyCode = 'Meta',
       zoomActivationKeyCode = 'Meta',
       snapToGrid = false,
-      snapGrid = [15, 15],
+      snapGrid = initSnapGrid,
       onlyRenderVisibleElements = false,
       selectNodesOnDrag = true,
       nodesDraggable,
@@ -179,7 +189,7 @@ const ReactFlow = forwardRef<ReactFlowRefType, ReactFlowProps>(
       minZoom,
       maxZoom,
       defaultZoom = 1,
-      defaultPosition = [0, 0],
+      defaultPosition = initDefaultPosition,
       translateExtent,
       preventScrolling = true,
       nodeExtent,
@@ -204,14 +214,17 @@ const ReactFlow = forwardRef<ReactFlowRefType, ReactFlowProps>(
       onEdgeUpdateStart,
       onEdgeUpdateEnd,
       edgeUpdaterRadius = 10,
-      nodeTypesId = '1',
-      edgeTypesId = '1',
+      onNodesChange,
+      onEdgesChange,
+      noDragClassName = 'nodrag',
+      noWheelClassName = 'nowheel',
+      noPanClassName = 'nopan',
       ...rest
     },
     ref
   ) => {
-    const nodeTypesParsed = useMemo(() => createNodeTypes(nodeTypes), [nodeTypesId]);
-    const edgeTypesParsed = useMemo(() => createEdgeTypes(edgeTypes), [edgeTypesId]);
+    const nodeTypesParsed = useMemo(() => createNodeTypes(nodeTypes), [nodeTypes]);
+    const edgeTypesParsed = useMemo(() => createEdgeTypes(edgeTypes), [edgeTypes]);
     const reactFlowClasses = cc(['react-flow', className]);
 
     return (
@@ -222,7 +235,8 @@ const ReactFlow = forwardRef<ReactFlowRefType, ReactFlowProps>(
             onMove={onMove}
             onMoveStart={onMoveStart}
             onMoveEnd={onMoveEnd}
-            onElementClick={onElementClick}
+            onNodeClick={onNodeClick}
+            onEdgeClick={onEdgeClick}
             onNodeMouseEnter={onNodeMouseEnter}
             onNodeMouseMove={onNodeMouseMove}
             onNodeMouseLeave={onNodeMouseLeave}
@@ -233,34 +247,18 @@ const ReactFlow = forwardRef<ReactFlowRefType, ReactFlowProps>(
             onNodeDragStop={onNodeDragStop}
             nodeTypes={nodeTypesParsed}
             edgeTypes={edgeTypesParsed}
-            connectionMode={connectionMode}
             connectionLineType={connectionLineType}
             connectionLineStyle={connectionLineStyle}
             connectionLineComponent={connectionLineComponent}
             selectionKeyCode={selectionKeyCode}
-            onElementsRemove={onElementsRemove}
             deleteKeyCode={deleteKeyCode}
             multiSelectionKeyCode={multiSelectionKeyCode}
             zoomActivationKeyCode={zoomActivationKeyCode}
-            onConnect={onConnect}
-            onConnectStart={onConnectStart}
-            onConnectStop={onConnectStop}
-            onConnectEnd={onConnectEnd}
-            snapToGrid={snapToGrid}
-            snapGrid={snapGrid}
             onlyRenderVisibleElements={onlyRenderVisibleElements}
-            nodesDraggable={nodesDraggable}
-            nodesConnectable={nodesConnectable}
-            elementsSelectable={elementsSelectable}
             selectNodesOnDrag={selectNodesOnDrag}
-            minZoom={minZoom}
-            maxZoom={maxZoom}
             defaultZoom={defaultZoom}
             defaultPosition={defaultPosition}
-            translateExtent={translateExtent}
             preventScrolling={preventScrolling}
-            nodeExtent={nodeExtent}
-            defaultMarkerColor={defaultMarkerColor}
             zoomOnScroll={zoomOnScroll}
             zoomOnPinch={zoomOnPinch}
             zoomOnDoubleClick={zoomOnDoubleClick}
@@ -284,8 +282,32 @@ const ReactFlow = forwardRef<ReactFlowRefType, ReactFlowProps>(
             onEdgeUpdateStart={onEdgeUpdateStart}
             onEdgeUpdateEnd={onEdgeUpdateEnd}
             edgeUpdaterRadius={edgeUpdaterRadius}
+            defaultMarkerColor={defaultMarkerColor}
+            noDragClassName={noDragClassName}
+            noWheelClassName={noWheelClassName}
+            noPanClassName={noPanClassName}
           />
-          <ElementUpdater elements={elements} />
+          <StoreUpdater
+            nodes={nodes}
+            edges={edges}
+            onConnect={onConnect}
+            onConnectStart={onConnectStart}
+            onConnectStop={onConnectStop}
+            onConnectEnd={onConnectEnd}
+            nodesDraggable={nodesDraggable}
+            nodesConnectable={nodesConnectable}
+            elementsSelectable={elementsSelectable}
+            minZoom={minZoom}
+            maxZoom={maxZoom}
+            nodeExtent={nodeExtent}
+            onNodesChange={onNodesChange}
+            onEdgesChange={onEdgesChange}
+            snapToGrid={snapToGrid}
+            snapGrid={snapGrid}
+            connectionMode={connectionMode}
+            translateExtent={translateExtent}
+          />
+
           {onSelectionChange && <SelectionListener onSelectionChange={onSelectionChange} />}
           {children}
         </Wrapper>
diff --git a/src/container/Viewport/index.tsx b/src/container/Viewport/index.tsx
new file mode 100644
index 00000000..590b7516
--- /dev/null
+++ b/src/container/Viewport/index.tsx
@@ -0,0 +1,25 @@
+import React, { ReactNode } from 'react';
+
+import { useStore } from '../../store';
+import { ReactFlowState } from '../../types';
+
+const selector = (s: ReactFlowState) => s.transform;
+
+type ViewportProps = {
+  children: ReactNode;
+};
+
+function Viewport({ children }: ViewportProps) {
+  const transform = useStore(selector);
+
+  return (
+    <div
+      className="react-flow__viewport react-flow__container"
+      style={{ transform: `translate(${transform[0]}px,${transform[1]}px) scale(${transform[2]})` }}
+    >
+      {children}
+    </div>
+  );
+}
+
+export default Viewport;
diff --git a/src/container/ZoomPane/index.tsx b/src/container/ZoomPane/index.tsx
index 5af2659b..d19207c1 100644
--- a/src/container/ZoomPane/index.tsx
+++ b/src/container/ZoomPane/index.tsx
@@ -1,12 +1,13 @@
 import React, { useEffect, useRef, ReactNode } from 'react';
 import { zoom, zoomIdentity } from 'd3-zoom';
 import { select, pointer } from 'd3-selection';
+import shallow from 'zustand/shallow';
 
 import { clamp } from '../../utils';
 import useKeyPress from '../../hooks/useKeyPress';
 import useResizeHandler from '../../hooks/useResizeHandler';
-import { useStoreState, useStoreActions, useStore } from '../../store/hooks';
-import { FlowTransform, TranslateExtent, PanOnScrollMode, KeyCode } from '../../types';
+import { useStore, useStoreApi } from '../../store';
+import { FlowTransform, PanOnScrollMode, KeyCode, ReactFlowState } from '../../types';
 
 interface ZoomPaneProps {
   selectionKeyPressed: boolean;
@@ -20,13 +21,14 @@ interface ZoomPaneProps {
   paneMoveable?: boolean;
   defaultPosition?: [number, number];
   defaultZoom?: number;
-  translateExtent?: TranslateExtent;
   onMove?: (flowTransform?: FlowTransform) => void;
   onMoveStart?: (flowTransform?: FlowTransform) => void;
   onMoveEnd?: (flowTransform?: FlowTransform) => void;
   zoomActivationKeyCode?: KeyCode;
   preventScrolling?: boolean;
   children: ReactNode;
+  noWheelClassName: string;
+  noPanClassName: string;
 }
 
 const viewChanged = (prevTransform: FlowTransform, eventTransform: any): boolean =>
@@ -40,7 +42,16 @@ const eventToFlowTransform = (eventTransform: any): FlowTransform => ({
   zoom: eventTransform.k,
 });
 
-const hasNoWheelClass = (event: any) => event.target.closest('.nowheel');
+const isWrappedWithClass = (event: any, className: string | undefined) => event.target.closest(`.${className}`);
+
+const selector = (s: ReactFlowState) => ({
+  d3Zoom: s.d3Zoom,
+  d3Selection: s.d3Selection,
+  d3ZoomHandler: s.d3ZoomHandler,
+
+  initD3Zoom: s.initD3Zoom,
+  updateTransform: s.updateTransform,
+});
 
 const ZoomPane = ({
   onMove,
@@ -57,36 +68,29 @@ const ZoomPane = ({
   paneMoveable = true,
   defaultPosition = [0, 0],
   defaultZoom = 1,
-  translateExtent,
   zoomActivationKeyCode,
   preventScrolling = true,
   children,
+  noWheelClassName,
+  noPanClassName,
 }: ZoomPaneProps) => {
+  const store = useStoreApi();
   const zoomPane = useRef<HTMLDivElement>(null);
   const prevTransform = useRef<FlowTransform>({ x: 0, y: 0, zoom: 0 });
-
-  const store = useStore();
-  const d3Zoom = useStoreState((s) => s.d3Zoom);
-  const d3Selection = useStoreState((s) => s.d3Selection);
-  const d3ZoomHandler = useStoreState((s) => s.d3ZoomHandler);
-
-  const initD3Zoom = useStoreActions((actions) => actions.initD3Zoom);
-  const updateTransform = useStoreActions((actions) => actions.updateTransform);
-
+  const { d3Zoom, d3Selection, d3ZoomHandler, initD3Zoom, updateTransform } = useStore(selector, shallow);
   const zoomActivationKeyPressed = useKeyPress(zoomActivationKeyCode);
 
   useResizeHandler(zoomPane);
 
   useEffect(() => {
     if (zoomPane.current) {
-      const state = store.getState();
-      const currentTranslateExtent = typeof translateExtent !== 'undefined' ? translateExtent : state.translateExtent;
-      const d3ZoomInstance = zoom().scaleExtent([state.minZoom, state.maxZoom]).translateExtent(currentTranslateExtent);
+      const { minZoom, maxZoom, translateExtent } = store.getState();
+      const d3ZoomInstance = zoom().scaleExtent([minZoom, maxZoom]).translateExtent(translateExtent);
       const selection = select(zoomPane.current as Element).call(d3ZoomInstance);
 
-      const clampedX = clamp(defaultPosition[0], currentTranslateExtent[0][0], currentTranslateExtent[1][0]);
-      const clampedY = clamp(defaultPosition[1], currentTranslateExtent[0][1], currentTranslateExtent[1][1]);
-      const clampedZoom = clamp(defaultZoom, state.minZoom, state.maxZoom);
+      const clampedX = clamp(defaultPosition[0], translateExtent[0][0], translateExtent[1][0]);
+      const clampedY = clamp(defaultPosition[1], translateExtent[0][1], translateExtent[1][1]);
+      const clampedZoom = clamp(defaultZoom, minZoom, maxZoom);
       const updatedTransform = zoomIdentity.translate(clampedX, clampedY).scale(clampedZoom);
 
       d3ZoomInstance.transform(selection, updatedTransform);
@@ -106,7 +110,7 @@ const ZoomPane = ({
       if (panOnScroll && !zoomActivationKeyPressed) {
         d3Selection
           .on('wheel', (event: any) => {
-            if (hasNoWheelClass(event)) {
+            if (isWrappedWithClass(event, noWheelClassName)) {
               return false;
             }
             event.preventDefault();
@@ -140,7 +144,7 @@ const ZoomPane = ({
       } else if (typeof d3ZoomHandler !== 'undefined') {
         d3Selection
           .on('wheel', (event: any) => {
-            if (!preventScrolling || hasNoWheelClass(event)) {
+            if (!preventScrolling || isWrappedWithClass(event, noWheelClassName)) {
               return null;
             }
 
@@ -158,6 +162,7 @@ const ZoomPane = ({
     zoomActivationKeyPressed,
     zoomOnPinch,
     preventScrolling,
+    noWheelClassName,
   ]);
 
   useEffect(() => {
@@ -230,22 +235,28 @@ const ZoomPane = ({
           return false;
         }
 
-        if (hasNoWheelClass(event) && event.type === 'wheel') {
+        // if the target element is inside the nowheel class, we prevent zooming
+        if (isWrappedWithClass(event, noWheelClassName) && event.type === 'wheel') {
+          return false;
+        }
+
+        // if the target element is inside the nopan class, we prevent panning
+        if (isWrappedWithClass(event, noPanClassName) && event.type !== 'wheel') {
           return false;
         }
 
         // when the target element is a node, we still allow zooming
-        if (
-          (event.target.closest('.react-flow__node') || event.target.closest('.react-flow__edge')) &&
-          event.type !== 'wheel'
-        ) {
-          return false;
-        }
+        // if (
+        //   (event.target.closest('.react-flow__node') || event.target.closest('.react-flow__edge')) &&
+        //   event.type !== 'wheel'
+        // ) {
+        //   return false;
+        // }
 
-        // when the target element is a node selection, we still allow zooming
-        if (event.target.closest('.react-flow__nodesselection') && event.type !== 'wheel') {
-          return false;
-        }
+        // // when the target element is a node selection, we still allow zooming
+        // if (event.target.closest('.react-flow__nodesselection') && event.type !== 'wheel') {
+        //   return false;
+        // }
 
         if (!zoomOnPinch && event.ctrlKey && event.type === 'wheel') {
           return false;
@@ -278,7 +289,7 @@ const ZoomPane = ({
   ]);
 
   return (
-    <div className="react-flow__renderer react-flow__zoompane" ref={zoomPane}>
+    <div className="react-flow__renderer react-flow__container" ref={zoomPane}>
       {children}
     </div>
   );
diff --git a/src/contexts/NodeIdContext.ts b/src/contexts/NodeIdContext.ts
index cc19227f..ba112d70 100644
--- a/src/contexts/NodeIdContext.ts
+++ b/src/contexts/NodeIdContext.ts
@@ -1,8 +1,6 @@
 import { createContext } from 'react';
 
-import { ElementId } from '../types';
-
-type ContextProps = ElementId | null;
+type ContextProps = string | null;
 
 export const NodeIdContext = createContext<Partial<ContextProps>>(null);
 export const Provider = NodeIdContext.Provider;
diff --git a/src/hooks/useGlobalKeyHandler.ts b/src/hooks/useGlobalKeyHandler.ts
index 340620f5..ead31a62 100644
--- a/src/hooks/useGlobalKeyHandler.ts
+++ b/src/hooks/useGlobalKeyHandler.ts
@@ -1,42 +1,55 @@
 import { useEffect } from 'react';
+import shallow from 'zustand/shallow';
 
-import { useStore, useStoreActions } from '../store/hooks';
+import { useStore, useStoreApi } from '../store';
 import useKeyPress from './useKeyPress';
-import { isNode, getConnectedEdges } from '../utils/graph';
-import { Elements, KeyCode, ElementId, FlowElement } from '../types';
+import { getConnectedEdges } from '../utils/graph';
+import { EdgeChange, KeyCode, NodeChange, ReactFlowState } from '../types';
 
 interface HookParams {
   deleteKeyCode: KeyCode;
   multiSelectionKeyCode: KeyCode;
-  onElementsRemove?: (elements: Elements) => void;
 }
 
-export default ({ deleteKeyCode, multiSelectionKeyCode, onElementsRemove }: HookParams): void => {
-  const store = useStore();
+const selector = (s: ReactFlowState) => ({
+  unsetNodesSelection: s.unsetNodesSelection,
+  setMultiSelectionActive: s.setMultiSelectionActive,
+  resetSelectedElements: s.resetSelectedElements,
+  onNodesChange: s.onNodesChange,
+  onEdgesChange: s.onEdgesChange,
+});
 
-  const unsetNodesSelection = useStoreActions((actions) => actions.unsetNodesSelection);
-  const setMultiSelectionActive = useStoreActions((actions) => actions.setMultiSelectionActive);
-  const resetSelectedElements = useStoreActions((actions) => actions.resetSelectedElements);
+export default ({ deleteKeyCode, multiSelectionKeyCode }: HookParams): void => {
+  const store = useStoreApi();
+  const { unsetNodesSelection, setMultiSelectionActive, resetSelectedElements, onNodesChange, onEdgesChange } =
+    useStore(selector, shallow);
 
   const deleteKeyPressed = useKeyPress(deleteKeyCode);
   const multiSelectionKeyPressed = useKeyPress(multiSelectionKeyCode);
 
   useEffect(() => {
-    const { edges, selectedElements } = store.getState();
+    const { nodeInternals, edges } = store.getState();
+    // @TODO: work with nodeInternals instead of converting it to an array
+    const nodes = Array.from(nodeInternals).map(([_, node]) => node);
+    const selectedNodes = nodes.filter((n) => n.selected);
+    const selectedEdges = edges.filter((e) => e.selected);
 
-    if (onElementsRemove && deleteKeyPressed && selectedElements) {
-      const selectedNodes = selectedElements.filter(isNode);
+    if (deleteKeyPressed && (selectedNodes || selectedEdges)) {
       const connectedEdges = getConnectedEdges(selectedNodes, edges);
-      const elementsToRemove = [...selectedElements, ...connectedEdges].reduce(
-        (res, item) => res.set(item.id, item),
-        new Map<ElementId, FlowElement>()
-      );
 
-      onElementsRemove(Array.from(elementsToRemove.values()));
+      const nodeChanges: NodeChange[] = selectedNodes.map((n) => ({ id: n.id, type: 'remove' }));
+      const edgeChanges: EdgeChange[] = [...selectedEdges, ...connectedEdges].map((e) => ({
+        id: e.id,
+        type: 'remove',
+      }));
+
+      onNodesChange?.(nodeChanges);
+      onEdgesChange?.(edgeChanges);
+
       unsetNodesSelection();
       resetSelectedElements();
     }
-  }, [deleteKeyPressed, onElementsRemove]);
+  }, [deleteKeyPressed, onNodesChange, onEdgesChange]);
 
   useEffect(() => {
     setMultiSelectionActive(multiSelectionKeyPressed);
diff --git a/src/hooks/useKeyPress.ts b/src/hooks/useKeyPress.ts
index 3b1b94d0..4095af3f 100644
--- a/src/hooks/useKeyPress.ts
+++ b/src/hooks/useKeyPress.ts
@@ -6,14 +6,16 @@ type Keys = Array<string>;
 type PressedKeys = Set<string>;
 type KeyOrCode = 'key' | 'code';
 export interface UseKeyPressOptions {
-  target: Document | HTMLElement | ShadowRoot;
+  target: Document | HTMLElement | ShadowRoot | null;
 }
 
+const doc = typeof document !== 'undefined' ? document : null;
+
 // the keycode can be a string 'a' or an array of strings ['a', 'a+d']
 // a string means a single key 'a' or a combination when '+' is used 'a+d'
 // an array means different possibilites. Explainer: ['a', 'd+s'] here the
 // user can use the single key 'a' or the combination 'd' + 's'
-export default (keyCode: KeyCode | null = null, options: UseKeyPressOptions = { target: document }): boolean => {
+export default (keyCode: KeyCode | null = null, options: UseKeyPressOptions = { target: doc }): boolean => {
   const [keyPressed, setKeyPressed] = useState(false);
 
   // we need to remember the pressed keys in order to support combinations
@@ -64,16 +66,16 @@ export default (keyCode: KeyCode | null = null, options: UseKeyPressOptions = {
         setKeyPressed(false);
       };
 
-      options.target.addEventListener('keydown', downHandler as EventListenerOrEventListenerObject);
-      options.target.addEventListener('keyup', upHandler as EventListenerOrEventListenerObject);
-      options.target.addEventListener('blur', resetHandler);
+      options?.target?.addEventListener('keydown', downHandler as EventListenerOrEventListenerObject);
+      options?.target?.addEventListener('keyup', upHandler as EventListenerOrEventListenerObject);
+      options?.target?.addEventListener('blur', resetHandler);
 
       return () => {
         pressedKeys.current.clear();
 
-        options.target.removeEventListener('keydown', downHandler as EventListenerOrEventListenerObject);
-        options.target.removeEventListener('keyup', upHandler as EventListenerOrEventListenerObject);
-        options.target.removeEventListener('blur', resetHandler);
+        options?.target?.removeEventListener('keydown', downHandler as EventListenerOrEventListenerObject);
+        options?.target?.removeEventListener('keyup', upHandler as EventListenerOrEventListenerObject);
+        options?.target?.removeEventListener('blur', resetHandler);
       };
     }
   }, [keyCode, setKeyPressed]);
diff --git a/src/hooks/useOnLoadHandler.ts b/src/hooks/useOnLoadHandler.ts
new file mode 100644
index 00000000..d75ba9f4
--- /dev/null
+++ b/src/hooks/useOnLoadHandler.ts
@@ -0,0 +1,63 @@
+import { useEffect, useRef } from 'react';
+
+import { pointToRendererPoint } from '../utils/graph';
+import { useStoreApi } from '../store';
+import useZoomPanHelper from '../hooks/useZoomPanHelper';
+import { OnLoad, XYPosition, Node, Edge, FlowExportObject } from '../types';
+
+function useOnLoadHandler(onLoad: OnLoad<any> | undefined) {
+  const isInitialized = useRef<boolean>(false);
+  const store = useStoreApi();
+  const { zoomIn, zoomOut, zoomTo, transform: setTransform, fitView, initialized } = useZoomPanHelper();
+
+  useEffect(() => {
+    if (!isInitialized.current && initialized) {
+      if (onLoad) {
+        const project = (position: XYPosition): XYPosition => {
+          const { transform, snapToGrid, snapGrid } = store.getState();
+          return pointToRendererPoint(position, transform, snapToGrid, snapGrid);
+        };
+
+        const getNodes = (): Node[] => {
+          const { nodeInternals } = store.getState();
+          // @TODO: work with nodeInternals instead of converting it to an array
+          const nodes = Array.from(nodeInternals).map(([_, node]) => node);
+          return nodes.map((n) => ({ ...n }));
+        };
+
+        const getEdges = (): Edge[] => {
+          const { edges = [] } = store.getState();
+          return edges.map((e) => ({ ...e }));
+        };
+
+        const toObject = (): FlowExportObject => {
+          const { nodeInternals, edges = [], transform } = store.getState();
+          // @TODO: work with nodeInternals instead of converting it to an array
+          const nodes = Array.from(nodeInternals).map(([_, node]) => node);
+          return {
+            nodes: nodes.map((n) => ({ ...n })),
+            edges: edges.map((e) => ({ ...e })),
+            position: [transform[0], transform[1]],
+            zoom: transform[2],
+          };
+        };
+
+        onLoad({
+          fitView: (params = { padding: 0.1 }) => fitView(params),
+          zoomIn,
+          zoomOut,
+          zoomTo,
+          setTransform,
+          project,
+          getNodes,
+          getEdges,
+          toObject,
+        });
+      }
+
+      isInitialized.current = true;
+    }
+  }, [onLoad, zoomIn, zoomOut, zoomTo, setTransform, fitView, initialized]);
+}
+
+export default useOnLoadHandler;
diff --git a/src/hooks/useResizeHandler.ts b/src/hooks/useResizeHandler.ts
index 164c7686..fc2c830b 100644
--- a/src/hooks/useResizeHandler.ts
+++ b/src/hooks/useResizeHandler.ts
@@ -1,10 +1,13 @@
 import { useEffect, MutableRefObject } from 'react';
-import { useStoreActions } from '../store/hooks';
 
+import { useStore } from '../store';
 import { getDimensions } from '../utils';
+import { ReactFlowState } from '../types';
 
-export default (rendererNode: MutableRefObject<HTMLDivElement | null>) => {
-  const updateSize = useStoreActions((actions) => actions.updateSize);
+const updateSizeSelector = (state: ReactFlowState) => state.updateSize;
+
+function useResizeHandler(rendererNode: MutableRefObject<HTMLDivElement | null>) {
+  const updateSize = useStore(updateSizeSelector);
 
   useEffect(() => {
     let resizeObserver: ResizeObserver;
@@ -39,4 +42,6 @@ export default (rendererNode: MutableRefObject<HTMLDivElement | null>) => {
       }
     };
   }, []);
-};
+}
+
+export default useResizeHandler;
diff --git a/src/hooks/useUpdateNodeInternals.ts b/src/hooks/useUpdateNodeInternals.ts
index 7d64af8a..27da895f 100644
--- a/src/hooks/useUpdateNodeInternals.ts
+++ b/src/hooks/useUpdateNodeInternals.ts
@@ -1,13 +1,15 @@
 import { useCallback } from 'react';
 
-import { useStoreActions } from '../store/hooks';
-import { ElementId, UpdateNodeInternals } from '../types';
+import { useStore } from '../store';
+import { UpdateNodeInternals, ReactFlowState } from '../types';
+
+const updateNodeDimsSelector = (state: ReactFlowState) => state.updateNodeDimensions;
 
 function useUpdateNodeInternals(): UpdateNodeInternals {
-  const updateNodeDimensions = useStoreActions((actions) => actions.updateNodeDimensions);
+  const updateNodeDimensions = useStore(updateNodeDimsSelector);
 
-  return useCallback<UpdateNodeInternals>((id: ElementId) => {
-    const nodeElement = document.querySelector(`.react-flow__node[data-id="${id}"]`);
+  return useCallback<UpdateNodeInternals>((id: string) => {
+    const nodeElement = document.querySelector(`.react-flow__node[data-id="${id}"]`) as HTMLDivElement;
 
     if (nodeElement) {
       updateNodeDimensions([{ id, nodeElement, forceUpdate: true }]);
diff --git a/src/hooks/useVisibleEdges.ts b/src/hooks/useVisibleEdges.ts
new file mode 100644
index 00000000..5ab023d5
--- /dev/null
+++ b/src/hooks/useVisibleEdges.ts
@@ -0,0 +1,72 @@
+import { useCallback } from 'react';
+
+import { useStore } from '../store';
+import { isEdgeVisible } from '../container/EdgeRenderer/utils';
+import { ReactFlowState, NodeInternals, Edge } from '../types';
+
+function groupEdgesByZLevel(edges: Edge[], nodeInternals: NodeInternals) {
+  let maxLevel = -1;
+
+  const levelLookup = edges.reduce<Record<string, Edge[]>>((tree, edge) => {
+    const z = edge.zIndex || Math.max(nodeInternals.get(edge.source)?.z || 0, nodeInternals.get(edge.target)?.z || 0);
+    if (tree[z]) {
+      tree[z].push(edge);
+    } else {
+      tree[z] = [edge];
+    }
+
+    maxLevel = z > maxLevel ? z : maxLevel;
+
+    return tree;
+  }, {});
+
+  return Object.entries(levelLookup).map(([key, edges]) => {
+    const level = +key;
+
+    return {
+      edges,
+      level,
+      isMaxLevel: level === maxLevel,
+    };
+  });
+}
+
+function useVisibleEdges(onlyRenderVisible: boolean, nodeInternals: NodeInternals) {
+  const edges = useStore(
+    useCallback(
+      (s: ReactFlowState) => {
+        if (!onlyRenderVisible) {
+          return s.edges;
+        }
+
+        return s.edges.filter((e) => {
+          const sourceNode = nodeInternals.get(e.source);
+          const targetNode = nodeInternals.get(e.target);
+
+          return (
+            sourceNode?.width &&
+            sourceNode?.height &&
+            targetNode?.width &&
+            targetNode?.height &&
+            isEdgeVisible({
+              sourcePos: sourceNode.position || { x: 0, y: 0 },
+              targetPos: targetNode.position || { x: 0, y: 0 },
+              sourceWidth: sourceNode.width,
+              sourceHeight: sourceNode.height,
+              targetWidth: targetNode.width,
+              targetHeight: targetNode.height,
+              width: s.width,
+              height: s.height,
+              transform: s.transform,
+            })
+          );
+        });
+      },
+      [onlyRenderVisible, nodeInternals]
+    )
+  );
+
+  return groupEdgesByZLevel(edges, nodeInternals);
+}
+
+export default useVisibleEdges;
diff --git a/src/hooks/useVisibleNodes.ts b/src/hooks/useVisibleNodes.ts
new file mode 100644
index 00000000..4252db59
--- /dev/null
+++ b/src/hooks/useVisibleNodes.ts
@@ -0,0 +1,24 @@
+import { useCallback } from 'react';
+
+import { useStore } from '../store';
+import { getNodesInside } from '../utils/graph';
+import { ReactFlowState } from '../types';
+
+function useVisibleNodes(onlyRenderVisible: boolean) {
+  const nodes = useStore(
+    useCallback(
+      (s: ReactFlowState) => {
+        // @TODO: work with nodeInternals instead of converting it to an array
+        const nodes = Array.from(s.nodeInternals).map(([_, node]) => node);
+        return onlyRenderVisible
+          ? getNodesInside(nodes, { x: 0, y: 0, width: s.width, height: s.height }, s.transform, true)
+          : nodes;
+      },
+      [onlyRenderVisible]
+    )
+  );
+
+  return nodes;
+}
+
+export default useVisibleNodes;
diff --git a/src/hooks/useZoomPanHelper.ts b/src/hooks/useZoomPanHelper.ts
index f5b91403..735ef4b6 100644
--- a/src/hooks/useZoomPanHelper.ts
+++ b/src/hooks/useZoomPanHelper.ts
@@ -1,9 +1,10 @@
 import { useMemo } from 'react';
 import { zoomIdentity } from 'd3-zoom';
+import shallow from 'zustand/shallow';
 
-import { useStoreState, useStore } from '../store/hooks';
+import { useStoreApi, useStore } from '../store';
 import { getRectOfNodes, pointToRendererPoint, getTransformForBounds } from '../utils/graph';
-import { FitViewParams, FlowTransform, ZoomPanHelperFunctions, Rect, XYPosition } from '../types';
+import { FitViewParams, FlowTransform, ZoomPanHelperFunctions, ReactFlowState, Rect, XYPosition } from '../types';
 
 const DEFAULT_PADDING = 0.1;
 
@@ -19,10 +20,14 @@ const initialZoomPanHelper: ZoomPanHelperFunctions = {
   initialized: false,
 };
 
+const selector = (s: ReactFlowState) => ({
+  d3Zoom: s.d3Zoom,
+  d3Selection: s.d3Selection,
+});
+
 const useZoomPanHelper = (): ZoomPanHelperFunctions => {
-  const store = useStore();
-  const d3Zoom = useStoreState((s) => s.d3Zoom);
-  const d3Selection = useStoreState((s) => s.d3Selection);
+  const store = useStoreApi();
+  const { d3Zoom, d3Selection } = useStore(selector, shallow);
 
   const zoomPanHelperFunctions = useMemo<ZoomPanHelperFunctions>(() => {
     if (d3Selection && d3Zoom) {
@@ -36,13 +41,14 @@ const useZoomPanHelper = (): ZoomPanHelperFunctions => {
           d3Zoom.transform(d3Selection, nextTransform);
         },
         fitView: (options: FitViewParams = { padding: DEFAULT_PADDING, includeHiddenNodes: false }) => {
-          const { nodes, width, height, minZoom, maxZoom } = store.getState();
-
+          const { nodeInternals, width, height, minZoom, maxZoom } = store.getState();
+          // @TODO: work with nodeInternals instead of converting it to an array
+          const nodes = Array.from(nodeInternals).map(([_, node]) => node);
           if (!nodes.length) {
             return;
           }
 
-          const bounds = getRectOfNodes(options.includeHiddenNodes ? nodes : nodes.filter((node) => !node.isHidden));
+          const bounds = getRectOfNodes(options.includeHiddenNodes ? nodes : nodes.filter((node) => !node.hidden));
           const [x, y, zoom] = getTransformForBounds(
             bounds,
             width,
diff --git a/src/index.ts b/src/index.ts
index 589e49c8..6aa465d3 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -1,3 +1,5 @@
+// import './wdyr';
+
 import ReactFlow from './container/ReactFlow';
 
 export default ReactFlow;
@@ -13,7 +15,6 @@ export { getMarkerEnd, getCenter as getEdgeCenter } from './components/Edges/uti
 export {
   isNode,
   isEdge,
-  removeElements,
   addEdge,
   getOutgoers,
   getIncomers,
@@ -22,11 +23,13 @@ export {
   getTransformForBounds,
   getRectOfNodes,
 } from './utils/graph';
+export { applyNodeChanges, applyEdgeChanges } from './utils/changes';
+
 export { default as useZoomPanHelper } from './hooks/useZoomPanHelper';
 export { default as useUpdateNodeInternals } from './hooks/useUpdateNodeInternals';
 
 export * from './additional-components';
-export * from './store/hooks';
+export { useStore, useStoreApi } from './store';
 export * from './types';
 
 export { ReactFlowProps } from './container/ReactFlow';
diff --git a/src/store/actions.ts b/src/store/actions.ts
deleted file mode 100644
index 0f6bbc47..00000000
--- a/src/store/actions.ts
+++ /dev/null
@@ -1,163 +0,0 @@
-import { createAction } from './utils';
-
-import {
-  Elements,
-  OnConnectEndFunc,
-  OnConnectFunc,
-  OnConnectStartFunc,
-  OnConnectStopFunc,
-  NodeDimensionUpdate,
-  NodePosUpdate,
-  NodeDiffUpdate,
-  XYPosition,
-  Transform,
-  Dimensions,
-  InitD3ZoomPayload,
-  TranslateExtent,
-  SetConnectionId,
-  SnapGrid,
-  ConnectionMode,
-  NodeExtent,
-} from '../types';
-
-import * as constants from './contants';
-
-export const setOnConnect = (onConnect: OnConnectFunc) =>
-  createAction(constants.SET_ON_CONNECT, {
-    onConnect,
-  });
-
-export const setOnConnectStart = (onConnectStart: OnConnectStartFunc) =>
-  createAction(constants.SET_ON_CONNECT_START, {
-    onConnectStart,
-  });
-
-export const setOnConnectStop = (onConnectStop: OnConnectStopFunc) =>
-  createAction(constants.SET_ON_CONNECT_STOP, {
-    onConnectStop,
-  });
-
-export const setOnConnectEnd = (onConnectEnd: OnConnectEndFunc) =>
-  createAction(constants.SET_ON_CONNECT_END, {
-    onConnectEnd,
-  });
-
-export const setElements = (elements: Elements) => createAction(constants.SET_ELEMENTS, elements);
-
-export const updateNodeDimensions = (updates: NodeDimensionUpdate[]) =>
-  createAction(constants.UPDATE_NODE_DIMENSIONS, updates);
-
-export const updateNodePos = (payload: NodePosUpdate) => createAction(constants.UPDATE_NODE_POS, payload);
-
-export const updateNodePosDiff = (payload: NodeDiffUpdate) => createAction(constants.UPDATE_NODE_POS_DIFF, payload);
-
-export const setUserSelection = (mousePos: XYPosition) => createAction(constants.SET_USER_SELECTION, mousePos);
-
-export const updateUserSelection = (mousePos: XYPosition) => createAction(constants.UPDATE_USER_SELECTION, mousePos);
-
-export const unsetUserSelection = () => createAction(constants.UNSET_USER_SELECTION);
-
-export const setSelection = (selectionActive: boolean) =>
-  createAction(constants.SET_SELECTION, {
-    selectionActive,
-  });
-
-export const unsetNodesSelection = () =>
-  createAction(constants.UNSET_NODES_SELECTION, {
-    nodesSelectionActive: false,
-  });
-
-export const resetSelectedElements = () =>
-  createAction(constants.RESET_SELECTED_ELEMENTS, {
-    selectedElements: null,
-  });
-
-export const setSelectedElements = (elements: Elements) => createAction(constants.SET_SELECTED_ELEMENTS, elements);
-
-export const addSelectedElements = (elements: Elements) => createAction(constants.ADD_SELECTED_ELEMENTS, elements);
-
-export const updateTransform = (transform: Transform) => createAction(constants.UPDATE_TRANSFORM, { transform });
-
-export const updateSize = (size: Dimensions) =>
-  createAction(constants.UPDATE_SIZE, {
-    width: size.width || 500,
-    height: size.height || 500,
-  });
-
-export const initD3Zoom = (payload: InitD3ZoomPayload) => createAction(constants.INIT_D3ZOOM, payload);
-
-export const setMinZoom = (minZoom: number) => createAction(constants.SET_MINZOOM, minZoom);
-
-export const setMaxZoom = (maxZoom: number) => createAction(constants.SET_MAXZOOM, maxZoom);
-
-export const setTranslateExtent = (translateExtent: TranslateExtent) =>
-  createAction(constants.SET_TRANSLATEEXTENT, translateExtent);
-
-export const setConnectionPosition = (connectionPosition: XYPosition) =>
-  createAction(constants.SET_CONNECTION_POSITION, { connectionPosition });
-
-export const setConnectionNodeId = (payload: SetConnectionId) => createAction(constants.SET_CONNECTION_NODEID, payload);
-
-export const setSnapToGrid = (snapToGrid: boolean) => createAction(constants.SET_SNAPTOGRID, { snapToGrid });
-
-export const setSnapGrid = (snapGrid: SnapGrid) => createAction(constants.SET_SNAPGRID, { snapGrid });
-
-export const setInteractive = (isInteractive: boolean) =>
-  createAction(constants.SET_INTERACTIVE, {
-    nodesDraggable: isInteractive,
-    nodesConnectable: isInteractive,
-    elementsSelectable: isInteractive,
-  });
-
-export const setNodesDraggable = (nodesDraggable: boolean) =>
-  createAction(constants.SET_NODES_DRAGGABLE, { nodesDraggable });
-
-export const setNodesConnectable = (nodesConnectable: boolean) =>
-  createAction(constants.SET_NODES_CONNECTABLE, { nodesConnectable });
-
-export const setElementsSelectable = (elementsSelectable: boolean) =>
-  createAction(constants.SET_ELEMENTS_SELECTABLE, { elementsSelectable });
-
-export const setMultiSelectionActive = (multiSelectionActive: boolean) =>
-  createAction(constants.SET_MULTI_SELECTION_ACTIVE, { multiSelectionActive });
-
-export const setConnectionMode = (connectionMode: ConnectionMode) =>
-  createAction(constants.SET_CONNECTION_MODE, { connectionMode });
-
-export const setNodeExtent = (nodeExtent: NodeExtent) => createAction(constants.SET_NODE_EXTENT, nodeExtent);
-
-export type ReactFlowAction = ReturnType<
-  | typeof setOnConnect
-  | typeof setOnConnectStart
-  | typeof setOnConnectStop
-  | typeof setOnConnectEnd
-  | typeof setElements
-  | typeof updateNodeDimensions
-  | typeof updateNodePos
-  | typeof updateNodePosDiff
-  | typeof setUserSelection
-  | typeof updateUserSelection
-  | typeof unsetUserSelection
-  | typeof setSelection
-  | typeof unsetNodesSelection
-  | typeof resetSelectedElements
-  | typeof setSelectedElements
-  | typeof addSelectedElements
-  | typeof updateTransform
-  | typeof updateSize
-  | typeof initD3Zoom
-  | typeof setMinZoom
-  | typeof setMaxZoom
-  | typeof setTranslateExtent
-  | typeof setConnectionPosition
-  | typeof setConnectionNodeId
-  | typeof setSnapToGrid
-  | typeof setSnapGrid
-  | typeof setInteractive
-  | typeof setNodesDraggable
-  | typeof setNodesConnectable
-  | typeof setElementsSelectable
-  | typeof setMultiSelectionActive
-  | typeof setConnectionMode
-  | typeof setNodeExtent
->;
diff --git a/src/store/configure-store.ts b/src/store/configure-store.ts
deleted file mode 100644
index 6a33bc10..00000000
--- a/src/store/configure-store.ts
+++ /dev/null
@@ -1,10 +0,0 @@
-import { createStore, Store } from 'redux';
-
-import { ReactFlowState } from '../types';
-import { ReactFlowAction } from './actions';
-import reactFlowReducer from './reducer';
-
-export default function configureStore(preloadedState: ReactFlowState): Store<ReactFlowState, ReactFlowAction> {
-  const store = createStore(reactFlowReducer, preloadedState);
-  return store;
-}
diff --git a/src/store/contants.ts b/src/store/contants.ts
deleted file mode 100644
index 389432df..00000000
--- a/src/store/contants.ts
+++ /dev/null
@@ -1,33 +0,0 @@
-export const SET_ON_CONNECT = 'SET_ON_CONNECT';
-export const SET_ON_CONNECT_START = 'SET_ON_CONNECT_START';
-export const SET_ON_CONNECT_STOP = 'SET_ON_CONNECT_STOP';
-export const SET_ON_CONNECT_END = 'SET_ON_CONNECT_END';
-export const SET_ELEMENTS = 'SET_ELEMENTS';
-export const UPDATE_NODE_DIMENSIONS = 'UPDATE_NODE_DIMENSIONS';
-export const UPDATE_NODE_POS = 'UPDATE_NODE_POS';
-export const UPDATE_NODE_POS_DIFF = 'UPDATE_NODE_POS_DIFF';
-export const SET_USER_SELECTION = 'SET_USER_SELECTION';
-export const UPDATE_USER_SELECTION = 'UPDATE_USER_SELECTION';
-export const UNSET_USER_SELECTION = 'UNSET_USER_SELECTION';
-export const SET_SELECTION = 'SET_SELECTION';
-export const UNSET_NODES_SELECTION = 'UNSET_NODES_SELECTION';
-export const SET_SELECTED_ELEMENTS = 'SET_SELECTED_ELEMENTS';
-export const RESET_SELECTED_ELEMENTS = 'RESET_SELECTED_ELEMENTS';
-export const ADD_SELECTED_ELEMENTS = 'ADD_SELECTED_ELEMENTS';
-export const UPDATE_TRANSFORM = 'UPDATE_TRANSFORM';
-export const UPDATE_SIZE = 'UPDATE_SIZE';
-export const INIT_D3ZOOM = 'INIT_D3ZOOM';
-export const SET_MINZOOM = 'SET_MINZOOM';
-export const SET_MAXZOOM = 'SET_MAXZOOM';
-export const SET_TRANSLATEEXTENT = 'SET_TRANSLATEEXTENT';
-export const SET_CONNECTION_POSITION = 'SET_CONNECTION_POSITION';
-export const SET_CONNECTION_NODEID = 'SET_CONNECTION_NODEID';
-export const SET_SNAPTOGRID = 'SET_SNAPTOGRID';
-export const SET_SNAPGRID = 'SET_SNAPGRID';
-export const SET_INTERACTIVE = 'SET_INTERACTIVE';
-export const SET_NODES_DRAGGABLE = 'SET_NODES_DRAGGABLE';
-export const SET_NODES_CONNECTABLE = 'SET_NODES_CONNECTABLE';
-export const SET_ELEMENTS_SELECTABLE = 'SET_ELEMENTS_SELECTABLE';
-export const SET_MULTI_SELECTION_ACTIVE = 'SET_MULTI_SELECTION_ACTIVE';
-export const SET_CONNECTION_MODE = 'SET_CONNECTION_MODE';
-export const SET_NODE_EXTENT = 'SET_NODE_EXTENT';
diff --git a/src/store/hooks.ts b/src/store/hooks.ts
deleted file mode 100644
index 4ee7affc..00000000
--- a/src/store/hooks.ts
+++ /dev/null
@@ -1,48 +0,0 @@
-import { bindActionCreators, Store, ActionCreator, ActionCreatorsMapObject } from 'redux';
-import {
-  useStore as useStoreRedux,
-  useSelector,
-  useDispatch as reduxUseDispatch,
-  TypedUseSelectorHook,
-} from 'react-redux';
-import { useMemo } from 'react';
-
-import { ReactFlowDispatch } from './index';
-import * as actions from './actions';
-import { ReactFlowAction } from './actions';
-import { ReactFlowState } from '../types';
-
-export const useTypedSelector: TypedUseSelectorHook<ReactFlowState> = useSelector;
-
-export type ActionCreatorSelector<Action> = (acts: typeof actions) => ActionCreator<Action>;
-export type ActionMapObjectSelector<Action> = (acts: typeof actions) => ActionCreatorsMapObject<Action>;
-export type ActionSelector<Action> = (acts: typeof actions) => ActionCreatorsMapObject<Action> | ActionCreator<Action>;
-
-export function useStoreActions<Action extends ReactFlowAction>(
-  actionSelector: ActionCreatorSelector<Action>
-): ActionCreator<Action>;
-
-export function useStoreActions<Action extends ReactFlowAction>(
-  actionSelector: ActionMapObjectSelector<Action>
-): ActionCreatorsMapObject<Action>;
-
-export function useStoreActions<Action extends ReactFlowAction>(actionSelector: ActionSelector<Action>) {
-  const dispatch: ReactFlowDispatch = reduxUseDispatch();
-  const currAction = actionSelector(actions);
-
-  const action = useMemo(() => {
-    // this looks weird but required if both ActionSelector and ActionMapObjectSelector are supported
-    return typeof currAction === 'function'
-      ? bindActionCreators(currAction, dispatch)
-      : bindActionCreators(currAction, dispatch);
-  }, [dispatch, currAction]);
-
-  return action;
-}
-
-export const useStoreState = useTypedSelector;
-export const useStore = (): Store<ReactFlowState, ReactFlowAction> => {
-  const store = useStoreRedux<ReactFlowState, ReactFlowAction>();
-  return store;
-};
-export const useDispatch: ReactFlowDispatch = reduxUseDispatch;
diff --git a/src/store/index.ts b/src/store/index.ts
index 6640d9e0..54c718ad 100644
--- a/src/store/index.ts
+++ b/src/store/index.ts
@@ -1,63 +1,371 @@
-import configureStore from './configure-store';
+import create from 'zustand';
+import createContext from 'zustand/context';
 
-import { ReactFlowState, ConnectionMode } from '../types';
+import { clampPosition, getDimensions } from '../utils';
+import {
+  ReactFlowState,
+  ConnectionMode,
+  Node,
+  Edge,
+  NodeDimensionUpdate,
+  NodeDiffUpdate,
+  InitD3ZoomPayload,
+  CoordinateExtent,
+  OnConnect,
+  OnConnectStart,
+  OnConnectStop,
+  OnConnectEnd,
+  SetConnectionId,
+  SnapGrid,
+  NodeChange,
+  OnNodesChange,
+  OnEdgesChange,
+  EdgeChange,
+  NodeDimensionChange,
+  Transform,
+  Dimensions,
+  XYPosition,
+} from '../types';
+import { isNode, isEdge, getRectOfNodes, getNodesInside, getConnectedEdges } from '../utils/graph';
+import { getHandleBounds } from '../components/Nodes/utils';
+import { createNodeInternals, createNodeOrEdgeSelectionChange } from './utils';
 
-export const initialState: ReactFlowState = {
-  width: 0,
-  height: 0,
-  transform: [0, 0, 1],
-  nodes: [],
-  edges: [],
-  selectedElements: null,
-  selectedNodesBbox: { x: 0, y: 0, width: 0, height: 0 },
+const { Provider, useStore, useStoreApi } = createContext<ReactFlowState>();
 
-  d3Zoom: null,
-  d3Selection: null,
-  d3ZoomHandler: undefined,
-  minZoom: 0.5,
-  maxZoom: 2,
-  translateExtent: [
-    [Number.NEGATIVE_INFINITY, Number.NEGATIVE_INFINITY],
-    [Number.POSITIVE_INFINITY, Number.POSITIVE_INFINITY],
-  ],
+const infiniteExtent: CoordinateExtent = [
+  [Number.NEGATIVE_INFINITY, Number.NEGATIVE_INFINITY],
+  [Number.POSITIVE_INFINITY, Number.POSITIVE_INFINITY],
+];
 
-  nodeExtent: [
-    [Number.NEGATIVE_INFINITY, Number.NEGATIVE_INFINITY],
-    [Number.POSITIVE_INFINITY, Number.POSITIVE_INFINITY],
-  ],
-
-  nodesSelectionActive: false,
-  selectionActive: false,
-
-  userSelectionRect: {
-    startX: 0,
-    startY: 0,
-    x: 0,
-    y: 0,
+const createStore = () =>
+  create<ReactFlowState>((set, get) => ({
     width: 0,
     height: 0,
-    draw: false,
-  },
-  connectionNodeId: null,
-  connectionHandleId: null,
-  connectionHandleType: 'source',
-  connectionPosition: { x: 0, y: 0 },
-  connectionMode: ConnectionMode.Strict,
+    transform: [0, 0, 1],
+    edges: [],
+    onNodesChange: null,
+    onEdgesChange: null,
+    selectedNodesBbox: { x: 0, y: 0, width: 0, height: 0 },
+    d3Zoom: null,
+    d3Selection: null,
+    d3ZoomHandler: undefined,
+    minZoom: 0.5,
+    maxZoom: 2,
+    translateExtent: infiniteExtent,
+    nodeExtent: infiniteExtent,
+    nodesSelectionActive: false,
+    selectionActive: false,
+    userSelectionRect: {
+      startX: 0,
+      startY: 0,
+      x: 0,
+      y: 0,
+      width: 0,
+      height: 0,
+      draw: false,
+    },
+    connectionNodeId: null,
+    connectionHandleId: null,
+    connectionHandleType: 'source',
+    connectionPosition: { x: 0, y: 0 },
+    connectionMode: ConnectionMode.Strict,
 
-  snapGrid: [15, 15],
-  snapToGrid: false,
+    snapGrid: [15, 15],
+    snapToGrid: false,
 
-  nodesDraggable: true,
-  nodesConnectable: true,
-  elementsSelectable: true,
+    nodesDraggable: true,
+    nodesConnectable: true,
+    elementsSelectable: true,
 
-  multiSelectionActive: false,
+    multiSelectionActive: false,
 
-  reactFlowVersion: typeof __REACT_FLOW_VERSION__ !== 'undefined' ? __REACT_FLOW_VERSION__ : '-',
-};
+    reactFlowVersion: typeof __REACT_FLOW_VERSION__ !== 'undefined' ? __REACT_FLOW_VERSION__ : '-',
 
-const store = configureStore(initialState);
+    nodeInternals: new Map(),
 
-export type ReactFlowDispatch = typeof store.dispatch;
+    setNodes: (nodes: Node[]) => {
+      const nodeInternals = createNodeInternals(nodes, get().nodeInternals);
 
-export default store;
+      set({ nodeInternals });
+    },
+    setEdges: (edges: Edge[]) => {
+      set({ edges });
+    },
+    updateNodeDimensions: (updates: NodeDimensionUpdate[]) => {
+      const { onNodesChange, transform, nodeInternals } = get();
+
+      const nodesToChange: NodeChange[] = updates.reduce<NodeChange[]>((res, update) => {
+        const node = nodeInternals.get(update.id);
+
+        if (node) {
+          const dimensions = getDimensions(update.nodeElement);
+          const doUpdate = !!(
+            dimensions.width &&
+            dimensions.height &&
+            (node.width !== dimensions.width || node.height !== dimensions.height || update.forceUpdate)
+          );
+
+          if (doUpdate) {
+            const handleBounds = getHandleBounds(update.nodeElement, transform[2]);
+            nodeInternals.set(node.id, {
+              ...node,
+              handleBounds,
+              ...dimensions,
+            });
+
+            const change = {
+              id: node.id,
+              type: 'dimensions',
+              dimensions,
+            } as NodeChange;
+            res.push(change);
+          }
+        }
+
+        return res;
+      }, []);
+
+      set({ nodeInternals: new Map(nodeInternals) });
+
+      if (nodesToChange?.length > 0) {
+        onNodesChange?.(nodesToChange);
+      }
+    },
+    updateNodePosition: ({ id, diff, dragging }: NodeDiffUpdate) => {
+      const { onNodesChange, nodeExtent, nodeInternals } = get();
+
+      if (onNodesChange) {
+        const nodes = Array.from(nodeInternals);
+        const matchingNodes = nodes.filter(([_, n]) => !!(n.selected || n.id === id));
+        if (matchingNodes?.length) {
+          onNodesChange(
+            matchingNodes?.map(([_, node]) => {
+              const change: NodeDimensionChange = {
+                id: node.id,
+                type: 'dimensions',
+                dragging: !!dragging,
+              };
+
+              if (diff) {
+                let currentExtent = nodeExtent || node.extent;
+
+                if (node.extent === 'parent' && node.parentNode && node.width && node.height) {
+                  const parent = nodeInternals.get(node.parentNode);
+                  currentExtent =
+                    parent?.width && parent?.height
+                      ? [
+                          [0, 0],
+                          [parent.width - node.width, parent.height - node.height],
+                        ]
+                      : currentExtent;
+                }
+
+                change.position = currentExtent
+                  ? clampPosition(
+                      {
+                        x: node.position.x + diff.x,
+                        y: node.position.y + diff.y,
+                      },
+                      currentExtent
+                    )
+                  : { x: node.position.x + diff.x, y: node.position.y + diff.y };
+              }
+
+              return change;
+            })
+          );
+        }
+      }
+    },
+    setUserSelection: (mousePos: XYPosition) => {
+      set({
+        selectionActive: true,
+        userSelectionRect: {
+          width: 0,
+          height: 0,
+          startX: mousePos.x,
+          startY: mousePos.y,
+          x: mousePos.x,
+          y: mousePos.y,
+          draw: true,
+        },
+      });
+    },
+    updateUserSelection: (mousePos: XYPosition) => {
+      const { userSelectionRect, nodeInternals, edges, transform, onNodesChange, onEdgesChange } = get();
+      const startX = userSelectionRect.startX ?? 0;
+      const startY = userSelectionRect.startY ?? 0;
+
+      const nextUserSelectRect = {
+        ...userSelectionRect,
+        x: mousePos.x < startX ? mousePos.x : userSelectionRect.x,
+        y: mousePos.y < startY ? mousePos.y : userSelectionRect.y,
+        width: Math.abs(mousePos.x - startX),
+        height: Math.abs(mousePos.y - startY),
+      };
+
+      // @TODO: work with nodeInternals instead of converting it to an array
+      const nodes = Array.from(nodeInternals).map(([_, node]) => node);
+      const selectedNodes = getNodesInside(nodes, nextUserSelectRect, transform, false, true);
+      const selectedEdgeIds = getConnectedEdges(selectedNodes, edges).map((e) => e.id);
+      const selectedNodeIds = selectedNodes.map((n) => n.id);
+
+      onNodesChange?.(
+        nodes.map((n) => createNodeOrEdgeSelectionChange(selectedNodeIds.includes(n.id))(n)) as NodeChange[]
+      );
+      onEdgesChange?.(
+        edges.map((e) => createNodeOrEdgeSelectionChange(selectedEdgeIds.includes(e.id))(e)) as EdgeChange[]
+      );
+
+      set({
+        userSelectionRect: nextUserSelectRect,
+      });
+    },
+    unsetUserSelection: () => {
+      const { userSelectionRect, nodeInternals } = get();
+      // @TODO: work with nodeInternals instead of converting it to an array
+      const nodes = Array.from(nodeInternals).map(([_, node]) => node);
+      const selectedNodes = nodes.filter((node) => node.selected);
+
+      const stateUpdate = {
+        selectionActive: false,
+        userSelectionRect: {
+          ...userSelectionRect,
+          draw: false,
+        },
+        selectedNodesBbox: { x: 0, y: 0, width: 0, height: 0 },
+        nodesSelectionActive: false,
+      };
+
+      if (selectedNodes && selectedNodes.length > 0) {
+        const selectedNodesBbox = getRectOfNodes(selectedNodes);
+        stateUpdate.selectedNodesBbox = selectedNodesBbox;
+        stateUpdate.nodesSelectionActive = true;
+      }
+
+      set(stateUpdate);
+    },
+    addSelectedElements: (selectedElementsArr: Array<Node | Edge>) => {
+      const { multiSelectionActive, onNodesChange, onEdgesChange, nodeInternals, edges } = get();
+      // @TODO: work with nodeInternals instead of converting it to an array
+      const nodes = Array.from(nodeInternals).map(([_, node]) => node);
+      let changedNodes;
+      let changedEdges;
+
+      if (multiSelectionActive) {
+        changedNodes = selectedElementsArr.filter(isNode).map(createNodeOrEdgeSelectionChange(true));
+        changedEdges = selectedElementsArr.filter(isEdge).map(createNodeOrEdgeSelectionChange(true));
+      } else {
+        changedNodes = nodes.map((node) =>
+          createNodeOrEdgeSelectionChange(selectedElementsArr.some((e) => e.id === node.id))(node)
+        );
+        changedEdges = edges.map((edge) =>
+          createNodeOrEdgeSelectionChange(selectedElementsArr.some((e) => e.id === edge.id))(edge)
+        );
+      }
+
+      if (changedNodes.length) {
+        onNodesChange?.(changedNodes as NodeChange[]);
+      }
+
+      if (changedEdges.length) {
+        onEdgesChange?.(changedEdges as EdgeChange[]);
+      }
+    },
+    unselectNodesAndEdges: () => {
+      const { nodeInternals, edges, onNodesChange, onEdgesChange } = get();
+      // @TODO: work with nodeInternals instead of converting it to an array
+      const nodes = Array.from(nodeInternals).map(([_, node]) => node);
+
+      const nodesToUnselect = nodes.map((n) => {
+        n.selected = false;
+        return createNodeOrEdgeSelectionChange(false)(n);
+      }) as NodeChange[];
+      const edgesToUnselect = edges.map(createNodeOrEdgeSelectionChange(false)) as EdgeChange[];
+
+      if (nodesToUnselect.length) {
+        onNodesChange?.(nodesToUnselect);
+      }
+      if (edgesToUnselect.length) {
+        onEdgesChange?.(edgesToUnselect);
+      }
+    },
+    initD3Zoom: ({ d3Zoom, d3Selection, d3ZoomHandler, transform }: InitD3ZoomPayload) =>
+      set({
+        d3Zoom,
+        d3Selection,
+        d3ZoomHandler,
+        transform,
+      }),
+    setMinZoom: (minZoom: number) => {
+      const { d3Zoom, maxZoom } = get();
+      d3Zoom?.scaleExtent([minZoom, maxZoom]);
+
+      set({ minZoom });
+    },
+    setMaxZoom: (maxZoom: number) => {
+      const { d3Zoom, minZoom } = get();
+      d3Zoom?.scaleExtent([minZoom, maxZoom]);
+
+      set({ maxZoom });
+    },
+    setTranslateExtent: (translateExtent: CoordinateExtent) => {
+      const { d3Zoom } = get();
+      d3Zoom?.translateExtent(translateExtent);
+
+      set({ translateExtent });
+    },
+
+    resetSelectedElements: () => {
+      const { nodeInternals, edges, onNodesChange, onEdgesChange } = get();
+      // @TODO: work with nodeInternals instead of converting it to an array
+      const nodes = Array.from(nodeInternals).map(([_, node]) => node);
+      const nodesToUnselect = nodes.filter((e) => e.selected).map(createNodeOrEdgeSelectionChange(false));
+      const edgesToUnselect = edges.filter((e) => e.selected).map(createNodeOrEdgeSelectionChange(false));
+
+      if (nodesToUnselect.length) {
+        onNodesChange?.(nodesToUnselect as NodeChange[]);
+      }
+      if (edgesToUnselect.length) {
+        onEdgesChange?.(edgesToUnselect as EdgeChange[]);
+      }
+    },
+    setNodeExtent: (nodeExtent: CoordinateExtent) => {
+      const { nodeInternals } = get();
+
+      nodeInternals.forEach((node) => {
+        node.positionAbsolute = clampPosition(node.position, nodeExtent);
+      });
+
+      set({
+        nodeExtent,
+        nodeInternals: new Map(nodeInternals),
+      });
+    },
+    unsetNodesSelection: () => set({ nodesSelectionActive: false }),
+    updateTransform: (transform: Transform) => set({ transform }),
+    updateSize: (size: Dimensions) => set({ width: size.width || 500, height: size.height || 500 }),
+    setOnConnect: (onConnect: OnConnect) => set({ onConnect }),
+    setOnConnectStart: (onConnectStart: OnConnectStart) => set({ onConnectStart }),
+    setOnConnectStop: (onConnectStop: OnConnectStop) => set({ onConnectStop }),
+    setOnConnectEnd: (onConnectEnd: OnConnectEnd) => set({ onConnectEnd }),
+    setConnectionPosition: (connectionPosition: XYPosition) => set({ connectionPosition }),
+    setConnectionNodeId: (params: SetConnectionId) => set({ ...params }),
+    setSnapToGrid: (snapToGrid: boolean) => set({ snapToGrid }),
+    setSnapGrid: (snapGrid: SnapGrid) => set({ snapGrid }),
+    setInteractive: (isInteractive: boolean) =>
+      set({
+        nodesDraggable: isInteractive,
+        nodesConnectable: isInteractive,
+        elementsSelectable: isInteractive,
+      }),
+    setNodesDraggable: (nodesDraggable: boolean) => set({ nodesDraggable }),
+    setNodesConnectable: (nodesConnectable: boolean) => set({ nodesConnectable }),
+    setElementsSelectable: (elementsSelectable: boolean) => set({ elementsSelectable }),
+    setMultiSelectionActive: (multiSelectionActive: boolean) => set({ multiSelectionActive }),
+    setConnectionMode: (connectionMode: ConnectionMode) => set({ connectionMode }),
+    setOnNodesChange: (onNodesChange: OnNodesChange) => set({ onNodesChange }),
+    setOnEdgesChange: (onEdgesChange: OnEdgesChange) => set({ onEdgesChange }),
+  }));
+
+export { Provider, useStore, createStore, useStoreApi };
diff --git a/src/store/reducer.ts b/src/store/reducer.ts
deleted file mode 100644
index d9b9fbb4..00000000
--- a/src/store/reducer.ts
+++ /dev/null
@@ -1,341 +0,0 @@
-import isEqual from 'fast-deep-equal';
-
-import { clampPosition, getDimensions } from '../utils';
-import {
-  getNodesInside,
-  getConnectedEdges,
-  getRectOfNodes,
-  isNode,
-  isEdge,
-  parseNode,
-  parseEdge,
-} from '../utils/graph';
-import { getHandleBounds } from '../components/Nodes/utils';
-
-import { ReactFlowState, Node, XYPosition, Edge } from '../types';
-import * as constants from './contants';
-import { ReactFlowAction } from './actions';
-
-import { initialState } from './index';
-
-type NextElements = {
-  nextNodes: Node[];
-  nextEdges: Edge[];
-};
-
-export default function reactFlowReducer(state = initialState, action: ReactFlowAction): ReactFlowState {
-  switch (action.type) {
-    case constants.SET_ELEMENTS: {
-      const propElements = action.payload;
-      const nextElements: NextElements = {
-        nextNodes: [],
-        nextEdges: [],
-      };
-      const { nextNodes, nextEdges } = propElements.reduce((res, propElement): NextElements => {
-        if (isNode(propElement)) {
-          const storeNode = state.nodes.find((node) => node.id === propElement.id);
-
-          if (storeNode) {
-            const updatedNode: Node = {
-              ...storeNode,
-              ...propElement,
-            };
-
-            if (storeNode.position.x !== propElement.position.x || storeNode.position.y !== propElement.position.y) {
-              updatedNode.__rf.position = propElement.position;
-            }
-
-            if (typeof propElement.type !== 'undefined' && propElement.type !== storeNode.type) {
-              // we reset the elements dimensions here in order to force a re-calculation of the bounds.
-              // When the type of a node changes it is possible that the number or positions of handles changes too.
-              updatedNode.__rf.width = null;
-            }
-
-            res.nextNodes.push(updatedNode);
-          } else {
-            res.nextNodes.push(parseNode(propElement, state.nodeExtent));
-          }
-        } else if (isEdge(propElement)) {
-          const storeEdge = state.edges.find((se) => se.id === propElement.id);
-
-          if (storeEdge) {
-            res.nextEdges.push({
-              ...storeEdge,
-              ...propElement,
-            });
-          } else {
-            res.nextEdges.push(parseEdge(propElement));
-          }
-        }
-
-        return res;
-      }, nextElements);
-
-      return { ...state, nodes: nextNodes, edges: nextEdges };
-    }
-    case constants.UPDATE_NODE_DIMENSIONS: {
-      const updatedNodes = state.nodes.map((node) => {
-        const update = action.payload.find((u) => u.id === node.id);
-        if (update) {
-          const dimensions = getDimensions(update.nodeElement);
-          const doUpdate =
-            dimensions.width &&
-            dimensions.height &&
-            (node.__rf.width !== dimensions.width || node.__rf.height !== dimensions.height || update.forceUpdate);
-
-          if (doUpdate) {
-            const handleBounds = getHandleBounds(update.nodeElement, state.transform[2]);
-
-            return {
-              ...node,
-              __rf: {
-                ...node.__rf,
-                ...dimensions,
-                handleBounds,
-              },
-            };
-          }
-        }
-
-        return node;
-      });
-
-      return {
-        ...state,
-        nodes: updatedNodes,
-      };
-    }
-    case constants.UPDATE_NODE_POS: {
-      const { id, pos } = action.payload;
-      let position: XYPosition = pos;
-
-      if (state.snapToGrid) {
-        const [gridSizeX, gridSizeY] = state.snapGrid;
-        position = {
-          x: gridSizeX * Math.round(pos.x / gridSizeX),
-          y: gridSizeY * Math.round(pos.y / gridSizeY),
-        };
-      }
-
-      const nextNodes = state.nodes.map((node) => {
-        if (node.id === id) {
-          return {
-            ...node,
-            __rf: {
-              ...node.__rf,
-              position,
-            },
-          };
-        }
-
-        return node;
-      });
-
-      return { ...state, nodes: nextNodes };
-    }
-    case constants.UPDATE_NODE_POS_DIFF: {
-      const { id, diff, isDragging } = action.payload;
-
-      const nextNodes = state.nodes.map((node) => {
-        if (id === node.id || state.selectedElements?.find((sNode) => sNode.id === node.id)) {
-          const updatedNode = {
-            ...node,
-            __rf: {
-              ...node.__rf,
-              isDragging,
-            },
-          };
-
-          if (diff) {
-            updatedNode.__rf.position = {
-              x: node.__rf.position.x + diff.x,
-              y: node.__rf.position.y + diff.y,
-            };
-          }
-
-          return updatedNode;
-        }
-
-        return node;
-      });
-
-      return { ...state, nodes: nextNodes };
-    }
-    case constants.SET_USER_SELECTION: {
-      const mousePos = action.payload;
-
-      return {
-        ...state,
-        selectionActive: true,
-        userSelectionRect: {
-          width: 0,
-          height: 0,
-          startX: mousePos.x,
-          startY: mousePos.y,
-          x: mousePos.x,
-          y: mousePos.y,
-          draw: true,
-        },
-      };
-    }
-    case constants.UPDATE_USER_SELECTION: {
-      const mousePos = action.payload;
-      const startX = state.userSelectionRect.startX ?? 0;
-      const startY = state.userSelectionRect.startY ?? 0;
-
-      const nextUserSelectRect = {
-        ...state.userSelectionRect,
-        x: mousePos.x < startX ? mousePos.x : state.userSelectionRect.x,
-        y: mousePos.y < startY ? mousePos.y : state.userSelectionRect.y,
-        width: Math.abs(mousePos.x - startX),
-        height: Math.abs(mousePos.y - startY),
-      };
-
-      const selectedNodes = getNodesInside(state.nodes, nextUserSelectRect, state.transform, false, true);
-      const selectedEdges = getConnectedEdges(selectedNodes, state.edges);
-
-      const nextSelectedElements = [...selectedNodes, ...selectedEdges];
-      const selectedElementsChanged = !isEqual(nextSelectedElements, state.selectedElements);
-      const selectedElementsUpdate = selectedElementsChanged
-        ? {
-            selectedElements: nextSelectedElements.length > 0 ? nextSelectedElements : null,
-          }
-        : {};
-
-      return {
-        ...state,
-        ...selectedElementsUpdate,
-        userSelectionRect: nextUserSelectRect,
-      };
-    }
-    case constants.UNSET_USER_SELECTION: {
-      const selectedNodes = state.selectedElements?.filter((node) => isNode(node) && node.__rf) as Node[];
-
-      const stateUpdate = {
-        ...state,
-        selectionActive: false,
-        userSelectionRect: {
-          ...state.userSelectionRect,
-          draw: false,
-        },
-      };
-
-      if (!selectedNodes || selectedNodes.length === 0) {
-        stateUpdate.selectedElements = null;
-        stateUpdate.nodesSelectionActive = false;
-      } else {
-        const selectedNodesBbox = getRectOfNodes(selectedNodes);
-        stateUpdate.selectedNodesBbox = selectedNodesBbox;
-        stateUpdate.nodesSelectionActive = true;
-      }
-
-      return stateUpdate;
-    }
-    case constants.SET_SELECTED_ELEMENTS: {
-      const elements = action.payload;
-      const selectedElementsArr = Array.isArray(elements) ? elements : [elements];
-      const selectedElementsUpdated = !isEqual(selectedElementsArr, state.selectedElements);
-      const selectedElements = selectedElementsUpdated ? selectedElementsArr : state.selectedElements;
-
-      return {
-        ...state,
-        selectedElements,
-      };
-    }
-    case constants.ADD_SELECTED_ELEMENTS: {
-      const { multiSelectionActive, selectedElements } = state;
-      const elements = action.payload;
-      const selectedElementsArr = Array.isArray(elements) ? elements : [elements];
-
-      let nextElements = selectedElementsArr;
-
-      if (multiSelectionActive) {
-        nextElements = selectedElements ? [...selectedElements, ...selectedElementsArr] : selectedElementsArr;
-      }
-
-      const selectedElementsUpdated = !isEqual(nextElements, state.selectedElements);
-      const nextSelectedElements = selectedElementsUpdated ? nextElements : state.selectedElements;
-
-      return { ...state, selectedElements: nextSelectedElements };
-    }
-    case constants.INIT_D3ZOOM: {
-      const { d3Zoom, d3Selection, d3ZoomHandler, transform } = action.payload;
-
-      return {
-        ...state,
-        d3Zoom,
-        d3Selection,
-        d3ZoomHandler,
-        transform,
-      };
-    }
-    case constants.SET_MINZOOM: {
-      const minZoom = action.payload;
-
-      state.d3Zoom?.scaleExtent([minZoom, state.maxZoom]);
-
-      return {
-        ...state,
-        minZoom,
-      };
-    }
-
-    case constants.SET_MAXZOOM: {
-      const maxZoom = action.payload;
-
-      state.d3Zoom?.scaleExtent([state.minZoom, maxZoom]);
-
-      return {
-        ...state,
-        maxZoom,
-      };
-    }
-    case constants.SET_TRANSLATEEXTENT: {
-      const translateExtent = action.payload;
-
-      state.d3Zoom?.translateExtent(translateExtent);
-
-      return {
-        ...state,
-        translateExtent,
-      };
-    }
-    case constants.SET_NODE_EXTENT: {
-      const nodeExtent = action.payload;
-      return {
-        ...state,
-        nodeExtent,
-        nodes: state.nodes.map((node) => {
-          return {
-            ...node,
-            __rf: {
-              ...node.__rf,
-              position: clampPosition(node.__rf.position, nodeExtent),
-            },
-          };
-        }),
-      };
-    }
-    case constants.SET_ON_CONNECT:
-    case constants.SET_ON_CONNECT_START:
-    case constants.SET_ON_CONNECT_STOP:
-    case constants.SET_ON_CONNECT_END:
-    case constants.RESET_SELECTED_ELEMENTS:
-    case constants.UNSET_NODES_SELECTION:
-    case constants.UPDATE_TRANSFORM:
-    case constants.UPDATE_SIZE:
-    case constants.SET_CONNECTION_POSITION:
-    case constants.SET_CONNECTION_NODEID:
-    case constants.SET_SNAPTOGRID:
-    case constants.SET_SNAPGRID:
-    case constants.SET_INTERACTIVE:
-    case constants.SET_NODES_DRAGGABLE:
-    case constants.SET_NODES_CONNECTABLE:
-    case constants.SET_ELEMENTS_SELECTABLE:
-    case constants.SET_MULTI_SELECTION_ACTIVE:
-    case constants.SET_CONNECTION_MODE:
-      return { ...state, ...action.payload };
-    default:
-      return state;
-  }
-}
diff --git a/src/store/utils.ts b/src/store/utils.ts
index c4bfbaf3..8247df6b 100644
--- a/src/store/utils.ts
+++ b/src/store/utils.ts
@@ -1,5 +1,112 @@
-export function createAction<T extends string>(type: T): { type: T };
-export function createAction<T extends string, P extends any>(type: T, payload: P): { type: T; payload: P };
-export function createAction(type: string, payload?: any) {
-  return { type, payload };
+import { Node, Edge, NodeInternals, NodeInternalsItem, XYZPosition } from '../types';
+import { isNumeric } from '../utils';
+
+type ParentNodes = Record<string, boolean>;
+
+function calculateXYZPosition(
+  node: NodeInternalsItem,
+  nodeInternals: NodeInternals,
+  parentNodes: ParentNodes,
+  result: XYZPosition
+): XYZPosition {
+  if (!node.parentNode) {
+    return result;
+  }
+  const parentNode = nodeInternals.get(node.parentNode)!;
+
+  return calculateXYZPosition(parentNode, nodeInternals, parentNodes, {
+    x: (result.x ?? 0) + (parentNode.position?.x ?? 0),
+    y: (result.y ?? 0) + (parentNode.position?.y ?? 0),
+    z: parentNode.z > node.z ? parentNode.z : node.z,
+  });
 }
+
+// function createTree(items: Node[]): any {
+//   const rootItems = [];
+//   const lookup: Record<any, any> = {};
+
+//   for (const item of items) {
+//     const parentId = item.parentNode;
+
+//     if (!lookup[item.id]) {
+//       lookup[item.id] = { childNodes: [], z: 0 };
+//     }
+
+//     lookup[item.id] = {
+//       node: item,
+//       childNodes: lookup[item.id].childNodes,
+//       z: lookup[item.id].z,
+//     };
+
+//     const treeItem = lookup[item.id];
+
+//     if (!parentId) {
+//       rootItems.push(treeItem);
+//     } else {
+//       if (!lookup[parentId]) {
+//         lookup[parentId] = { childNodes: [], z: 0 };
+//       }
+
+//       lookup[parentId].childNodes.push({ ...treeItem, z: lookup[parentId].z + 1 });
+//     }
+//   }
+
+//   console.log(lookup);
+
+//   return rootItems;
+// }
+
+export function createNodeInternals(nodes: Node[], nodeInternals: NodeInternals): NodeInternals {
+  const nextNodeInternals = new Map<string, NodeInternalsItem>();
+  const parentNodes: ParentNodes = {};
+
+  nodes.forEach((node) => {
+    const z = isNumeric(node.zIndex) ? node.zIndex : node.dragging || node.selected ? 1000 : 0;
+    const internals: NodeInternalsItem = {
+      ...nodeInternals.get(node.id),
+      ...node,
+      positionAbsolute: {
+        x: node.position.x,
+        y: node.position.y,
+      },
+      z,
+    };
+    if (node.parentNode) {
+      internals.parentNode = node.parentNode;
+      parentNodes[node.parentNode] = true;
+    }
+    nextNodeInternals.set(node.id, internals);
+  });
+
+  nextNodeInternals.forEach((node) => {
+    if (node.parentNode && !nextNodeInternals.has(node.parentNode)) {
+      throw new Error(`Parent node ${node.parentNode} not found`);
+    }
+
+    if (node.parentNode || parentNodes[node.id]) {
+      const { x, y, z } = calculateXYZPosition(node, nextNodeInternals, parentNodes, {
+        ...node.position,
+        z: node.z,
+      });
+
+      node.positionAbsolute = {
+        x,
+        y,
+      };
+
+      node.z = z;
+
+      if (parentNodes[node.id]) {
+        node.isParent = true;
+      }
+    }
+  });
+
+  return nextNodeInternals;
+}
+
+export const createNodeOrEdgeSelectionChange = (selected: boolean) => (item: Node | Edge) => ({
+  id: item.id,
+  type: 'select',
+  selected,
+});
diff --git a/src/style.css b/src/style.css
index 327cfbef..30290001 100644
--- a/src/style.css
+++ b/src/style.css
@@ -5,12 +5,10 @@
   overflow: hidden;
 }
 
-.react-flow__renderer,
-.react-flow__pane,
-.react-flow__selectionpane {
+.react-flow__container {
+  position: absolute;
   width: 100%;
   height: 100%;
-  position: absolute;
   top: 0;
   left: 0;
 }
@@ -19,6 +17,12 @@
   z-index: 1;
 }
 
+.react-flow__viewport {
+  transform-origin: 0 0;
+  z-index: 2;
+  pointer-events: none;
+}
+
 .react-flow__renderer {
   z-index: 4;
 }
@@ -27,18 +31,9 @@
   z-index: 5;
 }
 
-.react-flow__selection {
-  position: absolute;
-  top: 0;
-  left: 0;
-}
-
 .react-flow__edges {
-  position: absolute;
-  top: 0;
-  left: 0;
   pointer-events: none;
-  z-index: 2;
+  overflow: visible;
 }
 
 .react-flow__edge {
@@ -82,12 +77,8 @@
 }
 
 .react-flow__nodes {
-  position: absolute;
-  width: 100%;
-  height: 100%;
   pointer-events: none;
   transform-origin: 0 0;
-  z-index: 3;
 }
 
 .react-flow__node {
@@ -95,15 +86,11 @@
   user-select: none;
   pointer-events: all;
   transform-origin: 0 0;
+  box-sizing: border-box;
 }
 
 .react-flow__nodesselection {
   z-index: 3;
-  position: absolute;
-  width: 100%;
-  height: 100%;
-  top: 0;
-  left: 0;
   transform-origin: left top;
   pointer-events: none;
 
@@ -154,14 +141,6 @@
 
 /* additional components */
 
-.react-flow__background {
-  position: absolute;
-  top: 0;
-  left: 0;
-  width: 100%;
-  height: 100%;
-}
-
 .react-flow__controls {
   position: absolute;
   z-index: 5;
diff --git a/src/types/changes.ts b/src/types/changes.ts
new file mode 100644
index 00000000..fd29a477
--- /dev/null
+++ b/src/types/changes.ts
@@ -0,0 +1,28 @@
+import { XYPosition, Dimensions } from './utils';
+import { NodeHandleBounds } from './nodes';
+
+export type NodeDimensionChange = {
+  id: string;
+  type: 'dimensions';
+  dimensions?: Dimensions;
+  position?: XYPosition;
+  handleBounds?: NodeHandleBounds;
+  dragging?: boolean;
+};
+
+export type NodeSelectionChange = {
+  id: string;
+  type: 'select';
+  selected: boolean;
+};
+
+export type NodeRemoveChange = {
+  id: string;
+  type: 'remove';
+};
+
+export type NodeChange = NodeDimensionChange | NodeSelectionChange | NodeRemoveChange;
+
+export type EdgeSelectionChange = NodeSelectionChange;
+export type EdgeRemoveChange = NodeRemoveChange;
+export type EdgeChange = EdgeSelectionChange | EdgeRemoveChange;
diff --git a/src/types/edges.ts b/src/types/edges.ts
new file mode 100644
index 00000000..21ea21e2
--- /dev/null
+++ b/src/types/edges.ts
@@ -0,0 +1,159 @@
+import { CSSProperties, ReactNode, HTMLAttributes } from 'react';
+
+import { Position } from './utils';
+import { OnConnect, Connection } from './general';
+import { HandleElement } from './handles';
+import { Node } from './nodes';
+
+// interface for the user edge items
+export interface Edge<T = any> {
+  id: string;
+  type?: string;
+  source: string;
+  target: string;
+  sourceHandle?: string | null;
+  targetHandle?: string | null;
+  label?: string | ReactNode;
+  labelStyle?: CSSProperties;
+  labelShowBg?: boolean;
+  labelBgStyle?: CSSProperties;
+  labelBgPadding?: [number, number];
+  labelBgBorderRadius?: number;
+  style?: CSSProperties;
+  animated?: boolean;
+  hidden?: boolean;
+  data?: T;
+  className?: string;
+  sourceNode?: Node;
+  targetNode?: Node;
+  selected?: boolean;
+  markerStart?: EdgeMarkerType;
+  markerEnd?: EdgeMarkerType;
+  zIndex?: number;
+}
+
+// props that get passed to a custom edge
+export interface EdgeProps<T = any> {
+  id: string;
+  source: string;
+  target: string;
+  sourceX: number;
+  sourceY: number;
+  targetX: number;
+  targetY: number;
+  selected?: boolean;
+  animated?: boolean;
+  sourcePosition: Position;
+  targetPosition: Position;
+  label?: string | ReactNode;
+  labelStyle?: CSSProperties;
+  labelShowBg?: boolean;
+  labelBgStyle?: CSSProperties;
+  labelBgPadding?: [number, number];
+  labelBgBorderRadius?: number;
+  style?: CSSProperties;
+  data?: T;
+  sourceHandleId?: string | null;
+  targetHandleId?: string | null;
+  markerStart?: string;
+  markerEnd?: string;
+}
+
+export type EdgeMouseHandler = (event: React.MouseEvent, edge: Edge) => void;
+
+export interface WrapEdgeProps<T = any> {
+  id: string;
+  className?: string;
+  type: string;
+  data?: T;
+  onClick?: EdgeMouseHandler;
+  onEdgeDoubleClick?: EdgeMouseHandler;
+  selected: boolean;
+  animated?: boolean;
+  label?: string | ReactNode;
+  labelStyle?: CSSProperties;
+  labelShowBg?: boolean;
+  labelBgStyle?: CSSProperties;
+  labelBgPadding?: [number, number];
+  labelBgBorderRadius?: number;
+  style?: CSSProperties;
+  source: string;
+  target: string;
+  sourceHandleId: string | null;
+  targetHandleId: string | null;
+  sourceX: number;
+  sourceY: number;
+  targetX: number;
+  targetY: number;
+  sourcePosition: Position;
+  targetPosition: Position;
+  elementsSelectable?: boolean;
+  hidden?: boolean;
+  handleEdgeUpdate: boolean;
+  onConnectEdge: OnConnect;
+  onContextMenu?: EdgeMouseHandler;
+  onMouseEnter?: EdgeMouseHandler;
+  onMouseMove?: EdgeMouseHandler;
+  onMouseLeave?: EdgeMouseHandler;
+  edgeUpdaterRadius?: number;
+  onEdgeUpdateStart?: EdgeMouseHandler;
+  onEdgeUpdateEnd?: (event: MouseEvent, edge: Edge) => void;
+  markerStart?: EdgeMarkerType;
+  markerEnd?: EdgeMarkerType;
+}
+
+export interface EdgeSmoothStepProps<T = any> extends EdgeProps<T> {
+  borderRadius?: number;
+}
+
+export interface EdgeTextProps extends HTMLAttributes<SVGElement> {
+  x: number;
+  y: number;
+  label?: string | ReactNode;
+  labelStyle?: CSSProperties;
+  labelShowBg?: boolean;
+  labelBgStyle?: CSSProperties;
+  labelBgPadding?: [number, number];
+  labelBgBorderRadius?: number;
+}
+
+export enum ConnectionLineType {
+  Bezier = 'default',
+  Straight = 'straight',
+  Step = 'step',
+  SmoothStep = 'smoothstep',
+}
+
+export type ConnectionLineComponentProps = {
+  sourceX: number;
+  sourceY: number;
+  sourcePosition?: Position;
+  targetX: number;
+  targetY: number;
+  targetPosition?: Position;
+  connectionLineStyle?: CSSProperties;
+  connectionLineType: ConnectionLineType;
+  sourceNode?: Node;
+  sourceHandle?: HandleElement;
+};
+
+export type ConnectionLineComponent = React.ComponentType<ConnectionLineComponentProps>;
+
+export type OnEdgeUpdateFunc<T = any> = (oldEdge: Edge<T>, newConnection: Connection) => void;
+
+export interface EdgeMarker {
+  type: ArrowHeadType;
+  color?: string;
+  width?: number;
+  height?: number;
+  units?: string;
+  orient?: string;
+  strokeWidth?: number;
+}
+
+export type EdgeMarkerType = string | EdgeMarker;
+
+export enum ArrowHeadType {
+  Arrow = 'arrow',
+  ArrowClosed = 'arrowclosed',
+}
diff --git a/src/types/general.ts b/src/types/general.ts
new file mode 100644
index 00000000..5cafafef
--- /dev/null
+++ b/src/types/general.ts
@@ -0,0 +1,216 @@
+import { MouseEvent as ReactMouseEvent, ReactNode } from 'react';
+import { Selection as D3Selection, ZoomBehavior } from 'd3';
+
+import { XYPosition, Rect, Transform, CoordinateExtent, Dimensions } from './utils';
+import { NodeChange, EdgeChange } from './changes';
+import { Node, NodeInternals, NodeDimensionUpdate, NodeDiffUpdate } from './nodes';
+import { Edge } from './edges';
+import { HandleType } from './handles';
+
+export type FlowElement<T = any> = Node<T> | Edge<T>;
+
+export type Elements<T = any> = Array<FlowElement<T>>;
+
+export type NodeTypesType = { [key: string]: ReactNode };
+export type EdgeTypesType = NodeTypesType;
+
+export type FitView = (fitViewOptions?: FitViewParams) => void;
+
+export type Project = (position: XYPosition) => XYPosition;
+
+export type ToObject<T = any> = () => FlowExportObject<T>;
+
+export type OnNodesChange = (nodes: NodeChange[]) => void;
+
+export type OnEdgesChange = (nodes: EdgeChange[]) => void;
+
+export interface SelectionRect extends Rect {
+  startX: number;
+  startY: number;
+  draw: boolean;
+}
+
+export type OnLoadParams<T = any> = {
+  zoomIn: () => void;
+  zoomOut: () => void;
+  zoomTo: (zoomLevel: number) => void;
+  fitView: FitView;
+  project: Project;
+  getNodes: () => Node<T>[];
+  getEdges: () => Edge<T>[];
+  setTransform: (transform: FlowTransform) => void;
+  toObject: ToObject<T>;
+};
+
+export type OnLoad<T = any> = (params: OnLoadParams<T>) => void;
+
+export interface Connection {
+  source: string | null;
+  target: string | null;
+  sourceHandle: string | null;
+  targetHandle: string | null;
+}
+
+export enum ConnectionMode {
+  Strict = 'strict',
+  Loose = 'loose',
+}
+
+export type FlowExportObject<T = any> = {
+  nodes: Node<T>[];
+  edges: Edge<T>[];
+  position: [number, number];
+  zoom: number;
+};
+
+export type OnConnect = (connection: Connection) => void;
+
+export type FitViewParams = {
+  padding?: number;
+  includeHiddenNodes?: boolean;
+  minZoom?: number;
+  maxZoom?: number;
+};
+
+export type OnConnectStartParams = {
+  nodeId: string | null;
+  handleId: string | null;
+  handleType: HandleType | null;
+};
+
+export type OnConnectStart = (event: ReactMouseEvent, params: OnConnectStartParams) => void;
+
+export type OnConnectStop = (event: MouseEvent) => void;
+
+export type OnConnectEnd = (event: MouseEvent) => void;
+
+export type SetConnectionId = {
+  connectionNodeId: string | null;
+  connectionHandleId: string | null;
+  connectionHandleType: HandleType | null;
+};
+
+export enum BackgroundVariant {
+  Lines = 'lines',
+  Dots = 'dots',
+}
+
+export type FlowTransform = {
+  x: number;
+  y: number;
+  zoom: number;
+};
+
+export type KeyCode = string | Array<string>;
+
+export type SnapGrid = [number, number];
+
+export enum PanOnScrollMode {
+  Free = 'free',
+  Vertical = 'vertical',
+  Horizontal = 'horizontal',
+}
+
+export interface ZoomPanHelperFunctions {
+  zoomIn: () => void;
+  zoomOut: () => void;
+  zoomTo: (zoomLevel: number) => void;
+  transform: (transform: FlowTransform) => void;
+  fitView: FitView;
+  setCenter: (x: number, y: number, zoom?: number) => void;
+  fitBounds: (bounds: Rect, padding?: number) => void;
+  project: (position: XYPosition) => XYPosition;
+  initialized: boolean;
+}
+
+export type InitD3ZoomPayload = {
+  d3Zoom: ZoomBehavior<Element, unknown>;
+  d3Selection: D3Selection<Element, unknown, null, undefined>;
+  d3ZoomHandler: ((this: Element, event: any, d: unknown) => void) | undefined;
+  transform: Transform;
+};
+
+export interface ReactFlowState {
+  width: number;
+  height: number;
+  transform: Transform;
+  nodeInternals: NodeInternals;
+  edges: Edge[];
+  selectedNodesBbox: Rect;
+  onNodesChange: OnNodesChange | null;
+  onEdgesChange: OnEdgesChange | null;
+
+  d3Zoom: ZoomBehavior<Element, unknown> | null;
+  d3Selection: D3Selection<Element, unknown, null, undefined> | null;
+  d3ZoomHandler: ((this: Element, event: any, d: unknown) => void) | undefined;
+  minZoom: number;
+  maxZoom: number;
+  translateExtent: CoordinateExtent;
+  nodeExtent: CoordinateExtent;
+
+  nodesSelectionActive: boolean;
+  selectionActive: boolean;
+
+  userSelectionRect: SelectionRect;
+
+  connectionNodeId: string | null;
+  connectionHandleId: string | null;
+  connectionHandleType: HandleType | null;
+  connectionPosition: XYPosition;
+  connectionMode: ConnectionMode;
+
+  snapToGrid: boolean;
+  snapGrid: SnapGrid;
+
+  nodesDraggable: boolean;
+  nodesConnectable: boolean;
+  elementsSelectable: boolean;
+
+  multiSelectionActive: boolean;
+
+  reactFlowVersion: string;
+
+  setNodes: (nodes: Node[]) => void;
+  setEdges: (edges: Edge[], nodes: Node[]) => void;
+  updateNodeDimensions: (updates: NodeDimensionUpdate[]) => void;
+  updateNodePosition: (update: NodeDiffUpdate) => void;
+  setUserSelection: (mousePos: XYPosition) => void;
+  updateUserSelection: (mousePos: XYPosition) => void;
+  unsetUserSelection: () => void;
+  unsetNodesSelection: () => void;
+  resetSelectedElements: () => void;
+  unselectNodesAndEdges: () => void;
+  addSelectedElements: (elements: Elements) => void;
+  updateTransform: (transform: Transform) => void;
+  updateSize: (size: Dimensions) => void;
+  initD3Zoom: (payload: InitD3ZoomPayload) => void;
+  setMinZoom: (minZoom: number) => void;
+  setMaxZoom: (maxZoom: number) => void;
+  setTranslateExtent: (translateExtent: CoordinateExtent) => void;
+  setNodeExtent: (nodeExtent: CoordinateExtent) => void;
+  setOnConnect: (onConnectFunction: OnConnect) => void;
+  setOnConnectStart: (onConnectFunction: OnConnectStart) => void;
+  setOnConnectStop: (onConnectFunction: OnConnectStop) => void;
+  setOnConnectEnd: (onConnectFunction: OnConnectEnd) => void;
+  setConnectionPosition: (connectionPosition: XYPosition) => void;
+  setConnectionNodeId: (payload: SetConnectionId) => void;
+  setSnapToGrid: (snapToGrid: boolean) => void;
+  setSnapGrid: (snapGrid: SnapGrid) => void;
+  setInteractive: (isInteractive: boolean) => void;
+  setNodesDraggable: (nodesDraggable: boolean) => void;
+  setNodesConnectable: (nodesConnectable: boolean) => void;
+  setElementsSelectable: (elementsSelectable: boolean) => void;
+  setMultiSelectionActive: (multiSelectionActive: boolean) => void;
+  setConnectionMode: (connectionMode: ConnectionMode) => void;
+  setOnNodesChange: (onNodesChange: OnNodesChange) => void;
+  setOnEdgesChange: (onEdgesChange: OnEdgesChange) => void;
+
+  onConnect?: OnConnect;
+  onConnectStart?: OnConnectStart;
+  onConnectStop?: OnConnectStop;
+  onConnectEnd?: OnConnectEnd;
+}
+
+export type UpdateNodeInternals = (nodeId: string) => void;
+
+export type OnSelectionChangeFunc = (params: { nodes: Node[]; edges: Edge[] }) => void;
diff --git a/src/types/handles.ts b/src/types/handles.ts
new file mode 100644
index 00000000..9114b360
--- /dev/null
+++ b/src/types/handles.ts
@@ -0,0 +1,18 @@
+import { XYPosition, Position, Dimensions } from './utils';
+import { OnConnect, Connection } from './general';
+
+export type HandleType = 'source' | 'target';
+
+export interface HandleElement extends XYPosition, Dimensions {
+  id?: string | null;
+  position: Position;
+}
+
+export interface HandleProps {
+  type: HandleType;
+  position: Position;
+  isConnectable?: boolean;
+  onConnect?: OnConnect;
+  isValidConnection?: (connection: Connection) => boolean;
+  id?: string;
+}
diff --git a/src/types/index.ts b/src/types/index.ts
index a03f847e..bf9e688a 100644
--- a/src/types/index.ts
+++ b/src/types/index.ts
@@ -1,459 +1,6 @@
-import React, { CSSProperties, MouseEvent as ReactMouseEvent, HTMLAttributes, ReactNode } from 'react';
-import { Selection as D3Selection, ZoomBehavior } from 'd3';
-
-export type ElementId = string;
-
-export type FlowElement<T = any> = Node<T> | Edge<T>;
-
-export type Elements<T = any> = Array<FlowElement<T>>;
-
-export type Transform = [number, number, number];
-
-export enum Position {
-  Left = 'left',
-  Top = 'top',
-  Right = 'right',
-  Bottom = 'bottom',
-}
-
-export interface XYPosition {
-  x: number;
-  y: number;
-}
-
-export interface Dimensions {
-  width: number;
-  height: number;
-}
-
-export interface Rect extends Dimensions, XYPosition {}
-
-export interface Box extends XYPosition {
-  x2: number;
-  y2: number;
-}
-
-export type SnapGrid = [number, number];
-
-export interface Node<T = any> {
-  id: ElementId;
-  position: XYPosition;
-  type?: string;
-  __rf?: any;
-  data?: T;
-  style?: CSSProperties;
-  className?: string;
-  targetPosition?: Position;
-  sourcePosition?: Position;
-  isHidden?: boolean;
-  draggable?: boolean;
-  selectable?: boolean;
-  connectable?: boolean;
-  dragHandle?: string;
-}
-
-export enum ArrowHeadType {
-  Arrow = 'arrow',
-  ArrowClosed = 'arrowclosed',
-}
-
-export interface EdgeMarker {
-  type: ArrowHeadType;
-  color?: string;
-  width?: number;
-  height?: number;
-  units?: string;
-  orient?: string;
-  strokeWidth?: number;
-}
-
-export type EdgeMarkerType = string | EdgeMarker;
-
-export interface Edge<T = any> {
-  id: ElementId;
-  type?: string;
-  source: ElementId;
-  target: ElementId;
-  sourceHandle?: ElementId | null;
-  targetHandle?: ElementId | null;
-  label?: string | ReactNode;
-  labelStyle?: CSSProperties;
-  labelShowBg?: boolean;
-  labelBgStyle?: CSSProperties;
-  labelBgPadding?: [number, number];
-  labelBgBorderRadius?: number;
-  style?: CSSProperties;
-  animated?: boolean;
-  isHidden?: boolean;
-  data?: T;
-  className?: string;
-  markerStart?: EdgeMarkerType;
-  markerEnd?: EdgeMarkerType;
-}
-
-export enum BackgroundVariant {
-  Lines = 'lines',
-  Dots = 'dots',
-}
-
-export type HandleType = 'source' | 'target';
-
-export type NodeTypesType = { [key: string]: ReactNode };
-
-export type EdgeTypesType = NodeTypesType;
-
-export interface SelectionRect extends Rect {
-  startX: number;
-  startY: number;
-  draw: boolean;
-}
-
-export interface WrapEdgeProps<T = any> {
-  id: ElementId;
-  className?: string;
-  type: string;
-  data?: T;
-  onClick?: (event: React.MouseEvent, edge: Edge) => void;
-  onEdgeDoubleClick?: (event: React.MouseEvent, edge: Edge) => void;
-  selected: boolean;
-  animated?: boolean;
-  label?: string | ReactNode;
-  labelStyle?: CSSProperties;
-  labelShowBg?: boolean;
-  labelBgStyle?: CSSProperties;
-  labelBgPadding?: [number, number];
-  labelBgBorderRadius?: number;
-  style?: CSSProperties;
-  source: ElementId;
-  target: ElementId;
-  sourceHandleId: ElementId | null;
-  targetHandleId: ElementId | null;
-  sourceX: number;
-  sourceY: number;
-  targetX: number;
-  targetY: number;
-  sourcePosition: Position;
-  targetPosition: Position;
-  elementsSelectable?: boolean;
-  isHidden?: boolean;
-  handleEdgeUpdate: boolean;
-  onConnectEdge: OnConnectFunc;
-  onContextMenu?: (event: React.MouseEvent, edge: Edge) => void;
-  onMouseEnter?: (event: React.MouseEvent, edge: Edge) => void;
-  onMouseMove?: (event: React.MouseEvent, edge: Edge) => void;
-  onMouseLeave?: (event: React.MouseEvent, edge: Edge) => void;
-  edgeUpdaterRadius?: number;
-  onEdgeUpdateStart?: (event: React.MouseEvent, edge: Edge) => void;
-  onEdgeUpdateEnd?: (event: MouseEvent, edge: Edge) => void;
-  markerStart?: EdgeMarkerType;
-  markerEnd?: EdgeMarkerType;
-}
-
-export interface EdgeProps<T = any> {
-  id: ElementId;
-  source: ElementId;
-  target: ElementId;
-  sourceX: number;
-  sourceY: number;
-  targetX: number;
-  targetY: number;
-  selected?: boolean;
-  animated?: boolean;
-  sourcePosition: Position;
-  targetPosition: Position;
-  label?: string | ReactNode;
-  labelStyle?: CSSProperties;
-  labelShowBg?: boolean;
-  labelBgStyle?: CSSProperties;
-  labelBgPadding?: [number, number];
-  labelBgBorderRadius?: number;
-  style?: CSSProperties;
-  data?: T;
-  sourceHandleId?: ElementId | null;
-  targetHandleId?: ElementId | null;
-  markerStart?: string;
-  markerEnd?: string;
-}
-export interface EdgeSmoothStepProps<T = any> extends EdgeProps<T> {
-  borderRadius?: number;
-}
-
-export interface EdgeTextProps extends HTMLAttributes<SVGElement> {
-  x: number;
-  y: number;
-  label?: string | ReactNode;
-  labelStyle?: CSSProperties;
-  labelShowBg?: boolean;
-  labelBgStyle?: CSSProperties;
-  labelBgPadding?: [number, number];
-  labelBgBorderRadius?: number;
-}
-
-export interface NodeProps<T = any> {
-  id: ElementId;
-  type: string;
-  data: T;
-  selected: boolean;
-  isConnectable: boolean;
-  xPos?: number;
-  yPos?: number;
-  targetPosition?: Position;
-  sourcePosition?: Position;
-  isDragging?: boolean;
-  dragHandle?: string;
-}
-
-export interface NodeComponentProps<T = any> {
-  id: ElementId;
-  type: string;
-  data: T;
-  selected?: boolean;
-  isConnectable: boolean;
-  transform?: Transform;
-  xPos?: number;
-  yPos?: number;
-  targetPosition?: Position;
-  sourcePosition?: Position;
-  onClick?: (node: Node) => void;
-  onNodeDoubleClick?: (node: Node) => void;
-  onMouseEnter?: (node: Node) => void;
-  onMouseMove?: (node: Node) => void;
-  onMouseLeave?: (node: Node) => void;
-  onContextMenu?: (node: Node) => void;
-  onNodeDragStart?: (node: Node) => void;
-  onNodeDrag?: (node: Node) => void;
-  onNodeDragStop?: (node: Node) => void;
-  style?: CSSProperties;
-  isDragging?: boolean;
-  dragHandle?: string;
-}
-
-export interface WrapNodeProps<T = any> {
-  id: ElementId;
-  type: string;
-  data: T;
-  selected: boolean;
-  scale: number;
-  xPos: number;
-  yPos: number;
-  isSelectable: boolean;
-  isDraggable: boolean;
-  isConnectable: boolean;
-  selectNodesOnDrag: boolean;
-  onClick?: (event: ReactMouseEvent, node: Node) => void;
-  onNodeDoubleClick?: (event: ReactMouseEvent, node: Node) => void;
-  onMouseEnter?: (event: ReactMouseEvent, node: Node) => void;
-  onMouseMove?: (event: ReactMouseEvent, node: Node) => void;
-  onMouseLeave?: (event: ReactMouseEvent, node: Node) => void;
-  onContextMenu?: (event: ReactMouseEvent, node: Node) => void;
-  onNodeDragStart?: (event: ReactMouseEvent, node: Node) => void;
-  onNodeDrag?: (event: ReactMouseEvent, node: Node) => void;
-  onNodeDragStop?: (event: ReactMouseEvent, node: Node) => void;
-  style?: CSSProperties;
-  className?: string;
-  sourcePosition?: Position;
-  targetPosition?: Position;
-  isHidden?: boolean;
-  isInitialized?: boolean;
-  snapToGrid?: boolean;
-  snapGrid?: SnapGrid;
-  isDragging?: boolean;
-  resizeObserver: ResizeObserver | null;
-  dragHandle?: string;
-}
-
-export type FitViewParams = {
-  padding?: number;
-  includeHiddenNodes?: boolean;
-  minZoom?: number;
-  maxZoom?: number;
-};
-
-export type FlowExportObject<T = any> = {
-  elements: Elements<T>;
-  position: [number, number];
-  zoom: number;
-};
-
-export type FitViewFunc = (fitViewOptions?: FitViewParams) => void;
-export type ProjectFunc = (position: XYPosition) => XYPosition;
-export type ToObjectFunc<T = any> = () => FlowExportObject<T>;
-
-export type OnLoadParams<T = any> = {
-  zoomIn: () => void;
-  zoomOut: () => void;
-  zoomTo: (zoomLevel: number) => void;
-  fitView: FitViewFunc;
-  project: ProjectFunc;
-  getElements: () => Elements<T>;
-  setTransform: (transform: FlowTransform) => void;
-  toObject: ToObjectFunc<T>;
-};
-
-export type OnLoadFunc<T = any> = (params: OnLoadParams<T>) => void;
-
-export interface Connection {
-  source: ElementId | null;
-  target: ElementId | null;
-  sourceHandle: ElementId | null;
-  targetHandle: ElementId | null;
-}
-
-export enum ConnectionMode {
-  Strict = 'strict',
-  Loose = 'loose',
-}
-
-export enum ConnectionLineType {
-  Bezier = 'default',
-  Straight = 'straight',
-  Step = 'step',
-  SmoothStep = 'smoothstep',
-}
-
-export type ConnectionLineComponentProps = {
-  sourceX: number;
-  sourceY: number;
-  sourcePosition?: Position;
-  targetX: number;
-  targetY: number;
-  targetPosition?: Position;
-  connectionLineStyle?: CSSProperties;
-  connectionLineType: ConnectionLineType;
-  sourceNode?: Node;
-  sourceHandle?: HandleElement;
-};
-
-export type ConnectionLineComponent = React.ComponentType<ConnectionLineComponentProps>;
-
-export type OnConnectFunc = (connection: Connection) => void;
-export type OnConnectStartParams = {
-  nodeId: ElementId | null;
-  handleId: ElementId | null;
-  handleType: HandleType | null;
-};
-export type OnConnectStartFunc = (event: ReactMouseEvent, params: OnConnectStartParams) => void;
-export type OnConnectStopFunc = (event: MouseEvent) => void;
-export type OnConnectEndFunc = (event: MouseEvent) => void;
-
-export type SetConnectionId = {
-  connectionNodeId: ElementId | null;
-  connectionHandleId: ElementId | null;
-  connectionHandleType: HandleType | null;
-};
-
-export interface HandleElement extends XYPosition, Dimensions {
-  id?: ElementId | null;
-  position: Position;
-}
-
-export interface HandleProps {
-  type: HandleType;
-  position: Position;
-  isConnectable?: boolean;
-  onConnect?: OnConnectFunc;
-  isValidConnection?: (connection: Connection) => boolean;
-  id?: ElementId;
-}
-
-export type NodePosUpdate = {
-  id: ElementId;
-  pos: XYPosition;
-};
-
-export type NodeDiffUpdate = {
-  id?: ElementId;
-  diff?: XYPosition;
-  isDragging?: boolean;
-};
-
-export type FlowTransform = {
-  x: number;
-  y: number;
-  zoom: number;
-};
-
-export type TranslateExtent = [[number, number], [number, number]];
-export type NodeExtent = TranslateExtent;
-
-export type KeyCode = string | Array<string>;
-
-export enum PanOnScrollMode {
-  Free = 'free',
-  Vertical = 'vertical',
-  Horizontal = 'horizontal',
-}
-
-export interface ZoomPanHelperFunctions {
-  zoomIn: () => void;
-  zoomOut: () => void;
-  zoomTo: (zoomLevel: number) => void;
-  transform: (transform: FlowTransform) => void;
-  fitView: FitViewFunc;
-  setCenter: (x: number, y: number, zoom?: number) => void;
-  fitBounds: (bounds: Rect, padding?: number) => void;
-  project: (position: XYPosition) => XYPosition;
-  initialized: boolean;
-}
-
-export type OnEdgeUpdateFunc<T = any> = (oldEdge: Edge<T>, newConnection: Connection) => void;
-
-export type NodeDimensionUpdate = {
-  id: ElementId;
-  nodeElement: HTMLDivElement;
-  forceUpdate?: boolean;
-};
-
-export type InitD3ZoomPayload = {
-  d3Zoom: ZoomBehavior<Element, unknown>;
-  d3Selection: D3Selection<Element, unknown, null, undefined>;
-  d3ZoomHandler: ((this: Element, event: any, d: unknown) => void) | undefined;
-  transform: Transform;
-};
-
-export interface ReactFlowState {
-  width: number;
-  height: number;
-  transform: Transform;
-  nodes: Node[];
-  edges: Edge[];
-  selectedElements: Elements | null;
-  selectedNodesBbox: Rect;
-
-  d3Zoom: ZoomBehavior<Element, unknown> | null;
-  d3Selection: D3Selection<Element, unknown, null, undefined> | null;
-  d3ZoomHandler: ((this: Element, event: any, d: unknown) => void) | undefined;
-  minZoom: number;
-  maxZoom: number;
-  translateExtent: TranslateExtent;
-  nodeExtent: NodeExtent;
-
-  nodesSelectionActive: boolean;
-  selectionActive: boolean;
-
-  userSelectionRect: SelectionRect;
-
-  connectionNodeId: ElementId | null;
-  connectionHandleId: ElementId | null;
-  connectionHandleType: HandleType | null;
-  connectionPosition: XYPosition;
-  connectionMode: ConnectionMode;
-
-  snapToGrid: boolean;
-  snapGrid: SnapGrid;
-
-  nodesDraggable: boolean;
-  nodesConnectable: boolean;
-  elementsSelectable: boolean;
-
-  multiSelectionActive: boolean;
-
-  reactFlowVersion: string;
-
-  onConnect?: OnConnectFunc;
-  onConnectStart?: OnConnectStartFunc;
-  onConnectStop?: OnConnectStopFunc;
-  onConnectEnd?: OnConnectEndFunc;
-}
-
-export type UpdateNodeInternals = (nodeId: ElementId) => void;
+export * from './general';
+export * from './nodes';
+export * from './edges';
+export * from './handles';
+export * from './changes';
+export * from './utils';
diff --git a/src/types/nodes.ts b/src/types/nodes.ts
new file mode 100644
index 00000000..76d0cd34
--- /dev/null
+++ b/src/types/nodes.ts
@@ -0,0 +1,113 @@
+import { CSSProperties, MouseEvent as ReactMouseEvent } from 'react';
+
+import { SnapGrid } from './general';
+import { XYPosition, Position, CoordinateExtent } from './utils';
+import { HandleElement } from './handles';
+
+// interface for the user node items
+export interface Node<T = any> {
+  id: string;
+  position: XYPosition;
+  type?: string;
+  data?: T;
+  style?: CSSProperties;
+  className?: string;
+  targetPosition?: Position;
+  sourcePosition?: Position;
+  hidden?: boolean;
+  selected?: boolean;
+  dragging?: boolean;
+  draggable?: boolean;
+  selectable?: boolean;
+  connectable?: boolean;
+  dragHandle?: string;
+  width?: number | null;
+  height?: number | null;
+  parentNode?: string;
+  zIndex?: number;
+  extent?: 'parent' | CoordinateExtent;
+}
+
+// props that get passed to a custom node
+export interface NodeProps<T = any> {
+  id: string;
+  type: string;
+  data: T;
+  selected: boolean;
+  isConnectable: boolean;
+  xPos: number;
+  yPos: number;
+  dragging: boolean;
+  zIndex: number;
+  targetPosition?: Position;
+  sourcePosition?: Position;
+  dragHandle?: string;
+}
+
+export type NodeMouseHandler = (event: ReactMouseEvent, node: Node) => void;
+
+export interface WrapNodeProps<T = any> {
+  id: string;
+  type: string;
+  data: T;
+  selected: boolean;
+  isConnectable: boolean;
+  scale: number;
+  xPos: number;
+  yPos: number;
+  width?: number | null;
+  height?: number | null;
+  isSelectable: boolean;
+  isDraggable: boolean;
+  selectNodesOnDrag: boolean;
+  onClick?: NodeMouseHandler;
+  onNodeDoubleClick?: NodeMouseHandler;
+  onMouseEnter?: NodeMouseHandler;
+  onMouseMove?: NodeMouseHandler;
+  onMouseLeave?: NodeMouseHandler;
+  onContextMenu?: NodeMouseHandler;
+  onNodeDragStart?: NodeMouseHandler;
+  onNodeDrag?: NodeMouseHandler;
+  onNodeDragStop?: NodeMouseHandler;
+  style?: CSSProperties;
+  className?: string;
+  sourcePosition?: Position;
+  targetPosition?: Position;
+  hidden?: boolean;
+  isInitialized?: boolean;
+  snapToGrid?: boolean;
+  snapGrid?: SnapGrid;
+  dragging: boolean;
+  resizeObserver: ResizeObserver | null;
+  dragHandle?: string;
+  zIndex: number;
+  isParent: boolean;
+  noPanClassName: string;
+  noDragClassName: string;
+}
+
+export type NodeHandleBounds = {
+  source: HandleElement[] | null;
+  target: HandleElement[] | null;
+};
+
+export type NodeDiffUpdate = {
+  id?: string;
+  diff?: XYPosition;
+  dragging?: boolean;
+};
+
+export type NodeDimensionUpdate = {
+  id: string;
+  nodeElement: HTMLDivElement;
+  forceUpdate?: boolean;
+};
+
+export type NodeInternalsItem = Node & {
+  positionAbsolute?: XYPosition;
+  handleBounds?: NodeHandleBounds;
+  z: number;
+  isParent?: boolean;
+};
+
+export type NodeInternals = Map<string, NodeInternalsItem>;
diff --git a/src/types/utils.ts b/src/types/utils.ts
new file mode 100644
index 00000000..eadce743
--- /dev/null
+++ b/src/types/utils.ts
@@ -0,0 +1,29 @@
+export enum Position {
+  Left = 'left',
+  Top = 'top',
+  Right = 'right',
+  Bottom = 'bottom',
+}
+
+export interface XYPosition {
+  x: number;
+  y: number;
+}
+
+export type XYZPosition = XYPosition & { z: number };
+
+export interface Dimensions {
+  width: number;
+  height: number;
+}
+
+export interface Rect extends Dimensions, XYPosition {}
+
+export interface Box extends XYPosition {
+  x2: number;
+  y2: number;
+}
+
+export type Transform = [number, number, number];
+
+export type CoordinateExtent = [[number, number], [number, number]];
diff --git a/src/utils/changes.ts b/src/utils/changes.ts
new file mode 100644
index 00000000..64743226
--- /dev/null
+++ b/src/utils/changes.ts
@@ -0,0 +1,51 @@
+import { Node, Edge, EdgeChange, NodeChange } from '../types';
+
+function applyChanges(changes: NodeChange[] | EdgeChange[], elements: any[]): any[] {
+  const initElements: any[] = [];
+
+  return elements.reduce((res: any[], item: any) => {
+    const currentChange = changes.find((c) => c.id === item.id);
+
+    if (currentChange) {
+      switch (currentChange.type) {
+        case 'select': {
+          res.push({ ...item, selected: currentChange.selected });
+          return res;
+        }
+        case 'dimensions': {
+          const updateItem = { ...item };
+
+          if (typeof currentChange.dimensions !== 'undefined') {
+            updateItem.width = currentChange.dimensions.width;
+            updateItem.height = currentChange.dimensions.height;
+          }
+
+          if (typeof currentChange.position !== 'undefined') {
+            updateItem.position = currentChange.position;
+          }
+
+          if (typeof currentChange.dragging !== 'undefined') {
+            updateItem.dragging = currentChange.dragging;
+          }
+
+          res.push(updateItem);
+          return res;
+        }
+        case 'remove': {
+          return res;
+        }
+      }
+    }
+
+    res.push(item);
+    return res;
+  }, initElements);
+}
+
+export function applyNodeChanges(changes: NodeChange[], nodes: Node[]): Node[] {
+  return applyChanges(changes, nodes) as Node[];
+}
+
+export function applyEdgeChanges(changes: EdgeChange[], edges: Edge[]): Edge[] {
+  return applyChanges(changes, edges) as Edge[];
+}
diff --git a/src/utils/graph.ts b/src/utils/graph.ts
index 2c6f4309..b575bbf7 100644
--- a/src/utils/graph.ts
+++ b/src/utils/graph.ts
@@ -1,22 +1,6 @@
-import { Store } from 'redux';
+import { boxToRect, clamp, getBoundsOfBoxes, rectToBox } from '../utils';
 
-import { clampPosition, clamp } from '../utils';
-
-import {
-  ElementId,
-  Node,
-  Edge,
-  Elements,
-  Transform,
-  XYPosition,
-  Rect,
-  Box,
-  Connection,
-  FlowExportObject,
-  ReactFlowState,
-  NodeExtent,
-  EdgeMarkerType,
-} from '../types';
+import { Node, Edge, Elements, Connection, EdgeMarkerType, Transform, XYPosition, Rect } from '../types';
 
 export const isEdge = (element: Node | Connection | Edge): element is Edge =>
   'id' in element && 'source' in element && 'target' in element;
@@ -24,38 +8,25 @@ export const isEdge = (element: Node | Connection | Edge): element is Edge =>
 export const isNode = (element: Node | Connection | Edge): element is Node =>
   'id' in element && !('source' in element) && !('target' in element);
 
-export const getOutgoers = (node: Node, elements: Elements): Node[] => {
+export const getOutgoers = (node: Node, nodes: Node[], edges: Edge[]): Node[] => {
   if (!isNode(node)) {
     return [];
   }
 
-  const outgoerIds = elements.filter((e) => isEdge(e) && e.source === node.id).map((e) => (e as Edge).target);
-  return elements.filter((e) => outgoerIds.includes(e.id)) as Node[];
+  const outgoerIds = edges.filter((e) => e.source === node.id).map((e) => e.target);
+  return nodes.filter((n) => outgoerIds.includes(n.id));
 };
 
-export const getIncomers = (node: Node, elements: Elements): Node[] => {
+export const getIncomers = (node: Node, nodes: Node[], edges: Edge[]): Node[] => {
   if (!isNode(node)) {
     return [];
   }
 
-  const incomersIds = elements.filter((e) => isEdge(e) && e.target === node.id).map((e) => (e as Edge).source);
-  return elements.filter((e) => incomersIds.includes(e.id)) as Node[];
+  const incomersIds = edges.filter((e) => e.target === node.id).map((e) => e.source);
+  return nodes.filter((n) => incomersIds.includes(n.id));
 };
 
-export const removeElements = (elementsToRemove: Elements, elements: Elements): Elements => {
-  const nodeIdsToRemove = elementsToRemove.map((n) => n.id);
-
-  return elements.filter((element) => {
-    const edgeElement = element as Edge;
-    return !(
-      nodeIdsToRemove.includes(element.id) ||
-      nodeIdsToRemove.includes(edgeElement.target) ||
-      nodeIdsToRemove.includes(edgeElement.source)
-    );
-  });
-};
-
-const getEdgeId = ({ source, sourceHandle, target, targetHandle }: Connection): ElementId =>
+const getEdgeId = ({ source, sourceHandle, target, targetHandle }: Connection): string =>
   `reactflow__edge-${source}${sourceHandle}-${target}${targetHandle}`;
 
 export const getMarkerId = (marker: EdgeMarkerType | undefined): string => {
@@ -73,10 +44,9 @@ export const getMarkerId = (marker: EdgeMarkerType | undefined): string => {
     .join('&');
 };
 
-const connectionExists = (edge: Edge, elements: Elements) => {
-  return elements.some(
+const connectionExists = (edge: Edge, edges: Edge[]) => {
+  return edges.some(
     (el) =>
-      isEdge(el) &&
       el.source === edge.source &&
       el.target === edge.target &&
       (el.sourceHandle === edge.sourceHandle || (!el.sourceHandle && !edge.sourceHandle)) &&
@@ -84,10 +54,10 @@ const connectionExists = (edge: Edge, elements: Elements) => {
   );
 };
 
-export const addEdge = (edgeParams: Edge | Connection, elements: Elements): Elements => {
+export const addEdge = (edgeParams: Edge | Connection, edges: Edge[]): Edge[] => {
   if (!edgeParams.source || !edgeParams.target) {
     console.warn("Can't create edge. An edge needs a source and a target.");
-    return elements;
+    return edges;
   }
 
   let edge: Edge;
@@ -100,11 +70,11 @@ export const addEdge = (edgeParams: Edge | Connection, elements: Elements): Elem
     } as Edge;
   }
 
-  if (connectionExists(edge, elements)) {
-    return elements;
+  if (connectionExists(edge, edges)) {
+    return edges;
   }
 
-  return elements.concat(edge);
+  return edges.concat(edge);
 };
 
 export const updateEdge = (oldEdge: Edge, newConnection: Connection, elements: Elements): Elements => {
@@ -154,80 +124,16 @@ export const pointToRendererPoint = (
   return position;
 };
 
-export const onLoadProject = (currentStore: Store<ReactFlowState>) => {
-  return (position: XYPosition): XYPosition => {
-    const { transform, snapToGrid, snapGrid } = currentStore.getState();
-
-    return pointToRendererPoint(position, transform, snapToGrid, snapGrid);
-  };
-};
-
-export const parseNode = (node: Node, nodeExtent: NodeExtent): Node => {
-  return {
-    ...node,
-    id: node.id.toString(),
-    type: node.type || 'default',
-    __rf: {
-      position: clampPosition(node.position, nodeExtent),
-      width: null,
-      height: null,
-      handleBounds: {},
-      isDragging: false,
-    },
-  };
-};
-
-export const parseEdge = (edge: Edge): Edge => {
-  return {
-    ...edge,
-    source: edge.source.toString(),
-    target: edge.target.toString(),
-    sourceHandle: edge.sourceHandle ? edge.sourceHandle.toString() : null,
-    targetHandle: edge.targetHandle ? edge.targetHandle.toString() : null,
-    id: edge.id.toString(),
-    type: edge.type || 'default',
-  };
-};
-
-const getBoundsOfBoxes = (box1: Box, box2: Box): Box => ({
-  x: Math.min(box1.x, box2.x),
-  y: Math.min(box1.y, box2.y),
-  x2: Math.max(box1.x2, box2.x2),
-  y2: Math.max(box1.y2, box2.y2),
-});
-
-export const rectToBox = ({ x, y, width, height }: Rect): Box => ({
-  x,
-  y,
-  x2: x + width,
-  y2: y + height,
-});
-
-export const boxToRect = ({ x, y, x2, y2 }: Box): Rect => ({
-  x,
-  y,
-  width: x2 - x,
-  height: y2 - y,
-});
-
-export const getBoundsofRects = (rect1: Rect, rect2: Rect): Rect =>
-  boxToRect(getBoundsOfBoxes(rectToBox(rect1), rectToBox(rect2)));
-
 export const getRectOfNodes = (nodes: Node[]): Rect => {
   const box = nodes.reduce(
-    (currBox, { __rf: { position, width, height } = {} }) =>
-      getBoundsOfBoxes(currBox, rectToBox({ ...position, width, height })),
+    (currBox, { position, width, height }) =>
+      getBoundsOfBoxes(currBox, rectToBox({ ...position, width: width || 0, height: height || 0 })),
     { x: Infinity, y: Infinity, x2: -Infinity, y2: -Infinity }
   );
 
   return boxToRect(box);
 };
 
-export const graphPosToZoomedPos = ({ x, y }: XYPosition, [tx, ty, tScale]: Transform): XYPosition => ({
-  x: x * tScale + tx,
-  y: y * tScale + ty,
-});
-
 export const getNodesInside = (
   nodes: Node[],
   rect: Rect,
@@ -243,17 +149,23 @@ export const getNodesInside = (
     height: rect.height / tScale,
   });
 
-  return nodes.filter(({ selectable = true, __rf: { position, width, height, isDragging } }) => {
+  return nodes.filter(({ selectable = true, position, width, height, dragging }) => {
     if (excludeNonSelectableNodes && !selectable) {
       return false;
     }
 
-    const nBox = rectToBox({ ...position, width, height });
+    const nBox = rectToBox({ ...position, width: width || 0, height: height || 0 });
     const xOverlap = Math.max(0, Math.min(rBox.x2, nBox.x2) - Math.max(rBox.x, nBox.x));
     const yOverlap = Math.max(0, Math.min(rBox.y2, nBox.y2) - Math.max(rBox.y, nBox.y));
     const overlappingArea = Math.ceil(xOverlap * yOverlap);
 
-    if (width === null || height === null || isDragging) {
+    if (
+      typeof width === 'undefined' ||
+      typeof height === 'undefined' ||
+      width === null ||
+      height === null ||
+      dragging
+    ) {
       // nodes are initialized with width and height = null
       return true;
     }
@@ -262,7 +174,7 @@ export const getNodesInside = (
       return overlappingArea > 0;
     }
 
-    const area = width * height;
+    const area = (width || 0) * (height || 0);
 
     return overlappingArea >= area;
   });
@@ -274,40 +186,6 @@ export const getConnectedEdges = (nodes: Node[], edges: Edge[]): Edge[] => {
   return edges.filter((edge) => nodeIds.includes(edge.source) || nodeIds.includes(edge.target));
 };
 
-const parseElements = (nodes: Node[], edges: Edge[]): Elements => {
-  return [
-    ...nodes.map((node) => {
-      const n = { ...node };
-
-      n.position = n.__rf.position;
-
-      delete n.__rf;
-      return n;
-    }),
-    ...edges.map((e) => ({ ...e })),
-  ];
-};
-
-export const onLoadGetElements = (currentStore: Store<ReactFlowState>) => {
-  return (): Elements => {
-    const { nodes = [], edges = [] } = currentStore.getState();
-
-    return parseElements(nodes, edges);
-  };
-};
-
-export const onLoadToObject = (currentStore: Store<ReactFlowState>) => {
-  return (): FlowExportObject => {
-    const { nodes = [], edges = [], transform } = currentStore.getState();
-
-    return {
-      elements: parseElements(nodes, edges),
-      position: [transform[0], transform[1]],
-      zoom: transform[2],
-    };
-  };
-};
-
 export const getTransformForBounds = (
   bounds: Rect,
   width: number,
diff --git a/src/utils/index.ts b/src/utils/index.ts
index 29c3149f..926977cb 100644
--- a/src/utils/index.ts
+++ b/src/utils/index.ts
@@ -1,4 +1,4 @@
-import { Dimensions, XYPosition, NodeExtent } from '../types';
+import { Dimensions, XYPosition, CoordinateExtent, Box, Rect } from '../types';
 
 export const getDimensions = (node: HTMLDivElement): Dimensions => ({
   width: node.offsetWidth,
@@ -7,10 +7,36 @@ export const getDimensions = (node: HTMLDivElement): Dimensions => ({
 
 export const clamp = (val: number, min: number = 0, max: number = 1): number => Math.min(Math.max(val, min), max);
 
-export const clampPosition = (position: XYPosition, extent: NodeExtent) => ({
+export const clampPosition = (position: XYPosition, extent: CoordinateExtent) => ({
   x: clamp(position.x, extent[0][0], extent[1][0]),
   y: clamp(position.y, extent[0][1], extent[1][1]),
 });
 
 export const getHostForElement = (element: HTMLElement): Document | ShadowRoot =>
   (element.getRootNode?.() as Document | ShadowRoot) || window?.document;
+
+export const getBoundsOfBoxes = (box1: Box, box2: Box): Box => ({
+  x: Math.min(box1.x, box2.x),
+  y: Math.min(box1.y, box2.y),
+  x2: Math.max(box1.x2, box2.x2),
+  y2: Math.max(box1.y2, box2.y2),
+});
+
+export const rectToBox = ({ x, y, width, height }: Rect): Box => ({
+  x,
+  y,
+  x2: x + width,
+  y2: y + height,
+});
+
+export const boxToRect = ({ x, y, x2, y2 }: Box): Rect => ({
+  x,
+  y,
+  width: x2 - x,
+  height: y2 - y,
+});
+
+export const getBoundsofRects = (rect1: Rect, rect2: Rect): Rect =>
+  boxToRect(getBoundsOfBoxes(rectToBox(rect1), rectToBox(rect2)));
+
+export const isNumeric = (n: any): n is number => !isNaN(n) && isFinite(n);
diff --git a/src/wdyr.ts b/src/wdyr.ts
new file mode 100644
index 00000000..79be3097
--- /dev/null
+++ b/src/wdyr.ts
@@ -0,0 +1,10 @@
+/// <reference types="@welldone-software/why-did-you-render" />
+
+import React from 'react';
+
+if (process.env.NODE_ENV === 'development') {
+  const whyDidYouRender = require('@welldone-software/why-did-you-render');
+  whyDidYouRender(React, {
+    trackAllPureComponents: true,
+  });
+}