Simple IO type

IO.scala

object IO2a {

  /*
  The previous IO representation overflows the stack for some programs.
  The problem is that `run`
  */

  sealed trait IO[A] {
    def flatMap[B](f: A => IO[B]): IO[B] =
      FlatMap(this, (a: A) => Suspend { () => f(a) }) // note this returns immediately
    def map[B](f: A => B): IO[B] =
      flatMap(f andThen (Return(_)))
  }
  sealed trait Productive[A] extends IO[A]
  case class Return[A](a: A) extends Productive[A]
  case class Suspend[A](resume: () => IO[A]) extends Productive[A]
  case class FlatMap[A,B](sub: IO[A], k: A => Productive[B]) extends IO[B] // k must produce a `Productive`

  object IO extends Monad[IO] {
    def unit[A](a: => A): IO[A] = Return(a)
    def flatMap[A,B](a: IO[A])(f: A => IO[B]): IO[B] = a flatMap f
  }

  def printLine(s: String): IO[Unit] =
    Suspend(() => Return(println(s)))

  val p = IO.forever(printLine("Still going..."))

  val actions: Stream[IO[Unit]] =
    Stream.fill(100000)(printLine("Still going..."))
  val composite: IO[Unit] =
    actions.foldLeft(IO.unit(())) { (acc, a) => acc flatMap { _ => a } }

  @annotation.tailrec def run[A](io: IO[A]): A = io match {
    case Return(a) => a
    case Suspend(r) => run(r())
    case FlatMap(x, f) => x match {
      case Return(a) => run(f(a)) // since f: A => Productive[B], we are assured this returns in one step
      case Suspend(r) => run(FlatMap(r(), f))
      case FlatMap(y, g) =>
        run(y flatMap (a => g(a) flatMap f)) // g: A => Productive[B], so `g(a)` is assured to return in one step
    }
  }
}