Tree zipper

TreeZipper.scala

package test.zipper

object TreeZipper extends App {

  println("hello")

  // exp 2 * 3 + 4 * 6

  val _2 = Node("2", Nil)
  val _3 = Node("3", Nil)

  val _4 = Node("4", Nil)
  val _6 = Node("6", Nil)

  val mult1 = Node("*", List(_2, _3))
  val mult2 = Node("*", List(_4, _6))

  val plus = Node("+", List(mult1, mult2))

  val locOfPlus = Location(plus, Top)

  println(locOfPlus.tree)
  println()

  for (l1 <- locOfPlus.down;
       _2 <- l1.down;
       u2 <- _2.update(Node("5", Nil));
       up1 <- u2.up;
       up2 <- up1.up;
       dn1 <- up2.down;
       dn2 <- dn1.down) {
     println(dn2)
  }

}


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

sealed trait Context[+A]
case object Top extends Context[Nothing]
case class NodeContext[A](left: List[Tree[A]], parent: Location[A], right: List[Tree[A]]) extends Context[A]


trait ZipperNavigation[A] { self: Location[A] =>

  def left : Option[Location[A]] = ctx match {
    case NodeContext(Nil, parent, right) => None
    case NodeContext(h :: t, parent, right) => 
      Some(Location(h, NodeContext(t, parent, tree :: right)))
    case _ => None
  }

  def right : Option[Location[A]] = ctx match {
    case NodeContext(left, parent, Nil) => None
    case NodeContext(left, parent, h :: tail) => 
      Some(Location(h, NodeContext( tree :: left, parent, tail)))
    case _ => None
  }

  def down: Option[Location[A]] = this match {
    case Location(Node(_, Nil), ctx) => None
    case Location(Node(_, h :: t), ctx) => Some(Location(h, NodeContext(Nil, this, t)))
  }

  def up: Option[Location[A]] = this match {
    case Location(_, Top) => None
    case Location(t, NodeContext(l, Location(Node(a, _), Top), r)) => 
      Some(Location(Node(a, l ::: List(t) ::: r), Top))
    case Location(t, NodeContext(_, Location(Node(a, _), NodeContext(l, p, r)), _)) => 
      Some(Location(Node(a, l ::: List(t) ::: r), NodeContext(l.reverse, p, r)))
    case k => None
  }

}

trait ZipperMutator[A] { self : Location[A] =>

  def update(t: Tree[A]): Option[Location[A]] = this match {
    case Location(_, ctx) => Some(Location(t, ctx))
  }

  def insertLeft(newt : Tree[A]): Option[Location[A]] = this match {
    case Location(tree, Top) => None
    case Location(tree, NodeContext(l, p, r)) => 
      Some(Location(newt, NodeContext(l, p, tree :: r)))
  }

  def insertRight(newt : Tree[A]): Option[Location[A]] = this match {
    case Location(tree, Top) => None
    case Location(tree, NodeContext(l, p, r)) => 
      Some(Location(newt, NodeContext(l ::: List(tree), p, r)))
  }

  def delete: Option[Location[A]] = this match {
    case Location(tree, Top) => None
    case Location(tree, NodeContext(Nil, p, Nil)) => Some(p)
    case Location(tree, NodeContext(f :: l, p, Nil)) => 
      Some(Location(f, NodeContext(l, p, Nil)))
    case Location(tree, NodeContext(l, p, r :: rest)) => 
      Some(Location(r, NodeContext(l, p, rest)))
  }

}

case class Location[A](tree: Tree[A], ctx: Context[A]) extends ZipperNavigation[A] with ZipperMutator[A]