Cont.hs

{- cabal:
   build-depends: base
   build-depends: transformers
   ghc-options: -threaded

-}

import Control.Concurrent (myThreadId, threadDelay, forkOS)
import Control.Monad (forever,  void)
import Control.Monad.IO.Class (liftIO)
import Control.Monad.Trans.Cont (ContT(..), evalContT)
import Data.IORef (IORef, newIORef, writeIORef, readIORef, modifyIORef)

type Connection = IORef (Msg -> IO ())

type Msg = String

newConnection :: IO Connection
newConnection = do
  ref <- newIORef (const (pure ()))
  forkOS $ do
    myid <- myThreadId
    forever $ do
      f <- readIORef ref
      f $ "Hello from " ++ show myid
      threadDelay 1000000
  pure ref

handleMessage :: Connection -> (Msg -> IO ()) -> IO ()
handleMessage = writeIORef

doThingWithMessage msg = putStrLn $ "doThingWithMessage: " ++ msg

{-
main :: IO ()
main = do
  messageCount <- newIORef 0
  conn <- newConnection
  handleMessage conn (handler messageCount)
  forever $ do
    totalMessages <- readIORef messageCount
    putStrLn $ show totalMessages ++ " messages received"
    threadDelay 1000000

handler :: IORef Integer -> Msg -> IO ()
handler counter msg = do
  modifyIORef counter (+ 1)
  doThingWithMessage msg
-}

yield :: Connection -> ContT () IO Msg
yield conn = ContT $ \f -> handleMessage conn f

main :: IO ()
main = do
  let messageCount = 0
  conn <- newConnection
  evalContT $ void $ iterM messageCount $ handler conn
  forever $ threadDelay 1000000

iterM :: Monad m => a -> (a -> m a) -> m a
iterM a m = m a >>= flip iterM m

handler :: Connection -> Integer -> ContT () IO Integer
handler conn counter = do
  msg <- yield conn
  liftIO $ putStrLn $ show counter ++ " messages received"
  liftIO $ doThingWithMessage msg
  pure (counter + 1)