Fix DNElementRemover to remove the actual component first if possible

[simantics/district.git] / org.simantics.district.network.ui / src / org / simantics / district / network / ui / DNElementRemover.java

diff --git a/org.simantics.district.network.ui/src/org/simantics/district/network/ui/DNElementRemover.java b/org.simantics.district.network.ui/src/org/simantics/district/network/ui/DNElementRemover.java
index 06b31031f2efc794eca78c498aaaf5e0a1287e29..a32d5359ee10653e545761725e83baca179732e1 100644 (file)
--- a/org.simantics.district.network.ui/src/org/simantics/district/network/ui/DNElementRemover.java
+++ b/org.simantics.district.network.ui/src/org/simantics/district/network/ui/DNElementRemover.java
@@ -10,6 +10,7 @@ import org.simantics.district.network.DistrictNetworkUtil;
 import org.simantics.district.network.ontology.DistrictNetworkResource;
 import org.simantics.layer0.Layer0;
 import org.simantics.layer0.utils.binaryPredicates.OrderedSetElementsPredicate;
+import org.simantics.modeling.ModelingResources;
 import org.simantics.modeling.adapters.ElementRemover;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -71,9 +72,15 @@ public class DNElementRemover extends ElementRemover {
         Boolean trackChangesEnabled = DistrictNetworkUtil.trackChangesEnabled(graph, graph.getPossibleObject(resource, Layer0.getInstance(graph).PartOf));
         if (trackChangesEnabled) {
             DistrictNetworkResource DN = DistrictNetworkResource.getInstance(graph);
-            Resource mappedComponent = graph.getPossibleObject(resource, DN.MappedComponent);
-            if (mappedComponent != null) {
-                RemoverUtil.remove(graph, mappedComponent);
+            Resource element = graph.getPossibleObject(resource, DN.MappedComponent);
+            if (element != null) {
+                // get component
+                ModelingResources MOD = ModelingResources.getInstance(graph);
+                Resource component = graph.getPossibleObject(element, MOD.ElementToComponent);
+                if (component != null)
+                    RemoverUtil.remove(graph, component);
+                else
+                    RemoverUtil.remove(graph, element);
             }
         }
     }