Comprehension framework, upon which are (re)implemented, for, doseq, reducible/foldable for and reduce-based doseq

comprehensions.clj

;; I wrote this in the Eurostar on my way back from the last lambdanext.eu clojure course.

(ns comprehensions
  (:refer-clojure :exclude [for doseq])
  (:require [clojure.core.reducers :as r]))

;; borrowed from clojure.core
(defmacro ^{:private true} assert-args
  [& pairs]
  `(do (when-not ~(first pairs)
         (throw (IllegalArgumentException.
                  (str (first ~'&form) " requires " ~(second pairs) " in " ~'*ns* ":" (:line (meta ~'&form))))))
     ~(let [more (nnext pairs)]
        (when more
          (list* `assert-args more)))))

(defn emit-comprehension [&form {:keys [emit-other emit-inner]} seq-exprs body-expr]
  (assert-args
     (vector? seq-exprs) "a vector for its binding"
     (even? (count seq-exprs)) "an even number of forms in binding vector")
  (let [groups (reduce (fn [groups [k v]]
                         (if (keyword? k)
                           (conj (pop groups) (conj (peek groups) [k v]))
                           (conj groups [k v])))
                 [] (partition 2 seq-exprs))
        inner-group (peek groups)
        other-groups (pop groups)]
    (reduce emit-other (emit-inner body-expr inner-group) other-groups)))

