# 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.