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.

Answer 2

This should do the trick:

takeWhileM :: (Monad m) => (a -> Bool) -> [m a] -> m [a]
takeWhileM p [] = return []
takeWhileM p (m:ms) = do 
    x <- m
    if p x
      then liftM (x:) (takeWhileM p ms) 
      else return []

See sepp2k's answer for an explanation of why you are losing laziness. The Identity monad or the nonempty list monad, for example, wouldn't have this problem.

Answer 3

You can't mapM an infinite list of Maybes. mapM is map followed by sequence. Here is the definition of sequence:

sequence ms = foldr k (return []) ms
  where
    k m m' = do { x <- m; xs <- m'; return (x:xs) }

From this we see that sequence evaluates every monadic value in the list. Since it's an infinite list, this operation will not terminate.

EDIT:

luqui and Carl make a good point that this doesn't generalize to any monad. To see why it doesn't work for Maybe we need to look at the implementation of (>>=):

(>>=) m k = case m of
    Just x  -> k x
    Nothing -> Nothing

The important point here is that we do a case on m. This makes the m strict because we have to evaluate it to figure out how to continue execution. Note that we're not casing on x here, so it remains lazy.