(defn- do-mod [mod-pairs cont & {:keys [skip stop]}]
  (let [err (fn [& msg] (throw (IllegalArgumentException. ^String (apply str msg))))]
    (reduce 
      (fn [cont [k v]]
        (cond 
          (= k :let) `(let ~v ~cont)
          (= k :while) `(if ~v ~cont ~stop)
          (= k :when) `(if ~v ~cont ~skip)
          :else (err "Invalid 'for' keyword " k)))
      cont (reverse mod-pairs))))

(defmacro for
  "List comprehension. Takes a vector of one or more
   binding-form/collection-expr pairs, each followed by zero or more
   modifiers, and yields a lazy sequence of evaluations of expr.
   Collections are iterated in a nested fashion, rightmost fastest,
   and nested coll-exprs can refer to bindings created in prior
   binding-forms.  Supported modifiers are: :let [binding-form expr ...],
   :while test, :when test.

  (take 100 (for [x (range 100000000) y (range 1000000) :while (< y x)] [x y]))"
  {:added "1.0"}
  [seq-exprs body-expr]
  (let [emit-other
        (fn [sub-expr [bind expr & mod-pairs]]
          (let [giter (gensym "iter__")
                gxs (gensym "s__")]
            #_"not the inner-most loop"
            `((fn ~giter [~gxs]
                (lazy-seq
                  (loop [~gxs ~gxs]
                    (when-first [~bind ~gxs]
                      ~(do-mod mod-pairs
                         `(let [fs# (seq ~sub-expr)]
                            (if fs#
                              (concat fs# (~giter (rest ~gxs)))
                              (recur (rest ~gxs))))
                         :skip `(recur (rest ~gxs))
                         :stop nil)))))
               ~expr)))
        emit-inner 
        (fn [body-expr [bind expr & mod-pairs]]
          #_"inner-most loop"
          (let [giter (gensym "iter__")
                gxs (gensym "s__")
                gi (gensym "i__")
                gb (gensym "b__")]
            `((fn ~giter [~gxs]
                (lazy-seq
                  (loop [~gxs ~gxs]
                    (when-let [~gxs (seq ~gxs)]
                      (if (chunked-seq? ~gxs)
                        (let [c# (chunk-first ~gxs)
                              size# (int (count c#))
                              ~gb (chunk-buffer size#)]
                          (if (loop [~gi (int 0)]
                               (if (< ~gi size#)
                                 (let [~bind (.nth c# ~gi)]
                                   ~(do-mod mod-pairs
                                      `(do (chunk-append ~gb ~body-expr)
                                         (recur (unchecked-inc ~gi))) 
                                      :skip `(recur (unchecked-inc ~gi))
                                      :stop nil))
                                 true))
                            (chunk-cons
                              (chunk ~gb)
                             (~giter (chunk-rest ~gxs)))
                            (chunk-cons (chunk ~gb) nil)))
                        (let [~bind (first ~gxs)]
                          ~(do-mod mod-pairs
                             `(cons ~body-expr
                                (~giter (rest ~gxs)))
                             :skip `(recur (rest ~gxs))
                             :stop nil)))))))
               ~expr)))]
    (emit-comprehension &form
      {:emit-other emit-other :emit-inner emit-inner}
      seq-exprs body-expr)))

(defmacro doseq
  "Repeatedly executes body (presumably for side-effects) with
  bindings and filtering as provided by \"for\".  Does not retain
  the head of the sequence. Returns nil."
  {:added "1.0"}
  [seq-exprs & body-expr]
  (let [emit-loop
        (fn [body-expr [bind expr & mod-pairs]]
          #_"inner-most loop"
          (let [giter (gensym "iter__")
                gxs (gensym "s__")
                gi (gensym "i__")]
            `(loop [~gxs ~expr]
               (when-let [~gxs (seq ~gxs)]
                 (if (chunked-seq? ~gxs)
                   (let [c# (chunk-first ~gxs)
                         size# (int (count c#))]
                     (when (loop [~gi (int 0)]
                             (if (< ~gi size#)
                               (let [~bind (.nth c# ~gi)]
                                    ~(do-mod mod-pairs
                                       `(do ~@body-expr
                                          (recur (unchecked-inc ~gi))) 
                                       :skip `(recur (unchecked-inc ~gi))
                                       :stop nil))
                               true))
                       (recur (chunk-rest ~gxs)))
                     (let [~bind (first ~gxs)]
                       ~(do-mod mod-pairs
                          `(do  ~body-expr
                             (recur (rest ~gxs)))
                          :skip `(recur (rest ~gxs))
                          :stop nil))))))))]
    (emit-comprehension &form 
      {:emit-other emit-loop :emit-inner emit-loop}
      seq-exprs body-expr)))

(defn- reducer
  "Like clojure.core.reducers/reducer but if coll is a map then
   uses kv-reduce."
  ([coll xf]
     (reify
      clojure.core.protocols/CollReduce
      (coll-reduce [this f1]
        (r/reduce f1 (f1) this))
      (coll-reduce [_ f1 init]
        (r/reduce (xf f1) init coll)))))

(defn- folder
  "Like clojure.core.reducers/folder but if coll is a map then
   uses kv-reduce."
  ([coll xf]
     (reify
      clojure.core.protocols/CollReduce
      (coll-reduce [this f1]
        (r/reduce f1 (f1) this))
      (coll-reduce [_ f1 init]
        (r/reduce (xf f1) init coll))

      r/CollFold
      (coll-fold [_ n combinef reducef]
                 (r/coll-fold coll n combinef (xf reducef))))))

(defmacro rfor
  "Reducer comprehension, behaves like \"for\" but yields a reducible collection.
   Leverages kv-reduce when destructuring and iterating over a map."
  {:added "1.0"}
  [seq-exprs body-expr]
  (letfn [(emit-fn [form] 
            (fn [sub-expr [bind expr & mod-pairs]]
              (let [foldable (not-any? (comp #{:while} first) mod-pairs)
                    kv-able (and (vector? bind) (not-any? #{:as} bind)
                              (every? #(and (symbol? %) (not= % '&)) (take 2 bind)))
                    [kv-args kv-bind] 
                    (if kv-able
                      [(take 2 (concat bind (repeat `_#)))
                       (if (< 2 (count bind)) 
                         [(subvec bind 2) nil]
                         [])]
                      `[[k# v#] [~bind (clojure.lang.MapEntry. k# v#)]])
                    combiner (if kv-able
                               (if foldable `folder `reducer)
                               (if foldable `r/folder `r/reducer))
                    f (gensym "f__")
                    ret (gensym "ret__")
                    body (do-mod mod-pairs (form f ret sub-expr)
                           :skip ret
                           :stop `(reduced ~ret))]
                `(~combiner ~expr
                   (fn [~f]
                     (fn
                       ([] (~f))
                       ([~ret ~bind] ~body)
                       ([~ret ~@kv-args] (let ~kv-bind ~body))))))))]
    (emit-comprehension &form
      {:emit-other (emit-fn (partial list `r/reduce)) :emit-inner (emit-fn list)}
      seq-exprs body-expr)))

(defmacro rdoseq "doseq but based on reducers, leverages kv-reduce when iterating on maps."
  [bindings & body]
 `(reduce (constantly nil) (rfor ~bindings (do ~@body))))