Lazy Iterators on Scala

Exercise:

Implement proper (means lazy) iterator for data structure B-Tree.

Code (Scala):

import org.scalatest.FunSuite

class LazyIteratorTest extends FunSuite {

  trait LazyIt[T] {
    def iterator(): Iterator[T]
    protected def lazyIt[T](
      sq: Seq[LazyIt[T]]): Iterator[T] =
      sq match {
        case Seq(hd) => hd.iterator
        case Seq(hd, tl @ _*) =>
          hd.iterator ++ lazyIt(tl)
        case _ => Iterator.empty
      }
  }

  sealed trait BTree[K, V]
    extends LazyIt[(K, V)] {
    type T = (K, V)
    def map[R](f: T => R): Seq[R]
  }

  case class BTLeaf[K, V](vs: Seq[(K, V)])
    extends BTree[K, V] {
    override def toString =
      "leaf(" + shorten(showSeq(vs)) + ")"
    def iterator(): Iterator[T] = {
      println("* iterator() on " + this + " *")
      vs.iterator
    }
    def map[R](f: T => R): Seq[R] = vs.map(f)
  }

  case class BTNode[K, V](ts: Seq[BTree[K, V]])
    extends BTree[K, V] {
    override def toString =
      "node(" + shorten(showSeq(ts)) + ")"
    def iterator(): Iterator[T] = {
      println("* iterator() on " + this + " *")
      lazyIt(ts)
    }
    def map[R](f: T => R): Seq[R] =
      ts.flatMap(_.map(f))
  }

  def showSeq[T](sq: Seq[T]) = 
      sq.map(_.toString).mkString(",")
  def shorten[T](s: String) = 
      if (s.size <= 30) s else s.take(30) + "..."
  def time = System.currentTimeMillis()

  def leaf[K, V](vs: (K, V)*) = BTLeaf[K, V](vs)
  def node[K, V](ts: BTree[K, V]*) = BTNode[K, V](ts)

  def tree = node(
    node(leaf(1 -> "a", 2 -> "b", 3 -> "c"),
      leaf(4 -> "d", 5 -> "e", 6 -> "f"),
      leaf(7 -> "g", 8 -> "h", 9 -> "i")),
    node(leaf(10 -> "j", 11 -> "k", 12 -> "l"),
      leaf(13 -> "m", 14 -> "n", 15 -> "o"),
      leaf(16 -> "p", 17 -> "q", 18 -> "r")),
    node(leaf(19 -> "s", 20 -> "t", 21 -> "u"),
      leaf(22 -> "v", 23 -> "w", 24 -> "x"),
      leaf(25 -> "y", 26 -> "z", 27 -> "!")))

  test("lazy iterator") {
    println("warm up before...")
    tree.map(_.toString)
    println("do simple map")
    val t = time
    println(shorten(showSeq(
      tree.map(_.toString))))
    println("it took time: "
      + (time - t) + " ms")
    val lt = time
    println("get lazy iterator {")
    val it = tree.iterator()
    println("} got lazy iterator")
    println("do lazy take(1) {")
    val t1 = it.take(1)
    println("} lazy taken(1)")
    println("do lazy take(2) {")
    val t2 = it.take(2)
    println("} lazy taken(2)")
    println(t1.toList)
    println(t2.toList)
    println("do lazy foreach {")
    it.foreach { println }
    println("} done lazy foreach")
    println("lazily it took time: "
      + (time - lt) + " ms")
  }
}

Output:

warm up before...
do simple map
(1,a),(2,b),(3,c),(4,d),(5,e),...
it took time: 6 ms
get lazy iterator {
* iterator() on node(node(leaf((1,a),(2,b),(3,c)),l...) *
* iterator() on node(leaf((1,a),(2,b),(3,c)),leaf((...) *
* iterator() on leaf((1,a),(2,b),(3,c)) *
} got lazy iterator
do lazy take(1) {
} lazy taken(1)
do lazy take(2) {
} lazy taken(2)
List((1,a))
List((2,b), (3,c))
do lazy foreach {
* iterator() on leaf((4,d),(5,e),(6,f)) *
(4,d)
(5,e)
(6,f)
* iterator() on leaf((7,g),(8,h),(9,i)) *
(7,g)
(8,h)
(9,i)
* iterator() on node(leaf((10,j),(11,k),(12,l)),lea...) *
* iterator() on leaf((10,j),(11,k),(12,l)) *
(10,j)
(11,k)
(12,l)
* iterator() on leaf((13,m),(14,n),(15,o)) *
(13,m)
(14,n)
(15,o)
* iterator() on leaf((16,p),(17,q),(18,r)) *
(16,p)
(17,q)
(18,r)
* iterator() on node(leaf((19,s),(20,t),(21,u)),lea...) *
* iterator() on leaf((19,s),(20,t),(21,u)) *
(19,s)
(20,t)
(21,u)
* iterator() on leaf((22,v),(23,w),(24,x)) *
(22,v)
(23,w)
(24,x)
* iterator() on leaf((25,y),(26,z),(27,!)) *
(25,y)
(26,z)
(27,!)
} done lazy foreach
lazily it took time: 9 ms