ローカルな関数でメモ化再帰できるようにするマクロをClojureで

Clojureにはメモ化のために使うmemoizeという関数がある。

memoizeの使用例

user=> (def f (memoize (fn [n] (println n) (* n n))))
#'user/f
user=> (f 2)  ; 1回目の呼び出しではprintlnが実行される
2
4
user=> (f 2)  ; 2回めの呼び出しではキャッシュされた値が使われるため、printlnが実行されない、
4
user=> 

しかし、memoizeは意外と使い勝手が悪い。
メモ化を使いたいシチュエーションというのは再帰関数の場合が多い。いわゆるメモ化再帰というやつ。ところが、defを使ってトップレベルにメモ化した関数を定義してしまうと、メモ化に使われたキャッシュがいつまでも残ってメモリリークになる。加えて、最近のClojureでは再帰呼び出しが遅延束縛にならなくなったため、「メモ化されていない関数をトップレベルに定義しておき、メモ化して使いたい箇所だけmemoizeで囲んで使う」というのも意味がない。かと言って、letを使ってローカルに関数を定義しようとすると、今度はletのスコープが再帰関数の定義を許してくれない。
このあたりの悩みに直面する人はわりといるみたいで、以下のブログ記事でも確認できる。

• Clojureの関数のメモ化 >> memoize 1.1 と 1.2 ~ Y箱
• tnoda-clojure • メモ化再帰関数

この問題に対する解決策はいくつかある。
解決策のひとつは動的束縛を使う方法だ。メモ化して使う関数を^:dynamicを使って定義しておき、メモ化して使う箇所だけmemoizeで囲んで束縛し直す。この方法を使うと、たとえばフィボナッチ数を返す関数の定義は以下のようになる。

動的束縛を使ったfibの定義

(defn ^:dynamic fib [n]
  (if (<= n 1)
    1
    (+ (#'fib (- n 1)) (#'fib (- n 2)))))

(binding [fib (memoize fib)]
  (fib 10))

この方法の場合、再帰呼び出しはVar経由で行う必要があることに注意する。ここでの例では、fibの再帰呼び出しが#'fibになっている。これは上で触れた、再帰呼び出しが遅延束縛にならなくなったことに起因する。

もうひとつの解決策は、上の2つめのブログ記事で挙げられているように、letとatomを使って再帰関数を定義する方法だ。

atomを使ったfibの定義

(let [fib (atom nil)]
  (reset! fib
          (memoize (fn [n]
                     (if (<= n 1)
                       1
                       (+ (@fib (- n 1)) (@fib (- n 2)))))))
  (@fib 10))

他にもやり方は考えられるが、いずれの場合にしても通常の再帰関数の定義からやや手を加えてやる必要がある。しかし、パターンは見えているので、マクロを使えばそのパターンをマクロの裏へ追いやってしまうことができる。今回はマクロを使って、ふたつめに挙げた方法をletfnと似た見た目に仕立ててみよう。以下がマクロの実装だ。

letfn-memoized

;; このマクロではsymbol-macroletを使っているため、tools.macro(https://github.com/clojure/tools.macro)が必要。
(defmacro letfn-memoized [fn-defs & body]
  (let [fnames (map first fn-defs)
        gensyms (into {} (map #(vector % (gensym)) fnames))]
    `(let [~@(mapcat #(list (gensyms %) '(atom nil)) fnames)]
       (symbol-macrolet [~@(mapcat #(list % `@~(gensyms %)) fnames)]
         ~@(for [[fname args & fbody] fn-defs]
             `(reset! ~(gensyms fname) (memoize (fn ~args ~@fbody))))
         ~@body))))

letfn-memoizedマクロを使うと、さっきのfibの定義はこう書き換えられる。

letfn-memoizedを使ったfibの定義

(letfn-memoized
  [(fib [n]
     (if (<= n 1)
       1
       (+ (fib (- n 1)) (fib (- n 2)))))]
  (fib 10))

これは以下のように展開され、

letfn-memoizedの展開

(let [G__1131 (atom nil)]
  (symbol-macrolet [fib @G__1131]
    (reset! G__1131
            (memoize (fn [n]
                       (if (<= n 1)
                         1
                         (+ (fib (- n 1)) (fib (- n 2)))))))
    (fib 10)))

最終的には、上のatomを使ったfibの定義と同じように展開される。

letfn-memoizedによって、晴れてletfnと同じ構文でメモ化再帰できるローカルな関数が定義できるようになった。

##あとがき
とはいえ、letfn-memoizedは十分複雑なマクロで、よほどいろんな箇所でメモ化再帰使っているというような場合でなければ、わざわざ定義して使うのもわりに合わないだろう。普段は動的束縛を使った方法を使っておくのが吉だと思う。