diff --git a/example/SimpleGraph.js b/example/SimpleGraph.js
index 60de8e63..ea55a0ef 100644
--- a/example/SimpleGraph.js
+++ b/example/SimpleGraph.js
@@ -2,6 +2,24 @@ import React, { PureComponent } from 'react';
 import Graph from '../src';
 
 class App extends PureComponent {
+  constructor() {
+    super();
+
+    this.state = {
+      elements: [
+        { data: { id: '1', label: 'Tests', type: 'input' }, position: { x: 0, y: 0 } },
+        { data: { id: '2', label: 'This is a node This is a node This is a node This is a node' }, position: { x: 100, y: 100 } },
+        { data: { id: '3', label: 'This is a node' }, position: { x: 100, y: 200 }, style: { background: '#222', color: '#fff' } },
+        { data: { id: '4', label: 'nody nodes', type: 'output' }, position: { x: 50, y: 300 } },
+        { data: { id: '5', label: 'Another node', type: 'output' }, position: { x: 400, y: 300 } },
+        { data: { source: '1', target: '2' } },
+        { data: { source: '2', target: '3' } },
+        { data: { source: '3', target: '4' } },
+        { data: { source: '3', target: '5' } }
+      ]
+    }
+  }
+
   onLoad(graphInstance) {
     this.graphInstance = graphInstance;
 
@@ -22,22 +40,17 @@ class App extends PureComponent {
     this.graphInstance.fitView();
   }
 
-  render() {
-    const elements = [
-      { data: { id: '1', label: 'Tests', type: 'input' }, position: { x: 0, y: 0 } },
-      { data: { id: '2', label: 'This is a node This is a node This is a node This is a node' }, position: { x: 100, y: 100 } },
-      { data: { id: '3', label: 'This is a node' }, position: { x: 100, y: 200 }, style: { background: '#222', color: '#fff' } },
-      { data: { id: '4', label: 'nody nodes', type: 'output' }, position: { x: 50, y: 300 } },
-      { data: { id: '5', label: 'Another node', type: 'output' }, position: { x: 400, y: 300 } },
-      { data: { source: '1', target: '2' } },
-      { data: { source: '2', target: '3' } },
-      { data: { source: '3', target: '4' } },
-      { data: { source: '3', target: '5' } }
-    ];
+  onAdd() {
+    this.setState(prevState => ({
+      ...prevState,
+      elements: prevState.elements.concat({ data: { id: (prevState.elements.length + 1).toString(), label: 'added', type: 'input' }, position: { x: 50, y: 50 } })
+    }))
+  }
 
+  render() {
     return (
       <Graph
-        elements={elements}
+        elements={this.state.elements}
         onNodeClick={node => console.log(node)}
         style={{ width: '100%', height: '100%' }}
         onLoad={graphInstance => this.onLoad(graphInstance)}
@@ -50,6 +63,13 @@ class App extends PureComponent {
         >
           fit
         </button>
+        <button
+          type="button"
+          style={{ position: 'absolute', bottom: '10px', left: '10px' }}
+          onClick={() => this.onAdd()}
+        >
+          add
+        </button>
       </Graph>
     );
   }
diff --git a/package.json b/package.json
index e05e4044..ca0deb6d 100644
--- a/package.json
+++ b/package.json
@@ -3,6 +3,7 @@
   "version": "1.0.0",
   "module": "dist/ReactGraph.esm.js",
   "browser": "dist/ReactGraph.umd.js",
+  "main": "dist/ReactGraph.umd.js",
   "private": true,
   "dependencies": {
     "@emotion/core": "^10.0.14",
diff --git a/src/GraphContext/index.js b/src/GraphContext/index.js
index ac90d4b6..4a98b1df 100644
--- a/src/GraphContext/index.js
+++ b/src/GraphContext/index.js
@@ -1,6 +1,8 @@
-import React, { createContext, useState } from 'react';
+import React, { createContext, useState, useEffect } from 'react';
 import PropTypes from 'prop-types';
 import { zoomIdentity } from 'd3-zoom';
+import isEqual from 'lodash.isequal';
+
 import { getBoundingBox } from '../graph-utils';
 
 export const GraphContext = createContext({});
@@ -22,6 +24,19 @@ export const Provider = (props) => {
   const [edges, setEdges] = useState(initEdges);
   const [transform, setTransform] = useState({ x: 0, y: 0, k: 1 });
 
+  useEffect(() => {
+    const nodesChanged = !isEqual(nodes, props.nodes);
+    const edgesChanged = !isEqual(edges, props.edges);
+
+    if (nodesChanged) {
+      setNodes(props.nodes);
+    }
+
+    if (edgesChanged) {
+      setEdges(props.edges);
+    }
+  });
+
   const updateNodeData = (nodeId, updateData) => {
     const updatedNodes = nodes.map((n) => {
       if (n.data.id === nodeId) {
diff --git a/src/index.js b/src/index.js
index 0c61ba2f..d910e051 100644
--- a/src/index.js
+++ b/src/index.js
@@ -14,44 +14,14 @@ const GraphWrapper = styled.div`
 `;
 
 class ReactGraph extends PureComponent {
-  constructor(props) {
-    super();
-
-    const { elements } = props;
-
-    this.state = {
-      ...separateElements(elements)
-    };
-  }
-
-  componentDidUpdate(prevProps, prevState) {
-    const { elements } = this.props;
-    const { nodes, edges } = separateElements(elements);
-    const nodesChanged = !isEqual(nodes, prevState.nodes);
-    const edgesChanged = !isEqual(edges, prevState.edges);
-
-    if (!nodesChanged && !edgesChanged) {
-      return false;
-    }
-
-    if (nodesChanged) {
-      this.setState({ nodes });
-    }
-
-    if (edgesChanged) {
-      this.setState({ edges });
-    }
-  }
-
   render() {
     const {
-      style, onNodeClick, children, onLoad, onMove
-    } = this.props;
-    const { nodes, edges } = this.state;
+      style, onNodeClick, children, onLoad, onMove, elements
+		} = this.props;
 
     return (
       <GraphWrapper style={style}>
-        <Provider nodes={nodes} edges={edges} onNodeClick={onNodeClick}>
+        <Provider {...separateElements(elements)} onNodeClick={onNodeClick}>
           <GraphView onLoad={onLoad} onMove={onMove} />
           {children}
         </Provider>