1 files changed, 24 insertions, 10 deletions

diff --git a/plugins/org.eclipse.etrice.ui.behavior.fsm/src/org/eclipse/etrice/ui/behavior/fsm/support/DiagramUpdateFeature.java b/plugins/org.eclipse.etrice.ui.behavior.fsm/src/org/eclipse/etrice/ui/behavior/fsm/support/DiagramUpdateFeature.java
index 6d194add9..8cdd19284 100644
--- a/plugins/org.eclipse.etrice.ui.behavior.fsm/src/org/eclipse/etrice/ui/behavior/fsm/support/DiagramUpdateFeature.java
+++ b/plugins/org.eclipse.etrice.ui.behavior.fsm/src/org/eclipse/etrice/ui/behavior/fsm/support/DiagramUpdateFeature.java
@@ -14,9 +14,12 @@ package org.eclipse.etrice.ui.behavior.fsm.support;
 
 import java.util.ArrayList;
 
+import org.eclipse.etrice.core.fsm.fSM.FSMFactory;
 import org.eclipse.etrice.core.fsm.fSM.ModelComponent;
 import org.eclipse.etrice.core.fsm.fSM.StateGraph;
 import org.eclipse.etrice.ui.behavior.fsm.commands.StateGraphContext;
+import org.eclipse.etrice.ui.behavior.fsm.editor.AbstractFSMDiagramTypeProvider;
+import org.eclipse.etrice.ui.behavior.fsm.provider.GenModelProvider;
 import org.eclipse.etrice.ui.behavior.fsm.provider.IInjectorProvider;
 import org.eclipse.etrice.ui.behavior.fsm.support.util.DiagramEditingUtil;
 import org.eclipse.etrice.ui.behavior.fsm.support.util.FSMSupportUtil;
@@ -65,9 +68,7 @@ public class DiagramUpdateFeature extends AbstractUpdateFeature {
 	 */
 	@Override
 	public IReason updateNeeded(IUpdateContext context) {
-		ModelComponent mc = FSMSupportUtil.getInstance().getModelComponent(getDiagram());
-    	Injector injector = ((IInjectorProvider) getFeatureProvider()).getInjector();
-		StateGraphContext tree = StateGraphContext.createContextTree(mc, injector);
+		IStateGraphContext tree = createStateGraphContext();
 		
 		usedShapes.clear();
 		
@@ -110,9 +111,7 @@ public class DiagramUpdateFeature extends AbstractUpdateFeature {
 	 */
 	@Override
 	public boolean update(IUpdateContext context) {
-		ModelComponent mc = FSMSupportUtil.getInstance().getModelComponent(getDiagram());
-    	Injector injector = ((IInjectorProvider) getFeatureProvider()).getInjector();
-		StateGraphContext tree = StateGraphContext.createContextTree(mc, injector);
+		IStateGraphContext tree = createStateGraphContext();
 
 		StateGraph currentStateGraph = ContextSwitcher.getCurrentStateGraph(getDiagram());
 		
@@ -171,7 +170,7 @@ public class DiagramUpdateFeature extends AbstractUpdateFeature {
 	 * @param ctx
 	 * @return
 	 */
-	private IReason updateNeeded(StateGraphContext ctx) {
+	private IReason updateNeeded(IStateGraphContext ctx) {
 		StateGraph sg = ctx.getStateGraph();
 		ContainerShape cont = findStateGraphContainer(sg);
 		if (cont==null)
@@ -198,7 +197,7 @@ public class DiagramUpdateFeature extends AbstractUpdateFeature {
 		}
 		
 		// recursion
-		for (StateGraphContext child : ctx.getChildren()) {
+		for (IStateGraphContext child : ctx.getChildren()) {
 			IReason needUpdate = updateNeeded(child);
 			if (needUpdate.toBoolean())
 				return needUpdate;
@@ -211,7 +210,7 @@ public class DiagramUpdateFeature extends AbstractUpdateFeature {
 	 * @param ctx
 	 * @return
 	 */
-	private boolean update(StateGraphContext ctx) {
+	private boolean update(IStateGraphContext ctx) {
 		boolean changed = false;
 		
 		StateGraph sg = ctx.getStateGraph();
@@ -244,7 +243,7 @@ public class DiagramUpdateFeature extends AbstractUpdateFeature {
 		}
 		
 		// recursion
-		for (StateGraphContext child : ctx.getChildren()) {
+		for (IStateGraphContext child : ctx.getChildren()) {
 			if (update(child))
 				changed = true;
 		}
@@ -264,5 +263,20 @@ public class DiagramUpdateFeature extends AbstractUpdateFeature {
 		}
 		return null;
 	}
+	
+	private IStateGraphContext createStateGraphContext(){
+		ModelComponent mc = FSMSupportUtil.getInstance().getModelComponent(getDiagram());
+//    	Injector injector = ((IInjectorProvider) getFeatureProvider()).getInjector();
+//		StateGraphContext tree = StateGraphContext.createContextTree(mc, injector);
+		
+		// side effect on model, but diagram needs that for initial graph shape
+		if(mc.getStateMachine() == null || mc.getStateMachine().eIsProxy()){
+			mc.setStateMachine(FSMFactory.eINSTANCE.createStateGraph());
+		}
+		
+		IStateGraphContext tree = new GenModelStateGraphContext((AbstractFSMDiagramTypeProvider)getFeatureProvider().getDiagramTypeProvider());
+		
+		return tree;
+	}
 
 }