FreeA.scala

package jp.segfault.jdbc_sample.util.functional

import scala.language.higherKinds

import scalaz._
import scalaz.Scalaz._

sealed trait FreeA[F[_], +A] {

  import PartialFreeA._

  def flatMap[B](f: A => FreeA[F, B]): FreeA[F, B] = FreeA.Bind(this, f)

  def map[B](f: A => B): FreeA[F, B] = flatMap(f andThen FreeA.Pure.apply)

  def batch[I](implicit unkind: Unkind[F, I]): Seq[I] =
    this match {
      case FreeA.Bind(m, _) => m.batch
      case FreeA.Pure(_) => Seq()
      case FreeA.Send(a) => Seq(unkind(a))
      case FreeA.Apply(f, a) =>
        f.batch ++ a.batch
    }

  def flush[M[+_], O](results: Seq[O])(implicit bypass: Bypass[F, M, O]): PartialFreeA[F, M, A] =
  {
    def replace[A](next: FreeA[F, A]): State[Seq[O], PartialFreeA[F, M, A]] = next match {
      case FreeA.Pure(a) => State.state(Pure(a))
      case FreeA.Send(a) =>
        State[Seq[O], O](_ match { case h +: tl => tl -> h }) map (e => Done(bypass(e)(a)))

      case FreeA.Apply(f, a) =>
        for {
          f_ <- replace(f)
          a_ <- replace(a)
        } yield Apply(f_, a_)

      case FreeA.Bind(m, f) => replace(m).map(Bind(_, f))
    }
    val (Seq(), out) = replace(this) run results
    out
  }

  def convert[M[+_]](run: F ~> M)(implicit M: Monad[M]): M[A] = {
    this match {
      case FreeA.Pure(a) => M point a
      case FreeA.Send(a) => run(a)
      case FreeA.Bind(m, f) =>
        (m convert run) >>= (a => f(a) convert run)

      case FreeA.Apply(f, a) =>
        (a convert run) <*> (f convert run)
    }
  }

}

object FreeA {

  type With[F[_]] = { type f[+a] = FreeA[F, a] }

  def apply[F[_], A](a: => A): FreeA[F, A] = Pure(a)

  case class Send[F[_], A](get: F[A]) extends FreeA[F, A]

  case class Pure[F[_], A](get: A) extends FreeA[F, A]

  case class Bind[F[_], A, B](lhs: FreeA[F, A], rhs: A => FreeA[F, B]) extends FreeA[F, B]

  case class Apply[F[_], A, B](lhs: FreeA[F, A => B], rhs: FreeA[F, A]) extends FreeA[F, B] {
    def asBind: FreeA[F, B] =
      for {
        f <- lhs
        a <- rhs
      } yield f(a)
  }

  implicit def forMonad[F[_]] = new Monad[({ type f[+a] = FreeA[F, a] })#f] {

    def bind[A, B](fa: FreeA[F, A])(f: A => FreeA[F, B]): FreeA[F, B] = fa flatMap f

    def point[A](a: => A): FreeA[F, A] = FreeA(a)

    override def ap[A, B](fa: => FreeA[F, A])(f: => FreeA[F, A => B]): FreeA[F, B]  = Apply(f, fa)
  }

  implicit class Applies1[F[_], A, Z](val lhs: FreeA[F, A => Z]) extends AnyVal {
    def apply(rhs: FreeA[F, A]): FreeA[F, Z] = Apply[F, A, Z](lhs, rhs)
  }

  implicit class Applies2[F[_], A, B, Z](val lhs: FreeA[F, (A, B) => Z]) extends AnyVal {
    def apply(a: FreeA[F, A], b: FreeA[F, B]): FreeA[F, Z] = lhs.map(_.curried)(a)(b)
  }

  implicit class Applies3[F[_], A, B, C, Z](val lhs: FreeA[F, (A, B, C) => Z]) extends AnyVal {
    def apply(a: FreeA[F, A], b: FreeA[F, B], c: FreeA[F, C]): FreeA[F, Z] = lhs.map(_.curried)(a)(b)(c)
  }

}

sealed trait PartialFreeA[F[_], M[+_], +A] {

  import PartialFreeA._

  def get(implicit M: Monad[M]): M[\/[A, FreeA[F, A]]] = {
    this match {
      case Pure(a) => M.pure(-\/(a))
      case Bind(m, f) =>
        m.get.map(
          e => \/-(e.fold(f, _ flatMap f))
        )
      case Apply(f, a) =>
        for {
          f_ <- f.get
          a_ <- a.get
        } yield (f_, a_) match {
          case (\/-(f), \/-(a)) => \/-(f(a))
          case (-\/(f), \/-(a)) => \/-(FreeA(f)(a))
          case (\/-(f), -\/(a)) => \/-(f(FreeA(a)))
          case (-\/(f), -\/(a)) => -\/(f(a))
        }
      case Done(a) => a.map(-\/(_))
    }
  }

}

object PartialFreeA {

  def apply[F[_], M[+_], A](a: => A): PartialFreeA[F, M, A] = Pure(a)

  case class Pure[F[_], M[+_], A](get: A) extends PartialFreeA[F, M, A]

  case class Done[F[_], M[+_], A](get: M[A]) extends PartialFreeA[F, M, A]

  case class Bind[F[_], M[+_], A, B](lhs: PartialFreeA[F, M, A], rhs: A => FreeA[F, B]) extends PartialFreeA[F, M, B]

  case class Apply[F[_], M[+_], A, B](
    lhs: PartialFreeA[F, M, A => B],
    rhs: PartialFreeA[F, M, A]) extends PartialFreeA[F, M, B]
}

trait Unkind[F[_], I] {
  def apply[A](fa: F[A]): I
}

trait Bypass[F[_], M[+_], O] {
  def apply(out: O): F ~> M
}