How does foldr/build elimination work?

Or: How do Haskell "loops" (sometimes) go brrr?

For a declarative program optimization class this semester, I read the paper “Playing by the Rules,” which describes how GHC optimizes Haskell programs by rewriting terms into equivalent terms that ideally run more efficiently. These rewrites are defined by rules supplied by library authors or the developer.

The paper is a satisfying read—it’s very well written—but in particular it took me a very long time to understand one key optimization rule in the paper, called foldr/build elimination (a special case of short cut fusion). I assume it took such a long time because I am not so good at Haskell. This optimization was first introduced in another paper, “A Short Cut to Deforestation,” which is also well written and worth a read. Still, my path to understanding foldr/build elimination is different from the path outlined in that second paper, and I feel my perspective might be more approachable for people who are also not so good at Haskell.

The high level idea behind the foldr/build optimization is that sometimes we can avoid materializing lists in memory if those lists are produced (by build) and then immediately consumed by a foldr. For example, given a function down n that creates a list from n down to 1 (inclusive), foldr + 0 (down 5) computes the sum . As directly written, that computation requires materializing the intermediate list (5 : 4 : 3 : 2 : 1 : []). foldr/build elimination helps avoid materializing the intermediate in some circumstances; namely, when the list producer is written using build.

The first paper, “Playing by the Rules”, defines foldr, build, and the foldr/build rewrite rule as follows. (Read the definitions for build and the foldr/build rule, but don’t worry too much about understanding them; I will explain them after.)

foldr :: (a->b->b) -> b -> [a] -> b
foldr k z [] = z
foldr k z (x:xs) = k x (foldr k z xs)

build :: (forall b. (a->b->b) -> b -> b) -> [a]
build g = g (:) []

  forall k z (g::forall b.(a->b->b) -> b -> b) .
  foldr k z (build g) = g k z

Line 11—foldr k z (build g) = g k z—is the meat of the rewrite rule. Loosely, it says that a right associative fold (foldr) “cancels” with build. What does build do, you might ask? The Haskell docs say that build is “a list producer that can be fused with foldr.” Okay, so that’s also not very helpful. Also, for someone who is not fluent in Haskell, the type signature of build is not very enlightening. It accepts a function g with signature forall b. (a->b->b) -> b -> b. These mirror the arguments to foldr: the first argument of g, (a->b->b), looks like the binary operator, and the second argument looks like the initial element argument. But then g returns a value of type b, while build, which returns the result of g, returns a value of type [a] somehow? How does this typecheck? (More on this later.)

At this point I was lost. Here is the observation that sparked some understanding: building a list is in itself a right associative fold, where the binary operator is : (cons); and the initial element is [] (the empty list). For example, consider a function down x, which builds a list containing the values x down to 1. Evaluating down 5 is equivalent to:

down 5 (5 : (4 : (3 : (2 : (1 : [])))))

If I wanted to sum all the elements in down 5, I could call foldr + 0 (down 5). That would materialize the above list, and then fold over it immediately to compute a sum. But that would be slower than necessary; a faster program would just (somehow) replace : with + and replace [] with 0 inside down:

-- Faster sum of (down 5) that does not materialize an intermediate list
foldr + 0 (down 5)  (5 + (4 + (3 + (2 + (1 + 0)))))

This suggests that, to make some computations more efficient, down (and indeed, all functions that build lists) could have variants which additionally accept : and [] as arguments; that is, the variants would be generic over the “monoidal structure” of the fold. I will abbreviate these generic list building functions as GLBFs. By passing in different binary operators and initial elements to a GLBF, we could achieve different computations. For example, if we wanted a GLBF to return a materialized list, we could pass : and [] to it; if we wanted it to compute the sum of the elements, we could pass + and 0; if we wanted to compute the product, we could pass * and 1, and so on. Indeed, the paper does define such a variant, down':

down' 0 cons nil = nil
down' n cons nil = n `cons` (down' (n-1) cons nil)

We could then express our original down in terms of down' by specializing the latter for lists:

down n = down' n (:) []

But observe that the above definition for down is almost the same as the definition for build. So by creating a closure over down', we could instead define down in terms of build:

-- definition of build from above
build :: (forall b. (a->b->b) -> b -> b) -> [a]
build g = g (:) []

-- down, rewritten using build
down n = build (\cons nil -> down' n cons nil)

But what benefit does this get us? If build is just a wrapper function that invokes a GLBF to return a list, it doesn’t seem to be that useful—except that as a wrapper function, it’s extremely useful! build’s main purpose is to help GHC’s rewrite system identify opportunities to optimize foldr’s over a GLBF. That is, whenever the GHC rewrite system sees build, it knows that some GLBF is being called to materialize a list. If it sees that the list is immediately consumed by a call to foldr, it can remove both the build and foldr and instead call the GLBF directly, substituting the original arguments passed to foldr. In fact, that is exactly what the rewrite rule—foldr k z (build g) = g k z—says.

Without build, it would be much harder (and likely impossible!) for GHC to detect whether a function was indeed a GLBF so that it could perform this elimination. The following rewrite rule would likely be very difficult for GHC to find automatically, because it would have to reason about how down' uses k (cons) and z (nil) internally:

foldr k z (down' n (:) []) = down' n k z

So if all list building functions were defined without build, library authors would likely have to add one rewrite rule per (compatible) list-building function to get the same efficiency as foldr/build. “Standardizing” list building around build elegantly resolves this issue. A collorary is that if you are writing a library that produces big lists that people might fold on, if possible you should probably use build to build your list so that GHC can make your users’ code faster by performing foldr/build elimination. (But feel free to ignore this advice; I’ve written maybe 300 lines of Haskell in my life.)

Finally, why does g return b, but build returns [a]? (Recall the type of g is forall b. (a->b->b) -> b -> b).) The only way for this to typecheck is for b ≡ [a], and within build this is indeed the case: replacing b with [a], we see that : has the type (a->b->b) ≡ (a->[a]->[a]) and [] has the type b ≡ [a]. So the polymorphism forall b. is not necessary when g is called within build.

However, the polymorphism is necessary for foldr/build elimination. After elimination occurs, we get the term g k z, where g is a GLBF, and k is an arbitrary binary operator and z is an initial element of arbitrary type. Therefore, we do need g to also be polymorphic in its inputs and outputs so that the right-hand side of the rewrite is correctly typed. (Thanks to quchen for pointing this out in another article about build.)