From 1ca9fc78d897758fce86519bf1e1fe0fd76fd51a Mon Sep 17 00:00:00 2001
From: Mikhail Kuzmin <m.kuzmin@darkleaf.ru>
Date: Wed, 19 Feb 2020 22:44:06 +0400
Subject: [PATCH] memoize-multi

---
 README.md                               |  7 +++++++
 src/darkleaf/multidecorators.cljc       | 27 +++++++++++++++++++++----
 test/darkleaf/multidecorators_test.cljc | 18 +++++++++++++++++
 3 files changed, 48 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index 48b9750..64e71df 100644
--- a/README.md
+++ b/README.md
@@ -59,6 +59,13 @@ Like multimethods but multidecorators.
 (assert (= [] (func ::f)))
 ```
 
+## Memoization
+
+```clojure
+(defn -main []
+  (alter-var-root #'func md/memoize-multi))
+```
+
 ## Development
 
 ```
diff --git a/src/darkleaf/multidecorators.cljc b/src/darkleaf/multidecorators.cljc
index 427f83c..ee732d7 100644
--- a/src/darkleaf/multidecorators.cljc
+++ b/src/darkleaf/multidecorators.cljc
@@ -41,8 +41,27 @@
              f    (method @iregistry tag initial)]
          (apply f obj args))))))
 
+(defn memoize-multi [multi]
+  (case (:type (multi))
+    :memoized multi
+    :dynamic  (let [{:keys [iregistry
+                            dispatch
+                            initial]} (multi)
+                    registry          @iregistry
+                    mem-method        (memoize method)]
+                (fn
+                  ([] {:type     :memoized
+                       :registry registry
+                       :initial  initial
+                       :dispatch dispatch})
+                  ([obj & args]
+                   (let [tag (apply dispatch obj args)
+                         f   (mem-method registry tag initial)]
+                     (apply f obj args)))))))
+
 (defn ^{:style/indent :defn} decorate [multi tag decorator]
-  (let [state     (multi)
-        iregistry (:iregistry state)]
-    (swap! iregistry assoc tag decorator)
-    multi))
+  (case (:type (multi))
+    :dynamic (let [state     (multi)
+                   iregistry (:iregistry state)]
+               (swap! iregistry assoc tag decorator)
+               multi)))
diff --git a/test/darkleaf/multidecorators_test.cljc b/test/darkleaf/multidecorators_test.cljc
index c6a5505..a776f64 100644
--- a/test/darkleaf/multidecorators_test.cljc
+++ b/test/darkleaf/multidecorators_test.cljc
@@ -33,3 +33,21 @@
     (t/is (= [] (multi ::f)))
     (t/is (= [:a :b :c :d 's] (multi `s)))
     #?(:clj (t/is (= [:a :b :c :d :obj] (multi String))))))
+
+(t/deftest memoization
+  (let [multi (doto (md/multi identity (constantly []))
+                (md/decorate ::a (fn [super obj] (conj (super obj) :a)))
+                (md/decorate ::b (fn [super obj] (conj (super obj) :b)))
+                (md/decorate ::c (fn [super obj] (conj (super obj) :c)))
+                (md/decorate ::d (fn [super obj] (conj (super obj) :d)))
+                (md/decorate `s  (fn [super obj] (conj (super obj) 's)))
+                #?(:clj (md/decorate Object (fn [super obj] (conj (super obj) :obj)))))
+        mem-multi (md/memoize-multi multi)]
+    (doseq [_ (range 2)]
+      (t/is (= [:a] (mem-multi ::a)))
+      (t/is (= [:a :b] (mem-multi ::b)))
+      (t/is (= [:a :c] (mem-multi ::c)))
+      (t/is (= [:a :b :c :d] (mem-multi ::d)))
+      (t/is (= [] (mem-multi ::f)))
+      (t/is (= [:a :b :c :d 's] (mem-multi `s)))
+      #?(:clj (t/is (= [:a :b :c :d :obj] (mem-multi String)))))))