Haskell: iterate in State, how to force the behaviour I want?

Question

This is my first posting on SO, and I'm relatively new to Haskell, so please excuse any missteps or if my code is not idiomatic!

Consider the following two intuitive descriptions of: a, f(a), f(f(a))...

A. a list containing: a, the application of f to a, the application of f to that, the application of f to that...

B. a list containing, in the ith position, i nested applications of f to a.

My issue is that I got burnt trying to use the iterate function in Haskell to do A. My real application is a simulation, but the following contrived example highlights the problem.

import Control.Monad.State

example :: State Int [[String]]

step :: [String] -> State Int [String]
step l = do
         currentState <- get
         let result = if (currentState == 1)
                          then "foo":l
                          else "bar":l
         put (currentState + 1)
         return result

example = do
          sequence $ take 3 . iterate (>>= step) $ return []

With these definitions,

evalState example 1

results in:

[[],["foo"],["bar","bar"]]

Clearly, iterate does B, not A! Because the step function only ever adds something to the input list, step ["foo"] could not possibly result in ["bar", "bar"], no matter what the state!

Let me say that I do understand what is going on here, and also that - formally - the result is exactly "as it should be": step is a stateful function, so when f(a) comes up for evaluation as part of f(f(a)), it will be recomputed rather than taken from the second list item because the state has changed. I also realize I could avoid this in my real-life application by putting my accumulating list inside the state.

Nevertheless, there are two reasons for posting this.

First, the point is that iterate is frequently explained in a way that may potentially mislead a beginner to think it does A, when it actually does B. This includes Learn You A Haskell (which I otherwise found incredibly useful), but also post on SO (here and here, for example). In fact, the verbal explanation of iterate in LYAHFGG is almost exactly definition A above. So it might be useful to have a post on this that as a resource for other Haskell newbies who get a bug because of this and are searching for an explanation (so by all means do post more accurate, technical, better phrased, clarifications on the difference between A and B below).

Second, I would still be interested whether there is a function that will, in fact, do A! In other words, how can I, in the above stateful example, produce the list (with slight abuse of notation): [a, b = f(a), f(b), ...]? In other words, given

example2 = do
           firstResult <- step []
           secondResult <- step firstResult
           return $ [[], firstResult, secondResult]

for which

evalState example2 1

yields the desired result

[[],["foo"],["bar","foo"]]

How can I rewrite example2 using iterate?

On the beginners Haskell list, a related question regarding a memoizing version of iterate was posted. However, that query does not appear to have received an answer.

I'm not entirely sure laziness is really the problem in my application. Would a strict version of iterate do what I want? My own, naive, 'strict iterate' as below does not appear to make any difference.

iterate' f x = x : rest
               where previous = f x
                     rest = previous `seq` iterate f previous

Any insights on all of this would be much appreciated!

Answer 1

There is no difference between A. and B., they are the same thing by referential transparency.
The core of the problem seems to be that you're interpreting them in the context of execution of stateful computations. In that context, the analogue of A that you're expecting is
A': Produce a result list by 1. putting the result of the initial computation into the list, 2. determine the next computation from the previous result, execute it and append its result to the list, 3. goto 2.
The analogue of B is
B': Produce a list of computations (by iterating (>>= step)) and from that a result list by executing the computations one after the other.
For stateless computations, or when you pass the same initial state to all computations produced in B', the only difference is in efficiency, but if you're using sequence, each computation starts with a different state, so you get different results from A'. Breaking down your example, we have

actionList = take 3 $ iterate (>>= step) (return [])
           = [return [], return [] >>= step, return [] >>= step >>= step]

a list of actions (or monadic values) in State Int [String]. Now, when you apply sequence to that,

example = sequence actionList

you get an action that when executed, runs the first of these actions with the initial state, the second with the state as updated by the first, and the third with the state as updated by the second. To get the behaviour you expect, they all would have to be run with the same initial state.

Basically, a value of type State s v is a function of type s -> (v, s). iterate creates a list of functions, and sequence applys these functions, supplying different s arguments to them (each gets the s produced by the previous).

To get the desired behaviour, we must introduce a new combinator. Very nice, but usable in only very few Monads is

iterateM :: Monad m => (a -> m a) -> m a -> m [a]
iterateM step start = do
    first <- start
    rest <- iterateM step (step first)
    return (first:rest)

Which produces

Prelude Control.Monad Control.Monad.State> evalState (fmap (take 4) (iterateM step (return []))) 1
[[],["foo"],["bar","foo"],["bar","bar","foo"]]

But it works only in monads with sufficiently lazy (>>=), Control.Monad.State.Lazy is one of the few, Control.Monad.State.Strict is not. And even with C.M.S.Lazy, You cannot use the state after an iterateM, you have to put a new state before you can continue the computation. To get something usable with other monads, we could add a count parameter,

iterCountM :: (Monad m) => Int -> (a -> m a) -> m a -> m [a]
iterCountM 0 _ _ = return []
iterCountM k step start = do
    first <- start
    rest <- iterCountM (k-1) step (step fisrt)
    return (first:rest)

so we lose flexibility, but gain usability in more monads.