Add a Continuation data type

core/src/main/scala/cats/data/Continuation.scala

@@ -0,0 +1,108 @@
+package cats
+package data
+
+import java.io.Serializable
+import scala.annotation.tailrec
+
+sealed abstract class Continuation[O, +I] extends Serializable {
+  def map[I2](fn: I => I2): Continuation[O, I2] =
+    Continuation.Mapped(this, fn)
+
+  def flatMap[I2](fn: I => Continuation[O, I2]): Continuation[O, I2] =
+    Continuation.FlatMapped(this, fn)
+
+  final def apply(fn: I => O): O = {
+    import Continuation._
+    val re = new RunEnv[O]
+    re.eval(this, re.ScalaFn(fn)).get
+  }
+}
+
+object Continuation {
+  def pure[O] = new PureBuilder[O]
+  class PureBuilder[O] {
+    def apply[I](i: I): Continuation[O, I] = Const(i)
+  }
+  def from[I, O](fn: (I => O) => O): Continuation[O, I] = Cont(fn)
+  def from2[I1, I2, O](call: (Function2[I1, I2, O]) => O): Continuation[O, (I1, I2)] = from { fn: (((I1, I2)) => O) =>
+    call { (i1: I1, i2: I2) => fn((i1, i2)) }
+  }
+
+  implicit def catsDataContinuationMonad[O]: Monad[Continuation[O, ?]] = new Monad[Continuation[O, ?]] {
+    def pure[I](i: I): Continuation[O, I] = Const(i)
+    override def map[I1, I2](f: Continuation[O, I1])(fn: I1 => I2): Continuation[O, I2] = f.map(fn)
+    def flatMap[I1, I2](f: Continuation[O, I1])(fn1: I1 => Continuation[O, I2]): Continuation[O, I2] = f.flatMap(fn1)
+    /**
+     * flatMap is trampolined, BUT the stack depth grows proportional to the number of `from` calls
+     */
+    def tailRecM[A, B](a: A)(fn: A => Continuation[O, Either[A, B]]): Continuation[O, B] =
+      fn(a).flatMap {
+        case Right(b) => Const(b)
+        case Left(a) => tailRecM(a)(fn)
+      }
+  }
+
+  private case class Const[O, I](i: I) extends Continuation[O, I]
+
+  private case class Mapped[O, I1, I2](
+    c: Continuation[O, I1],
+    fn: I1 => I2) extends Continuation[O, I2]
+
+  private case class FlatMapped[O, I1, I2](
+    c: Continuation[O, I1],
+    fn: I1 => Continuation[O, I2]) extends Continuation[O, I2]
+
+  private case class Cont[O, I](runfn: (I => O) => O) extends Continuation[O, I]
+
+  /**
+   * We hide all our implementation in this class. Note there is only
+   * ever one output, so we put the type here.
+   */
+  private class RunEnv[O] {
+    sealed abstract class Result {
+      def get: O
+    }
+    case class Complete(get: O) extends Result
+    case class ContResult[I](runfn: (I => O) => O, fn: Fn[I]) extends Result {
+      def get: O = fn match {
+        case ScalaFn(sfn) => runfn(sfn)
+        case _ =>
+          // This is not stack safe, we could keep finding continuations
+          // the stack depth will scale as the number of inner continuations,
+          // not flatMaps (which you can see
+          runfn { i: I => eval(Const(i), fn).get }
+      }
+    }
+
+    /**
+     * This evaluates a continuation given an Fn
+     */
+    final def eval[I](c: Continuation[O, I], fn: Fn[I]): Result =
+      loop(c.asInstanceOf[Continuation[O, Any]], fn.asInstanceOf[Fn[Any]])
+
+    /**
+     * This is the tail recursive loop that partially evaluates until
+     * we reach the next continuation. Note, it compiles with the type below:
+     *
+     * final def loop[I](c: Continuation[O, I], fn: Fn[I]): Result = c match {
+     *
+     * but we erase I to Any so scala tailrec optimization can work
+     */
+    @tailrec
+    private def loop(c: Continuation[O, Any], fn: Fn[Any]): Result = c match {
+      case Cont(k) => ContResult(k, fn)
+      case Mapped(c, mfn) => loop(c, AndThen(mfn, fn))
+      case FlatMapped(c, next) => loop(c, RunCont(next, fn))
+      case Const(i) => fn match {
+        case ScalaFn(sfn) => Complete(sfn(i))
+        case RunCont(mkC, next0) => loop(mkC(i), next0)
+        case AndThen(sfn, next) => loop(Const(sfn(i)), next)
+      }
+    }
+
+    sealed abstract class Fn[-A]
+    case class ScalaFn[A](fn: A => O) extends Fn[A]
+    case class RunCont[A, B](fn: A => Continuation[O, B], next: Fn[B]) extends Fn[A]
+    case class AndThen[A, B](fn: A => B, next: Fn[B]) extends Fn[A]
+  }
+}

