Deriving Optimal List Comprehensions

Introduction

List comprehensions are a feature of various programming languages, like Haskell, Python and Erlang, in order to make nested traversals, and filterings of lists convenient to read and write. As a simple example, taking the difference of two lists can be expressed in Haskell as:

difference :: [a] -> [a] -> [a]
difference xs ys = [x | x <- xs, not (x `elem` ys)]

and the cartesian product of two lists can be expressed as

cartesianProduct :: [a] -> [b] -> [(a,b)]
cartesianProduct xs ys = [(x,y) | x <- xs, y <- ys]

This is all very well and good, but how can we translate this into a language without comprehensions? Say, for instance, we wished to add a list comprehension macro to Scheme[1]. (Throughout the rest of the post, I am assuming that the comprehensions are pure, and will be using a notation vaguely inspired by Haskell.)

A simple translation

Well, it's pretty easy to see how you would go about doing this. Nested traversals can be expressed simply by nesting uses of the map function, though this leads to a result of nested lists, rather than a single list. But we can use the concat function after all but the innermost call to map, to flatten it again.

Treating the innermost map specially leads to an additional rule, and it is simpler to treat it similar to the others, and modify the result expression to return a singleton list. Doing this also simplifies the way we express guards, which can now return the empty list in the false case. This is a common trick written up by Wadler in “How to replace failure by a list of successes” in 1985.

Putting this altogether, we can write these rules down as

                           [e|] ≝ [e]
[e| var <- generator, rest ...] ≝ concatMap (\var -> [e|rest ...]) generator
           [e| guard, rest ...] ≝ if guard then [e | rest ...] else []

or as a Scheme syntax-rules macro as

(import (only (srfi :1 lists) append-map))

