List comprehensions in eight lines of Clojure
List comprehensions are a syntactic construct that allow you to create new lists in an elegant and concise way.
Here’s a question, though: If your language did not support list comprehensions, what would it take to add them? For a language like Ruby, you may have to dig into the language implementation itself. If you’re using a Lisp, you can write a macro.
Lisp macros allows you to transform your code at compile time. Macro calls are run before the rest of your code during a process called macro expansion. You define macros almost like any other function, except that the arguments to a macro are unevaluated Lisp code. Because Lisp code is just data, you can operate on it like any other data.
This is how you can implement list comprehensions in eight lines of Clojure.
(defmacro list-comp [[binding seq-expr & bindings] body-expr] (cond (not binding) `(list ~body-expr) (= :when binding) `(when ~seq-expr (list-comp ~bindings ~body-expr)) :else `(mapcat (fn [~binding] (list-comp ~bindings ~body-expr)) ~seq-expr)))
I’ve been thinking about this macro for three days now, and it’s still terribly exciting to me. It’s the macro I wish I had seen when I was learning about them. It takes full advantage of Clojure at compile-time and uses both higher-order functions and recursion.
Let’s see what it can do.
;; copy lists (list-comp [x (range 10)] x) ;;=> (0 1 2 3 4 5 6 7 8 9) ;; double every number (list-comp [x (range 5)] (* x 2)) ;;=> (0 2 4 6 8) ;; filter out elements (list-comp [x (range 10) :when (odd? x)] x) ;;=> (1 3 5 7 9) ;; nested bindings (list-comp [x "abc" y [0 1 2]] [x y]) ;;=> ([\a 0] [\a 1] [\a 2] [\b 0] [\b 1] [\b 2] [\c 0] [\c 1] [\c 2]) ;; list permutations (defn permutations [xs] (if-not (seq xs) (list ()) (list-comp [x xs ys (permutations (list-comp [z xs :when (not= z x)] z))] (conj ys x)))) (permutations [1 2 3]) ;;=> ((1 2 3) (1 3 2) (2 1 3) (2 3 1) (3 1 2) (3 2 1))
At compile time, it is transformed into code that produces the correct result. It’s a pretty handy macro. We can see how it works by looking at some macro expansions.
At its core,
list-comp is just a recursive transformation. This is made
obvious by looking at the first macroexpansion of a comprehension with
(macroexpand '(list-comp [x "ab" y [0 1]] [x y])) ;; yields (mapcat (fn [x] (list-comp [y [0 1]] [x y])) "ab") ;;=> ([\a 0] [\a 1] [\b 0] [\b 1])
If we simply assume that the inner call to
list-comp does the right
thing – that is, it successfully binds
y to 0 and then 1 –
then it is clear that this expansion works. The call to
take care of binding
x properly, and then concatenate the results.
This will work with an arbitrary level of nesting.
What happens when we use the
(macroexpand '(list-comp [x "ab" :when (= x \a) y [0 1]] [x y])) ;; yields (mapcat (fn [x] (when (= x \a) (list-comp (y [0 1]) [x y]))) "ab") ;;=> ([\a 0] [\a 1])
If the condition is truthy, we execute the inner comprehension.
nil, which, when concatenated with some
other list, returns that other list.
However, the skeptical reader just noticed a possible wrench in the
:when wraps the recursive call, what happens when there is
no recursive call left? That is, doesn’t everything break if
last form to appear in the bindings?
Thankfully, it doesn’t. The base-case is not the case of a single binding pair, but instead the case of an empty set of bindings. Then, we just return the body’s result as a singleton list.
(macroexpand '(list-comp  :foo)) ;; yields (list :foo)
This means that our
:when clause can appear at the end of the bindings
(clojure.walk/macroexpand-all '(list-comp [x (range 10) :when (even? x)] x)) ;; yields (mapcat (fn [x] (when (even? x) (list x))) ;; the empty-binding case (range 10)) ;;=> (0 2 4 6 8)
And that’s all there is to it.
Edit Jan 2, 2013: I should have mentioned this, but for posterity’s sake, Clojure comes with list comprehensions already built in. See Clojure’s for macro.