feature(core): more multimethod related functions

5 files changed, 99 insertions, 18 deletions

diff --git a/.dir-locals.el b/.dir-locals.el
index 7c3e0a2..0d0d3cd 100644
--- a/.dir-locals.el
+++ b/.dir-locals.el
@@ -1,7 +1,9 @@
 ;;; Directory Local Variables
 ;;; For more information see (info "(emacs) Directory Variables")
 
-((fennel-mode . ((eval . (put 'deftest 'fennel-indent-function 'defun))
+((fennel-mode . ((eval . (put 'defmethod 'fennel-indent-function 'defun))
+                 (eval . (put 'defmulti 'bfennel-indent-function 'defun))
+                 (eval . (put 'deftest 'fennel-indent-function 'defun))
                  (eval . (put 'testing 'fennel-indent-function 'defun))
                  (eval . (put 'when-some 'fennel-indent-function 1))
                  (eval . (put 'if-some 'fennel-indent-function 1))
diff --git a/core.fnl b/core.fnl
index 26cedcb..2218706 100644
--- a/core.fnl
+++ b/core.fnl
@@ -423,4 +423,29 @@ found in the table."
        (set res not-found)))
    res))
 
+(fn* core.remove-method
+  [multifn dispatch-val]
+  (tset (. (getmetatable multifn) :multimethods) dispatch-val nil)
+  multifn)
+
+(fn* core.remove-all-methods
+  "Removes all of the methods of multimethod"
+  [multifn]
+  (let [mtable (. (getmetatable multifn) :multimethods)]
+    (each [k _ (pairs mtable)]
+      (tset mtable k nil))
+    multifn))
+
+(fn* core.methods
+  "Given a multimethod, returns a map of dispatch values -> dispatch fns"
+  [multifn]
+  (. (getmetatable multifn) :multimethods))
+
+(fn* core.get-method
+  "Given a multimethod and a dispatch value, returns the dispatch `fn'
+that would apply to that value, or `nil' if none apply and no default."
+  [multifn dispatch-val]
+  (or (. (getmetatable multifn) :multimethods dispatch-val)
+      (. (getmetatable multifn) :multimethods :default)))
+
 core
diff --git a/core_test.fnl b/core_test.fnl
index 25f9e1d..55710d1 100644
--- a/core_test.fnl
+++ b/core_test.fnl
@@ -1,5 +1,5 @@
 (import-macros {: fn*} :macros.fn)
-(import-macros {: into} :macros.core)
+(import-macros {: into : defmethod : defmulti} :macros.core)
 (import-macros {: assert-eq : assert-ne : assert* : testing : deftest} :test)
 
 (local
@@ -48,7 +48,11 @@
   : dec
   : assoc
   : get
-  : get-in}
+  : get-in
+  : get-method
+  : methods
+  : remove-method
+  : remove-all-methods}
  (require :core))
 
 (deftest equality
@@ -453,3 +457,47 @@
     (assert-eq (get-in t []) t)
     (assert* (not (pcall get-in)))
     (assert* (not (pcall get-in {})))))
+
+(deftest methods
+  (testing methods
+    (defmulti f identity)
+    (defmethod f :a [_] :a)
+    (defmethod f :b [_] :b)
+    (defmethod f :c [x] (* x x))
+    (assert-eq (methods f) (. (getmetatable f) :multimethods))
+    (assert* (not (pcall methods)))
+    (assert* (not (pcall methods f f))))
+
+  (testing get-method
+    (defmulti f identity)
+    (defmethod f :a [_] :a)
+    (defmethod f :b [_] :b)
+    (defmethod f :c [x] (* x x))
+    (assert-eq ((get-method f :a) 10) :a)
+    (assert-eq ((get-method f :b) 20) :b)
+    (assert-eq ((get-method f :c) 4) 16)
+    (assert* (not (pcall get-method)))
+    (assert* (not (pcall get-method f)))
+    (assert* (not (pcall get-method f :a :b))))
+
+  (testing remove-method
+    (defmulti f identity)
+    (defmethod f :a [_] :a)
+    (defmethod f :b [_] :b)
+    (remove-method f :a)
+    (assert-eq (get-method f :a) nil)
+    (defmethod f :default [_] :default)
+    (assert-eq (get-method f :a) (get-method f :default))
+    (assert* (not (pcall remove-method)))
+    (assert* (not (pcall remove-method f)))
+    (assert* (not (pcall remove-method f :a :b))))
+
+  (testing remove-all-methods
+    (defmulti f identity)
+    (defmethod f :a [_] :a)
+    (defmethod f :b [_] :b)
+    (defmethod f :default [_] :default)
+    (remove-all-methods f)
+    (assert-eq (methods f) {})
+    (assert* (not (pcall remove-all-methods)))
+    (assert* (not (pcall remove-all-methods f f)))))
diff --git a/macros/core.fnl b/macros/core.fnl
index bbff020..de6a1b0 100644
--- a/macros/core.fnl
+++ b/macros/core.fnl
@@ -152,21 +152,23 @@
   (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
-                  (let [dispatch-value# (,dispatch-fn ...)]
-                    ((or (. multimethods# dispatch-value#)
-                         (. multimethods# :default)
-                         (error (.. "No method in multimethod '"
-                                    ,(tostring name)
-                                    "' for dispatch value: "
-                                    dispatch-value#) 2)) ...)))
-                :multimethods multimethods#})))))
+    (if (in-scope? name)
+        nil
+        `(local ,name
+                (let [multimethods# {}]
+                  (setmetatable
+                   {}
+                   {:__call
+                    (fn [_# ...]
+                      ,docstring
+                      (let [dispatch-value# (,dispatch-fn ...)]
+                        ((or (. multimethods# dispatch-value#)
+                             (. multimethods# :default)
+                             (error (.. "No method in multimethod '"
+                                        ,(tostring name)
+                                        "' for dispatch value: "
+                                        dispatch-value#) 2)) ...)))
+                    :multimethods multimethods#}))))))
 
 (fn* core.defmethod
   [multifn dispatch-val & fn-tail]
diff --git a/macros_test.fnl b/macros_test.fnl
index 246b639..f8c72fa 100644
--- a/macros_test.fnl
+++ b/macros_test.fnl
@@ -68,6 +68,10 @@
     (assert-eq (if-some [a nil] a :nothing) :nothing)))
 
 (deftest multimethods
+  (testing defmulti
+    (defmulti x (fn [x] x))
+    (assert-eq (defmulti x (fn [x] (+ x 1))) nil))
+
   (testing defmethod
     (defmulti fac identity)
     (defmethod fac 0 [_] 1)