n-ary tree zipper

Zipper.scala

sealed trait TreeData
// ...

// could be made generic in data
case class Node(data: TreeData, children: List[Node])

sealed trait ZipCtx
case object Top extends ZipCtx
case class At(revLeft: List[Node], right: List[Node], ctx: ZipCtx, updata: TreeData) extends ZipCtx

// ZipTree(root, Top) is the start
case class ZipTree(n: Node, ctx: ZipCtx) {

  private def toOption[A, B](a: A)(pf: PartialFunction[A, B]): Option[B] =
    if (pf isDefinedAt a) Some(pf(a))
    else None

  def left = toOption(ctx) {
    case At(l :: xs, r, ctx, d) => ZipTree(l, At(xs, n :: r, ctx, d))
  }

  def right = toOption(ctx) {
    case At(l, r :: xs, ctx, d) => ZipTree(r, At(n :: l, xs, ctx, d))
  }

  def down = toOption(n.children) {
    case c :: cs => ZipTree(c, At(Nil, cs, ctx, n.data))
  }

  def up = toOption(ctx) {
    case At(l, r, upctx, d) => ZipTree(Node(d, (n :: l) reverse_::: r), upctx)
  }
}

object ZipTree {
  private def bindApply(f: ZipTree => Option[ZipTree])(loc: Option[ZipTree]) = loc flatMap f

  val left = bindApply(_.left) _
  val right = bindApply(_.right) _

  def up = bindApply(_.up) _

  def down = bindApply(_.down) _

  def walk(start: Option[ZipTree])(path: Seq[Option[ZipTree] => Option[ZipTree]]) =
    (start /: path) {
      case (loc, f) => f(loc)
    }

  def rightOf(tree: ZipTree): Stream[ZipTree] =
    tree.right map (r => Stream.cons(r, rightOf(r))) getOrElse Stream.empty

  def children(zt: ZipTree): Stream[ZipTree] = zt.down map {
    leftmost =>
      Stream.cons(leftmost, rightOf(leftmost))
  } getOrElse Stream.empty

  def childrenKeyed(zt: ZipTree, s: String): Stream[ZipTree] =
    children(zt) filter (_.n.data._1 == s)
}