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 (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 (:) []
{-# RULES
"foldr/build"
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 zSo 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.)