rewrite function arities that are checked once

module-check/src/main/clojure/clojure/core/typed/check/fn.clj

@@ -21,11 +21,13 @@
    :post [(-> % u/expr-type r/TCResult?)
           (vector? (::t/cmethods %))]}
   ;(prn "check-fn" methods)
-  (let [cmethods (fn-methods/check-fn-methods 
-                   methods
-                   (r/ret-t expected)
-                   :self-name (cu/fn-self-name fexpr))]
+  (let [{:keys [methods cmethods]}
+        (fn-methods/check-fn-methods 
+          methods
+          (r/ret-t expected)
+          :self-name (cu/fn-self-name fexpr))]
     (assoc fexpr
+           :methods methods
            ::t/cmethods cmethods
            u/expr-type  (r/ret (r/ret-t expected)
                                (fo/-FS fl/-top fl/-bot)))))

module-check/src/main/clojure/clojure/core/typed/check/fn_methods.clj

@@ -5,6 +5,7 @@
             [clojure.core.typed.contract-utils :as con]
             [clojure.core.typed.ast-utils :as ast-u]
             [clojure.core.typed.type-ctors :as c]
+            [clojure.core.typed.utils :as u]
             [clojure.core.typed.check.utils :as cu]
             [clojure.core.typed.errors :as err]
             [clojure.core.typed.parse-unparse :as prs]
@@ -20,7 +21,7 @@
                              r/FnIntersection?))
 
 (def method? (some-fn ast-u/fn-method? ast-u/deftype-method?))
-(def methods? (con/every-c? method?))
+(def methods? (con/vec-c? method?))
 
 (def opt-map? (con/hmap-c? (con/optional :recur-target-fn) ifn?
                            (con/optional :validate-expected-fn) ifn?
@@ -117,7 +118,7 @@
         _ (when validate-expected-fn
             (validate-expected-fn fin))
         ;collect all inferred Functions
-        cmethods
+        cmethodss
         (lex/with-locals (when-let [name self-name] ;self calls
                            {name exp})
           ;scope type variables from polymorphic type in body
@@ -130,21 +131,59 @@
             (dvar-env/with-dotted-mappings (case poly?
                                              :PolyDots {(-> inst-frees last r/F-original-name) (last inst-frees)}
                                              {})
-              (vec
-                (mapcat (fn [f]
-                          {:pre [(r/Function? f)]
-                           ;returns a collection of fn-method's
-                           :post [(methods? %)]}
-                          (check-Function
-                            mthods
-                            f
-                            {:recur-target-fn recur-target-fn}))
-                        (:types fin))))))]
-    cmethods))
+              (let [;; map from function type to a set of matching methods (integers).
+                    fn-matches
+                    (into {}
+                          (map (fn [t]
+                                 (let [ms (into #{}
+                                                (comp
+                                                  (map-indexed (fn [i m]
+                                                                 (when (expected-for-method m t)
+                                                                   i)))
+                                                  (keep identity))
+                                                mthods)]
+                                   ;; it is a type error if no matching methods are found.
+                                   (when (empty? ms)
+                                     (binding [vs/*current-expr* (impl/impl-case
+                                                                   :clojure (first mthods)
+                                                                   ; fn-method is not printable in cljs
+                                                                   :cljs vs/*current-expr*)
+                                               vs/*current-env* (or (:env (first mthods)) vs/*current-env*)]
+                                       (prs/with-unparse-ns (cu/expr-ns (first mthods))
+                                         (err/tc-delayed-error (str "No matching arities: " (prs/unparse-type t))))))
+                                   [t ms])))
+                          (:types fin))]
+                ;; if a method occurs more than once in the entire map, it will be
+                ;; checked twice, so we disable rewriting for that method.
+                (into []
+                      (map-indexed
+                        (fn [i m]
+                          (let [expecteds (keep 
+                                            (fn [[t es]]
+                                              (when (es i)
+                                                t))
+                                            fn-matches)]
+                            (u/rewrite-when (== 1 (count expecteds))
+                              (mapv (comp
+                                      :cmethod
+                                      #(fn-method1/check-fn-method1 m % :recur-target-fn recur-target-fn))
+                                    expecteds)))))
+                      mthods)))))]
+    {:methods (mapv 
+                (fn [cmethods m]
+                  (if (== 1 (count cmethods))
+                    (nth cmethods 0)
+                    m))
+                cmethodss
+                mthods)
+     :cmethods (into []
+                     (mapcat identity)
+                     cmethodss)}))
 
 (defn function-types [expected]
   {:pre [(r/Type? expected)]
-   :post [(every? function-type? %)]}
+   :post [(and (every? function-type? %)
+               (vector? %))]}
   (let [exp (c/fully-resolve-type expected)
         ts (filterv function-type?
                     (if (r/Union? exp)
@@ -164,13 +203,27 @@
   {:pre [(r/Type? expected)
          ((every-pred methods? seq) mthods)
          (opt-map? opt)]
-   :post [(methods? %)]}
+   :post [((con/hmap-c?
+             :methods (con/vec-c? method?)
+             :cmethods (con/vec-c? method?))
+           %)]}
   ;(prn "check-fn-methods")
   (let [ts (function-types expected)]
-    (if (empty? ts)
+    (cond
+      (empty? ts)
       (prs/with-unparse-ns (cu/expr-ns (first mthods))
         (err/tc-delayed-error (str (prs/unparse-type expected) " is not a function type")
-                              :return []))
-      (vec
-        (mapcat (fn [t] (check-fni t mthods opt))
-                ts)))))
+                              :return {:methods mthods
+                                       :cmethods []}))
+
+      (== 1 (count ts))
+      (check-fni (nth ts 0) mthods opt)
+
+      ;; disable rewriting in case we recheck a method arity
+      :else
+      (binding [vs/*can-rewrite* nil]
+        {:methods mthods
+         :cmethods
+         (into []
+               (mapcat (fn [t] (check-fni t mthods opt)))
+               ts)}))))