Laziness of (>>=) in folding

Question

Consider the following 2 expressions in Haskell:

foldl' (>>=) Nothing (repeat (\y -> Just (y+1)))
foldM (\x y -> if x==0 then Nothing else Just (x+y)) (-10) (repeat 1)

The first one takes forever, because it's trying to evaluate the infinite expression

...(((Nothing >>= f) >>= f) >>=f)...

and Haskell will just try to evaluate it inside out.

The second expression, however, gives Nothing right away. I've always thought foldM was just doing fold using (>>=), but then it would run into the same problem. So it's doing something more clever here - once it hits Nothing it knows to stop. How does foldM actually work?

Answer 1

foldM can't be implemented using foldl . It needs the power of foldr to be able to stop short. Before we get there, here's a version without anything fancy.

foldM f b [] = return b
foldM f b (x : xs) = f b x >>= \q -> foldM f q xs

We can transform this into a version that uses foldr . First we flip it around:

foldM f b0 xs = foldM' xs b0 where
  foldM' [] b = return b
  foldM' (x : xs) b = f b x >>= foldM' xs

Then move the last argument over:

  foldM' [] = return
  foldM' (x : xs) = \b -> f b x >>= foldM' xs

And then recognize the foldr pattern:

  foldM' = foldr go return where
    go x r = \b -> f b x >>= r

Finally, we can inline foldM' and move b back to the left:

foldM f b0 xs = foldr go return xs b0 where
  go x r b = f b x >>= r

This same general approach works for all sorts of situations where you want to pass an accumulator from left to right within a right fold. You first shift the accumulator all the way over to the right so you can use foldr to build a function that takes an accumulator, instead of trying to build the final result directly. Joachim Breitner did a lot of work to create the Call Arity compiler analysis for GHC 7.10 that helps GHC optimize functions written this way. The main reason to want to do so is that it allows them to participate in the GHC list libraries' fusion framework.

Answer 2

One way to define foldl in terms of foldr is:

foldl f z xn = foldr (\ x g y -> g (f y x)) id xn z

It's probably worth working out why that is for yourself. It can be re-written using >>> from Control.Arrow as

foldl f z xn = foldr (>>>) id (map (flip f) xn) z

The monadic equivalent of >>> is

f >=> g = \ x -> f x >>= \ y -> g y

which allows us to guess that foldM might be

foldM f z xn = foldr (>=>) return (map (flip f) xn) z

which turns out to be the correct definition. It can be re-written using foldr/map as

foldM f z xn = foldr (\ x g y -> f y x >>= g) return xn z