FreeParallel.scala

package cats.free

import cats.free.Free.{FlatMapped, Pure, Suspend}
import cats.{Applicative, Monad, Parallel, ~>}
import cats.implicits._

class FreeParallel[S[_], A](val value: Free[FreeApplicative[S, ?], A]) extends AnyVal

object FreeParallel {
  def apply[S[_], A](fp: Free[FreeApplicative[S, ? ], A]): FreeParallel[S, A] = new FreeParallel(fp)

  def parRunAp[M[_], F[_], A](ma: FreeApplicative[M, A])(implicit M: Monad[M], P: Parallel[M, F]): M[A] = {
    implicit val F: Applicative[F] = P.applicative
    P.sequential.apply(ma.compile(P.parallel(M)).fold(F))
  }


  def foldMap[M[_]: Monad, F[_], S[_], A](fp: FreeParallel[S, A])(f: S ~> M)(implicit P: Parallel[M, F]): M[A] = {
    val parInterpreter: FreeApplicative[M, ?] ~> M =
      λ[FreeApplicative[M, ?] ~> M](parRunAp(_))

    val liftedF: FreeApplicative[S, ?] ~> FreeApplicative[M, ?] =
      λ[FreeApplicative[S, ?] ~> FreeApplicative[M, ?]](_.compile(f))

    fp.value.foldMap(liftedF andThen parInterpreter)
  }

  def fold[M[_]: Monad, F[_], A](fp: FreeParallel[M, A])(implicit P: Parallel[M, F]): M[A] = {
    val parInterpreter: FreeApplicative[M, ?] ~> M =
      λ[FreeApplicative[M, ?] ~> M](parRunAp(_))

    fp.value.foldMap(parInterpreter)
  }

  def liftF[F[_], A](value: F[A]): FreeParallel[F, A] =
    FreeParallel(Free.liftF(FreeApplicative.lift(value)))

  def pure[S[_], A](a: A): FreeParallel[S, A] =
    FreeParallel(Free.pure(a))

  def compile[S[_], T[_], A](fp: FreeParallel[S, A])(f: S ~> T): FreeParallel[T, A] = {
    val liftedF: FreeApplicative[S, ?] ~> FreeApplicative[T, ?] =
      λ[FreeApplicative[S, ?] ~> FreeApplicative[T, ?]](_.compile(f))

    FreeParallel(fp.value.compile(liftedF))
  }

  implicit def catsFreeMonadForFreeParallel[S[_]]: Monad[FreeParallel[S, ?]] =
    new Monad[FreeParallel[S, ?]] {

      implicit val M: Monad[Free[FreeApplicative[S, ? ], ?]] =
        Free.catsFreeMonadForFree[FreeApplicative[S, ?]]

      override def flatMap[A, B](fa: FreeParallel[S, A])(f: (A) => FreeParallel[S, B]) =
        FreeParallel(Monad[Free[FreeApplicative[S, ?], ?]].flatMap(fa.value)(f andThen(_.value)))

      override def tailRecM[A, B](a: A)(f: (A) => FreeParallel[S, Either[A, B]]) =
        FreeParallel(Monad[Free[FreeApplicative[S, ?], ?]].tailRecM(a)(f andThen(_.value)))

      override def pure[A](a: A): FreeParallel[S, A] =
        FreeParallel(Monad[Free[FreeApplicative[S, ?], ?]].pure(a))
    }

  class ParFreeParallel[S[_], A](val value: FreeParallel[S, A])

