Working with the State monad in Haskell

Question

I have been learning some Haskell little by little and am (slowly) working on understanding the State monad, attempting to write a function that repeats a State computation until the state meets some boolean test and collecting the returned values in a list for the overall result. I finally succeeded with this:

collectUntil :: (s -> Bool) -> State s a -> State s [a]
collectUntil f s = do s0 <- get
                      let (a,s') = runState s s0
                      put s'
                      if (f s') then return [a] else liftM (a:) $ collectUntil f s

so that

simpleState = state (\x -> (x,x+1))

*Main> evalState (collectUntil (>10) simpleState) 0
[0,1,2,3,4,5,6,7,8,9,10]

Is this a reasonable function for this task, or is there a more idiomatic way?

Answer 1

You are making exactly the same mistakes that I made when I first started writing monadic code - making it way too complicated, overusing liftM and underusing >>= (equivalently, underusing the <- notation).

Ideally, you shouldn't have to mention runState or evalState inside the state monad at all. The functionality you want is as follows:

• Read the current state
• If it satisfies the predicate f, then return
• If not, then run the computation s and add its result to the output

You can do this quite directly as:

collectUntil f comp = do
    s <- get                              -- Get the current state
    if f s then return []                 -- If it satisfies predicate, return
           else do                        -- Otherwise...
               x  <- comp                 -- Perform the computation s
               xs <- collectUntil f comp  -- Perform the rest of the computation
               return (x:xs)              -- Collect the results and return them

Note that you can nest do statements if they are part of the same monad! This is very useful - it allows you to branch within one do block, as long as both branches of the if statement lead to something of the same monadic type.

The inferred type for this function is:

collectUntil :: MonadState t m => (t -> Bool) -> m a -> m [a]

If you wish, you can specialise that to the State s type, although you don't have to:

collectUntil :: (s -> Bool) -> State s a -> State s [a]

It might even be preferable to keep the more general state, in case you want to use a different monad later.

What's the intuition?

Whenever s is a stateful computation and you are inside the state monad, you can do

x <- s

and x will now have the result of the computation (as if you'd called evalState and fed in an initial state). If you ever need to check the state, you can do

s' <- get

and s' will have the value of the current state.