From 22ce2f5150bd6f012ee2fb30ea557d48d348ed00 Mon Sep 17 00:00:00 2001
From: Denis Tereshkin <denis@kasan.ws>
Date: Mon, 12 Jun 2023 11:16:04 +0700
Subject: [PATCH] FSM

---
 src/FSM.hs              | 47 +++++++++++++++++++++++++++++++++++++++++
 test/Spec.hs            |  4 +++-
 test/Test/FSM.hs        | 47 +++++++++++++++++++++++++++++++++++++++++
 transaq-connector.cabal | 10 +++++++++
 4 files changed, 107 insertions(+), 1 deletion(-)
 create mode 100644 src/FSM.hs
 create mode 100644 test/Test/FSM.hs

diff --git a/src/FSM.hs b/src/FSM.hs
new file mode 100644
index 0000000..0416f04
--- /dev/null
+++ b/src/FSM.hs
@@ -0,0 +1,47 @@
+
+module FSM
+  (
+    FSMState(..)
+  , FSMCallback(..)
+  , FSM(..)
+  , runFsm
+  , makeFsm
+  ) where
+
+import           Control.Concurrent.STM.TVar (TVar, newTVarIO, readTVarIO,
+                                              writeTVar)
+import           Control.Monad.Extra         (whileM)
+import           Control.Monad.IO.Class      (MonadIO (..))
+import           Control.Monad.STM           (atomically)
+import qualified Data.Map                    as M
+
+class FSMState a where
+  isTerminalState :: a -> Bool
+
+newtype FSMCallback m a = FSMCallback (m (Maybe a))
+
+data FSM a m = FSM
+  { fsmCurrentState :: TVar a
+  , fsmHandlers     :: M.Map a (FSMCallback m a)
+  }
+
+runFsm :: (MonadIO m, FSMState a, Ord a) => FSM a m -> m ()
+runFsm fsm = whileM $ do
+  currentState <- liftIO . readTVarIO $ fsmCurrentState fsm
+  case M.lookup currentState $ fsmHandlers fsm of
+    Just (FSMCallback cb) -> do
+      maybeNextState <- cb
+      case maybeNextState of
+        Just nextState -> do
+          liftIO . atomically $ writeTVar (fsmCurrentState fsm) nextState
+          pure True
+        Nothing -> pure (not . isTerminalState $ currentState)
+    Nothing -> pure False
+
+makeFsm :: (MonadIO m1,
+            MonadIO m,
+            FSMState a,
+            Ord a) => a -> [(a, FSMCallback m a)] -> m1 (FSM a m)
+makeFsm initialState handlers = do
+  currentState <- liftIO $ newTVarIO initialState
+  pure $ FSM currentState (M.fromList handlers)
diff --git a/test/Spec.hs b/test/Spec.hs
index d7aaee5..ea08f13 100644
--- a/test/Spec.hs
+++ b/test/Spec.hs
@@ -1,5 +1,6 @@
 import           Test.Tasty
 
+import qualified Test.FSM
 import qualified Test.TickTable
 
 main :: IO ()
@@ -7,4 +8,5 @@ main = defaultMain $ testGroup "Tests" [unitTests]
 
 unitTests :: TestTree
 unitTests = testGroup "Unit Tests"
-  [Test.TickTable.unitTests ]
+  [ Test.TickTable.unitTests
+  , Test.FSM.unitTests ]
diff --git a/test/Test/FSM.hs b/test/Test/FSM.hs
new file mode 100644
index 0000000..e285b93
--- /dev/null
+++ b/test/Test/FSM.hs
@@ -0,0 +1,47 @@
+
+module Test.FSM
+  (
+    unitTests
+  ) where
+
+import           Control.Concurrent.STM.TVar (readTVarIO)
+import           FSM
+import           System.Timeout              (timeout)
+import           Test.Tasty
+import           Test.Tasty.HUnit            (testCase, (@?=))
+
+unitTests :: TestTree
+unitTests = testGroup "FSM"
+  [ testSimpleStateChange
+  , testStateWithoutCallback ]
+
+data SimpleState =
+    StateInitial
+  | StateIntermediate
+  | StateFinal
+  deriving (Show, Eq, Ord)
+
+instance FSMState SimpleState where
+  isTerminalState StateFinal = True
+  isTerminalState _          = False
+
+testSimpleStateChange = testCase "Simple state change" $ do
+  fsm <- makeFsm StateInitial [(StateInitial, callbackChangeState StateIntermediate),
+                               (StateIntermediate, callbackChangeState StateFinal),
+                               (StateFinal, doNothing)]
+  timeout 100000 $ runFsm fsm
+  endState <- readTVarIO (fsmCurrentState fsm)
+  endState @?= StateFinal
+  where
+    callbackChangeState x = FSMCallback (pure $ Just x)
+    doNothing = FSMCallback (pure Nothing)
+
+testStateWithoutCallback = testCase "Test without callback" $ do
+  fsm <- makeFsm StateInitial [(StateInitial, callbackChangeState StateIntermediate),
+                               (StateFinal, doNothing)]
+  timeout 100000 $ runFsm fsm
+  endState <- readTVarIO (fsmCurrentState fsm)
+  endState @?= StateIntermediate
+  where
+    callbackChangeState x = FSMCallback (pure $ Just x)
+    doNothing = FSMCallback (pure Nothing)
diff --git a/transaq-connector.cabal b/transaq-connector.cabal
index bbdcd23..939aed6 100644
--- a/transaq-connector.cabal
+++ b/transaq-connector.cabal
@@ -26,6 +26,7 @@ executable transaq-connector
                      , TXML
                      , TXMLConnector 
                      , TickTable
+                     , FSM
   default-extensions:  OverloadedStrings
                      , MultiWayIf
   default-language:    Haskell2010
@@ -75,6 +76,12 @@ test-suite transaq-connector-test
   type:                exitcode-stdio-1.0
   hs-source-dirs:      test src
   main-is:             Spec.hs
+  other-modules:       Test.TickTable
+                     , Test.FSM
+
+                     , FSM
+                     , TickTable
+
   build-depends:       base
                      , containers
                      , libatrade
@@ -82,4 +89,7 @@ test-suite transaq-connector-test
                      , tasty
                      , tasty-hunit
                      , time
+                     , extra
+  default-extensions:  OverloadedStrings
+                     , MultiWayIf