Type independent gens for TypedPipe/Execution

Author: dieu

scalding-core/src/test/scala/com/twitter/scalding/ExecutionOptimizationRulesTest.scala

@@ -8,6 +8,8 @@ import cascading.tuple.{Fields, Tuple}
 import com.stripe.dagon.{Dag, Rule}
 import com.twitter.maple.tap.MemorySourceTap
 import com.twitter.scalding.typed.TypedPipeGen
+import com.twitter.scalding.typed.gen
+import com.twitter.scalding.typed.gen.{ExecutionGen, TypeGen, TypeWith}
 import java.io.{InputStream, OutputStream}
 import java.util.UUID
 import org.scalacheck.{Arbitrary, Gen}
@@ -98,9 +100,19 @@ class ExecutionOptimizationRulesTest extends FunSuite with PropertyChecks {
   def write(pipe: Gen[TypedPipe[Int]]): Gen[Execution[TypedPipe[Int]]] =
     pipe.map(_.writeThrough(new MemorySource[Int]()))
 
+  def write(pipe: Gen[TypedPipe[TypeWith[TypeGen]#Type]], withT: TypeWith[TypeGen]): Gen[Execution[TypedPipe[TypeWith[TypeGen]#Type]]] = {
+    implicit val tc = withT.evidence.tupleConverter
+    pipe.map(_.writeThrough(new MemorySource()))
+  }
+
   val mappedOrFlatMapped =
     Gen.oneOf(mapped(pipe), flatMapped(pipe))
 
+  val stdZippedWrites =
+    Gen.zip(gen.TypedPipeGen.pipeOfStd, gen.TypedPipeGen.pipeOfStd).flatMap { case ((tl, l), (tr, r)) =>
+      zipped(write(l, tl), write(r, tr))
+    }
+
   val zippedWrites =
     zipped(write(TypedPipeGen.genWithIterableSources), write(TypedPipeGen.genWithIterableSources))
 
@@ -175,15 +187,15 @@ class ExecutionOptimizationRulesTest extends FunSuite with PropertyChecks {
   }
 
   test("randomly generated executions trees are invertible") {
-    forAll(genExec) { exec =>
+    forAll(ExecutionGen.executionOfStd) { exec =>
       invert(exec)
     }
   }
 
   test("optimization rules are reproducible") {
     implicit val generatorDrivenConfig: PropertyCheckConfiguration = PropertyCheckConfiguration(minSuccessful = 500)
 
-    forAll(genExec, genRule) { (exec, rule) =>
+    forAll(ExecutionGen.executionOfStd, genRule) { (exec, rule) =>
       val optimized = ExecutionOptimizationRules.apply(exec, rule)
       val optimized2 = ExecutionOptimizationRules.apply(exec, rule)
       assert(optimized == optimized2)
@@ -193,7 +205,7 @@ class ExecutionOptimizationRulesTest extends FunSuite with PropertyChecks {
   test("standard rules are reproducible") {
     implicit val generatorDrivenConfig: PropertyCheckConfiguration = PropertyCheckConfiguration(minSuccessful = 500)
 
-    forAll(genExec) { exec =>
+    forAll(ExecutionGen.executionOfStd) { exec =>
       val optimized = ExecutionOptimizationRules.stdOptimizations(exec)
       val optimized2 = ExecutionOptimizationRules.stdOptimizations(exec)
       assert(optimized == optimized2)
@@ -217,7 +229,7 @@ class ExecutionOptimizationRulesTest extends FunSuite with PropertyChecks {
   }
 
   test("zip of writes merged") {
-    forAll(zippedWrites) { e =>
+    forAll(stdZippedWrites) { e =>
       val opt = ExecutionOptimizationRules.apply(e, ZipWrite)
 
       assert(e.isInstanceOf[Execution.Zipped[_, _]])

scalding-core/src/test/scala/com/twitter/scalding/typed/gen/ExecutionGen.scala

@@ -0,0 +1,61 @@
+package com.twitter.scalding.typed.gen
+
+import com.twitter.scalding.Execution
+import com.twitter.scalding.typed.TypedPipe
+import org.scalacheck.{Cogen, Gen}
+
+object ExecutionGen {
+  import TypedPipeGen._
+
+  private[this] def cogen(t: TypeWith[TypeGen]): Cogen[TypedPipe[t.Type]] =
+    Cogen[TypedPipe[t.Type]] { pipe: TypedPipe[t.Type] =>
+      pipe.hashCode.toLong
+    }
+
+  def executionOfStd: Gen[Execution[TypedPipe[TypeWith[TypeGen]#Type]]] =
+    TypeGen.std.flatMap { t =>
+      executionOf(t)(TypeGen.std)
+    }
+
+  def executionOf(a: TypeWith[TypeGen])(implicit tg: Gen[TypeWith[TypeGen]]): Gen[Execution[TypedPipe[a.Type]]] =
+    Gen.delay(
+      Gen.frequency(
+        5 -> genFrom(a),
+        1 -> genMap(a),
+        1 -> genFlatMap(a),
+        1 -> tg.flatMap { t =>
+          Gen.oneOf(
+            genZipped(a, t).map(_.map(_._1)),
+            genZipped(t, a).map(_.map(_._2))
+          )
+        }
+      )
+    )
+
+  def genFrom(a: TypeWith[TypeGen])(implicit tg: Gen[TypeWith[TypeGen]]): Gen[Execution[TypedPipe[a.Type]]] =
+    pipeOf(a).map(Execution.from(_))
+
+  def genMap(a: TypeWith[TypeGen])(implicit tg: Gen[TypeWith[TypeGen]]): Gen[Execution[TypedPipe[a.Type]]] =
+    tg.flatMap { t =>
+      executionOf(t).flatMap { exec =>
+        Gen.function1(pipeOf(a))(cogen(t)).map { f =>
+          exec.map(f)
+        }
+      }
+    }
+
+  def genFlatMap(a: TypeWith[TypeGen])(implicit tg: Gen[TypeWith[TypeGen]]): Gen[Execution[TypedPipe[a.Type]]] =
+    tg.flatMap { t =>
+      executionOf(t).flatMap { exec =>
+        Gen.function1(executionOf(a))(cogen(t)).map { f =>
+          exec.flatMap(f)
+        }
+      }
+    }
+
+  def genZipped(l: TypeWith[TypeGen], r: TypeWith[TypeGen])(implicit tg: Gen[TypeWith[TypeGen]]): Gen[Execution[(TypedPipe[l.Type], TypedPipe[r.Type])]] =
+    for {
+      le <- executionOf(l)
+      re <- executionOf(r)
+    } yield le.zip(re)
+}

scalding-core/src/test/scala/com/twitter/scalding/typed/gen/StdGen.scala

@@ -0,0 +1,29 @@
+package com.twitter.scalding.typed.gen
+
+import org.scalacheck.Gen
+import org.scalacheck.Gen._
+
+object StdGen {
+  implicit def num[A](implicit n: Numeric[A], c: Choose[A]): Gen[A] =
+    Gen.sized(s => c.choose(n.zero, n.max(n.fromInt(s), n.zero)))
+
+  implicit val stringGen: Gen[String] = Gen.alphaStr
+
+  implicit val charGen: Gen[Char] = Gen.alphaChar
+
+  implicit val booleanGen: Gen[Boolean] = Gen.oneOf(true, false)
+
+  implicit val unitGen: Gen[Unit] = Gen.const(())
+
+  implicit val byteGen: Gen[Byte] = num[Byte]
+
+  implicit val shortGen: Gen[Short] = num[Short]
+
+  implicit val intGen: Gen[Int] = num[Int]
+
+  implicit val longGen: Gen[Long] = num[Long]
+
+  implicit val floatGen: Gen[Float] = num[Float]
+
+  implicit val doubleGen: Gen[Double] = num[Double]
+}

scalding-core/src/test/scala/com/twitter/scalding/typed/gen/StdSemigroup.scala

@@ -0,0 +1,17 @@
+package com.twitter.scalding.typed.gen
+
+import com.twitter.algebird.Semigroup
+
+object StdSemigroup {
+  implicit val byteGroup: Semigroup[Byte] = Semigroup.from { case (l, r) =>
+    implicitly[Numeric[Byte]].plus(l, r)
+  }
+
+  implicit val charGroup: Semigroup[Char] = Semigroup.from { case (l, r) =>
+    implicitly[Numeric[Char]].plus(l, r)
+  }
+
+  implicit val stringGroup: Semigroup[String] = Semigroup.from { case (l, r) =>
+    l + r
+  }
+}

scalding-core/src/test/scala/com/twitter/scalding/typed/gen/TypeGen.scala

@@ -0,0 +1,59 @@
+package com.twitter.scalding.typed.gen
+
+
+import com.twitter.algebird.Semigroup
+import com.twitter.scalding.TupleConverter
+import org.scalacheck.{Cogen, Gen}
+
+trait TypeGen[A] {
+  val gen: Gen[A]
+  val cogen: Cogen[A]
+  val ordering: Ordering[A]
+  val semigroup: Semigroup[A]
+  val tupleConverter: TupleConverter[A]
+}
+
+object TypeGen {
+  import StdGen._
+  import StdSemigroup._
+
+  def apply[A](g: Gen[A], c: Cogen[A], o: Ordering[A], s: Semigroup[A], t: TupleConverter[A]): TypeGen[A] =
+    new TypeGen[A] {
+      val gen: Gen[A] = g
+      val cogen: Cogen[A] = c
+      val ordering: Ordering[A] = o
+      val semigroup: Semigroup[A] = s
+      val tupleConverter: TupleConverter[A] = t
+    }
+
+  def apply[A, B](a: TypeGen[A], b: TypeGen[B]): TypeGen[(A, B)] =
+    new TypeGen[(A, B)] {
+      val gen: Gen[(A, B)] = Gen.zip(a.gen, b.gen)
+      val cogen: Cogen[(A, B)] = Cogen.tuple2(a.cogen, b.cogen)
+      val ordering: Ordering[(A, B)] = Ordering.Tuple2(a.ordering, b.ordering)
+      val semigroup: Semigroup[(A, B)] = Semigroup.semigroup2(a.semigroup, b.semigroup)
+      val tupleConverter: TupleConverter[(A, B)] =
+        TupleConverter.build(a.tupleConverter.arity + b.tupleConverter.arity) { te =>
+          val ta = a.tupleConverter.apply(te)
+          val tb = b.tupleConverter.apply(te)
+          (ta, tb)
+        }
+    }
+
+  implicit def typeGen[A: Gen: Cogen: Ordering: Semigroup: TupleConverter]: TypeGen[A] =
+    TypeGen(implicitly, implicitly, implicitly, implicitly, implicitly)
+
+  implicit val std: Gen[TypeWith[TypeGen]] =
+    Gen.oneOf(
+      TypeWith[Unit, TypeGen],
+      TypeWith[Boolean, TypeGen],
+      TypeWith[Byte, TypeGen],
+      TypeWith[Char, TypeGen],
+      TypeWith[Short, TypeGen],
+      TypeWith[Int, TypeGen],
+      TypeWith[Long, TypeGen],
+      TypeWith[Float, TypeGen],
+      TypeWith[Double, TypeGen],
+      TypeWith[String, TypeGen]
+    )
+}

scalding-core/src/test/scala/com/twitter/scalding/typed/gen/TypeWith.scala

@@ -0,0 +1,16 @@
+package com.twitter.scalding.typed.gen
+
+sealed abstract class TypeWith[+Ev[_]] {
+  type Type
+  def evidence: Ev[Type]
+}
+
+object TypeWith {
+  type Aux[A, Ev[_]] = TypeWith[Ev] { type Type = A }
+
+  def apply[A, Ev[_]](implicit eva: Ev[A]): Aux[A, Ev] =
+    new TypeWith[Ev] {
+      type Type = A
+      def evidence: Ev[Type] = eva
+    }
+}