streams.scala

package streams

import scala.annotation.tailrec

sealed trait Stream[+A] {
  def fold[B](z: B)(f: (A, => B) => B): B = {
    def inner(as: Stream[A]) = as match {
      case Cons(h, t) => f(h(), t().fold(z)(f))
      case _ => z
    }
    inner(this)
  }

  def raise: Stream[Option[A]] =
    Stream.unfold(this) {
      case Cons(h, t) => Some((Some(h()), t()))
      case _ => Some((None, Empty))
    }

  def startsWith[B >: A](as: Stream[B]): Boolean =
    Stream.zipAllWith(as, this)({
      case (Some(a), Some(b)) => a == b
      case (Some(_), None) => false
      case _ => true
    }).fold(true)(_ && _)

  def toList: List[A] = {
    @tailrec
    def loop(as: Stream[A], acc: List[A]): List[A] = as match {
      case Empty => acc
      case Cons(h, t) => loop(t(), h() :: acc)
    }
    loop(this, List()).reverse
  }

  override def toString: String = toList.toString
}

case object Empty extends Stream[Nothing]
case class Cons[+A](h : () => A, t: () => Stream[A]) extends Stream[A]

object Stream {
  private def cons[A](hd: => A, tl: => Stream[A]): Stream[A] = {
    lazy val head = hd
    lazy val tail = tl
    Cons(() => head, () => tail)
  }

  private def empty[A]: Stream[A] = Empty

  def unfold[A, S](z: S)(f: S => Option[(A, S)]): Stream[A] =
    f(z) match {
      case Some((a, s)) => Stream.cons(a, unfold(s)(f))
      case _ => Stream.empty
    }

  def zipAllWith[A, B, C](as: Stream[A], bs: Stream[B])(f: (Option[A], Option[B]) => C): Stream[C] =
    unfold((as.raise, bs.raise))({
      case (Cons(h1, t1), Cons(h2, t2)) => (h1(), h2()) match {
        case (None, None) => None
        case (a, b) => Some((f(a, b), (t1(), t2())))
      }
      case _ => None
    })

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


object Streams {
  def main(args: Array[String]): Unit = {
    println(Stream() startsWith Stream())
    println(Stream(1, 2, 3, 4) startsWith Stream())
    println(Stream(1, 2, 3, 4) startsWith Stream(1, 2))
    println(Stream(1, 2, 3, 4) startsWith Stream(3, 4))
    println(Stream(1, 2, 3, 4) startsWith Stream(1, 2, 4))
    println(Stream(1, 2, 3, 4) startsWith Stream(1, 2, 3, 4, 5))
  }
}