Continuation.scala

package cats
package data

import scala.annotation.tailrec

sealed abstract class Continuation[O, +I] {
  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._
    //@tailrec
    def go(e: Either[((Any => O) => O, Fn[Any, O]), O]): O = e match {
      case Right(o) => o
      case Left((cont, ScalaFn(sfn))) => cont(sfn)
      case Left((cont, fn)) =>
        // This is not stack safe, we could keep finding continuations
        cont { i: Any => go(run(Right(i), fn)) }
    }
    go(run(Left(this), ScalaFn(fn).asInstanceOf[Fn[Any, O]]))
  }
}

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)

  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]

  @tailrec
  private def run[O](c: Either[Continuation[O, Any], Any], fn: Fn[Any, O]): Either[((Any => O) => O, Fn[Any, O]), O] = c match {
  //private def run[O, I](c: Either[Continuation[O, I], I], fn: Fn[I, O]): O = c match {
    case Left(Const(i)) => run(Right(i), fn)
    case Left(Cont(runfn)) => Left((runfn, fn))
    case Left(Mapped(c, n)) => run(Left(c), AndThen(ScalaFn(n), fn))
    case Left(FlatMapped(c, next)) => run(Left(c), RunCont(next, fn))
    case r@Right(i) =>
      fn match {
        case ScalaFn(fn) => Right(fn(i))
        case r@RunCont(_, _) =>
          run(Left(r.fn(i)).asInstanceOf[Either[Continuation[O, Any], Any]], r.next.asInstanceOf[Fn[Any, O]]) // scala type inference goofs up here
        case AndThen(ScalaFn(f), next) => run(Right(f(i)), next)
        case AndThen(r@RunCont(_, _), next) =>
          run(Left(r.fn(i)).asInstanceOf[Either[Continuation[O, Any], Any]], AndThen(r.next, next).asInstanceOf[Fn[Any, O]]) // scala type inference goofs up here
        case at@AndThen(_, _) =>
          /**
           * Note this method compiles with the following signature, but not
           * as a tailrec method, which scala does not support when the types
           * are not identical on each loop.
           * def rotateRight[X, Y, Z](s: AndThen[X, Y, Z]): Fn[X, Z] = s.fn match {
           */
          @tailrec
          def rotateRight[Z](s: AndThen[Any, Any, Z]): Fn[Any, Z] = s.fn match {
          //def rotateRight[X, Y, Z](s: AndThen[X, Y, Z]): Fn[X, Z] = s.fn match {
            case f: ScalaFn[_, _] => s
            case f: RunCont[_, _, _] => s
            case AndThen(a, b) => rotateRight(AndThen(a, AndThen(b, s.next)))
          }
          run(r, rotateRight(at))
      }
  }

  sealed abstract class Fn[-A, +B]
  case class ScalaFn[A, B](fn: A => B) extends Fn[A, B]
  case class RunCont[I1, I2, O](fn: I1 => Continuation[O, I2], next: Fn[I2, O]) extends Fn[I1, O]
  case class AndThen[A, B, C](fn: Fn[A, B], next: Fn[B, C]) extends Fn[A, C]
}

// http://www.di.ens.fr/~pouzet/cours/systeme/bib/koen-concurrency-poor-man-jpf93.pdf
// poor man's concurrency monad
sealed abstract class Action[M[_]]
object Action {
  type C[M[_], T] = Continuation[Action[M], T]

  private case class Stop[M[_]]() extends Action[M]
  private case class Fork[M[_]](a: Action[M], b: Action[M]) extends Action[M]
  private case class Atom[M[_]](act: M[Action[M]]) extends Action[M]

  def action[M[_], T](c: C[M, T]): Action[M] =
    c(_ => Stop[M]())

  def atom[M[_], T](m: M[T])(implicit M: Functor[M]): C[M, T] =
    Continuation.from[T, Action[M]] { next => Atom(M.map(m)(next)) }

  def stop[M[_]]: C[M, Nothing] =
    Continuation.from[Nothing, Action[M]] { next => Stop() }

  def par[M[_], T](a: C[M, T], b: C[M, T]): C[M, T] =
    Continuation.from[T, Action[M]] { next => Fork[M](a(next), b(next)) }

  def fork[M[_], T](a: C[M, T]): C[M, Unit] =
    Continuation.from[Unit, Action[M]] { next => Fork[M](action(a), next(())) }

  def run[M[_]: Monad](action: Action[M]): M[Unit] = runAll(Vector(action))

  def runAll[M[_]](action: Vector[Action[M]])(implicit M: Monad[M]): M[Unit] = {
    def step(as: Vector[Action[M]]): M[Either[Vector[Action[M]], Unit]] =
      as match {
        case Vector() => M.pure(Right(()))
        case Stop() +: tail => M.pure(Left(tail))
        case Fork(a, b) +: tail => M.pure(Left(tail :+ a :+ b))
        case Atom(ma) +: tail => M.map(ma) { a => Left(as :+ a) }
      }
    M.tailRecM(action)(step)
  }
}