Simple chat with bots using FreeT monad transformers for both bots and environment.

chat.hs

{-# LANGUAGE TypeFamilies, ExistentialQuantification, FlexibleInstances #-}
module Main where
 
import System.IO (isEOF, hFlush, stdout)
 
import Data.Char (toLower, isDigit)
import Data.Maybe (isNothing)
 
import Control.Monad.Trans.Free
import Control.Monad.IO.Class (MonadIO, liftIO)
import Control.Monad (forever, when, unless, void)
 
import Control.Monad.Reader
 
import Control.Concurrent.STM.TChan
import Control.Concurrent.STM (atomically)
import Control.Concurrent (forkIO, threadDelay)
 
-- kind of helpers
-- FIXME: actually I wanted to use something like
-- `liftOp_` from layers package in order to use
-- `forkIO` in any MonadIO monad, but I failed
-- to install the package using ghc-7.4
class MonadToIO m where
    toIO :: m a -> IO a
 
instance MonadToIO IO where
    toIO = id
 
-- | Tear down trough a free monad transformer using iteration.
-- Should be in Control.Monad.Trans.Free
iter :: (Functor f, Monad m) => (f (m a) -> m a) -> FreeT f m a -> m a
iter f (FreeT m) = do
    val <- m
    case fmap (iter f) val of
        Pure x -> return x
        Free y -> f y
 
-- | Abstract syntax for the environment.
data EnvF node x
    -- ^ Create new channel of type `a`.
    = forall a. ENewChan   (NodeChan node a -> x)
    -- ^ Duplicate a channel of type `a`.
    | forall a. EDupChan   (NodeChan node a) (NodeChan node a -> x)
    -- ^ Spawn a new node computation in parallel.
    | forall r. ESpawnNode (node r) x
    -- ^ Execute a node computation.
    | forall r. EExecNode  (node r) x
 
-- Unfortunately, -XDeriveFunctor does not work with existential types.
instance Functor (EnvF node) where
    fmap f (ENewChan g)     = ENewChan (f . g)
    fmap f (EDupChan c g)   = EDupChan c (f . g)
    fmap f (ESpawnNode n x) = ESpawnNode n (f x)
    fmap f (EExecNode n x)  = EExecNode n (f x)
 
-- | Environment free monad transformer.
type EnvT node = FreeT (EnvF node)
 
-- | Environment monad class.
class (Monad m) => MonadEnv m where
    -- | Environment node constructors.
    type EnvNode m :: * -> *
    -- | New channel command.
    newChan   :: m (NodeChan (EnvNode m) a)
    -- | Duplicate channel command.
    dupChan   :: NodeChan (EnvNode m) a -> m (NodeChan (EnvNode m) a)
    -- | Spawn a separate node computation command.
    spawnNode :: EnvNode m a -> m ()
    -- | Execute a node computation command.
    execNode  :: EnvNode m a -> m ()
 
instance (Monad m) => MonadEnv (FreeT (EnvF node) m) where
    type EnvNode (FreeT (EnvF node) m) = node
    newChan     = liftF $ ENewChan id
    dupChan   c = liftF $ EDupChan c id
    spawnNode n = liftF $ ESpawnNode n ()
    execNode  n = liftF $ EExecNode n ()
 
-- | Abstract syntax for nodes with channels of type `chan`.
data NodeF chan x
    -- ^ Send a message to a channel of type `a`.
    = forall a. NSend    (chan a) a x
    -- ^ Receive a message from a channel of type `a`.
    | forall a. NRecv    (chan a) (a -> x)
    -- ^ Check if a channel of type `a` is empty.
    | forall a. NIsEmpty (chan a) (Bool -> x)
    -- ^ Do something else.
 
-- Unfortunately, -XDeriveFunctor does not work with existential types.
instance Functor (NodeF chan) where
    fmap f (NSend c a x)  = NSend c a (f x)
    fmap f (NRecv c g)    = NRecv c (f . g)
    fmap f (NIsEmpty c g) = NIsEmpty c (f . g)
 
-- | Node free monad transformer.
type NodeT chan = FreeT (NodeF chan)
 
-- | Node monad class.
class Monad m => MonadNode m where
    -- | Node's channel constructor.
    type NodeChan m :: * -> *
    -- | Send command.
    send :: NodeChan m a -> a -> m ()
    -- | Receive command.
    recv :: NodeChan m a -> m a
    -- | Empty check command.
    isEmpty :: NodeChan m a -> m Bool
 
instance Monad m => MonadNode (FreeT (NodeF chan) m) where
    type NodeChan (FreeT (NodeF chan) m) = chan
    send c x  = liftF $ NSend c x ()
    recv c    = liftF $ NRecv c id
    isEmpty c = liftF $ NIsEmpty c id
 
-- | Execute node as a concurrent thread with @TChan@ channels.
stmNode :: (MonadIO m) => NodeT TChan m r -> m r
stmNode = iter stmNodeF
    where
        stmNodeF (NSend chan msg m) = (liftIO . atomically $ writeTChan chan msg) >>  m
        stmNodeF (NRecv chan m)     = (liftIO . atomically $ readTChan chan) >>= m
        stmNodeF (NIsEmpty chan m)  = (liftIO . atomically $ isEmptyTChan chan) >>= m
 
-- | Execute environment with nodes as concurrent threads with @TChan@ channels.
stmEnv :: (MonadIO m, MonadToIO m, Functor m) => EnvT (NodeT TChan m) m r -> m r
stmEnv = iter stmEnvF
    where
        stmEnvF (ENewChan m)        = (liftIO . atomically $ newTChan) >>= m
        stmEnvF (EDupChan chan m)   = (liftIO . atomically $ dupTChan chan) >>= m
        stmEnvF (ESpawnNode n m)    = (liftIO . forkIO . toIO . void $ stmNode n) >> m
        stmEnvF (EExecNode n m)     = stmNode n >> m
 
-- | Observable environment of a chat bot.
data ChatBotE chan = ChatBotE
    { chatInputChan  :: chan String     -- ^ input channel
    , chatOutputChan :: chan String     -- ^ output channel
    }
 
-- | Chat bot monad transformer.
type ChatBotT m = ReaderT (ChatBotE (NodeChan m)) m
 
instance (MonadNode m) => MonadNode (ReaderT r m) where
    type NodeChan (ReaderT r m) = NodeChan m
    send c  = lift . send c
    recv    = lift . recv
    isEmpty = lift . isEmpty
 
-- | Receive a message from an input channel.
chatRecv :: (MonadNode m) => ChatBotT m String
chatRecv = do
    cin <- asks chatInputChan
    recv cin
 
-- | Send a message to an output channel.
chatSend :: (MonadNode m) => String -> ChatBotT m ()
chatSend msg = do
    cout <- asks chatOutputChan
    send cout msg
 
-- | Check if input channel is empty.
chatIsEmpty :: (MonadNode m) => ChatBotT m Bool
chatIsEmpty = do
    cin <- asks chatInputChan
    isEmpty cin
 
-- | Simple echo bot.
echoBot :: (MonadNode m) => ChatBotT m r
echoBot = forever $ do
    msg <- chatRecv
    chatSend $ "echo: " ++ msg
 
-- | Echo bot that sleeps before responding.
slowEchoBot :: (MonadIO m, MonadNode m) => ChatBotT m r
slowEchoBot = forever $ do
    msg <- chatRecv
    liftIO $ threadDelay $ 2 * 10^6 -- wait 2 secs
    chatSend $ "slow: " ++ msg
 
-- | Bot that responds with "Hello!" message for various greeting messages.
helloBot :: (MonadNode m) => ChatBotT m r
helloBot = forever $ do
    msg <- chatRecv
    when (isHello msg) $ do
        chatSend "hellobot: Hello!"
    where
        isHello :: String -> Bool
        isHello = (`elem` ["hello", "hi", "hi there", "hello world"]) . map toLower
 
-- | Bot that squares numbers.
squareBot :: (MonadNode m) => ChatBotT m r
squareBot = forever $ do
    msg <- chatRecv
    when (all isDigit msg) $ do
        chatSend $ "square: " ++ show (read msg ^ 2 :: Integer)
 
-- | User interface for a chat room.
userNode :: (MonadIO m, MonadNode m) => ChatBotT m ()
userNode = do
    msg <- liftIO prompt    -- ask user for input
    case msg of
        Nothing -> return ()
        Just s -> do
            chatSend s                  -- send message
            liftIO $ threadDelay $ 10^5 -- wait 0.1 sec for immediate answers
            printAnswers                -- print available answers
            userNode
    where
        -- print all available messages
        printAnswers :: (MonadIO m, MonadNode m) => ChatBotT m ()
        printAnswers = do
            e <- chatIsEmpty
            unless e $ do
                ans <- chatRecv
                liftIO $ putStrLn ans
                printAnswers
 
        -- ask user for an input
        prompt :: IO (Maybe String)
        prompt = do
            putStr "you: "
            hFlush stdout
            eof <- isEOF
            case eof of
                True -> return Nothing
                False -> do
                    s <- getLine
                    return $ Just s
 
-- | Multi-agent environment with bots and user node.
chatRoom :: (MonadIO n, MonadNode n, Monad m) => [ChatBotT n a] -> ChatBotT n () -> EnvT n m ()
chatRoom bots user = do
    cin  <- newChan     -- user input (bots' messages)
    cout <- newChan     -- user output (user's messages)
    let botE  = ChatBotE cout cin
        userE = ChatBotE cin cout
    -- spawn bots (cin in bot's output, cout is bot's input)
    mapM_ (spawnBot botE) bots
    -- run user interface
    execNode $ runReaderT user userE
    where
        -- spawn a bot
        spawnBot :: (Monad m) => ChatBotE (NodeChan n) -> ChatBotT n a -> EnvT n m ()
        spawnBot e bot = do
            -- duplicate input channel to enable multiple readers
            cin <- dupChan $ chatInputChan e
            spawnNode $ runReaderT bot e{chatInputChan = cin}
 
-- | main
main :: IO ()
main = stmEnv $ chatRoom bots userNode
    where
        bots = [echoBot, slowEchoBot, helloBot, squareBot]