How to walk an AST with tail recursion in Clojure

Question

I have an ANTLR3 AST which I need to traverse using a post-order, depth-first traversal which I have implemented as roughly the following Clojure:

(defn walk-tree [^CommonTree node]
  (if (zero? (.getChildCount node))
    (read-string (.getText node))
    (execute-node
      (map-action node)
      (map walk-tree (.getChildren node)))))))

I would like to convert this to tail recursion using loop...recur, but I haven't been able to figure out how to effectively use an explicit stack to do this since I need a post-order traversal.

Answer 1

Instead of producing a tail recursive solution which traverses the tree and visits each node, you could produce a lazy sequence of the depth first traversal using the tree-seq function and then get the text out of each object in the traversal. Lazy sequences never blow the stack because they store all the state required to produce the next item in the sequence in the heap. They are very often used instead of recursive solutions like this where loop and recur are more diffacult.

I don't know what your tree looks like though a typical answer would look something like this. You would need to play with the "Has Children" "list of children" functions

(map #(.getText %) ;; Has Children?      List of Children    Input Tree
     (tree-seq #(> (.getChildCount #) 0) #(.getChildren %) my-antlr-ast))

If tree-seq does not suit your needs there are other ways to produce a lazy sequence from a tree. Look at the zipper library next.

Answer 2

As you mention, the only way to implement this using tail recursion is to switch to using an explicit stack. One possible approach is to convert the tree structure into a stack structure that is essentially a Reverse Polish notation representation of the tree (using a loop and an intermediate stack to accomplish this). You would then use another loop to traverse the stack and build up your result.

Here's a sample program that I wrote to accomplish this, using the Java code at postorder using tail recursion as an inspiration.

(def op-map {'+ +, '- -, '* *, '/ /})

;; Convert the tree to a linear, postfix notation stack
(defn build-traversal [tree]
  (loop [stack [tree] traversal []]
    (if (empty? stack)
      traversal
      (let [e (peek stack)
            s (pop stack)]
        (if (seq? e)
          (recur (into s (rest e)) 
                 (conj traversal {:op (first e) :count (count (rest e))}))
          (recur s (conj traversal {:arg e})))))))

;; Pop the last n items off the stack, returning a vector with the remaining
;; stack and a list of the last n items in the order they were added to
;; the stack
(defn pop-n [stack n]
  (loop [i n s stack t '()]
    (if (= i 0)
      [s t]
      (recur (dec i) (pop s) (conj t (peek s))))))

;; Evaluate the operations in a depth-first manner, using a temporary stack
;; to hold intermediate results.
(defn eval-traversal [traversal]
  (loop [op-stack traversal arg-stack []]
    (if (empty? op-stack)
      (peek arg-stack)
      (let [o (peek op-stack)
            s (pop op-stack)]
        (if-let [a (:arg o)]
          (recur s (conj arg-stack a))
          (let [[args op-args] (pop-n arg-stack (:count o))]
            (recur s (conj args (apply (op-map (:op o)) op-args)))))))))

(defn eval-tree [tree] (-> tree build-traversal eval-traversal))

You can call it as such:

user> (def t '(* (+ 1 2) (- 4 1 2) (/ 6 3)))
#'user/t
user> (eval-tree t)
6

I leave it as an exercise to the reader to convert this to work with a Antlr AST structure ;)

Answer 3

I'm not skilled up on clojure, but I think I understand what you're looking for.

Here's some pseudocode. The stack here in my pseudocode looks like a stateful object, but it's quite feasible to use an immutable one instead. It uses something like O(depth of tree * max children per node) heap.

walk_tree(TreeNode node) {
    stack = new Stack<Pair<TreeNode, Boolean>>();
    push(Pair(node, True), stack)
    walk_tree_aux(stack);
}
walk_tree_aux(Stack<Pair<TreeNode, Boolean>> stack) { -- this should be tail-recursive
    if stack is empty, return;
    let (topnode, topflag) = pop(stack);
    if (topflag is true) {
        push Pair(topnode, False) onto stack);
        for each child of topnode, in reverse order:
            push(Pair(child, True)) onto stack
        walk_tree_aux(stack);
    } else { -- topflag is false
        process(topnode)
        walk_tree_aux(stack);
    }
}