Combining multiple states in StateT

Question

I am writing a program that runs as a daemon. To create the daemon, the user supplies a set of implementations for each of the required classes (one of them is a database) All of these classes have functions have type signatures of the form StateT s IO a, but s is different for each class.

Suppose each of the classes follows this pattern:

import Control.Monad (liftM)
import Control.Monad.State (StateT(..), get)

class Hammer h where
  driveNail :: StateT h IO ()

data ClawHammer = MkClawHammer Int -- the real implementation is more complex

instance Hammer ClawHammer where
  driveNail = return () -- the real implementation is more complex

-- Plus additional classes for wrenches, screwdrivers, etc.

Now I can define a record that represents the implementation chosen by the user for each "slot".

data MultiTool h = MultiTool {
    hammer :: h
    -- Plus additional fields for wrenches, screwdrivers, etc.
  }

And the daemon does most of its work in the StateT (MultiTool h ...) IO () monad.

Now, since the multitool contains a hammer, I can use it in any situation where a hammer is needed. In other words, the MultiTool type can implement any of the classes it contains, if I write code like this:

stateMap :: Monad m => (s -> t) -> (t -> s) -> StateT s m a -> StateT t m a
stateMap f g (StateT h) = StateT $ liftM (fmap f) . h . g

withHammer :: StateT h IO () -> StateT (MultiTool h) IO ()
withHammer runProgram = do
  t <- get
  stateMap (\h -> t {hammer=h}) hammer runProgram

instance Hammer h => Hammer (MultiTool h) where
  driveNail = withHammer driveNail

But the implementations of withHammer, withWrench, withScrewdriver, etc. are basically identical. It would be nice to be able to write something like this...

--withMember accessor runProgram = do
--  u <- get
--  stateMap (\h -> u {accessor=h}) accessor runProgram

-- instance Hammer h => Hammer (MultiTool h) where
--   driveNail = withMember hammer driveNail

But of course that won't compile.

I suspect my solution is too object-oriented. Is there a better way? Monad transformers, maybe? Thank you in advance for any suggestions.

Answer 1

If you want to go with a large global state like in your case, then what you want to use is lenses, as suggested by Ben. I too recommend Edward Kmett's lens library. However, there is another, perhaps nicer way.

Servers have the property that the program runs continuously and performs the same operation over a state space. The trouble starts when you want to modularize your server, in which case you want more than just some global state. You want modules to have their own state.

Let's think of a module as something that transforms a Request to a Response:

Module :: (Request -> m Response) -> Module m

Now if it has some state, then this state becomes noticable in that the module might give a different answer the next time. There are a number of ways to do this, for example the following:

Module :: s -> ((Request, s) -> m (Response s)) -> Module m

But a much nicer and equivalent way to express this is the following constructor (we will build a type around it soon):

Module :: (Request -> m (Response, Module m)) -> Module m

This module maps a request to a response, but along the way also returns a new version of itself. Let's go further and make requests and responses polymorphic:

Module :: (a -> m (b, Module m a b)) -> Module m a b

Now if the output type of a module matches another module's input type, then you can compose them like regular functions. This composition is associative and has a polymorphic identity. This sounds a lot like a category, and in fact it is! It is a category, an applicative functor and an arrow.

newtype Module m a b =
    Module (a -> m (b, Module m a b))

instance (Monad m) => Applicative (Module m a)
instance (Monad m) => Arrow (Module m)
instance (Monad m) => Category (Module m)
instance (Monad m) => Functor (Module m a)

We can now compose two modules that have their own individual local state without even knowing about it! But that's not sufficient. We want more. How about modules that can be switched among? Let's extend our little module system such that modules can actually choose not to give an answer:

newtype Module m a b =
    Module (a -> m (Maybe b, Module m a b))

This allows another form of composition that is orthogonal to (.): Now our type is also a family of Alternative functors:

instance (Monad m) => Alternative (Module m a)

Now a module can choose whether to respond to a request, and if not, the next module will be tried. Simple. You have just reinvented the wire category. =)

Of course you don't need to reinvent this. The Netwire library implements this design pattern and comes with a large library of predefined "modules" (called wires). See the Control.Wire module for a tutorial.

Answer 2

