Search code examples
haskellstatemonadswriterreader

Lifting a computation from the State monad into the RWS monad


I structure a computation around the use of the RWS (Reader+Writer+State) monad:

newtype Problem a = Problem { unProblem :: RWS MyEnv MyLog MyState a }
                    deriving ({- lots of typeclasses -})

The computation is built step by step by assembling elementary computations of the form

foo :: a -> Problem b

Sometimes, however, the subcomputations do not need the full power of the RWS monad. For instance, consider

bar :: c -> State MyState d

I would like to use bar as a part of a larger computation in the context of the Problem monad. I can see three ways of doing this, none of which seems very elegant to me.

  1. Manually unwrap the the State computation and rewrap it in the RWS monad:

    baz :: a -> RWS MyEnv MyLog MyState c
    baz x = do temp <- foo x
               initialState <- get
               let (finalResult, finalState) = runState (bar temp) initialState
               put finalState
               return finalResult
    
  2. Modify the type signature of bar by lifting it into the Problem monad. This has the downside that the new type signature does not explicitly promise that bar is independent of MyEnv and logs nothing to MyLog.

  3. Replace the RWS monad by an explicit ReaderT MyEnv WriterT MyLog State MyState monad stack. This allows me to concisely lift.lift the bar subcomputation into the full monad; however, this trick will not work e.g. for a subcomputation of the form c -> Reader MyEnv d.

Is there a cleaner way to compose foo and bar? I have a hunch that a few clever definitions of type-class instances might do the trick, but I can't see exactly how to proceed.


Solution

  • I assume you are using mtl (if you aren't, consider doing so - the libraries are mostly compatible, except for what follows). You can derive instances of MonadReader MyEnv, MonadWriter MyLog, and MonadState MyState. Then, you can use these to generalize your functions over any monad stack that has these sort of constraints.

    {-# LANGUAGE GeneralizedNewtypeDeriving, MultiParamTypeClasses, FlexibleContexts #-}
    
    import Control.Monad.RWS
    import Control.Monad.Reader
    import Control.Monad.Writer
    import Control.Monad.State
    
    newtype Problem a = Problem { unProblem :: RWS MyEnv MyLog MyState a }
                        deriving (Functor, Applicative, Monad,
                                  MonadReader MyEnv, MonadWriter MyLog, MonadState MyState)
    

    From your example, maybe bar only needs to know that there is some state MyState, so you can give it signature

    bar :: MonadState MyState m => c -> m d
    bar = ...
    

    Then, even for foo, which might need the full RWS capabilities, you could write

    foo :: (MonadState MyState m, MonadReader MyEnv m, MonadWriter MyLog m) => a -> m b
    foo = ...
    

    Then, you can mix and match these to your liking:

    baz :: Problem ()
    baz = foo 2 >> bar "hi" >> return ()
    

    Now, why is this useful in general? It boils down to the flexibility of your monad "stack". If tomorrow you decide that you don't actually need RWS but only State, you can refactor Problem accordingly and bar (which only ever needed state in the first place) will continue to work without any changes.

    In general, I try to avoid using WriterT, StateT, RWST etc inside auxiliary functions like foo or bar. Opt to keep your code there as generic and implementation-independent as possible using the typeclasses MonadWriter, MonadState, MonadReader, etc. Then, you only have to use WriterT, StateT, RWST once inside your code: when you actually "run" your monad.

    Side note about transformers

    If you are using transformers, none of this works. That isn't necessarily a bad thing: mtl, by virtue of always being able to "find" components (like state or writer) has some problems.