feature(macros): implement simple multimethods

2 files changed, 52 insertions, 2 deletions

diff --git a/macros/core.fnl b/macros/core.fnl
index be0c1bb..d107ad6 100644
--- a/macros/core.fnl
+++ b/macros/core.fnl
@@ -138,4 +138,33 @@
              :else (error "expected table as first argument"))
            to#))))
 
+(fn first [tbl]
+  (. tbl 1))
+
+(fn rest [tbl]
+  [(unpack tbl 2)])
+
+(fn string? [x]
+  (= (type x) :string))
+
+(fn* core.defmulti
+  [name & opts]
+  (let [docstring (if (string? (first opts)) (first opts))
+        opts (if docstring (rest opts) opts)
+        dispatch-fn (first opts)]
+    `(local ,name
+            (let [multimethods# {}]
+              (setmetatable {} {:__call
+                                (fn [_# ...]
+                                  ,docstring
+                                  ((or (. multimethods# (,dispatch-fn ...))
+                                       (. multimethods# :default)) ...))
+                                :multimethods multimethods#})))))
+
+(fn* core.defmethod
+  [multifn dispatch-val & fn-tail]
+  `(tset (. (getmetatable ,multifn) :multimethods)
+         ,dispatch-val
+         (fn ,(unpack fn-tail))))
+
 core
diff --git a/macros_test.fnl b/macros_test.fnl
index 850b181..246b639 100644
--- a/macros_test.fnl
+++ b/macros_test.fnl
@@ -1,6 +1,6 @@
-(import-macros {: if-let : when-let : if-some : when-some : into} :macros.core)
+(import-macros {: if-let : when-let : if-some : when-some : into : defmethod : defmulti} :macros.core)
 (import-macros {: assert-eq : assert-ne : assert* : testing : deftest} :test)
-(local {: eq?} (require :core)) ;; required for testing
+(local {: eq? : identity} (require :core)) ;; required for testing
 
 (deftest into
   (testing into
@@ -66,3 +66,24 @@
     (assert-eq (if-some [a [1 2 3]] a :nothing) [1 2 3])
     (assert-eq (if-some [a false] a :nothing) false)
     (assert-eq (if-some [a nil] a :nothing) :nothing)))
+
+(deftest multimethods
+  (testing defmethod
+    (defmulti fac identity)
+    (defmethod fac 0 [_] 1)
+    (defmethod fac :default [x] (* x (fac (- x 1))))
+    (assert-eq (fac 42) 7538058755741581312)
+
+    (defmulti send-data (fn [protocol data] protocol))
+    (defmethod send-data :http [protocol data] (.. data " will be sent over HTTP"))
+    (defmethod send-data :icap [protocol data] (.. data " will be sent over ICAP"))
+    (assert-eq (send-data :http 42) "42 will be sent over HTTP")
+    (assert-eq (send-data :icap 42) "42 will be sent over ICAP")
+
+    (defmulti send-message (fn [message] (. message :protocol)))
+    (defmethod send-message :http [message] (.. "sending " (. message :message) " over HTTP"))
+    (defmethod send-message :icap [message] (.. "sending " (. message :message) " over ICAP"))
+    (assert-eq (send-message {:protocol :http :message "ваыв"})
+               "sending ваыв over HTTP")
+    (assert-eq (send-message {:protocol :icap :message 42})
+               "sending 42 over ICAP")))