diff --git a/AlgebraicDataflowArchitectureModel/src/application/editor/stages/DependencyModelingStage.java b/AlgebraicDataflowArchitectureModel/src/application/editor/stages/DependencyModelingStage.java
index 30516bf..8cb75de 100644
--- a/AlgebraicDataflowArchitectureModel/src/application/editor/stages/DependencyModelingStage.java
+++ b/AlgebraicDataflowArchitectureModel/src/application/editor/stages/DependencyModelingStage.java
@@ -212,10 +212,10 @@
                 geo.getX(), geo.getY(), geo.getWidth(), geo.getHeight(), newStyle);
 
 
-        ResourceNode currentNode = findResourceNodeByCell(oldCell);
-        if (currentNode != null) {
-            resNodeCells.put(currentNode, newCell);
-        }
+        StatefulObjectNode parentNode = findResourceNodeByCell(oldCell);
+        if( parentNode != null){
+            resNodeCells.put(parentNode.getResource(), newCell);
+        };
 
         int childCount = oldCell.getChildCount();
         for (int i = 0; i < childCount; i++) {
@@ -223,16 +223,9 @@
             mxCell oldChild  = (mxCell) oldCell.getChildAt(i);
             mxCell newChild  = cloneResourceCellHierarchy(oldChild, newCell, resNodeCells);
 
-            ResourceNode parentNode = currentNode;
-            ResourceNode childNode  = findResourceNodeByCell(oldChild);
+            StatefulObjectNode childNode  = findResourceNodeByCell(oldChild);
 
-            if (parentNode != null && childNode != null) {
-
-                ObjectNode parentObj = new StatefulObjectNode(parentNode);
-                ObjectNode childObj  = new StatefulObjectNode(childNode);
-
-                addDependencyGraph(parentObj, childObj);
-            }
+            addDependencyGraph(parentNode, childNode);
         }
 
         return newCell;
@@ -242,16 +235,16 @@
      * Helper method to retrieve the corresponding ResourceNode from a cell.
      * (Performs reverse lookup from resNodeToCell)
      */
-    private ResourceNode findResourceNodeByCell(mxCell cell) {
+    private StatefulObjectNode findResourceNodeByCell(mxCell cell) {
         for (Map.Entry<ResourceNode, mxCell> entry : resNodeToCell.entrySet()) {
             if (entry.getValue() == cell) {
-                return entry.getKey();
+                return new StatefulObjectNode(entry.getKey());
             }
         }
         return null;
     }
 
-    private void addDependencyGraph(ObjectNode parentNode, ObjectNode childNode){
+    private void addDependencyGraph(StatefulObjectNode parentNode, StatefulObjectNode childNode){
         dependencyGraph.addEdge(new DependencyEdge(parentNode, childNode));
     }