Add multi-methods to Basilisp (#222)

* Add multi-methods to Basilisp

* Simplify call logic for MultiFunction

Author: chrisrink10

src/basilisp/core/__init__.lpy

@@ -1351,3 +1351,51 @@
   [pattern s]
   (lazy-re-seq (seq (re/finditer pattern s))))
 
+;;;;;;;;;;;;;;;;;;
+;; Multimethods ;;
+;;;;;;;;;;;;;;;;;;
+
+(import basilisp.lang.multifn)
+
+(defmacro defmulti
+  "Define a new multimethod with the dispatch function."
+  [name & body]
+  (let [doc         (when (string? (first body))
+                      (first body))
+        body        (if doc
+                      (rest body)
+                      body)
+        dispatch-fn (first body)
+        opts        (apply hash-map (rest body))]
+    `(def ~name (basilisp.lang.multifn/MultiFunction ~name
+                                                     ~dispatch-fn
+                                                     (or (:default opts) :default)))))
+
+(defmacro defmethod
+  "Add a new method to the multi-function which responds to dispatch-val."
+  [multifn dispatch-val & fn-tail]
+  `(.add-method multi-fn dispatch-val (fn ~@fn-tail)))
+
+(defn methods
+  "Return a map of dispatch values to methods for the given multi function."
+  [multifn]
+  (.-methods multifn))
+
+(defn get-method
+  "Return the method which would respond to dispatch-val or nil if no method
+  exists for dispatch-val."
+  [multifn dispatch-val]
+  (.get-method multifn dispatch-val))
+
+(defn remove-method
+  "Remove the method which responds to dispatch-val, if it exists. Return the
+  multi function."
+  [multifn dispatch-val]
+  (.remove-method multifn dispatch-val)
+  multifn)
+
+(defn remove-all-methods
+  "Remove all method for this multi-function. Return the multi function."
+  [multifn dispatch-val]
+  (.remove-all-methods multifn)
+  multifn)

src/basilisp/lang/multifn.py

@@ -0,0 +1,78 @@
+from typing import TypeVar, Generic, Callable, Any, Optional
+
+import basilisp.lang.atom as atom
+import basilisp.lang.map as lmap
+import basilisp.lang.symbol as sym
+from basilisp.util import Maybe
+
+T = TypeVar('T')
+DispatchFunction = Callable[..., T]
+Method = Callable[..., Any]
+
+
+class MultiFunction(Generic[T]):
+    __slots__ = ('_name', '_default', '_dispatch', '_methods')
+
+    def __init__(self, name: sym.Symbol, dispatch: DispatchFunction, default: T) -> None:
+        self._name = name  # pylint:disable=assigning-non-slot
+        self._default = default  # pylint:disable=assigning-non-slot
+        self._dispatch = dispatch  # pylint:disable=assigning-non-slot
+        self._methods: atom.Atom = atom.Atom(lmap.Map.empty())  # pylint:disable=assigning-non-slot
+
+    def __call__(self, *args, **kwargs):
+        key = self._dispatch(*args, **kwargs)
+        method_cache = self.methods
+        method = Maybe(method_cache.entry(key, None)).or_else(
+            lambda: method_cache.entry(self._default, None))
+        if method:
+            return method(*args, **kwargs)
+        raise NotImplementedError
+
+    @staticmethod
+    def __add_method(m: lmap.Map, key: T, method: Method) -> lmap.Map:
+        """Swap the methods atom to include method with key."""
+        return m.assoc(key, method)
+
+    def add_method(self, key: T, method: Method) -> None:
+        """Add a new method to this function which will respond for
+        key returned from the dispatch function."""
+        self._methods.swap(MultiFunction.__add_method, key, method)
+
+    def get_method(self, key: T) -> Optional[Method]:
+        """Return the method which would handle this dispatch key or
+        None if no method defined for this key and no default."""
+        method_cache = self.methods
+        # The 'type: ignore' comment below silences a spurious MyPy error
+        # about having a return statement in a method which does not return.
+        return Maybe(method_cache.entry(key, None)).or_else(
+            lambda: method_cache.entry(self._default, None))  # type: ignore
+
+    @staticmethod
+    def __remove_method(m: lmap.Map, key: T) -> lmap.Map:
+        """Swap the methods atom to remove method with key."""
+        return m.dissoc(key)
+
+    def remove_method(self, key: T) -> Optional[Method]:
+        """Remove the method defined for this key and return it."""
+        method = self.methods.entry(key, None)
+        if method:
+            self._methods.swap(MultiFunction.__remove_method, key)
+        return method
+
+    def remove_all_methods(self) -> None:
+        """Remove all methods defined for this multi-function."""
+        self._methods.reset(lmap.Map.empty())
+
+    @property
+    def default(self) -> T:
+        return self._default
+
+    @property
+    def methods(self) -> lmap.Map:
+        return self._methods.deref()
+
+
+def multifn(dispatch: DispatchFunction, default=None) -> MultiFunction[T]:
+    """Decorator function which can be used to make Python multi functions."""
+    name = sym.symbol(dispatch.__qualname__, ns=dispatch.__module__)
+    return MultiFunction(name, dispatch, default)

tests/basilisp/multifn_test.py

@@ -0,0 +1,59 @@
+import pytest
+
+import basilisp.lang.keyword as kw
+import basilisp.lang.map as lmap
+import basilisp.lang.multifn as multifn
+import basilisp.lang.symbol as sym
+
+
+def test_multi_function():
+    def dispatch(v) -> kw.Keyword:
+        if v == "i":
+            return kw.keyword('a')
+        elif v == 'ii':
+            return kw.keyword('b')
+        return kw.keyword('default')
+
+    def fn_a(v) -> str:
+        return "1"
+
+    def fn_b(v) -> str:
+        return "2"
+
+    def fn_default(v) -> str:
+        return "BLAH"
+
+    f = multifn.MultiFunction(sym.symbol('test-fn'), dispatch, kw.keyword('default'))
+    f.add_method(kw.keyword('a'), fn_a)
+    f.add_method(kw.keyword('b'), fn_b)
+    f.add_method(kw.keyword('default'), fn_default)
+
+    assert lmap.map({kw.keyword('a'): fn_a,
+                     kw.keyword('b'): fn_b,
+                     kw.keyword('default'): fn_default}) == f.methods
+
+    assert kw.keyword('default') == f.default
+
+    assert fn_a is f.get_method(kw.keyword('a'))
+    assert fn_b is f.get_method(kw.keyword('b'))
+    assert fn_default is f.get_method(kw.keyword('default'))
+    assert fn_default is f.get_method(kw.keyword('other'))
+
+    assert "1" == f("i")
+    assert "2" == f("ii")
+    assert "BLAH" == f("iii")
+    assert "BLAH" == f("whatever")
+
+    f.remove_method(kw.keyword('b'))
+
+    assert "1" == f("i")
+    assert "BLAH" == f("ii")
+    assert "BLAH" == f("iii")
+    assert "BLAH" == f("whatever")
+
+    f.remove_all_methods()
+
+    assert lmap.Map.empty() == f.methods
+
+    with pytest.raises(NotImplementedError):
+        f('blah')