Support defmethod argcount validation

.clj-kondo/config.edn

@@ -21,7 +21,8 @@
   :warn-on-reflection                    {:level :warning}
 
   :unresolved-symbol
-  {}
+  {:exclude
+   [(clojure.test/is [macroexpansion-spec-error?])]}
 
   :consistent-alias
   {:aliases

.github/workflows/tests.yml

@@ -121,4 +121,5 @@ jobs:
   codespell:
     runs-on: ubuntu-20.04
     steps:
+      - uses: actions/checkout@v3
       - uses: codespell-project/actions-codespell@master

README.md

@@ -517,6 +517,58 @@ dispatch value form of any `defmethod` forms at macroexpansion time:
 
 This is a great way to make sure people use your multimethods correctly and catch errors right away.
 
+#### `:defmethod-arities`
+
+A set of allowed/required arities that `defmethod` forms are allowed to have. `defmethod` forms must have arities that
+match *all* of the specified `:defmethod-arities`, and all of its arities must be allowed by `:defmethod-arities`:
+
+```clj
+(m/defmulti ^:private mf
+  {:arglists '([x]), :defmethod-arities #{1}}
+  keyword)
+
+(m/defmethod mf :x [x] x)
+;; => ok
+
+(m/defmethod mf :x ([x] x) ([x y] x y))
+;; => error: {:arities {:disallowed #{2}}}
+
+(m/defmethod mf :x [x y] x y)
+;; => error: {:arities {:required #{1}}}
+
+
+(m/defmethod mf :x [x y] x)
+;; => error: {:arities {:required #{1 [:>= 3]}, :disallowed #{2}}}
+```
+
+`:defmethod-arities` must be a set of either integers or `[:> n]` forms to represent arities with `&` rest
+arguments, e.g. `[:>= 3]` to mean an arity of three *or-more* arguments:
+
+```clj
+;; methods must both a 1-arity and a 3+-arity
+(m/defmulti ^:private mf
+  {:arglists '([x] [x y z & more]), :defmethod-arities #{1 [:>= 3]}}
+  keyword)
+
+(m/defmethod mf :x ([x] x) ([x y z & more] x))
+;; => ok
+```
+
+When rest-argument arities are used, Methodical is smart enough to allow them when appropriate even if they do not
+specifically match an arity specified in `:defmethod-arities`:
+
+```clj
+(m/defmulti ^:private mf
+  {:arglists '([x y z & more]), :defmethod-arities #{[:>= 3]}}
+  keyword)
+
+(m/defmethod mf :x
+  ([a b c] x)
+  ([a b c d] x)
+  ([a b c d & more] x))
+;; => ok, because everything required by [:>= 3] is covered, and everything present is allowed by [:>= 3]
+```
+
 ### Debugging
 
 Methodical offers debugging facilities so you can see what's going on under the hood, such as the `trace` utility:
@@ -528,7 +580,7 @@ and the `describe` utility, which outputs Markdown-formatted documentation, for
 
 ![Describe](assets/describe.png)
 
-This extra information is automatically generated and appened to a multimethod's docstring whenever methods or
+This extra information is automatically generated and appended to a multimethod's docstring whenever methods or
 preferences are added or removed.
 
 Methodical multimethods also implement `datafy`:

src/methodical/macros.clj

@@ -1,11 +1,12 @@
 (ns methodical.macros
-  "Methodical versions of vanilla Clojure `defmulti` and [[defmethod]] macros."
+  "Methodical versions of vanilla Clojure [[defmulti]] and [[defmethod]] macros."
   (:refer-clojure :exclude [defmulti defmethod])
   (:require
    [clojure.spec.alpha :as s]
    [clojure.string :as str]
    [methodical.impl :as impl]
    [methodical.interface :as i]
+   [methodical.macros.validate-arities :as validate-arities]
    [methodical.util :as u])
   (:import
    (methodical.impl.standard StandardMultiFn)))
@@ -16,13 +17,29 @@
   (s/alt :arity-1 :clojure.core.specs.alpha/params+body
          :arity-n (s/+ (s/spec :clojure.core.specs.alpha/params+body))))
 
