04 June 2010

Loop-Unrolling Macros in Clojure

I've been fooling around with clojure a bit lately. Something fun: a dotimes-like macro, but it completely unrolls your loop at compile time:
(defmacro dotimes-unroll
  "Just like dotimes, but completely unrolls the given loop.  The
number of iterations must read as a number."
  [[i n] & body]
  (assert (number? n))
  (if (= n 0)
    `nil
    `(let [~i ~(- n 1)]
       (dotimes-unroll [~i ~(- n 1)] ~@body)
       ~@body)))
Or, unroll by a given number of repetitions:
(defmacro dotimes-partial-unroll
  "Like dotimes, but partially unrolls the loop.  The number of
  repetitions attempted is the first argument to the macro, which must
  read as a literal number."
  [nstep [i nexpr] & body]
  (assert (number? nstep))
  `(let [n# (int ~nexpr)
         nmain# (quot n# ~nstep)
         nextra# (rem n# ~nstep)
         extrainc# (* nmain# ~nstep)]
     (loop [iouter# (int 0)]
       (if (>= iouter# nmain#)
         (dotimes [j# nextra#]
           (let [~i (int (+ j# extrainc#))]
             ~@body))
         (let [inc# (int (* iouter# ~nstep))]
           (dotimes-unroll [j# ~nstep]
             (let [~i (int (+ j# inc#))]
               ~@body))
           (recur (int (+ iouter# 1))))))))
I have no idea whether (or to what extent) this sort of loop unrolling is done by the JVM. So, these macros may be completely useless and redundant, but they were fun to write.