How to takeWhile elements in a list wrapped in a monad

Question

Got a little puzzle I was wondering if you could help me clarify.

Let's define a function that returns a list:

let f = replicate 3

What we want to do is map this function to an infinite list, concatenate the results, and then take only things that match a predicate.

takeWhile (< 3) $ concatMap f [1..]

Great! That returns [1,1,1,2,2,2], which is what I want.

Now, I want to do something similar, but the function f now wraps its results in a Monad. In my usecase, this is the IO monad, but this works for discussing my problem:

let f' x = Just $ replicate 3 x

To map and concat, I can use:

fmap concat $ mapM f' [1..5]

That returns: Just [1,1,1,2,2,2,3,3,3,4,4,4,5,5,5]

If I want to use takeWhile, this still works:

fmap (takeWhile (< 3) . concat) $ mapM f' [1..5]

Which returns: Just [1,1,1,2,2,2]. Great!

But, if I make the list over which I map an infinite list this does not do what I expected:

fmap (takeWhile (< 3) . concat) $ mapM f' [1..]

Seems like the takeWhile is never happening. Somehow, I'm not getting the lazy computation I was expecting. I’m a bit lost.

Answer 1

The problem isn't that fmap + takeWhile doesn't work with infinite lists wrapped in a monad. The problem is that mapM can't produce an infinite list (at least not in the Maybe monad).

Think about it: If f' returns Nothing for any item in the list, mapM has to return Nothing. However mapM can't know whether that will happen until it has called f' on all items in the list. So it needs to iterate through the whole list before it knows whether the result is Nothing or Just. Obviously that's a problem with infinite lists.