how to implement walk/postwalk traversal using clojure.zip

Question

I can't figure out how to implement the clojure.walk/postwalk function using clojure.zip:

(clojure.walk/postwalk
   #(do (println %)
        %)
   [1
    [2 [3 4 5]]
    [6 [7 8]]])

outputs:

1
2
3
4
5
[3 4 5]
[2 [3 4 5]]
6
7
8
[7 8]
[6 [7 8]]
[1 [2 [3 4 5]] [6 [7 8]]]

Answer 1

(defn postwalk [f loc]
  (let [loc (if-some [loc (z/down loc)]
              (loop [loc loc]
                (let [loc (postwalk f loc)]
                  (if-some [loc (z/right loc)]
                    (recur loc)
                    (z/up loc))))
              loc)]
    (z/replace loc (f (z/node loc)))))

=> (postwalk #(doto % prn) (z/vector-zip [1 [2 [3 4 5]] [6 [7 8]]]))
1
2
3
4
5
[3 4 5]
[2 [3 4 5]]
6
7
8
[7 8]
[6 [7 8]]
[1 [2 [3 4 5]] [6 [7 8]]]

Edit: for prewalk, just perform the z/replace before going down.

(defn prewalk [f loc]
  (let [loc (z/replace loc (f (z/node loc)))]
    (if-some [loc (z/down loc)]
      (loop [loc loc]
        (let [loc (prewalk f loc)]
          (if-some [loc (z/right loc)]
            (recur loc)
            (z/up loc))))
      loc)))

Answer 2

I believe there is a more idiomatic way to implement post order traversal that preserves zipper style navigation:

(defn post-zip
  [loc]
  ;; start from the deepest left child
  (loop [loc loc]
    (if (z/branch? loc)
      (recur (z/down loc))
      loc)))

(defn post-next
  [loc]
  (if-let [sib (z/right loc)]
    ;; if we have a right sibling, move to it's deepest left child
    (post-zip sib)
    ;; If there is no right sibling move up if possible
    (if-let [parent (z/up loc)]
      parent
      ;; otherwise we are done
      [(z/node loc) :end])))

This allows you to make conditional modifications in the same way you would with a normal zipper (ie you don't always have to z/replace , sometimes you can just call post-next with no change to the current node).

The implementation of postwalk given these helper functions becomes quite trivial:

(defn postwalk [f zipper]
  (loop [loc (post-zip zipper)]
    (if (z/end? loc)
      (z/node loc)
      (recur (post-next (z/replace loc (f (z/node loc))))))))

This solution also has the advantage of not introducing recursion that could overflow the stack for large trees.