clojure zippers, returning a path that can be used with update-in or get-in, has alternate versions of update-in and assoc-in which work with clojure sets

clojure_zippers.clj

(require '[clojure.zip :as z])

(defn update-in*
  "Like update-in, but also works with nested sets."
  [m ks f & args]
  (let [up (fn up [m ks f args]
             (let [[k & ks] ks]
               (if ks
                 (if (set? m)
                   (conj (disj m k) (up (get m k) ks f args))
                   (assoc m k (up (get m k) ks f args)))
                 (if (set? m)
                   (conj (disj m k) (apply f (get m k) args))
                   (assoc m k (apply f (get m k) args))))))]
    (up m ks f args)))

(defn assoc-in*
  "Like assoc-in, but also works with nested sets."
  [m ks v]
  (update-in* m ks (constantly v)))

(defn has-children? [x]
  (boolean (or (map? x)
               (or (seq (:children x))
                   (coll? x)))))

(defn map-vec-zipper [m]
  (let [child-or-branch? (fn [x] (when (and (map? x)
                                            (seq (:children x)))
                                   (seq (:children x))))]
    (z/zipper child-or-branch?
              child-or-branch?
              (fn [node children]
                (if (has-children? node)
                  (assoc node :children (into (empty (:children node)) children))
                  (into (empty node) children)))
              m)))

(defn zip-next-seq
  "Given a clojure.zip zipper location loc return a lazy sequence of all
  clojure.zip/next locations from loc."
  [loc]
  (if (z/end? loc)
    ()
    (lazy-seq (cons loc (zip-next-seq (z/next loc))))))

(defn find-node [zip-seq k v]
  (some (fn [x]
          (let [n (z/node x)]
            (and (map? n)
                 (= (get n k) v)
                 x)))
        zip-seq))

(def test-tree {:id "123"
                :type :group
                :children [{:id "abc"
                            :type :rule}
                           {:id "456"
                            :type :group
                            :children #{{:id "foo"
                                         :type :rule}
                                        {:id "fart"
                                         :type :rule}
                                        {:id "bar"
                                         :type :group
                                         :children [{:id "flerg"}
                                                    {:id "qux"}]}}}]})
(defn path-to [rules id]
  "Provided a rules datastructure, returns the path to the rule where
  the key and value match. See example below."
  (let [find-node (fn [id zip-seq]
                    (some (fn [x]
                            (let [n (z/node x)]
                              (and (map? n)
                                   (= (get n :id) id)
                                   x)))
                          zip-seq))
        zipped-seq (zip-next-seq (map-vec-zipper rules))
        loc (find-node id zipped-seq)]
    (loop [l loc
           path '()]
      (if-some [up (z/up l)]
        (let [up-node (z/node up)
              idx (count (z/lefts l))]
          (cond
            (and (has-children? up-node)
                 (set? (:children up-node)))
            (recur up (concat [:children (z/node l)] path))

            (has-children? up-node)
            (recur up (concat [:children idx] path))

            :otherwise (recur up (cons idx path))))
        (vec (seq path))))))

(get-in {:a {:b {:c #{{:d {:e :f}}}}}}
        [:a :b :c {:d {:e :f}} :d])
;; => {:e :f}

(update-in* {:a {:b {:c #{{:d {:e :f}}}}}}
            [:a :b :c {:d {:e :f}} :d]
            assoc :g 1)
;; => {:a {:b {:c #{{:d {:e :f, :g 1}}}}}}

(let [p (path-to test-tree "bar")]
  (update-in* test-tree
              (conj p :children)
              conj
              {:id "new item"}))
;; => {:id "123",
;;     :type :group,
;;     :children
;;     [{:id "abc", :type :rule}
;;      {:id "456",
;;       :type :group,
;;       :children
;;       #{{:id "bar",
;;          :type :group,
;;          :children [{:id "flerg"} {:id "qux"} {:id "new item"}]}
;;         {:id "foo", :type :rule} {:id "fart", :type :rule}}}]}

(path-to test-tree "qux")
;; => [:children
;;     1
;;     :children
;;     {:id "bar", :type :group, :children [{:id "flerg"} {:id "qux"}]}
;;     :children
;;     1]