depth-first indexing any Clojure form using clojure.walk

Question

Given the following tree (or any other form in Clojure including maps and vectors):

'( (a b) (c d) )

I would like to generate a map in Clojure that indexes each sub-form according to a depth-first traversal of the entire form and also provides a vector (or list) of the indices of the form's children (if any).

0 -> a []
1 -> b []
2 -> (a b) [0 1]
3 -> c []
4 -> d []
5 -> (c d) [3 4]
6 -> ( (a b) (c d) ) [2 5]

I have so far only managed to use clojure.walk to produce the first part (indexing the subforms) but I am baffled as to how to generate the indices of the children as well. My code is appended at the end and produces:

user=> (depthFirstIndexing '( (a b) (c d) ))
{6 ((a b) (c d)), 5 (c d), 4 d, 3 c, 2 (a b), 1 b, 0 a}

So the indexes to the sub-forms are generated correctly according to depth-first traversal but I don't see how I can obtain the indices of the children of every sub-form. I tried to use the zippers module but I couldn't see how to perform a depth-first traversal to collect the indices.

half-way there code

(use 'clojure.walk)
(defn depthFirstIndexing [aform]
  (let [counter       (atom -1)
        idxToSubform  (atom {})
        ]
    (postwalk (fn [x]
                (def idx (swap! counter inc))
                (swap! idxToSubform assoc idx x)
                x)
    aform)
  @idxToSubform))

Answer 1

A walk is recursive and does not provide for an accumulator argument, which is why you have had to resort to updating atoms.

A zipper is iterative, so you can carry along other information without breaking a functional pattern.

The natural depth-first traversal is a pre-order traversal, but you are after a post-order, so this requires a little extra work.

Here is a post-order traversal using zippers:

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

(defn dfs-post-order-traversal [zipper]
  (loop [loc zipper, a []] 
    (cond 
      (z/end? loc) 
        (conj a (z/node loc))
      (z/branch? loc) 
        (recur (z/next loc) a)
      :else
        (recur 
          (z/next loc) 
          (into 
            (conj a (z/node loc)) 
            (reverse 
              (drop 
                ((fnil count [nil]) (z/path (z/next loc))) 
                (z/path loc))))))))

And the test case:

(dfs-post-order-traversal (z/seq-zip '((a b) (c d))))
=> [a b (a b) c d (c d) ((a b) (c d))]

Now to finish off your request, we need to map tree locations back to their indices:

(defn dfs-post-order-indexing [branch? children root]
  (let [pot (dfs-post-order-traversal (z/zipper branch? children conj root))
        m (zipmap pot (range))]
    (for [n pot] [(m n) n (if (branch? n) (map m (children n)) (list))])))

(dfs-post-order-indexing seq? identity '((a b) (c d)))
=>  ([0 a ()]
     [1 b ()]
     [2 (a b) (0 1)]
     [3 c ()]
     [4 d ()]
     [5 (c d) (3 4)]
     [6 ((a b) (c d)) (2 5)])

Something more exotic:

(dfs-post-order-indexing coll? seq [{:a :b :c :d} :e [:f [:g '(:h :i)]]])
=>  ([0 :a ()]
     [1 :b ()]
     [2 [:a :b] (0 1)]
     [3 :c ()]
     [4 :d ()]
     [5 [:c :d] (3 4)]
     [6 {:a :b, :c :d} (2 5)]
     [7 :e ()]
     [8 :f ()]
     [9 :g ()]
     [10 :h ()]
     [11 :i ()]
     [12 (:h :i) (10 11)]
     [13 [:g (:h :i)] (9 12)]
     [14 [:f [:g (:h :i)]] (8 13)]
     [15 [{:a :b, :c :d} :e [:f [:g (:h :i)]]] (6 7 14)])

Answer 2

(use '[clojure.walk :only (postwalk)])
(use '[clojure.set :only (map-invert)])

(defn idx [col]
  (let [m (map vector 
               (range) 
               (let [v (atom [])]
                 (postwalk (fn [f] (swap! v conj f) f) col) 
                 @v))
        rm (map-invert m)]
    (into {} (map (fn [[i e]]
                    [i [e (if (sequential? e) 
                            (mapv rm e)
                            [])]])
                  m))))

(idx '((a b) (c d)))
=> {0 [a []],
    1 [b []],
    2 [(a b) [0 1]],
    3 [c []],
    4 [d []],
    5 [(c d) [3 4]],
    6 [((a b) (c d)) [2 5]]}