ScalaCheck Arbitrary instance for a Directed Acyclic Graph, using memoized generators to share previously generated values

gistfile1.scala

  /**
   * Arbitrary instance for a CompNode: Load, ParallelDo, GroupByKey, Combine, Op or Return
   */
  import scalaz.Scalaz._
  import Gen._
  
  /**
   * we're using a sized generator where the size of the generated data will be the depth of the graph
   * This allows to control the generation phase for ScalaCheck properties with minSize and maxSize
   */   
  implicit lazy val arbitraryCompNode: Arbitrary[CompNode] = Arbitrary(sized(depth => genDComp(depth)))
  
  def genDComp(depth: Int = 1): Gen[CompNode] = lzy(Gen.frequency((3, genLoad(depth)),
                                                                  (4, genParallelDo(depth)),
                                                                  (4, genGroupByKey(depth)),
                                                                  (3, genMaterialize(depth)),
                                                                  (3, genCombine(depth)),
                                                                  (5, genFlatten(depth)),
                                                                  (2, genOp(depth)),
                                                                  (2, genReturn(depth))))


  def genLoad       (depth: Int = 1) = oneOf(load, load)
  def genReturn     (depth: Int = 1) = oneOf(rt, rt)
  def genParallelDo (depth: Int = 1) = if (depth <= 1) value(parallelDo(load)) else memo(genDComp(depth - 1) map (parallelDo _))
  def genFlatten    (depth: Int = 1) = if (depth <= 1) value(flatten(load)   ) else memo(choose(1, 3).flatMap(n => listOfN(n, genDComp(depth - 1))).map(l => flatten(l:_*)))
  def genCombine    (depth: Int = 1) = if (depth <= 1) value(cb(load)        ) else memo(genDComp(depth - 1) map (cb _))
  def genOp         (depth: Int = 1) = if (depth <= 1) value(op(load, load)  ) else memo((genDComp(depth - 1) |@| genDComp(depth - 1))((op _)))
  def genMaterialize(depth: Int = 1) = if (depth <= 1) value(mt(load)        ) else memo(genDComp(depth - 1) map (mt _))
  def genGroupByKey (depth: Int = 1) = if (depth <= 1) value(gbk(load)       ) else memo(genDComp(depth - 1) map (gbk _))

  /**
   * allow the use of the applicative syntax to build generators:
   *
   * (gen1 |@| gen2 |@| gen3)(function(_,_,_))
   *
   * It's especially useful when building generators for case classes
   *
   * import scalaz.Scalaz._
   * case class MyClass(a: A, b: B, c: C)
   *
   * val MyCaseClassGenerator: Gen[MyCaseClass] = (genA |@| genB |@| genC)(MyCaseClass)
   *
   */
  implicit def genIsApply: Apply[Gen] = new Apply[Gen] {
    def ap[A, B](fa: => Gen[A])(f: => Gen[(A) => B]) = fa.map2(f)((v, function) => function(v))
    def map[A, B](fa: Gen[A])(f: (A) => B) = fa map f
  }
   
  /**
   * @return a new generator memoizing the previous values and returning old ones x % of the time
   *         this allows to share existing generated data between other generators
   */
  def memo[T](g: Gen[T], ratio: Int = 30): Gen[T] = {
    lazy val previousValues = new scala.collection.mutable.HashSet[T]
    def memoizeValue(v: T) = { previousValues.add(v); v }
    for {
      v           <- g
      usePrevious <- choose(0, 100)
      n           <- choose(0, previousValues.size)
    } yield {
      if (usePrevious <= ratio && previousValues.nonEmpty)
        if (n == previousValues.size) previousValues.toList(n - 1) else previousValues.toList(n)
      else memoizeValue(v)
    }
  }