  private def apWithInner[F[_]: Applicative, A, B](freeFF: Free[F, (A) => B])(freeFA: Free[F, A]): Free[F, B] = {

    def twoFlatMapped[C, D](fc: Free[F, C], fd: Free[F, D], fffc: C => Free[F, A => B], fffd: D => Free[F, A]): FlatMapped[F, B, (C, D)] =
      FlatMapped((fc, fd).tupled, cd => apWithInner(fffc(cd._1))(fffd(cd._2)))

    def suspendFlatMapped[C](ff: F[A => B], c: Free[F, C], f: C => Free[F, A]): FlatMapped[F, B, C] =
      FlatMapped(c, c => apWithInner(Free.liftF(ff))(f(c)))

    def flatMappedSuspend[C](fa: F[A], c: Free[F, C], ff: C => Free[F, A => B]): FlatMapped[F, B, C] =
      FlatMapped(c, c => apWithInner(ff(c))(Free.liftF(fa)))

    (freeFF, freeFA) match {
      case (Pure(f), Pure(a)) => Pure(f(a))
      case (Suspend(ff), Pure(a)) => Suspend(Applicative[F].ap(ff)(Applicative[F].pure(a)))
      case (Pure(f), Suspend(fa)) => Suspend(Applicative[F].ap(Applicative[F].pure(f))(fa))
      case (Suspend(ff), Suspend(fa)) => Suspend(Applicative[F].ap(ff)(fa))
      case (Pure(f), FlatMapped(c, fb)) => FlatMapped(c, fb.map(_.map(f)))
      case (FlatMapped(c, fff), Pure(a)) => FlatMapped(c, fff.map(fff => fff.map(f => f(a))))
      case (Suspend(ff), FlatMapped(c, ffa)) => suspendFlatMapped(ff, c, ffa)
      case (FlatMapped(c, f), Suspend(fa)) => flatMappedSuspend(fa, c, f)
      case (FlatMapped(ffc, fcffab), FlatMapped(ffd, fcfa)) => twoFlatMapped(ffc, ffd, fcffab, fcfa)
    }
  }


  implicit def catsFreeApplicativeForParFreeParallel[S[_]]: Applicative[ParFreeParallel[S, ?]] =
    new Applicative[ParFreeParallel[S, ?]] {
      def pure[A](a: A): ParFreeParallel[S, A] =
        new ParFreeParallel(Applicative[FreeParallel[S, ?]].pure(a))

      def ap[A, B](ff: ParFreeParallel[S, A => B])(fa: ParFreeParallel[S, A]): ParFreeParallel[S, B] =
        new ParFreeParallel(
          new FreeParallel(
            apWithInner(ff.value.value)(fa.value.value)
          )
        )
    }

  implicit def catsFreeParallelForFreeParallel[S[_]]: Parallel[FreeParallel[S, ?], ParFreeParallel[S, ?]] =
    new Parallel[FreeParallel[S, ?], ParFreeParallel[S, ?]] {
      def applicative: Applicative[ParFreeParallel[S, ?]] = catsFreeApplicativeForParFreeParallel

      def sequential(implicit M: Monad[FreeParallel[S, ?]]): ParFreeParallel[S, ?] ~> FreeParallel[S, ?] =
        λ[ParFreeParallel[S, ?] ~> FreeParallel[S, ?]](_.value)

      def parallel(implicit M: Monad[FreeParallel[S, ?]]): FreeParallel[S, ?] ~> ParFreeParallel[S, ?] =
        λ[FreeParallel[S, ?] ~> ParFreeParallel[S, ?]](fp => new ParFreeParallel(fp))
    }
}

object FreeParTest {

  import FreeParallel._
  import cats.implicits._

  sealed trait Alg[A]
  case class GetS(i: Int) extends Alg[String]
  case object GetI extends Alg[Int]

  type FreeAlg[A] = FreeParallel[Alg, A]

  def getS(i: Int): FreeAlg[String] =
    liftF(GetS(i))

  def getI: FreeAlg[Int] =
    liftF(GetI)

  def prog: FreeAlg[Int] = for {
    i <- getI
    xs <- List(getS(i), getS(2)).sequence
    s <- getS(3)
    i2 <- getI
  } yield (xs.toString + i2 + s + i).length


  def parProg: FreeAlg[Int] = for {
    i <- getI
    xs <- List(getS(i), getS(2)).parSequence
    s <- getS(3)
    i2 <- getI
  } yield (xs.toString + i2 + s + i).length


  def eitherCompiler: Alg ~> Either[String, ?]  =
    new (Alg ~> Either[String, ?]) {


      def apply[A](fa: Alg[A]): Either[String, A] =
        fa match {
          case GetS(i) => Left("Hello!" + i)
          case GetI => Right(42)
        }
    }

  val result = FreeParallel.foldMap(prog)(eitherCompiler)
  // scala.util.Either[String,Int] = Left(Hello!42)

  val resultPar = FreeParallel.foldMap(parProg)(eitherCompiler)
  // scala.util.Either[String,Int] = Left(Hello!42Hello!2)


}