              (recur acc coll k more)))))
 
 (defn- counts-apply
-  "Apply F to values of COLL flattened by K, then return unflattened/updated results."
+  "Apply F to values of COLL flattened by K, then return unflattened/updated results.
+
+    (counts-apply [{:a [{:b 1} {:b 2}], :c 2}
+                   {:a {:b 3}, :c 4}]
+      :a #(update-in % [:b] (partial * 2)))
+
+      -> [{:a [{:b 2} {:b 4}], :c 2}
+          {:a {:b 3}, :c 4}]"
   [coll k f]
   (let [counts (counts-of coll k)
         new-vals (-> coll

              (map var-get)
              (filter (u/fn-> type (= :korma.core/Entity)))
              (filter :hydration-keys)
              (mapcat (fn [{:keys [batched-hydration-keys] :as entity}]
                        (assert (and (set? batched-hydration-keys) (every? keyword? batched-hydration-keys))
                                (str ":hydration-keys should be a set of keywords. In: " entity))
                        (map (u/rpartial vector entity)
                             batched-hydration-keys)))
              (into {}))))

(def ^:private batched-hydration-keys

;; At a high level, these functions let you aggressively flatten a sequence of maps by a key
;; so you can apply some function across it, and then unflatten that sequence.
;;
;; 1.  Get a value that can be used to unflatten a sequence later with `counts-of`.
;; 2.  Flatten the sequence with `counts-flatten`
;; 3.  Modify the flattened sequence as needed
;; 4.  Unflatten the sequence by calling `counts-unflatten` with the modified sequence value from step 1
;; 5.  `map merge` the original sequence and the unflattened sequence.
;;
;; For your convenience `counts-apply` combines these steps for you.

(defn- counts-of
  "Return a sequence of counts / keywords that can be used to unflatten
   COLL later.

    (counts-of [{:a [{:b 1} {:b 2}], :c 2}
                {:a {:b 3}, :c 4}] :a)
      -> [2 :atom]

   For each `x` in COLL, return:

   *  `(count (k x))` if `(k x)` is sequential
   *  `:atom`         if `(k x)` is otherwise non-nil
