├── .github ├── env │ └── action.yml └── workflows │ ├── tag.yml │ ├── go.yml │ └── staticcheck-action.yml ├── .gitignore ├── go.mod ├── perf_test.go ├── state_action.go ├── LICENSE ├── go.sum ├── GFSM_UML.md ├── gfsm.go ├── sm_builder.go ├── gfsm_test.go ├── README.md ├── examples └── two-phase-commit │ └── main.go └── cmd └── gfsm_uml └── main.go /.github/env/action.yml: -------------------------------------------------------------------------------- 1 | name: prepare-env 2 | 3 | runs: 4 | using: composite 5 | steps: 6 | - name: Setup Go 7 | uses: actions/setup-go@v5 8 | with: 9 | go-version: 1.23.x 10 | - name: Generate 11 | shell: bash 12 | run: | 13 | go install tool 14 | 15 | go generate ./... 16 | -------------------------------------------------------------------------------- /.github/workflows/tag.yml: -------------------------------------------------------------------------------- 1 | name: CI 2 | on: 3 | push: 4 | tags: 5 | - 'v[0-9]+.[0-9]+.[0-9]+' 6 | 7 | jobs: 8 | publish: 9 | name: "pkg.go.dev publishing" 10 | runs-on: ubuntu-latest 11 | steps: 12 | - name: Publishing new version 13 | run: | 14 | curl https://sum.golang.org/lookup/github.com/astavonin/gfsm@${{ github.ref_name }} 15 | -------------------------------------------------------------------------------- /.github/workflows/go.yml: -------------------------------------------------------------------------------- 1 | name: main 2 | 3 | on: 4 | push: 5 | branches: [ "main" ] 6 | pull_request: 7 | branches: [ "main" ] 8 | 9 | jobs: 10 | build: 11 | runs-on: ubuntu-latest 12 | steps: 13 | - uses: actions/checkout@v4 14 | - name: prepare-env 15 | uses: ./.github/env 16 | - name: Build 17 | run: go build -v ./... 18 | - name: Test 19 | run: go test -v ./... 20 | -------------------------------------------------------------------------------- /.github/workflows/staticcheck-action.yml: -------------------------------------------------------------------------------- 1 | name: "CI: staticcheck" 2 | on: ["push", "pull_request"] 3 | 4 | jobs: 5 | ci: 6 | name: "Run CI" 7 | runs-on: ubuntu-latest 8 | steps: 9 | - uses: actions/checkout@v2 10 | with: 11 | fetch-depth: 1 12 | - name: prepare-env 13 | uses: ./.github/env 14 | - uses: dominikh/staticcheck-action@v1 15 | with: 16 | version: "latest" 17 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # If you prefer the allow list template instead of the deny list, see community template: 2 | # https://github.com/github/gitignore/blob/main/community/Golang/Go.AllowList.gitignore 3 | # 4 | # Binaries for programs and plugins 5 | *.exe 6 | *.exe~ 7 | *.dll 8 | *.so 9 | *.dylib 10 | 11 | # Test binary, built with `go test -c` 12 | *.test 13 | 14 | # Output of the go coverage tool, specifically when used with LiteIDE 15 | *.out 16 | 17 | # Dependency directories (remove the comment below to include it) 18 | # vendor/ 19 | 20 | # Go workspace file 21 | go.work 22 | 23 | # IDEs 24 | .idea 25 | 26 | # stringer generated files 27 | *_string.go 28 | 29 | # OS-specific 30 | .DS_Store 31 | -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module github.com/astavonin/gfsm 2 | 3 | go 1.24 4 | 5 | tool ( 6 | github.com/astavonin/gfsm/cmd/gfsm_uml 7 | golang.org/x/tools/cmd/stringer 8 | ) 9 | 10 | require ( 11 | github.com/google/uuid v1.6.0 12 | github.com/stretchr/testify v1.9.0 13 | ) 14 | 15 | require ( 16 | github.com/davecgh/go-spew v1.1.1 // indirect 17 | github.com/kr/pretty v0.3.1 // indirect 18 | github.com/pmezard/go-difflib v1.0.0 // indirect 19 | github.com/rogpeppe/go-internal v1.12.0 // indirect 20 | golang.org/x/mod v0.23.0 // indirect 21 | golang.org/x/sync v0.11.0 // indirect 22 | golang.org/x/tools v0.30.0 // indirect 23 | gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15 // indirect 24 | gopkg.in/yaml.v3 v3.0.1 // indirect 25 | ) 26 | -------------------------------------------------------------------------------- /perf_test.go: -------------------------------------------------------------------------------- 1 | package gfsm 2 | 3 | import "testing" 4 | 5 | func BenchmarkSliceAccess(b *testing.B) { 6 | for i := 0; i < b.N; i++ { 7 | _ = sliceTest() 8 | } 9 | } 10 | 11 | func BenchmarkMapAccess(b *testing.B) { 12 | for i := 0; i < b.N; i++ { 13 | _ = mapTest() 14 | } 15 | } 16 | 17 | func sliceTest() bool { 18 | transitions := map[StartStopSM]struct{}{ 19 | Stop: {}, 20 | Start: {}, 21 | InProgress: {}, 22 | } 23 | _, found := transitions[Stop] 24 | return found 25 | } 26 | 27 | func mapTest() bool { 28 | transitions := []StartStopSM{ 29 | Start, Stop, InProgress, 30 | } 31 | found := false 32 | for _, transition := range transitions { 33 | if Stop == transition { 34 | found = true 35 | break 36 | } 37 | } 38 | return found 39 | } 40 | -------------------------------------------------------------------------------- /state_action.go: -------------------------------------------------------------------------------- 1 | package gfsm 2 | 3 | // EventContext is an abstraction that represent any data that user passes to the current state for execution 4 | // in StateMachineHandler.ProcessEvent(...) call. The data will be forwarded as StateAction.Execute(...) argument. 5 | type EventContext interface{} 6 | 7 | // StateMachineContext is an abstraction that represent any data that user passes to the state machine. The data will 8 | // be forwarded as StateAction.OnEnter amd OnExit arguments. 9 | type StateMachineContext interface{} 10 | 11 | // StateAction is the interface which each state must implement. 12 | type StateAction[StateIdentifier comparable] interface { 13 | // OnEnter will be called once on the state entering. 14 | OnEnter(smCtx StateMachineContext) 15 | // OnExit will be called once on the state exiting. 16 | OnExit(smCtx StateMachineContext) 17 | // Execute is the call that state machine routes to the current state from StateMachineHandler.ProcessEvent(...) 18 | Execute(smCtx StateMachineContext, eventCtx EventContext) StateIdentifier 19 | } 20 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | BSD 3-Clause License 2 | 3 | Copyright (c) 2024, Alexander Stavonin 4 | 5 | Redistribution and use in source and binary forms, with or without 6 | modification, are permitted provided that the following conditions are met: 7 | 8 | 1. Redistributions of source code must retain the above copyright notice, this 9 | list of conditions and the following disclaimer. 10 | 11 | 2. Redistributions in binary form must reproduce the above copyright notice, 12 | this list of conditions and the following disclaimer in the documentation 13 | and/or other materials provided with the distribution. 14 | 15 | 3. Neither the name of the copyright holder nor the names of its 16 | contributors may be used to endorse or promote products derived from 17 | this software without specific prior written permission. 18 | 19 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 20 | AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 21 | IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 22 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 23 | FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 24 | DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 25 | SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 26 | CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 27 | OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 28 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 29 | -------------------------------------------------------------------------------- /go.sum: -------------------------------------------------------------------------------- 1 | github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= 2 | github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= 3 | github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= 4 | github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI= 5 | github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= 6 | github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= 7 | github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= 8 | github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= 9 | github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= 10 | github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= 11 | github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= 12 | github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA= 13 | github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= 14 | github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= 15 | github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs= 16 | github.com/rogpeppe/go-internal v1.12.0 h1:exVL4IDcn6na9z1rAb56Vxr+CgyK3nn3O+epU5NdKM8= 17 | github.com/rogpeppe/go-internal v1.12.0/go.mod h1:E+RYuTGaKKdloAfM02xzb0FW3Paa99yedzYV+kq4uf4= 18 | github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg= 19 | github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= 20 | golang.org/x/mod v0.23.0 h1:Zb7khfcRGKk+kqfxFaP5tZqCnDZMjC5VtUBs87Hr6QM= 21 | golang.org/x/mod v0.23.0/go.mod h1:6SkKJ3Xj0I0BrPOZoBy3bdMptDDU9oJrpohJ3eWZ1fY= 22 | golang.org/x/sync v0.11.0 h1:GGz8+XQP4FvTTrjZPzNKTMFtSXH80RAzG+5ghFPgK9w= 23 | golang.org/x/sync v0.11.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk= 24 | golang.org/x/tools v0.30.0 h1:BgcpHewrV5AUp2G9MebG4XPFI1E2W41zU1SaqVA9vJY= 25 | golang.org/x/tools v0.30.0/go.mod h1:c347cR/OJfw5TI+GfX7RUPNMdDRRbjvYTS0jPyvsVtY= 26 | gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= 27 | gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15 h1:YR8cESwS4TdDjEe65xsg0ogRM/Nc3DYOhEAlW+xobZo= 28 | gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= 29 | gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= 30 | gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= 31 | -------------------------------------------------------------------------------- /GFSM_UML.md: -------------------------------------------------------------------------------- 1 | # gfsm_uml 2 | 3 | **gfsm_uml** is a code generator tool that extracts state machine definitions from your Go source code and generates corresponding state diagrams. It supports multiple state machine definitions per file using a custom naming function (`SetSMName`), and it can output diagrams in either Mermaid or PlantUML format. 4 | 5 | ## Features 6 | 7 | - **Automatic Diagram Generation:** Scans your code for state machine builder chains and generates diagrams automatically. 8 | - **Custom Naming:** Use `SetSMName` in your builder chain to uniquely identify each state machine. 9 | - **Multiple Formats:** Output diagrams in either `Mermaid` or `PlantUML` format. 10 | - **Integration with go generate:** Easily integrate with your build process using go generate. 11 | 12 | ## Installation 13 | 14 | If you're working from within your module, you can install the tool locally. Since `gfsm_uml` is part of the `gfsm` repository (with the command-line app located in `cmd/gfsm_uml`), install it by running: 15 | 16 | ```bash 17 | go install github.com/astavonin/gfsm/cmd/gfsm_uml@latest 18 | ``` 19 | 20 | Make sure that your `$GOPATH/bin` (or your module-aware binary install location) is in your PATH so that you can invoke `gfsm_uml` directly. 21 | 22 | ## Usage 23 | 24 | ### Step 1. Annotate Your Code 25 | 26 | In the Go source file that contains your state machine definitions, add a go:generate directive at the top. For example: 27 | 28 | ```go 29 | //go:generate gfsm_uml -format=plantuml 30 | ``` 31 | 32 | Then, within your builder chain, use `SetSMName` to name your state machine. For example: 33 | 34 | ```go 35 | return NewBuilder[StartStopSM](). 36 | SetSMName("FooSM"). // Custom annotation for naming the state machine 37 | SetDefaultState(Start). 38 | RegisterState(Start, &StartState{}, []StartStopSM{Stop, InProgress}). 39 | RegisterState(Stop, &StopState{}, []StartStopSM{Start}). 40 | RegisterState(InProgress, &InProgressState{}, []StartStopSM{Stop}). 41 | Build() 42 | ``` 43 | 44 | ### Step 2. Run the Generator 45 | 46 | From the root of your project, run: 47 | 48 | ```bash 49 | go generate ./... 50 | ``` 51 | 52 | When invoked via `go generate`, `gfsm_uml` will read the file specified by the `GOFILE` environment variable, extract the builder chain, and generate a state diagram. If you specified PlantUML as the format (with `-format=plantuml`), the diagram will be written to a file named `FooSM.uml` (based on the name provided by `SetSMName`). 53 | 54 | ### Command Line Options 55 | 56 | The tool supports the following flag: 57 | 58 | - **`-format`**: Specifies the output diagram format. Valid options are: 59 | - `mermaid` (default) 60 | - `plantuml` 61 | 62 | For example, to generate a Mermaid diagram from a file: 63 | 64 | ```bash 65 | gfsm_uml -format=mermaid 66 | ``` 67 | 68 | ## Example 69 | 70 | Consider the following builder chain in your Go source file: 71 | 72 | ```go 73 | return NewBuilder[StartStopSM](). 74 | SetSMName("FooSM"). 75 | SetDefaultState(Start). 76 | RegisterState(Start, &StartState{}, []StartStopSM{Stop, InProgress}). 77 | RegisterState(Stop, &StopState{}, []StartStopSM{Start}). 78 | RegisterState(InProgress, &InProgressState{}, []StartStopSM{Stop}). 79 | Build() 80 | ``` 81 | 82 | When you run `go generate ./...`, the tool processes the file, extracts the state machine named `"FooSM"`, and generates a diagram file named `FooSM.uml` (if using PlantUML) or `FooSM.mermaid` (if using Mermaid). 83 | 84 | ## Troubleshooting 85 | 86 | - **No Diagram Generated:** 87 | Verify that your file contains valid state machine definitions (calls to both `SetSMName` and `RegisterState`). 88 | Run `go generate -x ./...` to see detailed output and debug any issues. 89 | 90 | - **Installation Issues:** 91 | Ensure that `gfsm_uml` is installed and available in your PATH. Use `go install` as described above. 92 | -------------------------------------------------------------------------------- /gfsm.go: -------------------------------------------------------------------------------- 1 | // Package gfsm implements basic state machine functionality. 2 | package gfsm 3 | 4 | import "fmt" 5 | 6 | var ( 7 | ErrNoValidTransition = fmt.Errorf("no valid transition") 8 | ) 9 | 10 | // Transitions represents all available transitions from the state. 11 | type Transitions[StateIdentifier comparable] map[StateIdentifier]struct{} 12 | 13 | // The state is a struct that is defined by the StateActions it can take, the Transitions it can make 14 | type state[StateIdentifier comparable] struct { 15 | action StateAction[StateIdentifier] 16 | transitions Transitions[StateIdentifier] 17 | } 18 | 19 | // StatesMap represent full state machine transactions and allows to verify path from any state to another. 20 | // It is a map of StateIdentifiers to state 21 | type StatesMap[StateIdentifier comparable] map[StateIdentifier]state[StateIdentifier] 22 | 23 | // StateMachineHandler is the main state machine interface. All manipulation with the state machine object shall 24 | // be performed using this interface. 25 | type StateMachineHandler[StateIdentifier comparable] interface { 26 | // Start is the first function that user MUST call before any further interactions with the state machine. 27 | // On Start call, state machine will switch to the defined default state, which must be specified during state 28 | // machine creation using StateMachineBuilder.SetDefaultState(...) call 29 | Start() 30 | 31 | // Stop call shutdowns the state machine. Any further State or ProcessEvent are not permitted on stopped 32 | // state machine. 33 | Stop() 34 | 35 | // State returns current state machine state 36 | State() StateIdentifier 37 | 38 | // ProcessEvent pass data to the sate machine for processing. The data will be forwarded to StateAction.Execute 39 | // method of the current state. If the event processing will lead to unexpected transaction, ProcessEvent call 40 | // will return ErrNoValidTransition error 41 | ProcessEvent(eventCtx EventContext) error 42 | 43 | // Reset will return the statemachine to its default state 44 | Reset() 45 | } 46 | 47 | type stateMachine[StateIdentifier comparable] struct { 48 | currentStateID StateIdentifier 49 | defaultStateID StateIdentifier 50 | states StatesMap[StateIdentifier] 51 | smCtx StateMachineContext 52 | name string 53 | } 54 | 55 | func (s *stateMachine[StateIdentifier]) Start() { 56 | state := s.states[s.currentStateID] 57 | state.action.OnEnter(s.smCtx) 58 | } 59 | 60 | func (s *stateMachine[StateIdentifier]) Stop() { 61 | state := s.states[s.currentStateID] 62 | state.action.OnExit(s.smCtx) 63 | } 64 | 65 | func (s *stateMachine[StateIdentifier]) State() StateIdentifier { 66 | return s.currentStateID 67 | } 68 | 69 | func (s *stateMachine[StateIdentifier]) ProcessEvent(eventCtx EventContext) error { 70 | currentState := s.states[s.currentStateID] 71 | nextStateID := currentState.action.Execute(s.smCtx, eventCtx) 72 | // do not need to change state 73 | if nextStateID == s.currentStateID { 74 | return nil 75 | } 76 | 77 | return s.ChangeState(nextStateID) 78 | } 79 | 80 | func (s *stateMachine[StateIdentifier]) ChangeState(nextStateID StateIdentifier) error { 81 | currentState := s.states[s.currentStateID] 82 | _, canSwitch := currentState.transitions[nextStateID] 83 | if !canSwitch { 84 | return fmt.Errorf("cannot switch from %v to %v: %w", s.currentStateID, nextStateID, ErrNoValidTransition) 85 | } 86 | s.currentStateID = nextStateID 87 | currentState.action.OnExit(s.smCtx) 88 | nextState := s.states[nextStateID] 89 | nextState.action.OnEnter(s.smCtx) 90 | 91 | return nil 92 | } 93 | 94 | func (s *stateMachine[StateIdentifier]) Reset() { 95 | currentState := s.states[s.currentStateID] 96 | currentState.action.OnExit(s.smCtx) 97 | 98 | defaultState := s.states[s.defaultStateID] 99 | defaultState.action.OnEnter(s.smCtx) 100 | 101 | s.currentStateID = s.defaultStateID 102 | } 103 | -------------------------------------------------------------------------------- /sm_builder.go: -------------------------------------------------------------------------------- 1 | package gfsm 2 | 3 | import "fmt" 4 | 5 | // StateMachineBuilder interface provides access to a builder that simplifies state machine creation. Builder usage is optional, 6 | // and state machine object can be created manually if needed. 7 | // Refer to newSmManual test for manual state machine creation or newSmWithBuilder as the alternative approach with builder. 8 | type StateMachineBuilder[StateIdentifier comparable] interface { 9 | // SetSMName provides optional name for the SM. Primary uses for debugging proposes in cases when an app has more than one state machine. 10 | SetSMName(smName string) StateMachineBuilder[StateIdentifier] 11 | // RegisterState call register one more state referenced by stateID with list of all valid transactions listed in transitions 12 | // and handler (action) into the state machine. 13 | RegisterState(stateID StateIdentifier, action StateAction[StateIdentifier], transitions []StateIdentifier) StateMachineBuilder[StateIdentifier] 14 | // SetDefaultState tells which state is the default for the state machine. Each state machine must have a default state. 15 | // On StateMachineHandler.Start() call, state machine will switch to the defined default state. 16 | SetDefaultState(stateID StateIdentifier) StateMachineBuilder[StateIdentifier] 17 | // SetSmContext is an optional call that allow to pass any context that is unique and persistent (but mutable) for each state machine. 18 | SetSmContext(ctx StateMachineContext) StateMachineBuilder[StateIdentifier] 19 | 20 | // Build is the final call that aggregates all the data from previous calls and creates new state machine. 21 | Build() StateMachineHandler[StateIdentifier] 22 | } 23 | 24 | // NewBuilder function generates StateMachineBuilder which simplifies state machine creation process. 25 | func NewBuilder[StateIdentifier comparable]() StateMachineBuilder[StateIdentifier] { 26 | return &stateMachineBuilder[StateIdentifier]{ 27 | hasState: false, 28 | sm: &stateMachine[StateIdentifier]{ 29 | states: StatesMap[StateIdentifier]{}, 30 | }, 31 | } 32 | } 33 | 34 | // stateMachineBuilder[StateIdentifier comparable] is an implementation for 35 | // the StateMachineBuilder[StateIdentifier comparable] interface 36 | type stateMachineBuilder[StateIdentifier comparable] struct { 37 | hasState bool 38 | hasDefaultState bool 39 | 40 | sm *stateMachine[StateIdentifier] 41 | } 42 | 43 | func (s *stateMachineBuilder[StateIdentifier]) Build() StateMachineHandler[StateIdentifier] { 44 | if !s.hasState || !s.hasDefaultState { 45 | panic("state machine is not properly initialised yet") 46 | } 47 | return s.sm 48 | } 49 | 50 | func (s *stateMachineBuilder[StateIdentifier]) RegisterState( 51 | stateID StateIdentifier, 52 | action StateAction[StateIdentifier], 53 | transitions []StateIdentifier) StateMachineBuilder[StateIdentifier] { 54 | 55 | _, ok := s.sm.states[stateID] 56 | if ok { 57 | panic(fmt.Sprintf("state %v is already registered", stateID)) 58 | } 59 | 60 | s.sm.states[stateID] = state[StateIdentifier]{ 61 | action: action, 62 | transitions: makeTransitions(transitions), 63 | } 64 | s.hasState = true 65 | 66 | return s 67 | } 68 | 69 | func makeTransitions[StateIdentifier comparable](transitions []StateIdentifier) Transitions[StateIdentifier] { 70 | trs := Transitions[StateIdentifier]{} 71 | for _, transition := range transitions { 72 | trs[transition] = struct{}{} 73 | } 74 | return trs 75 | } 76 | 77 | func (s *stateMachineBuilder[StateIdentifier]) SetDefaultState(stateID StateIdentifier) StateMachineBuilder[StateIdentifier] { 78 | s.sm.currentStateID = stateID 79 | s.sm.defaultStateID = stateID 80 | s.hasDefaultState = true 81 | 82 | return s 83 | } 84 | 85 | func (s *stateMachineBuilder[StateIdentifier]) SetSmContext(ctx StateMachineContext) StateMachineBuilder[StateIdentifier] { 86 | s.sm.smCtx = ctx 87 | 88 | return s 89 | } 90 | 91 | func (s *stateMachineBuilder[StateIdentifier]) SetSMName(smName string) StateMachineBuilder[StateIdentifier] { 92 | s.sm.name = smName 93 | 94 | return s 95 | } 96 | -------------------------------------------------------------------------------- /gfsm_test.go: -------------------------------------------------------------------------------- 1 | package gfsm 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/google/uuid" 7 | "github.com/stretchr/testify/assert" 8 | ) 9 | 10 | type StartStopSM int 11 | 12 | const ( 13 | Start StartStopSM = iota 14 | Stop 15 | InProgress 16 | ) 17 | 18 | type StartState struct { 19 | t *testing.T 20 | } 21 | 22 | func (s *StartState) OnEnter(smCtx StateMachineContext) { 23 | s.t = smCtx.(*aContext).t 24 | } 25 | 26 | func (s *StartState) OnExit(_ StateMachineContext) { 27 | } 28 | 29 | func (s *StartState) Execute(_ StateMachineContext, eventCtx EventContext) StartStopSM { 30 | assert.NotNil(s.t, eventCtx) 31 | _, ok := eventCtx.(StartData) 32 | assert.True(s.t, ok) 33 | 34 | return InProgress 35 | } 36 | 37 | type StopState struct { 38 | t *testing.T 39 | } 40 | 41 | func (s *StopState) OnEnter(smCtx StateMachineContext) { 42 | s.t = smCtx.(*aContext).t 43 | } 44 | 45 | func (s *StopState) OnExit(_ StateMachineContext) { 46 | } 47 | 48 | func (s *StopState) Execute(_ StateMachineContext, _ EventContext) StartStopSM { 49 | return Start 50 | } 51 | 52 | type InProgressState struct { 53 | t *testing.T 54 | } 55 | 56 | func (s *InProgressState) OnEnter(smCtx StateMachineContext) { 57 | s.t = smCtx.(*aContext).t 58 | } 59 | 60 | func (s *InProgressState) OnExit(_ StateMachineContext) { 61 | } 62 | 63 | func (s *InProgressState) Execute(_ StateMachineContext, eventCtx EventContext) StartStopSM { 64 | assert.NotNil(s.t, eventCtx) 65 | _, ok := eventCtx.(InProgressData) 66 | assert.True(s.t, ok) 67 | 68 | return InProgress 69 | } 70 | 71 | type StartData struct { 72 | id uuid.UUID 73 | } 74 | 75 | type InProgressData struct { 76 | } 77 | 78 | type aContext struct { 79 | t *testing.T 80 | } 81 | 82 | func newSmManual(t *testing.T) StateMachineHandler[StartStopSM] { 83 | return &stateMachine[StartStopSM]{ 84 | currentStateID: Start, 85 | states: StatesMap[StartStopSM]{ 86 | Start: state[StartStopSM]{ 87 | action: &StartState{}, 88 | transitions: Transitions[StartStopSM]{ 89 | Stop: struct{}{}, 90 | InProgress: struct{}{}, 91 | }, 92 | }, 93 | Stop: state[StartStopSM]{ 94 | action: &StopState{}, 95 | transitions: Transitions[StartStopSM]{ 96 | Start: struct{}{}, 97 | }, 98 | }, 99 | InProgress: state[StartStopSM]{ 100 | action: &InProgressState{}, 101 | transitions: Transitions[StartStopSM]{ 102 | Stop: struct{}{}, 103 | }, 104 | }, 105 | }, 106 | smCtx: &aContext{t: t}, 107 | } 108 | } 109 | 110 | func newSmWithBuilder(t *testing.T) StateMachineHandler[StartStopSM] { 111 | return NewBuilder[StartStopSM](). 112 | SetDefaultState(Start). 113 | SetSmContext(&aContext{t: t}). 114 | RegisterState(Start, &StartState{}, []StartStopSM{Stop, InProgress}). 115 | RegisterState(Stop, &StopState{}, []StartStopSM{Start}). 116 | RegisterState(InProgress, &InProgressState{}, []StartStopSM{Stop}). 117 | Build() 118 | } 119 | 120 | func TestStateMachine(t *testing.T) { 121 | var tests = []struct { 122 | sm StateMachineHandler[StartStopSM] 123 | testName string 124 | }{ 125 | {newSmManual(t), "manual SM creation"}, 126 | {newSmWithBuilder(t), "builder SM creation"}, 127 | } 128 | 129 | for _, test := range tests { 130 | t.Run(test.testName, func(t *testing.T) { 131 | sm := test.sm 132 | 133 | sm.Start() 134 | 135 | assert.Equal(t, sm.State(), Start) 136 | err := sm.ProcessEvent(StartData{id: uuid.New()}) 137 | assert.NoError(t, err) 138 | assert.Equal(t, sm.State(), InProgress) 139 | err = sm.ProcessEvent(InProgressData{}) 140 | assert.NoError(t, err) 141 | assert.Equal(t, sm.State(), InProgress) 142 | 143 | sm.Stop() 144 | }) 145 | } 146 | } 147 | 148 | func TestDoubleStateCreation(t *testing.T) { 149 | builder := NewBuilder[StartStopSM](). 150 | RegisterState(Stop, &StopState{}, []StartStopSM{Stop}) 151 | 152 | assert.Panics(t, func() { 153 | builder.RegisterState(Stop, &StopState{}, []StartStopSM{Stop}) 154 | }) 155 | } 156 | 157 | func TestResetStatMachine(t *testing.T) { 158 | sm := newSmWithBuilder(t) 159 | 160 | sm.Start() 161 | err := sm.ProcessEvent(StartData{id: uuid.New()}) 162 | assert.NoError(t, err) 163 | assert.Equal(t, sm.State(), InProgress) 164 | 165 | sm.Reset() 166 | assert.Equal(t, sm.State(), Start) 167 | 168 | sm.Stop() 169 | } 170 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # GFSM - simple and fast Finite State Machine for Go 2 | 3 | [![PkgGoDev](https://pkg.go.dev/badge/github.com/astavonin/gfsm)](https://pkg.go.dev/github.com/astavonin/gfsm) 4 | ![build-and-test](https://github.com/astavonin/gfsm/actions/workflows/go.yml/badge.svg) 5 | ![CI](https://github.com/astavonin/gfsm/actions/workflows/staticcheck-action.yml/badge.svg) 6 | 7 | The GFSM library provides a simple and **fast** implementation of a Finite State Machine (FSM) for Go. The library adopts C++-style approaches with a primary focus on speed and simplicity, which is the main difference from alternative Go FSM implementations like [looplab/fsm](https://github.com/looplab/fsm). You can also create a visual representation of an FSM using the [gfsm_uml](./GFSM_UML.md) generator. 8 | 9 | # How To use GFSM 10 | 11 | [Two Phase Commit protocol](https://en.wikipedia.org/wiki/Two-phase_commit_protocol) (TPC protocol) is an excellent example of Finite State Machine use case. 12 | 13 | *NOTE: Full example is in [examples/two-phase-commit/main.go]() folder.* 14 | 15 | Having the following State Machine: 16 | 17 | ```mermaid 18 | graph TD; 19 | Init-->Wait; 20 | Wait-->Abort; 21 | Wait-->Commit; 22 | ``` 23 | we should: 24 | 1. Enumerate all possible states 25 | 2. Describe each states 26 | 3. Describe a state machine 27 | 4. Run the state machine 28 | 29 | ## Enumerate all possible states 30 | 31 | The GFSM library can consume any `comparable` interface as a states enumerator, but the recommended approach is using `int`-based types. For example, we can define type `State` with all TPC protocol-related states. 32 | 33 | ```go 34 | type State int 35 | const ( 36 | Init State = iota 37 | Wait 38 | Abort 39 | Commit 40 | ) 41 | ``` 42 | 43 | ## Describe each states 44 | 45 | Each State is represented by its unique `stateID`, `action` handler, and list of possible `transitions`. Each state can have a unique handler that supports `StateAction` interface: 46 | 47 | ```go 48 | type StateAction[StateIdentifier comparable] interface { 49 | OnEnter(smCtx StateMachineContext) 50 | OnExit(smCtx StateMachineContext) 51 | Execute(smCtx StateMachineContext, eventCtx EventContext) StateIdentifier 52 | } 53 | ``` 54 | Where `OnEnter` and `OnExit` will be called once on the state entering or exiting respectfully, and `Execute` is the call that the state machine routes to the current state from `StateMachineHandler.ProcessEvent(...)`. 55 | 56 | ## Describe a state machine 57 | 58 | Any State Machine can have its unique context. For the TPC protocol, it's important to have information about the voting process, like the expected voter count and committed ID. 59 | ```go 60 | type coordinatorContext struct { 61 | commitID string 62 | partCnt int 63 | } 64 | ``` 65 | `StateMachineContext` is stored inside the constructed `StateMachineHandler` object and will be passed to each state on any calls (`OnEnter`, `OnExit`, and `Execute`). 66 | 67 | To create a new state machine, `StateMachineBuilder` should be used. To build a TPC protocol state machine, you can use the following approach: 68 | 69 | ```go 70 | sm := gfsm.NewBuilder[State](). 71 | SetDefaultState(Init). 72 | SetSmContext(&coordinatorContext{partCnt: 3}). 73 | RegisterState(Init, &initState{}, []State{Wait}). 74 | RegisterState(Wait, &waitState{}, []State{Abort, Commit}). 75 | RegisterState(Abort, &responseState{ 76 | keepResp: Abort, 77 | }, []State{Init}). 78 | RegisterState(Commit, &responseState{ 79 | keepResp: Commit, 80 | }, []State{Init}). 81 | Build() 82 | ``` 83 | ## Run the state machine 84 | 85 | After the construction, the state machine shall be executed and terminated on exit: 86 | 87 | ```go 88 | sm.Start() 89 | defer sm.Stop() 90 | ``` 91 | 92 | To process an event, the state machine object implements a `ProcessEvent(eventCtx EventContext) error` call, where `eventCtx` can be any data that the state will process internally as primary input data. 93 | 94 | ```go 95 | err := sm.ProcessEvent(commitRequest{"commit_1"}) 96 | ``` 97 | 98 | During the event processing, the state machine will call `Execute` with the passed `EventContext` and `StateMachineContext`. Based on these data, the state can decide either to keep the current state (`Init` in the example below) or make a switch (`Wait`) by returning the new expected state. 99 | 100 | ```go 101 | func (s *initState) Execute(smCtx gfsm.StateMachineContext, eventCtx gfsm.EventContext) State { 102 | cCtx := smCtx.(*coordinatorContext) 103 | req, ok := eventCtx.(commitRequest) 104 | if !ok { 105 | // ... 106 | return Init 107 | } 108 | // ... 109 | return Wait 110 | } 111 | ``` 112 | 113 | -------------------------------------------------------------------------------- /examples/two-phase-commit/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "fmt" 5 | 6 | gfsm2 "github.com/astavonin/gfsm" 7 | ) 8 | 9 | //go:generate stringer -type=State 10 | type State int 11 | 12 | const ( 13 | Init State = iota 14 | Wait 15 | Abort 16 | Commit 17 | ) 18 | 19 | type coordinatorContext struct { 20 | commitID string 21 | partCnt int 22 | } 23 | 24 | // ========= Init state handler ========= 25 | type initState struct { 26 | } 27 | 28 | type commitRequest struct { 29 | commitID string 30 | } 31 | 32 | func (s *initState) OnEnter(_ gfsm2.StateMachineContext) { 33 | } 34 | 35 | func (s *initState) OnExit(_ gfsm2.StateMachineContext) { 36 | } 37 | 38 | func (s *initState) Execute(smCtx gfsm2.StateMachineContext, eventCtx gfsm2.EventContext) State { 39 | cCtx := smCtx.(*coordinatorContext) 40 | req, ok := eventCtx.(commitRequest) 41 | if !ok { 42 | fmt.Printf("invalid request\n") 43 | // nothing to process, keeping state 44 | return Init 45 | } 46 | 47 | fmt.Printf("got request %s\n", req.commitID) 48 | // forwarding request to participants and switching the state 49 | //for i := 0; i < cCtx.votesCnt; i++ { 50 | // ... 51 | //} 52 | // and saving commit ID as the state machine context 53 | cCtx.commitID = req.commitID 54 | 55 | return Wait 56 | } 57 | 58 | // ========= Wait state handler ========= 59 | 60 | type waitState struct { 61 | votesCnt int 62 | } 63 | 64 | type commitVote struct { 65 | commit bool 66 | } 67 | 68 | func (s *waitState) OnEnter(_ gfsm2.StateMachineContext) { 69 | s.votesCnt = 0 70 | } 71 | 72 | func (s *waitState) OnExit(_ gfsm2.StateMachineContext) { 73 | } 74 | 75 | func (s *waitState) Execute(smCtx gfsm2.StateMachineContext, eventCtx gfsm2.EventContext) State { 76 | cCtx := smCtx.(*coordinatorContext) 77 | vote, ok := eventCtx.(commitVote) 78 | if !ok || !vote.commit { 79 | fmt.Printf("invalid vote or vote for commit %s was rejected\n", cCtx.commitID) 80 | // rejecting commit 81 | return Abort 82 | } 83 | 84 | fmt.Printf("one more commit confirmation for %s!\n", cCtx.commitID) 85 | s.votesCnt++ 86 | if s.votesCnt == cCtx.partCnt { 87 | // all votes were positive, committing 88 | fmt.Printf("all participants confirmed commit %s!\n", cCtx.commitID) 89 | return Commit 90 | } 91 | 92 | return Wait 93 | } 94 | 95 | // ========= Commit/Abort state handler ========= 96 | 97 | type responseState struct { 98 | votesCnt int 99 | keepResp State 100 | } 101 | 102 | func (s *responseState) OnEnter(smCtx gfsm2.StateMachineContext) { 103 | cCtx := smCtx.(*coordinatorContext) 104 | s.votesCnt = cCtx.partCnt 105 | fmt.Printf("committing %s\n", cCtx.commitID) 106 | //for i := 0; i < cCtx.votesCnt; i++ { 107 | // sending commit/abort message to each participant 108 | //} 109 | } 110 | 111 | func (s *responseState) OnExit(_ gfsm2.StateMachineContext) { 112 | } 113 | 114 | func (s *responseState) Execute(_ gfsm2.StateMachineContext, eventCtx gfsm2.EventContext) State { 115 | resp, ok := eventCtx.(commitVote) 116 | if !ok { 117 | fmt.Printf("invalid response\n") 118 | // nothing to process, keeping state 119 | return s.keepResp 120 | } 121 | if !resp.commit { 122 | // this is abnormal situation, as during the Commit/Abort phase participant can only confirm the commit. 123 | // we will need to resend commit message to the participant. 124 | // go resendCommit() 125 | return s.keepResp 126 | } 127 | 128 | s.votesCnt-- 129 | if s.votesCnt != 0 { 130 | return s.keepResp 131 | } 132 | return Init 133 | } 134 | 135 | //go:generate gfsm_uml -format=plantuml 136 | func main() { 137 | sm := gfsm2.NewBuilder[State](). 138 | SetSMName("TwoPhaseCommit"). 139 | SetDefaultState(Init). 140 | SetSmContext(&coordinatorContext{partCnt: 3}). 141 | RegisterState(Init, &initState{}, []State{Wait}). 142 | RegisterState(Wait, &waitState{}, []State{Abort, Commit}). 143 | RegisterState(Abort, &responseState{ 144 | keepResp: Abort, 145 | }, []State{Init}). 146 | RegisterState(Commit, &responseState{ 147 | keepResp: Commit, 148 | }, []State{Init}). 149 | Build() 150 | 151 | sm.Start() 152 | defer sm.Stop() 153 | 154 | fmt.Printf("SM state (pre commit request): %s\n", sm.State()) 155 | err := sm.ProcessEvent(commitRequest{"commit_1"}) 156 | if err != nil { 157 | fmt.Printf("SM state: %s\n", sm.State()) 158 | return 159 | } 160 | fmt.Printf("SM state (postcommit request): %s\n", sm.State()) 161 | 162 | for i := 0; i < 3; i++ { 163 | err := sm.ProcessEvent(commitVote{commit: true}) 164 | if err != nil { 165 | fmt.Printf("unable to vote: %v\n", err) 166 | return 167 | } 168 | fmt.Printf("SM state (voting): %s\n", sm.State()) 169 | } 170 | fmt.Printf("SM state (pre confirm): %s\n", sm.State()) 171 | for i := 0; i < 3; i++ { 172 | err := sm.ProcessEvent(commitVote{commit: true}) 173 | if err != nil { 174 | fmt.Printf("unable complete: %v\n", err) 175 | return 176 | } 177 | fmt.Printf("SM state (confirming): %s\n", sm.State()) 178 | } 179 | fmt.Printf("SM state (post confirm): %s\n", sm.State()) 180 | } 181 | -------------------------------------------------------------------------------- /cmd/gfsm_uml/main.go: -------------------------------------------------------------------------------- 1 | // run-new-generator.go 2 | // 3 | // Usage: go generate 4 | // In your source file, include a directive such as: 5 | // 6 | // //go:generate gfsm_uml -format=plantuml 7 | // 8 | // This generator reads the file specified by the GOFILE environment variable, 9 | // then for each state machine builder chain (identified by a terminating Build() 10 | // call), it extracts the SM name from a SetSMName call and collects all 11 | // RegisterState transitions. It then writes a diagram for each state machine 12 | // into a separate file. 13 | package main 14 | 15 | import ( 16 | "errors" 17 | "flag" 18 | "fmt" 19 | "go/ast" 20 | "go/parser" 21 | "go/token" 22 | "log" 23 | "os" 24 | "strings" 25 | ) 26 | 27 | const ( 28 | setSMNameCall = "SetSMName" 29 | registerStateCall = "RegisterState" 30 | buildCall = "Build" 31 | ) 32 | 33 | // Transition represents a state transition. 34 | type Transition struct { 35 | Source string 36 | Destinations []string 37 | } 38 | 39 | // StateMachine holds the name and all transitions for a state machine. 40 | type StateMachine struct { 41 | Name string 42 | Transitions []Transition 43 | } 44 | 45 | func main() { 46 | format := getFormat() 47 | filename, err := getOutputName() 48 | if err != nil { 49 | log.Fatalf("Failed to get output file name: %v", err) 50 | } 51 | 52 | machines, err := doParse(filename) 53 | if err != nil { 54 | log.Fatalf("Parse error: %s", err) 55 | } 56 | 57 | err = writeDiagram(machines, format, filename) 58 | if err != nil { 59 | log.Fatalf("Failed to write diagram: %v", err) 60 | } 61 | } 62 | 63 | func doParse(filename string) (map[string]StateMachine, error) { 64 | // Parse the file. 65 | fset := token.NewFileSet() 66 | node, err := parser.ParseFile(fset, filename, nil, parser.AllErrors) 67 | if err != nil { 68 | return nil, fmt.Errorf("failed to parse file %q: %v", filename, err) 69 | } 70 | 71 | // Map to hold state machines keyed by their name. 72 | machines := make(map[string]StateMachine) 73 | 74 | // Walk the AST to find builder chains ending with a call to Build(). 75 | ast.Inspect(node, func(n ast.Node) bool { 76 | call, ok := n.(*ast.CallExpr) 77 | if !ok { 78 | return true 79 | } 80 | 81 | // Look for calls to Build() which terminate a builder chain. 82 | sel, ok := call.Fun.(*ast.SelectorExpr) 83 | if !ok || sel.Sel.Name != buildCall { 84 | return true 85 | } 86 | 87 | // Extract the full call chain. 88 | chain := extractCallChain(call) 89 | 90 | // Look for the custom naming function and register state calls. 91 | smName, transitions := processChain(chain) 92 | if smName == "" { 93 | // If no SM name is provided, you might skip or assign a default name. 94 | smName = "Unnamed" 95 | } 96 | 97 | // Merge with any previously discovered machine of the same name. 98 | if existing, ok := machines[smName]; ok { 99 | existing.Transitions = append(existing.Transitions, transitions...) 100 | machines[smName] = existing 101 | } else { 102 | machines[smName] = StateMachine{ 103 | Name: smName, 104 | Transitions: transitions, 105 | } 106 | } 107 | 108 | return true 109 | }) 110 | return machines, nil 111 | } 112 | 113 | func getOutputName() (string, error) { 114 | // Use `GOFILE` environment variable (set by go generate). 115 | filename := os.Getenv("GOFILE") 116 | if filename == "" { 117 | return "", errors.New("GOFILE environment variable is not set") 118 | } 119 | return filename, nil 120 | } 121 | 122 | func getFormat() string { 123 | // Define flag for the output format ("mermaid" or "plantuml"). 124 | var format string 125 | flag.StringVar(&format, "format", "mermaid", "output format: mermaid or plantuml") 126 | flag.Parse() 127 | return format 128 | } 129 | 130 | func writeDiagram(machines map[string]StateMachine, format string, filename string) error { 131 | // Write each state machine's diagram to a file. 132 | for _, sm := range machines { 133 | var output string 134 | outFmt := strings.ToLower(format) 135 | switch outFmt { 136 | case "mermaid": 137 | output = buildMermaid(sm) 138 | case "plantuml": 139 | output = buildPlantUML(sm) 140 | default: 141 | return fmt.Errorf("unknown output format: %s", outFmt) 142 | } 143 | ext := ".mermaid" 144 | if outFmt == "plantuml" { 145 | ext = ".uml" 146 | } 147 | outFilename := sm.Name + ext 148 | err := os.WriteFile(outFilename, []byte(output), 0644) 149 | if err != nil { 150 | return fmt.Errorf("failed to write diagram to file %q: %v", outFilename, err) 151 | } 152 | log.Printf("Diagram for state machine %q written to %s\n\n", sm.Name, outFilename) 153 | } 154 | 155 | // If no state machines were found, log a message. 156 | if len(machines) == 0 { 157 | log.Println("No state machine definitions found in", filename) 158 | } 159 | return nil 160 | } 161 | 162 | // extractCallChain traverses the fluent API call chain starting at expr. 163 | // It returns a slice of CallExpr pointers in the order they were invoked. 164 | func extractCallChain(expr ast.Expr) []*ast.CallExpr { 165 | var chain []*ast.CallExpr 166 | current := expr 167 | for { 168 | call, ok := current.(*ast.CallExpr) 169 | if !ok { 170 | break 171 | } 172 | chain = append(chain, call) 173 | // Each call is a selector, e.g. previousCall.Method() 174 | sel, ok := call.Fun.(*ast.SelectorExpr) 175 | if !ok { 176 | break 177 | } 178 | current = sel.X 179 | } 180 | // Reverse the chain so that it is in left-to-right order. 181 | for i, j := 0, len(chain)-1; i < j; i, j = i+1, j-1 { 182 | chain[i], chain[j] = chain[j], chain[i] 183 | } 184 | return chain 185 | } 186 | 187 | // processChain looks through the call chain for a SetSMName call and RegisterState calls. 188 | // It returns the SM name (from SetSMName) and a slice of transitions. 189 | func processChain(chain []*ast.CallExpr) (string, []Transition) { 190 | var smName string 191 | var transitions []Transition 192 | 193 | // Iterate over each call in the chain. 194 | for _, callExpr := range chain { 195 | sel, ok := callExpr.Fun.(*ast.SelectorExpr) 196 | if !ok { 197 | continue 198 | } 199 | methodName := sel.Sel.Name 200 | 201 | switch methodName { 202 | case setSMNameCall: 203 | // Expect a single argument: a string literal with the SM name. 204 | if len(callExpr.Args) >= 1 { 205 | if lit, ok := callExpr.Args[0].(*ast.BasicLit); ok && lit.Kind == token.STRING { 206 | smName = strings.Trim(lit.Value, `"`) 207 | } 208 | } 209 | case registerStateCall: 210 | // Expect: RegisterState(source, stateInstance, []SM{dest1, dest2, ...}) 211 | if len(callExpr.Args) < 3 { 212 | continue 213 | } 214 | // First argument: source state (identifier). 215 | srcIdent, ok := callExpr.Args[0].(*ast.Ident) 216 | if !ok { 217 | continue 218 | } 219 | source := srcIdent.Name 220 | 221 | // Third argument: allowed transitions as a composite literal. 222 | compLit, ok := callExpr.Args[2].(*ast.CompositeLit) 223 | if !ok { 224 | continue 225 | } 226 | var dests []string 227 | for _, elt := range compLit.Elts { 228 | if ident, ok := elt.(*ast.Ident); ok { 229 | dests = append(dests, ident.Name) 230 | } 231 | } 232 | transitions = append(transitions, Transition{ 233 | Source: source, 234 | Destinations: dests, 235 | }) 236 | } 237 | } 238 | return smName, transitions 239 | } 240 | 241 | // buildMermaid generates a Mermaid state diagram for the state machine. 242 | func buildMermaid(sm StateMachine) string { 243 | var b strings.Builder 244 | b.WriteString("```mermaid\n") 245 | b.WriteString("stateDiagram-v2\n") 246 | for _, t := range sm.Transitions { 247 | for _, dest := range t.Destinations { 248 | b.WriteString(fmt.Sprintf(" %s --> %s\n", t.Source, dest)) 249 | } 250 | } 251 | b.WriteString("```\n") 252 | return b.String() 253 | } 254 | 255 | // buildPlantUML generates a PlantUML state diagram for the state machine. 256 | func buildPlantUML(sm StateMachine) string { 257 | var b strings.Builder 258 | b.WriteString("@startuml\n") 259 | for _, t := range sm.Transitions { 260 | for _, dest := range t.Destinations { 261 | b.WriteString(fmt.Sprintf("%s --> %s\n", t.Source, dest)) 262 | } 263 | } 264 | b.WriteString("@enduml\n") 265 | return b.String() 266 | } 267 | --------------------------------------------------------------------------------