State composition with ReaderT that can be fully mocked

Main.hs

{-# LANGUAGE TypeInType #-}

module Main where

import Control.Concurrent.STM
import Control.Exception.Base
import Control.Monad.Except
import Control.Monad.Identity
import Control.Monad.Reader
import Control.Monad.State
import Data.IORef
import Data.Kind
import Data.Proxy
import GHC.Records

data State1 = State1 { _name :: String, _age :: Int } deriving (Show, Eq)

data State2 = State2 { _address :: String, _phone :: Maybe String }
    deriving (Show, Eq)

data State3 = State3 { _fbUsername :: String } deriving (Show, Eq)

-- How to compose state in a real-world Haskell application?

-- What do we want?
-- 1. Read/write states from multiple nested states.
-- 2. Call abstract functions that access hidden references to do stuff.
-- 3. Handle errors.

-- What is the easiest and most flexible solution?
-- Use one ReaderT to one state which contains all of the nested
-- states as IORefs.

-- | Composed state and other hidden references needed for
-- the typeclass abstractions.
data Config = Config
    { _state1 :: IORef State1
    , _state2 :: TVar State2
    , _state3 :: TVar State3
    }

data MyError = NotFound
             | InvalidName String
             | Other String
             deriving (Show, Eq)
instance Exception MyError


-- | Typeclass abstractions over IORef code that can be
-- mocked with pure code.
class Ref c a s where
    getRef    :: c a -> s a
    putRef    :: c a -> a -> s ()
    modifyRef :: c a -> (a -> a) -> s ()

type family TransactionType (m :: * -> *) = (r :: * -> *) | r -> m
type instance TransactionType TVar = STM

class (Monad (TransactionType c), Ref c a s) => Transaction c a s where
    getRefTrans    :: Proxy s -> c a -> (TransactionType c) a
    putRefTrans    :: Proxy s -> c a -> a -> (TransactionType c) ()
    modifyRefTrans :: Proxy s -> c a -> (a -> a) -> (TransactionType c) ()
    runTransaction :: Proxy a -> (TransactionType c) r -> s r

class HasState1 ref c | c -> ref where
    getState1 :: c -> ref State1

class HasState2 ref c | c -> ref where
    getState2 :: c -> ref State2

class HasState3 ref c | c -> ref where
    getState3 :: c -> ref State3 

instance HasState1 IORef Config where getState1 = getField @"_state1"
instance HasState2 TVar Config where getState2 = getField @"_state2"
instance HasState3 TVar Config where getState3 = getField @"_state3"

-- | The ReaderT IO stack to be used in production.
newtype AppT a = AppT { unAppT :: ReaderT Config IO a }
    deriving (Functor, Applicative, Monad, MonadIO, MonadReader Config)

runAppT = runReaderT . unAppT

-- | Error instances for AppT use exceptions.
instance MonadError MyError AppT where
    throwError e = throw e
    catchError m f = do
        config <- ask
        liftIO $ catch (runAppT m config) (flip runAppT config . f)

instance Ref IORef t AppT where
    getRef = liftIO . readIORef
    putRef ref v = liftIO $ writeIORef ref v
    modifyRef ref v = liftIO $ modifyIORef' ref v

instance Ref TVar t AppT where
    getRef = liftIO . readTVarIO
    putRef ref v = liftIO $ atomically $ writeTVar ref v
    modifyRef ref v = liftIO $ atomically $ modifyTVar ref v

instance Transaction TVar t AppT where
    getRefTrans _ = readTVar
    putRefTrans _ = writeTVar
    modifyRefTrans _ = modifyTVar
    runTransaction _ = liftIO . atomically

data MockConfig = MockConfig
    { _state1 :: State1
    , _state2 :: State2
    , _state3 :: State3
    } deriving (Show, Eq)

data Accessor s = Accessor
    { _getter :: MockConfig -> s
    , _setter :: (s -> s) -> MockConfig -> MockConfig
    }

instance HasState1 Accessor MockConfig where
    getState1 _ = Accessor
        { _getter = \c -> getField @"_state1" c
        , _setter = \f c -> c { _state1 = f (getField @"_state1" c) }
        }

instance HasState2 Accessor MockConfig where
    getState2 _ = Accessor
        { _getter = \c -> getField @"_state2" c
        , _setter = \f c -> c { _state2 = f (getField @"_state2" c) }
        }

instance HasState3 Accessor MockConfig where
    getState3 _ = Accessor
        { _getter = \c -> getField @"_state3" c
        , _setter = \f c -> c { _state3 = f (getField @"_state3" c) }
        }

-- | The mocked version of AppT that is pure.
newtype MockT a = MockT
    { unMockT :: StateT MockConfig (ExceptT MyError Identity) a }
    deriving ( Functor, Applicative, Monad
             , MonadState MockConfig, MonadError MyError
             )

instance MonadReader MockConfig MockT where
    ask = get
    local f m = do
        s <- get
        put (f s)
        r <- m
        put s
        return r

runMockT m config = runIdentity
                  . runExceptT
                  . flip runStateT config
                  . unMockT
                  $ m

instance Ref Accessor t MockT where
    getRef accessor = (_getter accessor) <$> get
    putRef accessor v = modify ((_setter accessor) (const v))
    modifyRef accessor f = modify ((_setter accessor) f)

type instance TransactionType Accessor = MockT
instance Transaction Accessor t MockT where
    getRefTrans _ = getRef
    putRefTrans _ = putRef
    modifyRefTrans _ = modifyRef
    runTransaction _ = id

-- Generic functions that work with both AppT and MockT.

getPhone :: ( MonadError MyError m
            , MonadReader r m
            , HasState2 ref r
            , Ref ref State2 m
            )
         => m String
getPhone = 
    asks getState2 >>= getRef >>= \case
        State2{ _phone = Just phone } -> return phone
        _                             -> throwError NotFound

changeStatesAndGetAddr :: forall m a b r.
                          ( Monad m
                          , MonadError MyError m
                          , MonadReader r m
                          , HasState1 a r
                          , HasState2 b r
                          , HasState3 b r
                          , Ref a State1 m
                          , Transaction b State2 m
                          , Transaction b State3 m
                          )
                       => m (String, String)
changeStatesAndGetAddr = do
    state1Ref <- asks getState1
    modifyRef state1Ref (\s -> s { _age = _age s + 1 })
    state2 <- asks getState2 >>= getRef
    phone <- catchError getPhone $ pure . \case
        NotFound -> "Phone number not found"
        e        -> "Error: " ++ show e

    ref <- asks getState2
    ref2 <- asks getState3
    runTransaction (Proxy :: Proxy State2) $ do
        let p = Proxy :: Proxy m
        state2 <- getRefTrans p ref
        state3 <- getRefTrans p ref2
        putRefTrans p ref2 state3 { _fbUsername = "Zucc" }
        putRefTrans p ref state2 { _address = "FB Headquarters" }
        return ()
    return (_address state2, phone)

main :: IO ()
main = do
    let state1 = State1 { _name = "Bob", _age = 20 }
        state2 = State2 { _address = "123", _phone = Nothing }
        state3 = State3 { _fbUsername = "Bill Gates" }
    state1Ref <- newIORef state1
    state2Ref <- newTVarIO state2
    state3Ref <- newTVarIO state3
    let config = Config
            { _state1 = state1Ref
            , _state2 = state2Ref
            , _state3 = state3Ref
            }
        mockConfig = MockConfig
            { _state1 = state1
            , _state2 = state2
            , _state3 = state3
            }

    address' <- runAppT changeStatesAndGetAddr config
    let address'' = runMockT changeStatesAndGetAddr mockConfig
    print address'
    print address''