'supdate': Clojure's update with superpowers

supdate.clj

(ns utils.supdate
  "A macro for transforming maps (and other data structures) using a spec reflecting the
  schema of the value to transform.")

;; ------------------------------------------------------------------------------
;; Example Usage

(comment
  ;;;; nominal cases
  (supdate
    {:a 1 :b [{:c 2} {:c 3}]} ;; the data structure to transform
    {:a inc :b [{:c dec}]} ;; the transformation
    )
  => {:a 2, :b [{:c 1} {:c 2}]}

  ;; if the transform is a function: apply it to the value
  (supdate
    0
    inc)
  => 1

  ;; if it's a map: recursively transform the value for the given keys.
  (supdate
    {:a 1}
    {:a inc})
  => {:a 2}

  ;; if it's a vector with one (transform) element: transform each item in the collection.
  (supdate
    [1 2 3 4]
    [inc])
  => [2 3 4 5]

  ;; you can nest transforms arbitrarily.
  (supdate
    {:a {:b [{:c 1}]}}
    {:a {:b [{:c inc}]}})
  => {:a {:b [{:c 2}]}}

  ;; if a key is missing, no transform is applied
  (supdate
    {:a 1}
    {:a inc :b inc})
  => {:a 2}

  ;; a vector with 2 or more elements means 'chain the transforms'
  (supdate
    0
    [inc inc inc])
  => 3

  ;; a transform of `false` consists of dissoc'ing in the enclosing map
  (supdate
    {:a 1 :b 2}
    {:a false})
  => {:b 2}

  ;;;; Implementation notes
  ;; `supdate` is a macro which will attempt to perform static code analysis
  ;;  to generate low-level transformation code inline in order to achieve high performance,
  ;; and fall back to a dynamic implementation where that static code analysis fails
  ;; (see the supdate* function)

  )

;; ------------------------------------------------------------------------------
;; Implementation

(defn upd!*
  "A slightly modified version of update.
  If map `m` contains key `k`, will add 
  (f (get m k)) to the transient map `tm`."
  ;; we have to pass the original (non transient) map tm,
  ;; because it's the only one that supports the `contains?` operation.
  [m tm k f]
  (let [v (get tm k)]
    (if v
      (assoc! tm k (f v))
      (if (contains? m k)
        (assoc! tm k (f v))
        tm))))

(defn upd-dynamic!*
  "A version of upd! where we're not sure f is a function"
  [m tm k f]
  (let [v (get tm k)]
    (if v
      (if (false? f)
        (dissoc! tm k)
        (assoc! tm k (f v)))
      (if (contains? m k)
        (if (false? f)
          (dissoc! tm k)
          (assoc! tm k (f v)))
        tm))))

(defn supd-map*
  [f coll]
  (if (vector? coll)
    (mapv f coll)
    (map f coll)))

(defn supdate*
  [v transform]
  (cond
    (fn? transform)
    (transform v)

    (keyword? transform)
    (transform v)

    (map? transform)
    (when v
      (persistent!
        (reduce-kv
          (fn [tm k spec]
            (if (false? spec)
              (dissoc! tm k)
              (upd!* v tm k #(supdate* % spec))))
          (transient v) transform)))

    (and (vector? transform) (= (count transform) 1))
    (let [sspec (first transform)]
      (supd-map* #(supdate* % sspec) v))

    (sequential? transform)
    (reduce supdate* v transform)
    ))

(defn- static-transform?
  [t-form]
  (or (map? t-form) (vector? t-form) (false? t-form) (keyword? t-form)))

(defn- static-key?
  [key-form]
  (or (keyword? key-form) (string? key-form) (number? key-form) (#{true false} key-form)))

(defmacro supdate
  "'Super Update' - convenience for transforming all kinds of values, useful for format coercions etc.
Accepts a value `v` and a transform specification `transform` that represents a transformation to apply on v:
* if transform is a function or keyword, will apply it to v
* if transform is a map, will treat v as a map, and recursively modify the values of v for the keys transform supplies. If the transform value for a key is `false`, then the key is dissoc'ed.
* if transform is a vector with one element (a nested transform), will treat v as a collection an apply the nested transform to each element.
* if transform is a sequence, will apply each transform in the sequence in order."
  [v transform]
  (let [vsym (gensym "v")]
    `(let [~vsym ~v]
       ~(cond
          (map? transform)
          (let [tvsym (gensym "tv")]
            `(when ~vsym
               (as-> (transient ~vsym) ~tvsym
                 ~@(for [[k transform] transform]
                     (let [sk? (static-key? k)
                           ksym (if sk? k (gensym "k"))
                           form (cond
                                  (false? transform)
                                  `(dissoc! ~tvsym ~ksym)

                                  (static-transform? transform)
                                  `(upd!* ~vsym ~tvsym ~ksym (fn [v#] (supdate v# ~transform)))

                                  :dynamic
                                  `(let [spec# ~transform]
                                     (upd-dynamic!* ~vsym ~tvsym ~ksym (fn [v#] (supdate* v# spec#)))))]
                       (if sk? form `(let [~ksym ~k] ~form))))
                 (persistent! ~tvsym))))

          (keyword? transform)
          `(~transform ~vsym)

          (and (vector? transform) (= (count transform) 1))
          `(supd-map* (fn [e#] (supdate e# ~(first transform))) ~vsym)

          (and (vector? transform) (> (count transform) 1))
          `(as-> ~vsym ~vsym
             ~@(for [spec transform]
                 `(supdate ~vsym ~spec)))

          :else
          `(supdate* ~vsym ~transform)
          ))))