scala - How to make tree mapping tail-recursive?

Question

Suppose I have a tree data structure like this:

trait Node { val name: String }
case class BranchNode(name: String, children: List[Node]) extends Node
case class LeafNode(name: String) extends Node

Suppose also I've got a function to map over leaves:

def mapLeaves(root: Node, f: LeafNode => LeafNode): Node = root match {
  case ln: LeafNode => f(ln)
  case bn: BranchNode => BranchNode(bn.name, bn.children.map(ch => mapLeaves(ch, f)))
}

Now I am trying to make this function tail-recursive but having a hard time to figure out how to do it. I've read this answer but still don't know to make that binary tree solution work for a multiway tree.

How would you rewrite mapLeaves to make it tail-recursive?

Answer 1

"Call stack" and "recursion" are merely popular design patterns that later got incorporated into most programming languages (and thus became mostly "invisible"). There is nothing that prevents you from reimplementing both with heap data structures. So, here is "the obvious" 1960's TAOCP retro-style solution:

trait Node { val name: String }
case class BranchNode(name: String, children: List[Node]) extends Node
case class LeafNode(name: String) extends Node

def mapLeaves(root: Node, f: LeafNode => LeafNode): Node = {
  case class Frame(name: String, mapped: List[Node], todos: List[Node])
  @annotation.tailrec
  def step(stack: List[Frame]): Node = stack match {
    // "return / pop a stack-frame"
    case Frame(name, done, Nil) :: tail => {
      val ret = BranchNode(name, done.reverse)
      tail match {
        case Nil => ret
        case Frame(tn, td, tt) :: more => {
          step(Frame(tn, ret :: td, tt) :: more)
        }
      }
    }
    case Frame(name, done, x :: xs) :: tail => x match {
      // "recursion base"
      case l @ LeafNode(_) => step(Frame(name, f(l) :: done, xs) :: tail)
      // "recursive call"
      case BranchNode(n, cs) => step(Frame(n, Nil, cs) :: Frame(name, done, xs) :: tail)
    }
    case Nil => throw new Error("shouldn't happen")
  }
  root match {
    case l @ LeafNode(_) => f(l)
    case b @ BranchNode(n, cs) => step(List(Frame(n, Nil, cs)))
  }
}

The tail-recursive step function takes a reified stack with "stack frames". A "stack frame" stores the name of the branch node that is currently being processed, a list of child nodes that have already been processed, and the list of the remaining nodes that still must be processed later. This roughly corresponds to an actual stack frame of your recursive mapLeaves function.

With this data structure,

• returning from recursive calls corresponds to deconstructing a Frame object, and either returning the final result, or at least making the stack one frame shorter.
• recursive calls correspond to a step that prepends a Frame to the stack
• base case (invoking f on leaves) does not create or remove any frames

Once one understands how the usually invisible stack frames are represented explicitly, the translation is straightforward and mostly mechanical.

Example:

val example = BranchNode("x", List(
  BranchNode("y", List(
    LeafNode("a"),
    LeafNode("b")
  )),
  BranchNode("z", List(
    LeafNode("c"),
    BranchNode("v", List(
      LeafNode("d"),
      LeafNode("e")
    ))
  ))
))

println(mapLeaves(example, { case LeafNode(n) => LeafNode(n.toUpperCase) }))

Output (indented):

BranchNode(x,List(
  BranchNode(y,List(
    LeafNode(A),
    LeafNode(B)
  )),
  BranchNode(z, List(
    LeafNode(C),
    BranchNode(v,List(
      LeafNode(D),
      LeafNode(E)
    ))
  ))
))