Scala tree recursive fold method

Question

Given the following definition for a (not binary) tree:

sealed trait Tree[+A]
case class Node[A](value: A, children: List[Node[A]]) extends Tree[A]
object Tree {...}

I have written the following fold method:

def fold[A, B](t: Node[A])(f: A ⇒ B)(g: (B, List[B]) ⇒ B): B =
  g(f(t.value), t.children map (fold(_)(f)(g)))

that can be nicely used for (among other things) this map method:

def map[A, B](t: Node[A])(f: A ⇒ B): Node[B] =
  fold(t)(x ⇒ Node(f(x), List()))((x, y) ⇒ Node(x.value, y))

Question: can someone help me on how to write a tail recursive version of the above fold?

Answer 1

I believe you need a stack to do such a traversal, just as in imperative programming, there would be no natural way to write that without a recursive method. Of course, you can always manage the stack yourself, which moves it into the heap, and prevents stack overflows. Here is an example :

sealed trait Tree[+A]
case class Node[+A](value: A, children: List[Node[A]]) extends Tree[A]

case class StackFrame[+A,+B](
  value: A, 
  computedChildren: List[B], 
  remainingChildren: List[Node[A]])

def fold[A,B](t: Node[A])(f: A => B)(g: (B, List[B]) => B) : B = {

  def go(stack: List[StackFrame[A,B]]) : B = stack match {
    case StackFrame(v, cs, Nil) :: tail => 
      val folded = g(f(v), cs.reverse)
      tail match {
        case Nil => folded
        case StackFrame(vUp, csUp, remUp) :: rest => 
          go(StackFrame(vUp, folded::csUp, remUp)::rest)
      }
    case StackFrame(v, cs, nextChild :: others) :: tail =>
      go(
        StackFrame(nextChild.value, Nil, nextChild.children) ::
        StackFrame(v, cs, others) :: 
        tail)
    case Nil => sys.error("Should not go there")
  }

  go(StackFrame(t.value, Nil,  t.children) :: Nil)    
}

Note: I made Node covariant, not strictly necessary, but if it is not, you will need to be explicit in the type of a few Nil (e.g replace by List[X]()) in some place.

go it clearly tail recursive, but just because it manages the stack itself.

You may find a more principled and systematic technique (but not easy to grasp at first) based on continuations and trampolines, in this nice blog post.