Lazy-Step1.scala

///////////////////////////////////////////////////////////////////
// STEP 1 - Evaluation by need monad
///////////////////////////////////////////////////////////////////

import scala.annotation.unchecked.{ uncheckedVariance => uV }
final class Need[+A](private[this] var thunk: Need.Thunk[A @uV]) { A =>
  import Need._
  
  def value: A = thunk match {
    case Done(x) => x
    case x => 
      val result = x.value
      thunk = Done(result)
      result
  }
  
  def map[B](f: A => B): Need[B] = new Need(FlatMap[A, B](this, a => now(f(a))))
  def flatMap[B](f: A => Need[B]): Need[B] = new Need(FlatMap(this, f))
  def  *> [B](B: Need[B]): Need[B] = A.flatMap(_ => B)
  def <*  [B](B: Need[B]): Need[A] = A.flatMap(a => B.map(_ => a))
  def <*> [B](B: Need[B]): Need[(A, B)] = A.flatMap(a => B.map(b => (a, b)))
}
object Need {
  sealed trait Thunk[+A] {
    def value: A
  }
  final case class Done[A](value: A) extends Thunk[A]
  final case class FlatMap[Z, A](left: Need[Z], f: Z => Need[A]) extends Thunk[A] {
    def value = f(left.value).value
  }
  
  def apply[A](a: => A): Need[A] = new Need(FlatMap[Unit, A](unit, _ => now(a)))
  def now[A](a: A): Need[A] = new Need(Done(a))
  val unit: Need[Unit] = now(())
}

///////////////////////////////////////////////////////////////////
// STEP 2 - Typeclass for types with a Thunk
///////////////////////////////////////////////////////////////////

trait Delay[A] extends Any with Serializable {
  def delay(a: Need[A]): A
  def project(a: A): Need[A]
}
object Delay {
  def apply[A](implicit A: Delay[A]): Delay[A] = A

  implicit def string: Delay[String] = new Delay[String] {
    override def delay(a: Need[String]): String = a.value
    override def project(a: String): Need[String] = Need.now(a)
  }
  
  implicit def need[A]: Delay[Need[A]] = new Delay[Need[A]] {
    override def delay(a: Need[Need[A]]): Need[A] = a.flatMap(x => x)
    override def project(a: Need[A]): Need[Need[A]] = a.map(Need.now)
  }
}

def testDelayLaws[A](a: A)(implicit A: Delay[A]): Unit = {
  var c = 0
  def SE = { c += 1 }
  def assertSideEffect(): Unit = {
    assert(c == 1)
    c = 0
  }
  def assertNoSideEffect(): Unit = {
    assert(c == 0)
  }
  
  val x = A.delay(Need {SE; a})
  assertNoSideEffect()
  val y = A.project(x)
  assertNoSideEffect()
  y.value
  assertSideEffect()
  y.value
  assertNoSideEffect()
}

testDelayLaws[Need[Int]](Need.now(10))
// testDelayLaws[String]("a")
println("Done!")

///////////////////////////////////////////////////////////////////
// STEP 3 - Example type
///////////////////////////////////////////////////////////////////

sealed trait Maybe[+A] {
  import Maybe._
  def thunk: Need[Strict[A]]

  def map[B](f: A => B): Maybe[B] =
    Maybe.Thunk(this.thunk.map {
      case Maybe.None => Maybe.None
      case Maybe.Some(a) => Maybe.Some(f(a))
    })

  def flatMap[B](f: A => Maybe[B]): Maybe[B] =
    Maybe.Thunk(this.thunk.flatMap {
      case Maybe.None => Need.now(Maybe.None)
      case Maybe.Some(a) => f(a).thunk
    })
}
object Maybe {
  sealed trait Strict[+A] extends Maybe[A]
  final case class Thunk[+A](thunk: Need[Strict[A]]) extends Maybe[A]
  final case object None extends Strict[Nothing] {
    val thunk: Need[this.type] = Need.now(this)
  }
  final case class Some[A](get: A) extends Strict[A] {
    def thunk: Need[this.type] = Need.now(this)
  }

  implicit def delay[A]: Delay[Maybe[A]] = new Delay[Maybe[A]] {
    override def delay(a: Need[Maybe[A]]): Maybe[A] = Thunk(a.flatMap(_.thunk))
    override def project(a: Maybe[A]): Need[Maybe[A]] = a.thunk
  }
}

testDelayLaws[Maybe[Int]](Maybe.Some(10))

///////////////////////////////////////////////////////////////////
// STEP 4 - Typeclasses
///////////////////////////////////////////////////////////////////

trait Functor[F[_]] {
  def map[A, B](fa: F[A])(f: A => B): F[B]
}
trait Applicative[F[_]] extends Functor[F] {
  def pure[A](a: A): F[A]
  def zip[A, B](fa: F[A], fb: F[B]): F[(A, B)]
}
trait Monad[F[_]] extends Applicative[F] {
  def flatMap[A, B](fa: F[A])(f: A => F[B]): F[B]
}
trait MonadFix[F[_]] extends Monad[F] {
  def fix[A](f: Need[A] => F[A]): F[A]
}
trait Traverse[F[_]] extends Functor[F] {
  def traverse[G[_], A, B](fa: F[A])(f: A => G[B])(implicit G: Applicative[G]): Need[G[F[B]]] 
}

new MonadFix[Maybe] with Traverse[Maybe] {
  type F[A] = Maybe[A]
  
  def pure[A](a: A): F[A] = Maybe.Some(a)
  def zip[A, B](fa: F[A], fb: F[B]): F[(A, B)] = fa.flatMap(a => fb.map(b => (a, b)))
  def map[A, B](fa: F[A])(f: A => B): F[B] = fa.map(f)
  def flatMap[A, B](fa: F[A])(f: A => F[B]): F[B] = fa.flatMap(f)
  def fix[A](f: Need[A] => F[A]): F[A] = {
    def unJust(m: Maybe[A]): Need[A] = 
      m.thunk.map { case Maybe.Some(a) => a }
    lazy val a: Maybe[A] = f(unJust(a))
    a
  }
  def traverse[G[_], A, B](fa: F[A])(f: A => G[B])(implicit G: Applicative[G]): Need[G[F[B]]] = 
    fa.thunk.map {
      case Maybe.None => G.pure(Maybe.None : F[B])
      case Maybe.Some(a) => G.map(f(a))(Maybe.Some(_) : F[B])
    }
}