mpagenko | bb47bc2 | 2021-04-20 13:29:09 +0000 | [diff] [blame] | 1 | package fsm |
| 2 | |
| 3 | import ( |
| 4 | "bytes" |
| 5 | "fmt" |
| 6 | ) |
| 7 | |
| 8 | const highlightingColor = "#00AA00" |
| 9 | |
| 10 | // MermaidDiagramType the type of the mermaid diagram type |
| 11 | type MermaidDiagramType string |
| 12 | |
| 13 | const ( |
| 14 | // FlowChart the diagram type for output in flowchart style (https://mermaid-js.github.io/mermaid/#/flowchart) (including current state) |
| 15 | FlowChart MermaidDiagramType = "flowChart" |
| 16 | // StateDiagram the diagram type for output in stateDiagram style (https://mermaid-js.github.io/mermaid/#/stateDiagram) |
| 17 | StateDiagram MermaidDiagramType = "stateDiagram" |
| 18 | ) |
| 19 | |
| 20 | // VisualizeForMermaidWithGraphType outputs a visualization of a FSM in Mermaid format as specified by the graphType. |
| 21 | func VisualizeForMermaidWithGraphType(fsm *FSM, graphType MermaidDiagramType) (string, error) { |
| 22 | switch graphType { |
| 23 | case FlowChart: |
| 24 | return visualizeForMermaidAsFlowChart(fsm), nil |
| 25 | case StateDiagram: |
| 26 | return visualizeForMermaidAsStateDiagram(fsm), nil |
| 27 | default: |
| 28 | return "", fmt.Errorf("unknown MermaidDiagramType: %s", graphType) |
| 29 | } |
| 30 | } |
| 31 | |
| 32 | func visualizeForMermaidAsStateDiagram(fsm *FSM) string { |
| 33 | var buf bytes.Buffer |
| 34 | |
| 35 | sortedTransitionKeys := getSortedTransitionKeys(fsm.transitions) |
| 36 | |
| 37 | buf.WriteString("stateDiagram\n") |
| 38 | buf.WriteString(fmt.Sprintln(` [*] -->`, fsm.current)) |
| 39 | |
| 40 | for _, k := range sortedTransitionKeys { |
| 41 | v := fsm.transitions[k] |
| 42 | buf.WriteString(fmt.Sprintf(` %s --> %s: %s`, k.src, v, k.event)) |
| 43 | buf.WriteString("\n") |
| 44 | } |
| 45 | |
| 46 | return buf.String() |
| 47 | } |
| 48 | |
| 49 | // visualizeForMermaidAsFlowChart outputs a visualization of a FSM in Mermaid format (including highlighting of current state). |
| 50 | func visualizeForMermaidAsFlowChart(fsm *FSM) string { |
| 51 | var buf bytes.Buffer |
| 52 | |
| 53 | sortedTransitionKeys := getSortedTransitionKeys(fsm.transitions) |
| 54 | sortedStates, statesToIDMap := getSortedStates(fsm.transitions) |
| 55 | |
| 56 | writeFlowChartGraphType(&buf) |
| 57 | writeFlowChartStates(&buf, sortedStates, statesToIDMap) |
| 58 | writeFlowChartTransitions(&buf, fsm.transitions, sortedTransitionKeys, statesToIDMap) |
| 59 | writeFlowChartHightlightCurrent(&buf, fsm.current, statesToIDMap) |
| 60 | |
| 61 | return buf.String() |
| 62 | } |
| 63 | |
| 64 | func writeFlowChartGraphType(buf *bytes.Buffer) { |
| 65 | buf.WriteString("graph LR\n") |
| 66 | } |
| 67 | |
| 68 | func writeFlowChartStates(buf *bytes.Buffer, sortedStates []string, statesToIDMap map[string]string) { |
| 69 | for _, state := range sortedStates { |
| 70 | buf.WriteString(fmt.Sprintf(` %s[%s]`, statesToIDMap[state], state)) |
| 71 | buf.WriteString("\n") |
| 72 | } |
| 73 | |
| 74 | buf.WriteString("\n") |
| 75 | } |
| 76 | |
| 77 | func writeFlowChartTransitions(buf *bytes.Buffer, transitions map[eKey]string, sortedTransitionKeys []eKey, statesToIDMap map[string]string) { |
| 78 | for _, transition := range sortedTransitionKeys { |
| 79 | target := transitions[transition] |
| 80 | buf.WriteString(fmt.Sprintf(` %s --> |%s| %s`, statesToIDMap[transition.src], transition.event, statesToIDMap[target])) |
| 81 | buf.WriteString("\n") |
| 82 | } |
| 83 | buf.WriteString("\n") |
| 84 | } |
| 85 | |
| 86 | func writeFlowChartHightlightCurrent(buf *bytes.Buffer, current string, statesToIDMap map[string]string) { |
| 87 | buf.WriteString(fmt.Sprintf(` style %s fill:%s`, statesToIDMap[current], highlightingColor)) |
| 88 | buf.WriteString("\n") |
| 89 | } |