memoize.clj

;;; fast-memoize
;;; Lifted from https://github.com/ztellman/potemkin/blob/master/src/potemkin/utils.clj

(definline re-nil [x]
  `(let [x# ~x]
     (if (identical? ::nil x#) nil x#)))

(definline de-nil [x]
  `(let [x# ~x]
     (if (nil? x#) ::nil x#)))

(defmacro memoize-form [m f kf & args]
  `(let [k# (~kf ~@args)]
     (let [v# (.get ~m k#)]
       (if-not (nil? v#)
         (re-nil v#)
         (let [v# (de-nil (~f ~@args))]
           (or (.putIfAbsent ~m k# v#) v#))))))

(defn fast-memoize
  "A version of `memoize` which has equivalent behavior, but is faster.
   Takes an optional function kf that will be used to compute the cache/memoize key"
  ([f] (fast-memoize f hash))
  ([f kf]
   (let [m (ConcurrentHashMap.)]
     (fn
       ([]
        (memoize-form m f kf))
       ([x]
        (memoize-form m f kf x))
       ([x y]
        (memoize-form m f kf x y))
       ([x y z]
        (memoize-form m f kf x y z))
       ([x y z w]
        (memoize-form m f kf x y z w))
       ([x y z w u]
        (memoize-form m f kf x y z w u))
       ([x y z w u v]
        (memoize-form m f kf x y z w u v))
       ([x y z w u v & rest]
        (let [k (apply kf (list* x y z w u v rest))]
          (let [v (.get ^ConcurrentHashMap m k)]
            (if-not (nil? v)
              (re-nil v)
              (let [v (de-nil (apply f k))]
                (or (.putIfAbsent m k v) v))))))))))