Tips for more elegant code with monads?

Question

I finally got a hold on how to use monads (don't know if I understand them...), but my code is never very elegant. I guess is from a lack of grip on how all those functions on Control.Monad can really help. So I'd thought it would be nice to ask for tips on this in a particular piece of code using the state monad.

The goal of the code is to calculate many kinds of random walks, and it's something I'm trying to do before something more complicated. The problem is that I have two stateful computations at the same time, and I'd like to know how to compose them with elegance:

1. The function that updates the random number generator is something of type Seed -> (DeltaPosition, Seed)
2. The function that updates the position of the random walker is something of type DeltaPosition -> Position -> (Log, Position) (where Log is just some way for me to report what is the current position of the random walker).

What I've done is this:

I have a function to compose this two stateful computations:

composing :: (g -> (b, g)) -> (b -> s -> (v,s)) -> (s,g) -> (v, (s, g))
composing generate update (st1, gen1) = let (rnd, gen2) = generate gen1
                                            (val, st2)  = update rnd st1
                                        in (val, (st2, gen2))

and then I turn it into a function that compose states:

stateComposed :: State g b -> (b -> State s v) -> State (s,g) v
stateComposed rndmizer updater = let generate = runState rndmizer
                                     update x = runState $ updater x
                                 in  State $ composing generate update 

And then I have the simplest thing, for example, a random walker that will just sum a random number to its current position:

update :: Double -> State Double Double
update x = State (\y -> let z = x+y
                        in  (z,z))

generate :: State StdGen Double
generate = State random

rolling1 = stateComposed generate update 

and a function to do this repeatedly:

rollingN 1 = liftM (:[]) rolling1
rollingN n = liftM2 (:) rolling1 rollings
    where rollings = rollingN (n-1) 

And then, if I load this in ghci and run:

*Main> evalState (rollingN 5) (0,mkStdGen 0)
[0.9872770354820595,0.9882724161698186,1.9620425108498993,2.0923229488759123,2.296045158010918]

I get what I want, which is a list of the positions occupied by the random walker. But... I feel there must be a more elegant way to do this. I have two questions:

1. Can I rewrite those functions in a more "monadic" way, using clever functions from Control.Monad?

2. Is there a general pattern about combining states like this that can be used? Does this have something to do with monad transformers or something like that?

Answer 1

Update: I should have mentioned that there's actually a much nicer way to do this that doesn't require State or monads at all:

takeStep :: (Double, StdGen) -> (Double, StdGen)
takeStep (p, g) = let (d, g') = random g in (p + d, g')

takeSteps n = take n . tail . map fst $ iterate takeStep (0, mkStdGen 0)

It works as desired:

*Main> takeSteps 5
[0.9872770354820595,0.9882724161698186,1.9620425108498993,2.0923229488759123,2.296045158010918]

---

If you're not committed to the idea of "composing" two separate stateful computations, you can accomplish the same thing much more straightforwardly:

takeStep :: State (Double, StdGen) Double
takeStep = do
  (pos, gen) <- get
  let (delta, gen') = random gen
  let pos' = pos + delta
  put (pos', gen')
  return pos'

takeSteps n = evalState (replicateM n takeStep) (0, mkStdGen 0)

This produces the same output as your example:

*Main> takeSteps 5
[0.9872770354820595,0.9882724161698186,1.9620425108498993,2.0923229488759123,2.296045158010918]

This approach (doing all the state manipulation in a single monad instead of trying to compose a State A and State B) seems to me like the most elegant solution.

---

Update: To answer your question about using monad transformers to stack State monads: it's certainly possible. We can write the following, for example:

update' :: (Monad m) => Double -> StateT Double m Double
update' x = StateT $ \y -> let z = x + y in return (z, z)

generate' :: (Monad m) => StateT StdGen m Double
generate' = StateT $ return . random

takeStep' :: StateT Double (State StdGen) Double
takeStep' = update' =<< lift generate'

takeSteps' n = evalState (evalStateT (replicateM n takeStep') 0) $ mkStdGen 0

We could also do the stacking in the opposite order.

This version again produces the same output, but in my opinion the non-StateT version is a bit clearer.

Answer 2

The usual way to compose 2 monads (and the only way for most monads) is with monad transformers, but with different State monads you have more options. For example: you could use these functions:

leftState :: State a r -> State (a,b) r
leftState act = state $ \ ~(a,b) -> let
  (r,a') = runState act a
  in (r,(a',b))

rightState :: State b r -> State (a,b) r
rightState act = state $ \ ~(a,b) -> let
  (r,b') = runState act b
  in (r,(a,b'))