Support for subroutine init

Better detection of wait transitions
Support for RestoreStep op

files.go

@@ -54,6 +54,7 @@ func (p *File) parseAst(fileAst *ast.File) {
 
 const TypeStateUpdate = "StateUpdate"
 const GetInitStateForFunc = "GetInitStateFor"
+const GetSubroutineInitState = "GetSubroutineInitState"
 
 func (p *File) parseFuncDecl(fd *ast.FuncDecl) *MethodDecl {
 	if fd.Type.Results == nil {
@@ -72,8 +73,9 @@ func (p *File) parseFuncDecl(fd *ast.FuncDecl) *MethodDecl {
 	case fd.Name == nil:
 		return nil
 	default:
-		if fd.Name.Name == GetInitStateForFunc {
-			md = p.findFuncWith(fd.Type.Results.List, GetInitStateForFunc, "InitFunc")
+		switch fd.Name.Name {
+		case GetInitStateForFunc, GetSubroutineInitState:
+			md = p.findFuncWith(fd.Type.Results.List, fd.Name.Name, "InitFunc")
 		}
 		if md == nil {
 			return nil

methods.go

@@ -61,6 +61,7 @@ type MethodTransition struct {
 	HiddenPropagate string
 
 	InheritMigration bool
+	WaitTransition   bool
 
 	TransitionTo *MethodDecl
 	HiddenPropTo *MethodDecl

trace_build.go

@@ -117,12 +117,23 @@ func (p *ExecTrace) addContextOpTransition(su *StateUpdate, mt *MethodTransition
 		mt.Operation = `Repeat(` + p.shortenArgs(su.args, maxArgLen) + `)`
 		return true
 
+	case "RestoreStep":
+		mt.WaitTransition = true
+		mt.Operation = su.name
+
 	default:
 		if strings.HasPrefix(su.name, "Then") {
+			switch op := su.parent.name; {
+			case op == "Sleep":
+				mt.WaitTransition = true
+			case strings.HasPrefix(op, "Wait"):
+				mt.WaitTransition = true
+			}
+
 			mt.Operation, mt.DelayedStart = p.buildOperation(su.parent)
 		}
 
-		if strings.HasSuffix(su.name, "Ext") || strings.HasSuffix(su.name, "Step") {
+		if strings.HasSuffix(su.name, "Ext") {
 			break
 		}
 		if len(su.args) == 0 { // repeat and similar ops

writer.go

@@ -90,11 +90,9 @@ func (p *Writer) WriteDecl(d *SMDecl) {
 		for _, tr := range step.Transitions {
 			connIdx++
 
-			waitOperation := false
+			waitOperation := tr.WaitTransition
 
 			if tr.TransitionTo == nil || !tr.TransitionTo.IsSubroutine {
-				waitOperation = tr.Operation != ""
-
 				m := ""
 				switch {
 				case tr.Transition == "<stop>":