Stream implementation

gistfile1.scala

object Chap5 extends App {
  import Stream._

  def ones: Stream[Int] = cons(1, ones)
  def lin: Stream[Int] = cons(1, lin.map(_ + 1))
 
  println(ones.flatMap(_ => empty[Int])) // Stack overflow!
}

sealed trait Stream[A] {

  import Stream._

  def uncons[B](empty: => B, cons: (=> A, => Stream[A]) => B): B

  def head: Option[A] = uncons(None, (h, _) => Some(h))

  def tail: Option[Stream[A]] = uncons(None, (_, t) => Some(t))

  def toList: List[A] = foldRight[List[A]](Nil)(_ :: _)

  def take(n: Int): Stream[A] =
    uncons(empty, (h, t) => if (n > 0) cons(h, t.take(n - 1)) else empty)

  def foldRight[B](z: => B)(f: (=> A, => B) => B): B =
    uncons(z, (h, t) => f(h, t.foldRight(z)(f)))

  def map[B](f: A => B): Stream[B] =
    foldRight[Stream[B]](empty)((h, t) => cons(f(h), t))

  def concat(other: => Stream[A]): Stream[A] =
    foldRight(other)((a, as) => (cons(a, as)))

  def flatMap[B](f: A => Stream[B]): Stream[B] =
    foldRight[Stream[B]](empty)((h, t) => f(h).concat(t))

  //override def toString = "Stream(" + head.map(_ + ", ...").getOrElse("") + ")"

}

object Stream {

  def empty[A]: Stream[A] = new Stream[A] {
    def uncons[B](empty: => B, cons: (=> A, => Stream[A]) => B) = empty
  }

  def cons[A](hd: => A, tl: => Stream[A]): Stream[A] = new Stream[A] {
    def uncons[B](empty: => B, cons: (=> A, => Stream[A]) => B) = cons(hd, tl)
  }

  def id[A](as: Stream[A]): Stream[A] = as.uncons(empty, cons)
  def id2[A](as: Stream[A]): Stream[A] = as.foldRight[Stream[A]](empty)(cons)

  def apply[A](as: A*): Stream[A] =
    if (as.isEmpty) empty
    else cons(as.head, apply(as.tail: _*))

}