+(s/def :methodical.macros.defmulti.attr-map/dispatch-value-spec
+  ;; not sure how to validate this. Is there a predicate for something that can be used as a
+  ;; spec? [[clojure.spec.alpha/spec]] will happily turn random things into specs for us
+  any?)
+
+(s/def :methodical.macros.defmulti.attr-map/defmethod-arities
+  (s/nilable ::validate-arities/arities-set))
+
+(s/def :methodical.macros.defmulti/attr-map
+  (s/keys :opt-un [:methodical.macros.defmulti.attr-map/dispatch-value-spec
+                   :methodical.macros.defmulti.attr-map/defmethod-arities]))
+
 (s/def ::defmulti-args
-  (s/cat :name-symb   (every-pred symbol? (complement namespace))
-         :docstring   (s/? string?)
-         :attr-map    (s/? map?)
-         :dispatch-fn (s/? any?)
-         :options     (s/* (s/cat :k keyword?
-                                  :v any?))))
+  (s/& (s/cat :name-symb   (every-pred symbol? (complement namespace))
+              :docstring   (s/? string?)
+              :attr-map    (s/? map?)
+              :dispatch-fn (s/? any?)
+              :options     (s/* (s/cat :k keyword?
+                                       :v any?)))
+       ;; do a "soft" cut here where any map argument is always interpreted as an attribute map, and only later subject
+       ;; it to stricter validation. See my questions in Slack here
+       ;; https://clojurians.slack.com/archives/C1B1BB2Q3/p1662755403589769
+       (s/keys :opt-un [:methodical.macros.defmulti/attr-map])))
 
 (defn- emit-defmulti
   [name-symb dispatch-fn {:keys [hierarchy dispatcher combo method-table cache default-value]
@@ -49,7 +66,7 @@
        (let [impl# (impl/standard-multifn-impl ~combo ~dispatcher ~method-table)]
          (vary-meta (impl/multifn impl# ~mta ~cache) merge (meta (var ~name-symb)))))))
 
-(defn default-dispatch-value-spec
+(defn- default-dispatch-value-spec
   "A dispatch value as parsed to [[defmethod]] (i.e., not-yet-evaluated) can be ANYTHING other than the following two
   things:
 
@@ -78,10 +95,10 @@
           ...)))
      ```
 
-      but we are just UNFORTUNATELY going to have to throw up our hands and say we don't support it. The reason is in
-      the example above it's ambiguous whether this is a `:before` aux method with dispatch value `([toucan pigeon] i)`,
-      or a primary method with dispatch value `:before`. It's just impossible to tell what you meant. If you really want
-      to do something wacky like this, let-bind the dispatch value to a symbol or something.
+     but we are just UNFORTUNATELY going to have to throw up our hands and say we don't support it. The reason is in
+     the example above it's ambiguous whether this is a `:before` aux method with dispatch value `([toucan pigeon] i)`,
+     or a primary method with dispatch value `:before`. It's just impossible to tell what you meant. If you really want
+     to do something wacky like this, let-bind the dispatch value to a symbol or something.
 
   Note that if you specify a custom `:dispatch-value-spec` it overrides this spec. Hopefully your spec is stricter than
   this one is and it won't be a problem."
@@ -103,7 +120,7 @@
     (emit-defmulti name-symb dispatch-fn options)))
 
 (defmacro defmulti
-  "Creates a new Methodical multimethod named by a Var. Usage of this macro mimics usage of vanilla Clojure `defmulti`,
+  "Creates a new Methodical multimethod named by a Var. Usage of this macro mimics usage of [[clojure.core/defmulti]],
   and it can be used as a drop-in replacement; it does, however, support a larger set of options. Note the dispatch-fn
   is optional (if omitted, then identity will be used). In addition to the usual `:default` and `:hierarchy` options,
   you many specify:
@@ -132,14 +149,93 @@
      value -> methods. The default implementation is a pair of simple maps.
 
   The above options comprise the main constituent parts of a Methodical multimethod, and the majority of those parts
-  have several alternative implementations available in `methodical.impl`. Defining additional implementations is
-  straightforward as well: see `methodical.interface` for more details.
+  have several alternative implementations available in [[methodical.impl]]. Defining additional implementations is
+  straightforward as well: see [[methodical.interface]] for more details.
 
-  Other improvements over vanilla Clojure `defmulti`:
+  Other improvements over [[clojure.core/defmulti]]:
 
   * Evaluating the form a second time (e.g., when reloading a namespace) will *not* redefine the multimethod, unless
     you have modified its form -- unlike vanilla Clojure multimethods, which need to be unmapped from the namespace to
-    make such minor tweaks as changing the dispatch function."
+    make such minor tweaks as changing the dispatch function.
+
+  Attribute map options:
+
+  `defmulti` supports a few additional options in its attributes map that will be used to validate `defmethod` forms
+  during macroexpansion time. These are meant to help the users of your multimethods use them correctly by catching
+  mistakes right away rather than waiting for them to pull their hair out later wondering why a method they added isn't
+  getting called.
+
+  * `:dispatch-value-spec` -- a spec for the `defmethod` dispatch value:
+
+    ```clj
+    (m/defmulti mf
+      {:arglists '([x y]), :dispatch-value-spec (s/cat :x keyword?, :y int?)}
+      (fn [x y] [x y]))
+
+    (m/defmethod mf [:x 1]
+      [x y]
+      {:x x, :y y})
+    ;; => ok
+
+    (m/defmethod mf [:x]
+      [x y]
+      {:x x, :y y})
+    ;; failed: Insufficient input in: [0] at: [:args-for-method-type :primary :dispatch-value :y] [:x]
+    ```
+
+    Note that this spec is applied to the unevaluated arguments at macroexpansion time, not the actual evaluated values.
+    Note also that if you want to allow a `:default` method your spec will have to support it.
+
+  * `:defmethod-arities` -- a set of allowed/required arities that `defmethod` forms are allowed to have. `defmethod`
+    forms must have arities that match *all* of the specified `:defmethod-arities`, and all of its arities must be
+    allowed by `:defmethod-arities`:
+
+    ```clj
+    (m/defmulti ^:private mf
+      {:arglists '([x]), :defmethod-arities #{1}}
+      keyword)
+
+    (m/defmethod mf :x [x] x)
+    ;; => ok
+
+    (m/defmethod mf :x ([x] x) ([x y] x y))
+    ;; => error: {:arities {:disallowed #{2}}}
+
+    (m/defmethod mf :x [x y] x y)
+    ;; => error: {:arities {:required #{1}}}
+
+
+    (m/defmethod mf :x [x y] x)
+    ;; => error: {:arities {:required #{1 [:>= 3]}, :disallowed #{2}}}
+    ```
+
+    `:defmethod-arities` must be a set of either integers or `[:> n]` forms to represent arities with `&` rest
+    arguments, e.g. `[:>= 3]` to mean an arity of three *or-more* arguments:
+
+    ```clj
+    ;; methods must both a 1-arity and a 3+-arity
+    (m/defmulti ^:private mf
+      {:arglists '([x] [x y z & more]), :defmethod-arities #{1 [:>= 3]}}
+      keyword)
+
+    (m/defmethod mf :x ([x] x) ([x y z & more] x))
+    ;; => ok
+    ```
+
+    When rest-argument arities are used, Methodical is smart enough to allow them when appropriate even if they do not
+    specifically match an arity specified in `:defmethod-arities`:
+
+    ```clj
+    (m/defmulti ^:private mf
+      {:arglists '([x y z & more]), :defmethod-arities #{[:>= 3]}}
+      keyword)
+
+    (m/defmethod mf :x
+      ([a b c] x)
+      ([a b c d] x)
+      ([a b c d & more] x))
+    ;; => ok, because everything required by [:>= 3] is covered, and everything present is allowed by [:>= 3]
+    ```"
   {:arglists     '([name-symb docstring? attr-map? dispatch-fn?
                     & {:keys [hierarchy default-value prefers combo method-table cache]}]
                    [name-symb docstring? attr-map? & {:keys [dispatcher combo method-table cache]}])
@@ -230,15 +326,17 @@
         primary-methods-allowed? (contains? allowed-qualifiers nil)
         allowed-aux-qualifiers   (disj allowed-qualifiers nil)
         dispatch-value-spec      (or (some-> (get (meta multifn) :dispatch-value-spec) s/spec)
-                                     (default-dispatch-value-spec allowed-aux-qualifiers))]
+                                     (default-dispatch-value-spec allowed-aux-qualifiers))
+        arities-spec             (validate-arities/allowed-arities-fn-tail-spec (:defmethod-arities (meta multifn)))]
     (s/cat :args-for-method-type (s/alt :primary (if primary-methods-allowed?
                                                    (s/cat :dispatch-value dispatch-value-spec)
                                                    (constantly false))
                                         :aux     (s/cat :qualifier      allowed-aux-qualifiers
                                                         :dispatch-value dispatch-value-spec
                                                         :unique-key     (s/? (complement (some-fn string? sequential?)))))
            :docstring            (s/? string?)
-           :fn-tail              (s/& (s/+ any?) (s/nonconforming ::fn-tail)))))
+           :fn-tail              (s/& (s/+ any?)
+                                      (s/nonconforming (s/& ::fn-tail arities-spec))))))
 
 (defn- parse-defmethod-args [multifn args]
   (let [spec      (defmethod-args-spec multifn)

src/methodical/macros/validate_arities.clj

@@ -0,0 +1,111 @@
+(ns ^:no-doc methodical.macros.validate-arities
+  "Implementation for `defmethod` arities validation (see [#59](https://github.com/camsaul/methodical/issues/59)).
+
+  The arities sets passed around here are sets of numeric arity counts, or `[:> n]` for arity `n` that also accepts
+  varargs, i.e.
+
+  ```clj
+  (fn
+    ([x])
+    ([x y & more]))
+  ```
+
+  has the arities `#{1 [:>= 2]}`."
+  (:require
+   [clojure.data :as data]
+   [clojure.set :as set]
+   [clojure.spec.alpha :as s]))
+
+(defn- int-between-zero-and-twenty-inclusive? [n]
+  (and (integer? n)
+       (<= 0 n 20)))
+
+(s/def ::arities-set
+  (s/every
+   (s/or :int      int-between-zero-and-twenty-inclusive?
+         :>=-form  (s/spec (s/cat :>=  (partial = :>=)
+                                  :int int-between-zero-and-twenty-inclusive?)))
+   :kind set?
+   :min-count 1))
+
+(defn- arglist-arities
+  "Determine the arity sets for `arglists`. Deals with arglists already conformed using the
+  `:clojure.core.specs.alpha/param-list` spec.
+
+  ```clj
+  (arglist-arities (s/conform (s/+ :clojure.core.specs.alpha/param-list) '([] [x])))
+  => #{0 1}
+
+  (arglist-arities (s/conform (s/+ :clojure.core.specs.alpha/param-list) '([] [x] [x y z & more])))
+  => #{0 1 [:>= 3]}
+  ```"
+  [arglists]
+  (into #{}
+        (map (fn [{:keys [params var-params]}]
+               (if var-params
+                 [:>= (count params)]
+                 (count params))))
+        arglists))
+
+(defn- fn-tail-arities
+  "Determine arity sets for `fn-tails`. `fn-tails` should already be conformed using the
+  `:methodical.macros/fn-tail` spec or similar:
+
+  ```clj
+  (fn-tail-arities (s/conform :methodical.macros/fn-tail '[([x] x) ([x y & more] x)]))
+  =>
+  #{1 [:>= 2]}
+  ```"
+  [[arity-type x]]
+  (set (case arity-type
+         :arity-1 (arglist-arities [(:params x)])
+         :arity-n (arglist-arities (map :params x)))))
+
+(defn- expand-arity
+  "Given an arity like `1` or `[:> 3]` expand the arity into a flat set of all arities that can be used to invoke a
+  function tail with that specific arity, e.g. `[:> 3]` can be invoked with `3`, `4`, `5,` or so on arguments.
+
+  `clojure.lang.IFn/invoke` only supports distinct arities between 0 and 20, inclusive; any more than 20 arguments must
+  be invoked with `.applyTo`, so that's all we need to consider here; we'll use the keyword `:more` to represent > 20
+  arguments.
+
+  ```clj
+  (expand-arity 1)      => #{1}
+  (expand-arity [:> 3]) => #{3 4 5 6 7 8 9 ... :more}
+  ```"
+  [arity]
+  (if (integer? arity)
+    (sorted-set arity)
+    (let [[_ arity] arity]
+      (into #{:more} (range arity 21)))))
+
+(defn- expand-arities [arities]
+  (into #{} (mapcat expand-arity) arities))
+
+(defn- diff-arities
+  "Given a set of `required` arities (e.g., the `:defmethod-arities` in the `defmulti` metadata) and `actual`
+  arities (e.g. the function tail arities in a `defmethod` form, as generated by [[fn-tail-arities]]), return a map with
+  any `:required` arities that are missing and any `:disallowed` arities that are present. Returns `nil` if there are no
+  missing required arities or disallowed arities present."
+  [required actual]
+  (let [[missing disallowed] (data/diff (expand-arities required) (expand-arities actual))]
+    (not-empty
+     (into {} (for [[k orig expanded] [[:required required missing]
+                                       [:disallowed actual disallowed]]
+                    :when             (seq expanded)]
+                [k (set (for [arity orig
+                              :when (not-empty (set/intersection (expand-arity arity) expanded))]
+                          arity))])))))
+
+(defn allowed-arities-fn-tail-spec
+  "Create a spec for a function tail to make sure it has all of the `required-arities`, and all of its arities are
+  allowed."
+  [required-arities]
+  (if (empty? required-arities)
+    identity
+    (s/and (s/conformer (fn [fn-tail]
+                    (let [arities (fn-tail-arities fn-tail)
+                          diff    (diff-arities required-arities arities)]
+                      (when (seq diff)
+                        {:arities diff}))))
+           empty?)))