State refactor

akka-core/src/main/scala/actor/Fsm.scala

@@ -5,17 +5,17 @@ import se.scalablesolutions.akka.actor.{ActorRef, Scheduler, Actor}
 
 trait Fsm[S] { self: Actor =>
   
-  type State = scala.PartialFunction[Event, NextState]
+  type StateFunction = scala.PartialFunction[Event, State]
 
-  @volatile var currentState: NextState = initialState
+  @volatile var currentState: State = initialState
   @volatile var timeoutActor: Option[ActorRef] = None
 
-  def initialState: NextState
+  def initialState: State
 
-  def handleEvent: State = {
+  def handleEvent: StateFunction = {
     case event@Event(value,stateData) =>
       log.warning("No state for event with value %s - keeping current state %s", value, stateData)
-      NextState(currentState.state, stateData, currentState.timeout)
+      State(NextState, currentState.stateFunction, stateData, currentState.timeout)
   }
 
 
@@ -24,15 +24,30 @@ trait Fsm[S] { self: Actor =>
       timeoutActor = timeoutActor flatMap { ref => Scheduler.unschedule(ref); None }
 
       val event = Event(value, currentState.stateData)
-      currentState = (currentState.state orElse handleEvent).apply(event)
+      currentState = (currentState.stateFunction orElse handleEvent).apply(event)
 
       currentState.timeout.foreach{timeout =>
         timeoutActor = Some(Scheduler.scheduleOnce(this.self, StateTimeout, timeout, TimeUnit.MILLISECONDS))
       }
+
+      currentState match {
+        case State(Reply, _, _, _, reply) => reply.foreach(this.self.reply)
+      }
     }
   }
 
-  case class NextState(state: State, stateData: S, timeout: Option[Int] = None)
+  
+  case class State(stateEvent: StateEvent,
+                   stateFunction: StateFunction,
+                   stateData: S,
+                   timeout: Option[Int] = None,
+                   reply: Option[Any] =None)
+
   case class Event(event: Any, stateData: S)
+
+  sealed trait StateEvent
+  object NextState extends StateEvent
+  object Reply extends StateEvent
+
   object StateTimeout
 }

akka-core/src/test/scala/FsmActorSpec.scala

@@ -17,29 +17,29 @@ object FsmActorSpec {
              unlockedLatch: StandardLatch,
              lockedLatch: StandardLatch) extends Actor with Fsm[CodeState] {
 
-    def initialState = NextState(locked, CodeState("", "33221"))
+    def initialState = State(NextState, locked, CodeState("", "33221"))
 
-    def locked: State = {
+    def locked: StateFunction = {
       case Event(digit: Char, CodeState(soFar, code)) => {
            soFar + digit match {
              case incomplete if incomplete.length < code.length =>
-               NextState(locked, CodeState(incomplete, code))
+               State(NextState, locked, CodeState(incomplete, code))
              case codeTry if (codeTry == code) => {
                doUnlock
-               new NextState(open, CodeState("", code), Some(timeout))
+               new State(NextState, open, CodeState("", code), Some(timeout))
              }
              case wrong => {
                log.error("Wrong code %s", wrong)
-               NextState(locked, CodeState("", code))
+               State(NextState, locked, CodeState("", code))
              }
            }
       }
     }
 
-    def open: State = {
+    def open: StateFunction = {
       case Event(StateTimeout, stateData) => {
         doLock
-        NextState(locked, stateData)
+        State(NextState, locked, stateData)
       }
     }