AtomicState.hs

{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE RankNTypes #-}
module AtomicState
    ( getState, modifyState, forkThread, execAtomicStateT, AtomicStateT
    , readState, runReadOnly, AtomicReadT
    )
where

import Control.Monad.Reader
import Control.Concurrent
import Control.Concurrent.STM

newtype AtomicReadT r m a =
    AtomicReadT { unpackAtomicReadT :: AtomicStateT r m a }
    deriving (Monad, MonadIO, MonadTrans)

readState :: MonadIO m => AtomicReadT r m r
readState =
    AtomicReadT $ getState

runReadOnly :: MonadIO m => AtomicReadT r m r -> AtomicStateT r m r
runReadOnly readOnlyAction =
    unpackAtomicReadT readOnlyAction

newtype AtomicStateT s m a =
    AtomicStateT { unpackAtomicStateT :: ReaderT (TVar s) m a }
    deriving (Monad, MonadIO, MonadTrans)

getState :: MonadIO m => AtomicStateT s m s
getState =
    AtomicStateT $
    do var <- ask
       liftIO $ atomically $ readTVar var

modifyState :: MonadIO m => (s -> s) -> AtomicStateT s m ()
modifyState f =
    AtomicStateT $
    do var <- ask
       liftIO $ atomically $ modifyTVar var f

forkThread :: MonadIO m
           => AtomicStateT s m ()
           -> (forall a. m a -> IO a)
           -> AtomicStateT s m ThreadId
forkThread thread ioLift =
    AtomicStateT $
    do var <- ask
       liftIO $ forkIO $ ioLift (runReaderT (unpackAtomicStateT thread) var)

execAtomicStateT :: MonadIO m => AtomicStateT s m a -> s -> m (s, a)
execAtomicStateT action initSt =
    do var <- liftIO $ newTVarIO initSt
       res <- runReaderT (unpackAtomicStateT action) var
       st' <- liftIO $ atomically $ readTVar var
       return (st', res)