Pure-functional IO in Scala

IO.scala

sealed trait IOOperation[A] {
  def map[B](f: A => B): IOOperation[B] =
    this match {
      case PutChar(c, a) => 
        PutChar(c, f(a))
      case GetChar(g) =>
        GetChar(f compose g)
    }
}
case class PutChar[A](c: Char, a: A) extends IOOperation[A]
case class GetChar[A](f: Char => A) extends IOOperation[A]

sealed trait IO[A] {
  def map[B](f: A => B): IO[B] =
    this match {
      case Done(a) =>
        Done(f(a))
      case More(a) => 
        More(a map (_ map f))
    }

  def flatMap[B](f: A => IO[B]): IO[B] =
    this match {
      case Done(a) =>
        f(a)
      case More(a) => 
        More(a map (_ flatMap f))
    }
}
case class Done[A](a: A) extends IO[A]
case class More[A](a: IOOperation[IO[A]]) extends IO[A]

object IO {
  def putChar(c: Char): IO[Unit] =
    More(PutChar(c, Done(())))

  def getChar: IO[Char] =
    More(GetChar(Done(_)))
}

object Program {
  import IO._

  def program: IO[Unit] =
    for {
      c1 <- getChar
      c2 <- getChar
      _ <- putChar(c1)
      _ <- putChar(c2)
    } yield ()
}