In Clojure, is it possible to combine memoization and tail call optimization?

Question

In clojure, I would like to write a tail-recursive function that memoizes its intermediate results for subsequent calls.

[EDIT: this question has been rewritten using gcd as an example instead of factorial.]

The memoized gcd (greatest common divisor) could be implemented like this:

(def gcd (memoize (fn [a b] 
   (if (zero? b) 
       a 
       (recur b (mod a b)))) 

In this implementation, intermediate results are not memoized for subsequent calls. For example, in order to calculate gcd(9,6), gcd(6,3) is called as an intermediate result. However, gcd(6,3) is not stored in the cache of the memoized function because the recursion point of recur is the anonymous function that is not memoized.

Therefore, if after having called gcd(9,6), we call gcd(6,3) we won't benefit from the memoization.

The only solution I can think about will be to use mundane recursion (explicitely call gcd instead of recur) but then we will not benefit from Tail Call Optimization.

Bottom Line

Is there a way to achieve both:

1. Tail call optimization
2. Memoization of intermediate results for subsequent calls

Remarks

1. This question is similar to Combine memoization and tail-recursion. But all the answers there are related to F#. Here, I am looking for an answer in clojure.
2. This question has been left as an exercise for the reader by The Joy of Clojure (chap 12.4). You can consult the relevant page of the book at http://bit.ly/HkQrio.

Answer 1

in your case it's hard to show memoize do anything with factorial because the intermediate calls are unique, so I'll rewrite a somewhat contrived example assuming the point is to explore ways to avoid blowing the stack:

(defn stack-popper [n i] 
    (if (< i n) (* i (stack-popper n (inc i))) 1)) 

which can then get something out of a memoize:

(def stack-popper 
   (memoize (fn [n i] (if (< i n) (* i (stack-popper n (inc i))) 1))))

the general approaches to not blowing the stack are:

use tail calls

(def stack-popper 
    (memoize (fn [n acc] (if (> n 1) (recur (dec n) (* acc (dec n))) acc))))

use trampolines

(def stack-popper 
    (memoize (fn [n acc] 
        (if (> n 1) #(stack-popper (dec n) (* acc (dec n))) acc))))
(trampoline (stack-popper 4 1))

use a lazy sequence

(reduce * (range 1 4))

None of these work all the time, though I have yet to hit a case where none of them work. I almost always go for the lazy ones first because I find them to be most clojure like, then I head for tail calling with recur or tramplines

Answer 2

(defmacro memofn
  [name args & body]
  `(let [cache# (atom {})]
     (fn ~name [& args#]
       (let [update-cache!# (fn update-cache!# [state# args#]
                              (if-not (contains? state# args#)
                                (assoc state# args#
                                       (delay
                                         (let [~args args#]
                                           ~@body)))
                                state#))]
         (let [state# (swap! cache# update-cache!# args#)]
           (-> state# (get args#) deref))))))

This will allow a recursive definition of a memoized function, which also caches intermediate results. Usage:

(def fib (memofn fib [n]
           (case n
             1 1
             0 1
             (+ (fib (dec n)) (fib (- n 2))))))

Answer 3

(def gcd 
  (let [cache (atom {})]
    (fn [a b]
      @(or (@cache [a b])
         (let [p (promise)]
           (deliver p
             (loop [a a b b]
               (if-let [p2 (@cache [a b])]
                 @p2
                 (do
                   (swap! cache assoc [a b] p)
                   (if (zero? b) 
                     a 
                     (recur b (mod a b))))))))))))

There is some concurrency issues (double evaluation, the same problem as with memoize, but worse because of the promises) which may be fixed using @kotarak's advice.

Turning the above code into a macro is left as an exercise to the reader. (Fogus's note was imo tongue-in-cheek.)

Turning this into a macro is really a simple exercise in macrology, please remark that the body (the 3 last lines) remain unchanged.

Answer 4

This is factorial function implemented with anonymous recursion with tail call and memoization of intermediate results. The memoization is integrated with the function and a reference to shared buffer (implemented using Atom reference type) is passed by a lexical closure.

Since the factorial function operates on natural numbers and the arguments for succesive results are incremental, Vector seems more tailored data structure to store buffered results.

Instead of passing the result of a previous computation as an argument (accumulator) we're getting it from the buffer.

(def !                                            ; global variable referring to a function
  (let [m (atom [1 1 2 6 24])]                    ; buffer of results
    (fn [n]                                       ; factorial function definition
      (let [m-count (count @m)]                   ; number of results in a buffer
        (if (< n m-count)                         ; do we have buffered result for n?
          (nth @m n)                              ; · yes: return it
          (loop [cur m-count]                     ; · no: compute it recursively
            (let [r (*' (nth @m (dec cur)) cur)]  ; new result
              (swap! m assoc cur r)               ; store the result
              (if (= n cur)                       ; termination condition:
                r                                 ; · base case
                (recur (inc cur))))))))))         ; · recursive case

(time (do (! 8000) nil))  ; => "Elapsed time: 154.280516 msecs"
(time (do (! 8001) nil))  ; => "Elapsed time: 0.100222 msecs"
(time (do (! 7999) nil))  ; => "Elapsed time: 0.090444 msecs"
(time (do (! 7999) nil))  ; => "Elapsed time: 0.055873 msecs"

Answer 5

Using Clojure's recur you can write factorial using an accumulator that has no stack growth, and just memoize it:

(defn fact
  ([n]
     (fact n 1))
  ([n acc]
     (if (= 1 n)
       acc
       (recur (dec n)
              (* n acc)))))