tests/src/test/scala/cats/tests/ContinuationTests.scala

@@ -0,0 +1,50 @@
+package cats
+package tests
+
+import cats.data.Continuation
+import cats.laws.discipline._
+import org.scalacheck.{ Arbitrary, Gen }
+
+class ContinuationTests extends CatsSuite {
+
+  implicit def arbContFn[O, I](implicit I: Arbitrary[I], O: Arbitrary[O]): Arbitrary[(I => O) => O] = {
+    def call(i: I): (I => O) => O = { fn => fn(i) }
+    def const(o: O): (I => O) => O = { fn => o }
+
+    Arbitrary(Gen.oneOf(I.arbitrary.map(call(_)), O.arbitrary.map { o => const(o) }))
+  }
+  implicit def arbCont[O, I](implicit I: Arbitrary[I], fn: Arbitrary[(I => O) => O]): Arbitrary[Continuation[O, I]] =
+    Arbitrary(Gen.oneOf(I.arbitrary.map(Continuation.pure[O](_)),
+      fn.arbitrary.map(Continuation.from(_))))
+
+  def eq0[I, O: Eq](fn: I => O) = new Eq[Continuation[O, I]] {
+    def eqv(a: Continuation[O, I], b: Continuation[O, I]) =
+      Eq[O].eqv(a(fn), b(fn))
+  }
+  implicit val eqInt = eq0[Int, Int](x => x*x*x)
+  implicit val eqInt3 = eq0[(Int, Int, Int), Int] { case (x, y, z) => x*y*z }
+  implicit val iso = CartesianTests.Isomorphisms.invariant[Continuation[Int, ?]]
+
+  checkAll("Continuation[Int, ?]", MonadTests[Continuation[Int, ?]].monad[Int, Int, Int])
+  checkAll("Monad[Continuation[Int, ?]]", SerializableTests.serializable(Monad[Continuation[Int, ?]]))
+
+  test("continuations are called") {
+    val c = Continuation.from2[Int, Int, Int]((0 to 100).reduce(_))
+    assert(c { case (a, b) => a + b } == (0 to 100).reduce(_ + _))
+  }
+  test("flatMaps are sane") {
+    val exists = Continuation.from[Int, Boolean](List(1, 2, 3, 4).exists(_))
+    val forall = Continuation.from[Int, Boolean](List(1, 2, 3, 4).forall(_))
+    val c = for {
+      e <- exists
+      f <- forall
+    } yield (e, f)
+
+    // Does there exist a number that is <= all the numbers? yes.
+    assert(c { case (x, y) => (x <= y) } == true)
+    // Does there exist a number that is >= all the numbers? yes.
+    assert(c { case (x, y) => (x >= y) } == true)
+    // Does there exist a number that is > all the numbers? no.
+    assert(c { case (x, y) => (x == y) } == false)
+  }
+}