Monadic takeWhile

Say we have a list of actions:

printx :: (Show a) => a -> IO a
printx x = do
    print x
    return x

as :: [IO Int]
as = map printx [1..5]

We can then use sequence to execute the actions and collect the results:

*Main> sequence as
1
2
3
4
5
[1,2,3,4,5]

Note that the single-number lines are from print, and the last line is the list of values returned from the actions. A slightly less trivial action clarifies the distinction:

*Main> mapM (\x -> printx x >>= return . (*2)) [1..3]
1
2
3
[2,4,6]

Applying η-reduction yields (return . (*2) =<<) . printx—interesting but slightly ugly. Also note that mapM is defined in terms of sequence:

mapM f as = sequence (map f as)

Now say we want to stop executing these actions after a certain point:

*Main> takeWhile (< 3) $ sequence as

<interactive>:1:26:
    Couldn't match expected type `[a]' against inferred type `IO Int'
      Expected type: [[a]]
      Inferred type: [IO Int]
    In the first argument of `sequence', namely `as'
    In the second argument of `($)', namely `sequence as'

Unlike sequence, takeWhile is pure, so we have to lift it inside the monad:

*Main> liftM (takeWhile (< 3)) (sequence as)
1
2
3
4
5
[1,2]

The result is deceptive: although it is a subset, all the actions ran. Not a big deal in the case of printing to the standard output, but potentially disastrous if the canonical launchTheMissiles were lurking.

But isn't Haskell supposed to be lazy? To see why they all run, look at the definition of sequence:

sequence xs = foldr (liftM2 (:)) (return []) xs

liftM2 :: (Monad m) => (a -> b -> c) -> m a -> m b -> m c
liftM2 f m1 m2 = do
    x1 <- m1
    x2 <- m2
    return (f x1 x2)

We have no opportunity to prevent m2 because it runs unconditionally. We could redefine liftM2 as in the following:

sequenceWhile p xs = foldr (liftM2' (:)) (return []) xs
  where liftM2' f m1 m2 = do
          x1 <- m1
          guard $ p x1
          x2 <- m2
          return $ f x1 x2

But note the type:

*Main> :t sequenceWhile 
sequenceWhile :: (MonadPlus m) => (t -> Bool) -> [m t] -> m [t]

IO is not an instance of MonadPlus (a constraint due to guard), so we must be more deliberate:

sequenceWhile :: (Monad m) => (a -> Bool) -> [m a] -> m [a]
sequenceWhile p xs = foldr (liftM2' (:) []) (return []) xs
  where liftM2' f z m1 m2 = do
            x1 <- m1
            if p x1 then do x2 <- m2
                            return $ f x1 x2
                    else return z

Now we don't launch the missiles:

*Main> sequenceWhile (< 3) as
1
2
3
[1,2]

The definition also maintains laziness:

*Main> sequenceWhile (< 3) $ map printx [1..]
1
2
3
[1,2]

Monadic span would also work:

spanM :: (Monad m) => (a -> Bool) -> [m a] -> m ([a], [m a])
spanM _ [] = return ([], [])
spanM p (a:as) = do
  x <- a
  if p x then do (xs,as') <- spanM p as
                 return (x:xs, as')
         else return ([x], as)

Apply it as in the following example. Unlike span, spanM includes the failing value as the last element of the result list:

*Main> (xs,as') <- spanM (< 3) as
1
2
3
*Main> xs
[1,2,3]
*Main> sequence as'
4
5
[4,5]