(define-syntax comp
  (syntax-rules (<-)
    ((comp expression)
     (list expression))
    ((comp expr (var <- generator) rest ...)
     (append-map (lambda (var)
                   (comp expr rest ...))
                 generator))
    ((comp expr guard rest ...)
     (if guard
         (comp expr rest ...)
         '()))))

An optimal translation

For explaining how list comprehensions can be translated into a language without them, this is pretty nice, and the translation above is similar to the one given in the Haskell 2010 report, but it is not the best we can do in terms of efficiency. We build up, and then later destruct lists at every level, when we should be able to get away with only needing one cons operation per element in the resulting list. Now, a sufficiently smart compiler could fix this, but most compilers aren't sufficiently smart.

In Simon Peyton Jone's 1987 book, “The Implementation of Functional Programming Languages”(available online), Phil Wadler contributes a chapter on List Comprehensions, and gives such a transformation. A modified version of his rules[2] is

             [e|generators ...] ≝ {e|generators ...} []
                           {e|} ≝ (\s -> e : s)
{e| var <- generator, rest ...} ≝ (\s -> foldr (\var s' -> {e | rest ...} s') s generator)
           {e| guard, rest ...} ≝ (\s -> if guard then {e | rest ...} s else s)

or a syntax-rules implementation is

(define-syntax comp
  (syntax-rules (<-)
    ((comp expression generators ...)
     ((comp* expression generators ...) '()))))

(define-syntax comp*
  (syntax-rules (<-)
    ((comp expression)
     (lambda (s)
       (cons expression s)))
    ((comp expr (var <- generator) rest ...)
     (lambda (s)
       (fold-right (lambda (var s*)
                     ((comp* expr rest ...) s*))
                   s
                   generator)))
    ((comp expr guard rest ...)
     (lambda (s)
       (if guard
           ((comp* expr rest ...) s)
           s)))))

Explaining the translation

Now, if all you are doing is implementing these you might stop there[3], but if you are like me it isn't good enough to just see the answers. You want to know how to get there for yourself.

One approach is to work with the original definition, and to use various laws to transform it into the optimal version. However, to me, a different approach was more helpful. Specifically, I started with an optimal stateful version, and was able to gradually change that into the pure version.

The stateful version

So, what is the optimal stateful version? Well, it's quite obvious really. Traversing a generator is a for-each, guards are a use of if, and for the inner expression you turn this into a mutation of a comprehension-local variable.

             [e|generators ...] ≝ let state = [] in {e|generators ...}; state
                           {e|} ≝ state := state ++ [e]
{e| var <- generator, rest ...} ≝ foreach (\var -> {e|rest ...}) generator
           {e| guard, rest ...} ≝ if guard then {e | rest ...} s else ()

From Stateful to Pure

So, now that we have that, we need to find a way to obtain a pure version, based on it. The details are involved, so I won't be including them in this post. The high level idea should be enough to give you the flavour.

So, you have a stateful version of something, and you want to make it functional. An easy way, then, is to just use a state monad. The imperative foreach translates nicely to the mapM_ function on the state monad, and since we only want the returned state value, we call execState on the result.

             [e|generators ...] ≝ execState {e|generators ...} []
                           {e|} ≝ modify (++ [e])
{e| var <- generator, rest ...} ≝ mapM_ (\var -> {e|rest ...}) generator
           {e| guard, rest ...} ≝ if guard then {e | rest ...} else return ()

So far so good, but hang on, we're doing a left to right traversal on a list, updating a state parameter, and returning the state parameter at the end. Doesn't this sound familiar to you? If you are a functional programmer it should: it's a left fold. Making the state explicit as an argument, using foldl, and performing some simplification we get

             [e|generators ...] ≝ {e|generators ...} []
                           {e|} ≝ (\s -> s ++ [e])
{e| var <- generator, rest ...} ≝ (\s -> foldl (\s' var -> {e|rest ...} s') s generator
           {e| guard, rest ...} ≝ (\x -> if guard then {e | rest ...} s else s)

This is looking better, but we have a nasty foldl, whereas we use a foldr in the optimal version, and we're still appending to the end. Fortunately, it is known to many functional programmers that

foldl f a xs = foldr (flip f) a (reverse xs)

so we can just apply that law.

             [e|generators ...] ≝ {e|generators ...} []
                           {e|} ≝ (\s -> s ++ [e])
{e| var <- generator, rest ...} ≝ (\s -> foldr (\var s' -> {e|rest ...} s') s (reverse generator)
           {e| guard, rest ...} ≝ (\x -> if guard then {e | rest ...} s else s)

Almost Done! Now notice that we are traversing the generator expressions in reverse, and we are likewise building the result from the front to the end. Intuitively, we can just switch to not reversing the generator, and building from the end to the front, and we are done.

             [e|generators ...] ≝ {e|generators ...} []
                           {e|} ≝ (\s -> [e] ++ s)
{e| var <- generator, rest ...} ≝ (\s -> foldr (\var s' -> {e|rest ...} s') s generator)
           {e| guard, rest ...} ≝ (\x -> if guard then {e | rest ...} s else s)

Conclusion

So, that's an intuitive picture of how you get from the stateful version, to the pure version, and I don't blame you if you think I've been hand-wavy. I have. But fear not, I have worked out all the ugly details (or at least, I think I have), though I will save them for another post.

Before I leave you, a quick bit of the history. This post was actually a year in the making. In February of last year I came across a paper by Guy Lapalme (pdf) describing an implementation of list comprehensions as a Common Lisp defmacro. The paper wasn't much interesting in and of itself, but it pointed me to Wadler's rules. When I saw that, I felt that I had to try and work it out for myself.

The full disclosure is that when I attempted last year, I actually failed. Being an “eager” functional programmer, I came up with the above foldl version, and went from there to a version that built the result list in reverse and then did a giant reverse at the end. This isn't optimal, insofar as you pay an additional cons for each one in the result list, however it isn't bad. And at the time, in what was probably sour grapes, I claimed a “moral success”, since without lazy guarded recursion, the above version is likely to give a stack overflow for large lists.

Some talk of comprehensions on the guile IRC channel prompted me to revisit this last week, and procrastinate on posting it till this week. :)

Footnotes

[1]In practice, I don't recommend using the translation I give since, people have written other, more featureful solutions already. Common Lisp has the infamous LOOP macro, and for Scheme I recommend using foof-loop, written by Taylor Campbell. Foof-loop favours parallel traversal rather than nested, although there are nested extensions. Guile users with Guildhall installed can obtain foof-loop by guild install wak-foof-loop

[2] In research for this article, I noticed that the rules I give using foldr are similar to a set given in a note in the GHC source code. See the subsection “Foldr/Build desugaring of list comprehensions” in the file compiler/deSugar/DsListComp.lhs

[3] If you are using a Scheme implementation like Guile, which does some optimisation, you actually get some decent looking code. Check for yourself with the ,optimize meta-command, but of course you can go further and inline the definition of foldr in the macro.