Add Entry/Exit (#4)

README.md

@@ -44,6 +44,17 @@ f.Transition(
 )
 ```
 
+Functions call be called when entering or leaving a state :
+
+```go
+f.Entry(StateFoo, func() {
+    // do something	
+})
+f.Exit(StateFoo, func() {
+    // do something	
+})
+```
+
 ## Installation
 
 go get github.com/cocoonspace/fsm

fsm.go

@@ -44,6 +44,8 @@ func (t *transition) apply(fsm *FSM) {
 // FSM is a finite state machine.
 type FSM struct {
 	transitions []transition
+	entry       map[State]func()
+	exit        map[State]func()
 	current     State
 	initial     State
 	previous    int
@@ -53,13 +55,16 @@ type FSM struct {
 // New creates a new finite state machine having the specified initial state.
 func New(initial State) *FSM {
 	return &FSM{
+		entry:   map[State]func(){},
+		exit:    map[State]func(){},
 		current: initial,
 		initial: initial,
 	}
 }
 
 // Option defines a transition option.
 type Option func(*transition)
+
 type result int
 
 const (
@@ -69,6 +74,7 @@ const (
 )
 
 type optionCondition func(e Event, times int, fsm *FSM) result
+
 type optionAction func(*FSM)
 
 // Transition creates a new transition, usually having trigger On an Event, from a Src State, to a Dst State.
@@ -110,7 +116,16 @@ func On(e Event) Option {
 func Dst(s State) Option {
 	return func(t *transition) {
 		t.actions = append(t.actions, func(fsm *FSM) {
+			if fsm.current == s {
+				return
+			}
+			if fn, ok := fsm.exit[fsm.current]; ok {
+				fn()
+			}
 			fsm.current = s
+			if fn, ok := fsm.entry[fsm.current]; ok {
+				fn()
+			}
 		})
 	}
 }
@@ -174,6 +189,16 @@ func (f *FSM) Current() State {
 	return f.current
 }
 
+// Entry sets a func that will be called when entering a state.
+func (f *FSM) Entry(state State, fn func()) {
+	f.entry[state] = fn
+}
+
+// Exit sets a func that will be called when entering a state.
+func (f *FSM) Exit(state State, fn func()) {
+	f.exit[state] = fn
+}
+
 // Event send an Event to a machine, applying at most one transition.
 // true is returned if a transition is applied.
 func (f *FSM) Event(e Event) bool {

fsm_test.go

@@ -154,3 +154,38 @@ func ExampleTimes() {
 	_ = f.Event(EventFoo) // no transition
 	_ = f.Event(EventFoo) // transition to StateBar
 }
+
+func TestEntryExit(t *testing.T) {
+	f := fsm.New(StateFoo)
+	f.Transition(
+		fsm.On(EventFoo), fsm.Src(StateFoo),
+		fsm.Dst(StateBar),
+	)
+	f.Transition(
+		fsm.On(EventBar), fsm.Src(StateBar),
+		fsm.Dst(StateFoo),
+	)
+	entry, exit := false, false
+	f.Entry(StateBar, func() {
+		entry = true
+	})
+	f.Exit(StateBar, func() {
+		exit = true
+	})
+
+	_ = f.Event(EventFoo)
+	if !entry {
+		t.Error("Entry func has not been called")
+	}
+	if exit {
+		t.Error("Exit func has wrongly been called")
+	}
+	entry, exit = false, false
+	_ = f.Event(EventBar)
+	if entry {
+		t.Error("Entry func has wrongly been called")
+	}
+	if !exit {
+		t.Error("Exit func has not been called")
+	}
+}