Here's a concrete example of how to use lens like everybody else is talking about. In the following code example, Type1 is the local state (i.e. your hammer), and Type2 is the global state (i.e. your multitool). lens provides the zoom function which lets you run a localized state computation that zooms in on any field defined by a lens:

import Control.Lens
import Control.Monad.Trans.Class (lift)
import Control.Monad.Trans.State

data Type1 = Type1 {
    _field1 :: Int   ,
    _field2 :: Double}

field1 :: SimpleLens Type1 Int
field1 = lens _field1 (\x a -> x { _field1 = a})

field2 :: SimpleLens Type1 Double
field2 = lens _field2 (\x a -> x { _field2 = a})

data Type2 = Type2 {
    _type1  :: Type1 ,
    _field3 :: String}

type1 :: SimpleLens Type2 Type1
type1 = lens _type1 (\x a -> x { _type1 = a})

field3 :: SimpleLens Type2 String
field3 = lens _field3 (\x a -> x { _field3 = a})

localCode :: StateT Type1 IO ()
localCode = do
    field1 += 3
    field2 .= 5.0
    lift $ putStrLn "Done!"

globalCode :: StateT Type2 IO ()
globalCode = do
    f1 <- zoom type1 $ do
        localCode
        use field1
    field3 %= (++ show f1)
    f3 <- use field3
    lift $ putStrLn f3

main = runStateT globalCode (Type2 (Type1 9 4.0) "Hello: ")

zoom is not limited to immediate sub-fields of a type. Since lenses are composable, you can zoom as deep as you want in a single operation just by doing something like:

zoom (field1a . field2c . field3b . field4j) $ do ...

Answer 3

This sounds very much like an application of lenses.

Lenses are a specification of a sub-field of some data. The idea is you have some value toolLens and functions view and set so that view toolLens :: MultiTool h -> h fetches the tool and set toolLens :: MultiTool h -> h -> MultiTool h replaces it with a new value. Then you can easily define your withMember as a function just accepting a lens.

Lens technology has advanced a great deal recently, and they are now incredibly capable. The most powerful library around at the time of writing is Edward Kmett's lens library, which is a bit much to swallow, but pretty simple once you find the features you want. You can also search for more questions about lenses here on SO, e.g. Functional lenses which links to lenses, fclabels, data-accessor - which library for structure access and mutation is better, or the lenses tag.

Answer 4

I created a lensed extensible record library called data-diverse-lens which allows combining multiple ReaderT (or StateT) like this gist:

{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeApplications #-}

module Main where

import Control.Lens
import Control.Monad.Reader
import Control.Monad.State
import Data.Diverse.Lens
import Data.Semigroup

foo :: (MonadReader r m, HasItem' Int r, HasItem' String r) => m (Int, String)
foo = do
    i <- view (item' @Int) -- explicitly specify type
    s <- view item' -- type can also be inferred
    pure (i + 10, s <> "bar")

bar :: (MonadState s m, HasItem' Int s, HasItem' String s) => m ()
bar = do
    (item' @Int) %= (+10) -- explicitly specify type
    item' %= (<> "bar") -- type can also be inferred
    pure ()

main :: IO ()
main = do
    -- example of running ReaderT with multiple items
    (i, s) <- runReaderT foo ((2 :: Int) ./ "foo" ./ nil)
    putStrLn $ show i <> s -- prints out "12foobar"
    -- example of running StateT with multiple items
    is <- execStateT bar ((2 :: Int) ./ "foo" ./ nil)
    putStrLn $ show (view (item @Int) is) <> (view (item @String) is) -- prints out "12foobar"

Data.Has is a simpler library that does the same with tuples. Example from the library front page:

 {-# LANGUAGE FlexibleContexts #-}

 -- in some library code
 ...
 logInAnyReaderHasLogger :: (Has Logger r, MonadReader r m) => LogString -> m ()
 logInAnyReaderHasLogger s = asks getter >>= logWithLogger s

 queryInAnyReaderHasSQL :: (Has SqlBackEnd r, MonadReader r m) => Query -> m a
 queryInAnyReaderHasSQL q = asks getter >>= queryWithSQL q
 ...

 -- now you want to use these effects together
 ...
 logger <- initLogger  ...
 sql <- initSqlBackEnd ...

 (`runReader` (logger, sql)) $ do
       ...
       logInAnyReaderHasLogger ...
       ...
       x <- queryInAnyReaderHasSQL ...
       ...