Tail-recursive Tree fold using CPS

README.md

Tail-recursive Tree fold

Both the original and the proposed solution assume an associative and commutative combining function (g).

That bugged me so a colleague and I put this together.

TailRecursiveTreeFold.scala

object TailRecursiveTreeFold {
  sealed trait Tree[A]
  case class Leaf[A](get: A) extends Tree[A]
  case class Branch[A](left: Tree[A], right: Tree[A]) extends Tree[A]

  def fold[A, B](tree: Tree[A], f: A => B, g: (B, B) => B): B = {
    def lk(r: Tree[A], k: B => B): B => B = (lb) => go(r, rk(lb, k))
    def rk(lb: B, k: B => B): B => B = (rb) => k(g(lb, rb))

    @annotation.tailrec
    def go(t: Tree[A], k: B => B): B = t match {
      case Leaf(a) => k(f(a))
      case Branch(l, r) => go(l, lk(r, k))
    }

    go(tree, identity)
  }
}