Improving the implementation of distinct-except (and swapping the argument order).

Author: krisajenkins

src/yesql/types.clj

@@ -26,7 +26,7 @@
         query-args (mapv replace-question-mark-with-gensym args)]
     {:query-args    query-args
      :function-args (distinct query-args)
-     :display-args  (distinct-except args #{'?})}))
+     :display-args  (distinct-except #{'?} args)}))
 
 (defn- fn-symbol
   "Attach metadata (docstring/argument lists) to the given symbol."

src/yesql/util.clj

@@ -10,22 +10,17 @@
 
 (ann distinct-except
   (All [x]
-       [(Option (Seq x)) (Option (Set x)) -> (Option (Seq x))]))
+       [[x -> Boolean] (Option (Seq x)) -> (Option (Seq x))]))
 (defn distinct-except
-  "Same as distinct, but keeps duplicates from the exceptions set."
-  [coll exceptions]
-  {:pre [(coll? coll)
-         (set? exceptions)]}
-  (let [step (fn step [xs seen]
-               (lazy-seq
-                ((fn [[f :as xs] seen]
-                   (when-let [s (seq xs)]
-                     (if (and (contains? seen f)
-                              (not (exceptions f)))
-                       (recur (rest s) seen)
-                       (cons f (step (rest s) (conj seen f))))))
-                 xs seen)))]
-    (step coll #{})))
+  "Same as distinct, but keeps duplicates if they pass exception?"
+  [exception? [head & tail :as coll]]
+  (lazy-seq
+   (when head
+     (cons head
+           (distinct-except exception?
+                            (if (exception? head)
+                              tail
+                              (remove (partial = head) tail)))))))
 
 (ann whitespace?
   (IFn [nil -> false]

test/yesql/util_test.clj

@@ -6,11 +6,11 @@
 ;;; Test distinct-except
 (let [coll '[a b c a b]]
   (expect '[a b c a]
-          (distinct-except coll #{'a}))
+          (distinct-except #{'a} coll))
   (expect '[a b c b]
-          (distinct-except coll #{'b}))
+          (distinct-except #{'b} coll))
   (expect '[a b c]
-          (distinct-except coll #{'c})))
+          (distinct-except #{'c} coll)))
 
 ;;; Test underscores-to-dashes
 (do-template [input output]