Generators.scala

object Generators extends App {

  trait Tree

  case class Inner(left: Tree, right: Tree) extends Tree
  case class Leaf(x: Int) extends Tree

  trait Generator[T] {
    self =>

    def generate: T

    def map[S](f: T => S): Generator[S] = new Generator[S] {
      def generate: S = f(self.generate)
    }

    def flatMap[S](f: T => Generator[S]): Generator[S] = new Generator[S] {
      def generate: S = f(self.generate).generate
    }
  }

  implicit def intGen: Generator[Int] = new Generator[Int] {
    def generate: Int = new java.util.Random().nextInt()
  }

  implicit def boolGen: Generator[Boolean] = new Generator[Boolean] {
    def generate: Boolean = new java.util.Random().nextBoolean();
  }

  def treeGen(implicit intGen: Generator[Int], boolGen: Generator[Boolean]): Generator[Tree] =
    for {
      isLeaf <- boolGen
      tree <- {
        if (isLeaf) {
          for {
            x <- intGen
          } yield Leaf(x)
        } else {
          for {
            x <- treeGen(intGen, boolGen)
            y <- treeGen(intGen, boolGen)
          } yield Inner(x, y)
        }
      }
    } yield tree

  def boolArrGenerator(arr: Array[Boolean]): Generator[Boolean] = new Generator[Boolean] {
    var index = 0
    def generate: Boolean = {
      if (index < arr.length) {
        val bool = arr(index)
        index += 1
        return bool
      }
      true
    }
  }

  def onlyOnesIntGenerator: Generator[Int] = new Generator[Int] { def generate = 1 }

  val expectedTree = Inner(
    Inner(
      Leaf(1),
      Inner(
        Leaf(1),
        Leaf(1)
      )
    ),
    Leaf(1)
  )

  assert(
    expectedTree == treeGen(
      onlyOnesIntGenerator,
      boolArrGenerator(Array(false, false, true, false, true))
    ).generate
  )
}