visualizer_graphviz.go

package msm

import (
	"bytes"
	"fmt"
	"sort"
)

func getSortedTransitionKeys(fm *FSM) []TransitionKey {
	// we sort the key alphabetically to have a reproducible graph output
	sortedTransitionKeys := make([]TransitionKey, 0)
	for transition := range fm.Table {
		sortedTransitionKeys = append(sortedTransitionKeys, transition)
	}

	sort.Slice(sortedTransitionKeys, func(i, j int) bool {
		if sortedTransitionKeys[i].State == sortedTransitionKeys[j].State {
			return sortedTransitionKeys[i].Event.Name() < sortedTransitionKeys[j].Event.Name()
		}
		return sortedTransitionKeys[i].State < sortedTransitionKeys[j].State
	})

	return sortedTransitionKeys
}

func getSortedStates(fm *FSM) []string {
	// we sort the key alphabetically to have a reproducible graph output
	states := make(map[string]struct{})
	for src, transitions := range fm.Table {
		states[string(src.State)] = struct{}{}
		for _, transition := range transitions {
			states[string(transition.NewState)] = struct{}{}
		}
	}

	sortedStates := make([]string, 0)
	for state := range states {
		sortedStates = append(sortedStates, state)
	}
	sort.Slice(sortedStates, func(i, j int) bool {
		return sortedStates[i] < sortedStates[j]
	})

	return sortedStates
}

// Visualize outputs a visualization of a FSM in Graphviz format.
func Visualize(fm *FSM) string {
	var buf bytes.Buffer

	// we sort the key alphabetically to have a reproducible graph output
	sortedEKeys := getSortedTransitionKeys(fm)
	sortedStateKeys := getSortedStates(fm)

	buf.WriteString(`digraph fsm {\n`)

	for _, k := range sortedEKeys {
		v := fm.Table[k]
		for _, transition := range v {
			buf.WriteString(fmt.Sprintf(`    "%s" -> "%s" [ label = "%s" ];`, k.State, transition.NewState, k.Event))
			buf.WriteString("\n")
		}
	}
	buf.WriteString("\n")

	for _, k := range sortedStateKeys {
		buf.WriteString(fmt.Sprintf(`    "%s";`, k))
		buf.WriteString("\n")
	}

	buf.WriteString("}\n")

	return buf.String()
}