Nested folds

I’m finally getting around to reading Algorithm Design with Haskell (hereafter abbreviated as ADH), by Jeremy Gibbons and Richard Bird. I’ve had it for a while, and I have no excuse for waiting this long to read it, but anyway. I’m enjoying it so far, and wanted to share something I (indirectly) learned. I’m sure there are some who already know this, but I didn’t. I’ll share both the fun takeaway and then also the interesting, roundabout path I took to get there.

Composed folds are nested folds

Here’s the punchline:

• foldl :: (b -> a -> b) -> b -> [a] -> b
• foldl . foldl :: (b -> a -> b) -> b -> [[a]] -> b
• foldl . foldl . foldl :: (b -> a -> b) -> b -> [[[a]]] -> b

Actually, it’s a bit more general than this, since foldl works over any Foldable, not just lists: in fact, foldl . foldl . foldl can be used to fold any t1 (t2 (t3 a)) as long as t1, t2, and t3 are all instances of Foldable. For example, here is how we can add up all the integers contained in a Maybe (Tree [Int]):

λ> (foldl . foldl . foldl) (+) 0 (Just (Node [1,2,3] [Node [5,6] [], Node [7] [], Node [9,12] [Node [6] []]]))
51

We can make sense of this if we look at the type of foldl. Its type is

Foldable t => (b -> a -> b) -> b -> t a -> b

and we usually think of it as taking three arguments: a combining function of type b -> a -> b, an initial value of type b, and a structure to fold. But we can also think of it as a one-argument function. That is, it takes a function of type b -> a -> b and transforms it into a function of type b -> t a -> b:

foldl :: Foldable t => (b -> a -> b) -> (b -> t a -> b)

Due to the magic of currying this is equivalent, and the second set of parentheses above is redundant. However, with this change of perspective it is easy to see what’s going on: the result of foldl is a function of type b -> t a -> b, which has the right shape to be the argument of foldl again, but this time with t a in place of a, yielding

(b -> t a -> b) -> (b -> t2 (t a) -> b)

and so on.

What about foldr?

foldr :: Foldable t => (a -> b -> b) -> (b -> t a -> b)

The shapes of the input and output to foldr don’t quite match, but they will if we throw in an extra flip, which switches the arguments of a two-argument function. So we can either do

foldr . flip :: Foldable t => (b -> a -> b) -> (b -> t a -> b)

if we want to match the type of foldl, or

flip . foldr :: Foldable t => (a -> b -> b) -> (t a -> b -> b)

In any case, we can now iterate just as with foldl; for example

foldr . flip . foldr . flip . foldr :: (a -> b -> b) -> (b -> t1 (t2 (t3 a)) -> b)

(with some appropriate Foldable constraints thrown in as well).

My roundabout journey

So how did I come to realize this? Sometimes the journey is more interesting than the destination. The first chapter of ADH talks about the foldr fusion rule, which says that

h . foldr f e == foldr g (h e)

as long as h (f x y) == g x (h y) for all x and y. In other words, if we have a foldr followed by a function h, we can turn this into a single foldr (i.e. “fuse away” the h) as long as we can find an appropriate function g that satisfies the given criterion.

One of the exercises asks us to use foldr fusion to simplify

foldr f e . concat

which performs a fold over a nested list by first flattening the list and then doing a fold. You may wish to go try it yourself before reading on!

The solution

We can compute as follows, where g is some function we will need to define appropriately:

  foldr f e . concat
=                              { definition of concat }
  foldr f e . foldr (++) []
=                              { foldr fusion law, with h = foldr f e }
  foldr g (foldr f e [])
=                              { definition of foldr }
  foldr g e

According to the fusion condition, we need g to satisfy foldr f e (x ++ y) == g x (foldr f e y). This was already given as a previous exercise; I actually solved it by thinking about how foldr works and intuiting the right answer, but we can also calculate it using a second round of foldr fusion:

  g x (foldr f e y)
=
  foldr f e (x ++ y)
=                              { definition of (++) }
  foldr f e (foldr (:) y x)
=                              { foldr fusion }
  foldr f (foldr f e y) x

The last equation follows since foldr f e (a : b) = f a (foldr f e b) by definition of foldr. Hence, using z in place of foldr f e y, we have g x z = foldr f z x = flip (foldr f) x z, and so g = flip (foldr f), and we have

foldr f e . concat = foldr (flip (foldr f)) e

We can simplify even further: if we define nestedFold f e = foldr f e . concat = foldr (flip (foldr f)) e then we can eta-reduce to get nestedFold = foldr . flip . foldr.

When I derived this, my eyes bugged out and I started playing around with it, which is how I ended up figuring out the thing with foldl. Presumably, one could use the similar foldl fusion law to fuse foldl f e . concat and end up deriving the foldl . foldl result; I’ll leave that as an exercise for interested readers.