Free.scala

package scalaz

import annotation.tailrec
import Free._
// See explanation in comments on function1CovariantByName
import std.function.{function1Covariant => _, function1CovariantByName, _}
import std.tuple._

// TODO report compiler bug when this appears just above FreeInstances:
//      "java.lang.Error: typeConstructor inapplicable for <none>"
object Free extends FreeInstances with FreeFunctions {

  /** Return from the computation with the given value. */
  private[scalaz] case class Return[S[_], A](a: A) extends Free[S, A]

  /** Suspend the computation with the given suspension. */
  private[scalaz] case class Suspend[S[_], A](a: S[Free[S, A]]) extends Free[S, A]

  /** Call a subroutine and continue with the given function. */
  private sealed abstract case class Gosub[S[_], B]() extends Free[S, B] {
    type C
    val a: () => Free[S, C]
    val f: C => Free[S, B]
  }

  def gosub[S[_], A, B](a0: () => Free[S, A])(f0: A => Free[S, B]): Free[S, B] =
    new Gosub[S, B] {
      type C = A
      val a = a0
      val f = f0
    }

  /** A computation that can be stepped through, suspended, and paused */
  type Trampoline[A] = Free[Function0, A]

  /** A computation that produces values of type `A`, eventually resulting in a value of type `B`. */
  type Source[A, B] = Free[({type f[x] = (A, x)})#f, B]

  /** A computation that accepts values of type `A`, eventually resulting in a value of type `B`.
    * Note the similarity to an [[scalaz.iteratee.Iteratee]].
    */
  type Sink[A, B] = Free[({type f[x] = (=> A) => x})#f, B]

  /** A free monad over the free functor generated by `S` */
  type FreeC[S[_], A] = Free[({type f[x] = Coyoneda[S, x]})#f, A]
}

/** A free operational monad for some functor `S`. Binding is done using the heap instead of the stack,
  * allowing tail-call elimination. */
sealed abstract class Free[S[_], A] {
  final def map[B](f: A => B): Free[S, B] =
    flatMap(a => Return(f(a)))

  /** Alias for `flatMap` */
  final def >>=[B](f: A => Free[S, B]): Free[S, B] = this flatMap f

  /** Binds the given continuation to the result of this computation.
    * All left-associated binds are reassociated to the right. */
  final def flatMap[B](f: A => Free[S, B]): Free[S, B] = this match {
    case a @ Gosub() => gosub(a.a)(x => gosub(() => a.f(x))(f))
    case a           => gosub(() => a)(f)
  }

  /** Catamorphism. Run the first given function if Return, otherwise, the second given function. */
  final def fold[B](r: A => B, s: S[Free[S, A]] => B)(implicit S: Functor[S]): B =
    resume.fold(s, r)

  /** Evaluates a single layer of the free monad. */
  @tailrec final def resume(implicit S: Functor[S]): (S[Free[S, A]] \/ A) = this match {
    case Return(a)  => \/-(a)
    case Suspend(t) => -\/(t)
    case x @ Gosub() => x.a() match {
      case Return(a)   => x.f(a).resume
      case Suspend(t)  => -\/(S.map(t)(_ flatMap x.f))
      case y @ Gosub() => y.a().flatMap(z => y.f(z) flatMap x.f).resume
    }
  }

  /** Changes the suspension functor by the given natural transformation. */
  final def mapSuspension[T[_]](f: S ~> T)(implicit S: Functor[S], T: Functor[T]): Free[T, A] =
    resume match {
      case -\/(s) => Suspend(f(S.map(s)(((_: Free[S, A]) mapSuspension f))))
      case \/-(r) => Return(r)
    }

  /** Modifies the first suspension with the given natural transformation. */
  final def mapFirstSuspension(f: S ~> S)(implicit S: Functor[S]): Free[S, A] = resume match {
    case -\/(s) => Suspend(f(s))
    case \/-(r) => Return(r)
  }

  /** Applies a function `f` to a value in this monad and a corresponding value in the dual comonad, annihilating both. */
  final def zapWith[G[_], B, C](bs: Cofree[G, B])(f: (A, B) => C)(implicit S: Functor[S], G: Functor[G], d: Zap[S, G]): C =
    Zap.monadComonadZap.zapWith(this, bs)(f)

  /** Applies a function in a comonad to the corresponding value in this monad, annihilating both. */
  final def zap[G[_], B](fs: Cofree[G, A => B])(implicit S: Functor[S], G: Functor[G], d: Zap[S, G]): B =
    zapWith(fs)((a, f) => f(a))

  /** Runs a single step, using a function that extracts the resumption from its suspension functor. */
  final def bounce(f: S[Free[S, A]] => Free[S, A])(implicit S: Functor[S]): Free[S, A] = resume match {
    case -\/(s) => f(s)
    case \/-(r) => Return(r)
  }

  /** Runs to completion, using a function that extracts the resumption from its suspension functor. */
  final def go(f: S[Free[S, A]] => Free[S, A])(implicit S: Functor[S]): A = {
    @tailrec def go2(t: Free[S, A]): A = t.resume match {
      case -\/(s) => go2(f(s))
      case \/-(r) => r
    }
    go2(this)
  }

  /**
   * Runs to completion, using a function that maps the resumption from `S` to a monad `M`.
   * @since 7.0.1
   */
  final def runM[M[_]](f: S[Free[S, A]] => M[Free[S, A]])(implicit S: Functor[S], M: Monad[M]): M[A] = {
    def runM2(t: Free[S, A]): M[A] = t.resume match {
      case -\/(s) => Monad[M].bind(f(s))(runM2)
      case \/-(r) => Monad[M].pure(r)
    }
    runM2(this)
  }

  /**
   * Catamorphism for `Free`.
   * Runs to completion, mapping the suspension with the given transformation at each step and
   * accumulating into the monad `M`.
   */
  final def foldMap[M[_]](f: S ~> M)(implicit S: Functor[S], M: Monad[M]): M[A] =
    this.resume match {
      case -\/(s) => Monad[M].bind(f(s))(_.foldMap(f))
      case \/-(r) => Monad[M].pure(r)
    }

  import Id._

  /**
   * Folds this free recursion to the right using the given natural transformations.
   */
  final def foldRight[G[_]](z: Id ~> G)(f: ({type λ[α] = S[G[α]]})#λ ~> G)(implicit S: Functor[S]): G[A] =
    this.resume match {
      case -\/(s) => f(S.map(s)(_.foldRight(z)(f)))
      case \/-(r) => z(r)
    }

  /** Runs to completion, allowing the resumption function to thread an arbitrary state of type `B`. */
  final def foldRun[B](b: B)(f: (B, S[Free[S, A]]) => (B, Free[S, A]))(implicit S: Functor[S]): (B, A) = {
    @tailrec def foldRun2(t: Free[S, A], z: B): (B, A) = t.resume match {
      case -\/(s) =>
        val (b1, s1) = f(z, s)
        foldRun2(s1, b1)
      case \/-(r) => (z, r)
    }
    foldRun2(this, b)
  }

  /** Runs a trampoline all the way to the end, tail-recursively. */
  def run(implicit ev: Free[S, A] =:= Trampoline[A]): A =
    ev(this).go(_())

  /** Interleave this computation with another, combining the results with the given function. */
  def zipWith[B, C](tb: Free[S, B])(f: (A, B) => C)(implicit S: Functor[S]): Free[S, C] = {
    (resume, tb.resume) match {
      case (-\/(a), -\/(b)) => Suspend(S.map(a)(x => Suspend(S.map(b)(y => x.zipWith(y)(f)))))
      case (-\/(a), \/-(b)) => Suspend(S.map(a)(x => x.zipWith(Return(b))(f)))
      case (\/-(a), -\/(b)) => Suspend(S.map(b)(y => Return(a).zipWith(y)(f)))
      case (\/-(a), \/-(b)) => Return(f(a, b))
    }
  }

  /** Runs a `Source` all the way to the end, tail-recursively, collecting the produced values. */
  def collect[B](implicit ev: Free[S, A] =:= Source[B, A]): (Vector[B], A) = {
    @tailrec def go(c: Source[B, A], v: Vector[B] = Vector()): (Vector[B], A) =
      c.resume match {
        case -\/((b, cont)) => go(cont, v :+ b)
        case \/-(r)         => (v, r)
      }
    go(ev(this))
  }

  /** Drive this `Source` with the given Sink. */
  def drive[E, B](sink: Sink[Option[E], B])(implicit ev: Free[S, A] =:= Source[E, A]): (A, B) = {
    @tailrec def go(src: Source[E, A], snk: Sink[Option[E], B]): (A, B) =
      (src.resume, snk.resume) match {
        case (-\/((e, c)), -\/(f)) => go(c, f(Some(e)))
        case (-\/((e, c)), \/-(y)) => go(c, Sink.sinkMonad[Option[E]].pure(y))
        case (\/-(x), -\/(f))      => go(Source.sourceMonad[E].pure(x), f(None))
        case (\/-(x), \/-(y))      => (x, y)
      }
    go(ev(this), sink)
  }

  /** Feed the given stream to this `Source`. */
  def feed[E](ss: Stream[E])(implicit ev: Free[S, A] =:= Sink[E, A]): A = {
    @tailrec def go(snk: Sink[E, A], rest: Stream[E]): A = (rest, snk.resume) match {
      case (x #:: xs, -\/(f)) => go(f(x), xs)
      case (Stream(), -\/(f)) => go(f(sys.error("No more values.")), Stream())
      case (_, \/-(r))        => r
    }
    go(ev(this), ss)
  }

  /** Feed the given source to this `Sink`. */
  def drain[E, B](source: Source[E, B])(implicit ev: Free[S, A] =:= Sink[E, A]): (A, B) = {
    @tailrec def go(src: Source[E, B], snk: Sink[E, A]): (A, B) = (src.resume, snk.resume) match {
      case (-\/((e, c)), -\/(f)) => go(c, f(e))
      case (-\/((e, c)), \/-(y)) => go(c, Sink.sinkMonad[E].pure(y))
      case (\/-(x), -\/(f))      => sys.error("Not enough values in source.")
      case (\/-(x), \/-(y))      => (y, x)
    }
    go(source, ev(this))
  }
}

object Trampoline extends TrampolineInstances {

  def done[A](a: A): Trampoline[A] =
    Free.Return[Function0,A](a)

  def delay[A](a: => A): Trampoline[A] =
    suspend(done(a))

  def suspend[A](a: => Trampoline[A]): Trampoline[A] =
    Free.Suspend[Function0, A](() => a)
}

sealed trait TrampolineInstances {
  implicit val trampolineInstance: Monad[Trampoline] with Comonad[Trampoline] =
    new Monad[Trampoline] with Comonad[Trampoline] {
      override def point[A](a: => A) = return_[Function0, A](a)
      def bind[A, B](ta: Trampoline[A])(f: A => Trampoline[B]) = ta flatMap f
      def copoint[A](fa: Trampoline[A]) = fa.run
      def cobind[A, B](fa: Trampoline[A])(f: Trampoline[A] => B) = return_(f(fa))
      override def cojoin[A](fa: Trampoline[A]) = Return(fa)
    }
}

object Sink extends SinkInstances

sealed trait SinkInstances {
  implicit def sinkMonad[S]: Monad[({type f[x] = Sink[S, x]})#f] =
    new Monad[({type f[x] = Sink[S, x]})#f] {
      def point[A](a: => A) =
        Suspend[({type f[x] = (=> S) => x})#f, A](s =>
         Return[({type f[x] = (=> S) => x})#f, A](a))
      def bind[A, B](s: Sink[S, A])(f: A => Sink[S, B]) = s flatMap f
    }
}

object Source extends SourceInstances

sealed trait SourceInstances {
  implicit def sourceMonad[S]: Monad[({type f[x] = Source[S, x]})#f] =
    new Monad[({type f[x] = Source[S, x]})#f] {
      override def point[A](a: => A) = Return[({type f[x] = (S, x)})#f, A](a)
      def bind[A, B](s: Source[S, A])(f: A => Source[S, B]) = s flatMap f
    }
}

sealed abstract class FreeInstances3 {
  implicit def freeFoldable[F[_]: Foldable: Functor]: Foldable[({type λ[α] = Free[F, α]})#λ] =
    new FreeFoldable[F] {
      def F = implicitly
      def F0 = implicitly
    }
}

sealed abstract class FreeInstances2 extends FreeInstances3 {
  implicit def freeFoldable1[F[_]: Foldable1: Functor]: Foldable1[({type λ[α] = Free[F, α]})#λ] =
    new FreeFoldable1[F] {
      def F = implicitly
      def F0 = implicitly
    }
}

sealed abstract class FreeInstances1 extends FreeInstances2 {
  implicit def freeTraverse[F[_]: Traverse]: Traverse[({type λ[α] = Free[F, α]})#λ] =
    new FreeTraverse[F] {
      def F = implicitly
    }
}

sealed abstract class FreeInstances0 extends FreeInstances1 {
  implicit def freeTraverse1[F[_]: Traverse1]: Traverse1[({type λ[α] = Free[F, α]})#λ] =
    new FreeTraverse1[F] {
      def F = implicitly
    }
}

// Trampoline, Sink, and Source are type aliases. We need to add their type class instances
// to Free to be part of the implicit scope.
sealed abstract class FreeInstances extends FreeInstances0 with TrampolineInstances with SinkInstances with SourceInstances {
  implicit def freeMonad[S[_]:Functor]: Monad[({type f[x] = Free[S, x]})#f] =
    new Monad[({type f[x] = Free[S, x]})#f] {
      def point[A](a: => A) = Return(a)
      override def map[A, B](fa: Free[S, A])(f: A => B) = fa map f
      def bind[A, B](a: Free[S, A])(f: A => Free[S, B]) = a flatMap f
    }
}

trait FreeFunctions {
  /** Collapse a trampoline to a single step. */
  def reset[A](r: Trampoline[A]): Trampoline[A] = { val a = r.run; return_(a) }

  /** Suspend the given computation in a single step. */
  def return_[S[_], A](value: => A)(implicit S: Applicative[S]): Free[S, A] =
    Suspend[S, A](S.point(Return[S, A](value)))

  /** Return the given value in the free monad. */
  def point[S[_], A](value: => A): Free[S, A] = Return[S, A](value)

  /** Alias for `point` */
  def pure[S[_], A](value: => A): Free[S, A] = point(value)

  def suspend[S[_], A](value: => Free[S, A])(implicit S: Applicative[S]): Free[S, A] =
    Suspend[S, A](S.point(value))

  /** Suspends a value within a functor in a single step. */
  def liftF[S[_], A](value: => S[A])(implicit S: Functor[S]): Free[S, A] =
    Suspend(S.map(value)(Return[S, A]))

  /** A version of `liftF` that infers the nested type constructor. */
  def liftFU[MA](value: => MA)(implicit MA: Unapply[Functor, MA]): Free[MA.M, MA.A] =
    liftF(MA(value))(MA.TC)

  /** A free monad over a free functor of `S`. */
  def liftFC[S[_], A](s: S[A]): FreeC[S, A] =
    liftFU(Coyoneda lift s)

  /** Interpret a free monad over a free functor of `S` via natural transformation to monad `M`. */
  def runFC[S[_], M[_], A](sa: FreeC[S, A])(interp: S ~> M)(implicit M: Monad[M]): M[A] =
    sa.foldMap[M](new (({type λ[α] = Coyoneda[S, α]})#λ ~> M) {
      def apply[A](cy: Coyoneda[S, A]): M[A] =
        M.map(interp(cy.fi))(cy.k)
      })

  /** A trampoline step that doesn't do anything. */
  def pause: Trampoline[Unit] =
    return_(())

  /** A source that produces the given value. */
  def produce[A](a: A): Source[A, Unit] =
    Suspend[({type f[x] = (A, x)})#f, Unit](a -> Return[({type f[x] = (A, x)})#f, Unit](()))

  /** A sink that waits for a single value and returns it. */
  def await[A]: Sink[A, A] =
    Suspend[({type f[x] = (=> A) => x})#f, A](a => Return[({type f[x] = (=> A) => x})#f, A](a))
}

private sealed trait FreeFoldable[F[_]] extends Foldable[({type λ[α] = Free[F, α]})#λ] {
  def F: Foldable[F]
  implicit def F0: Functor[F]

  override final def foldMap[A, B: Monoid](fa: Free[F, A])(f: A => B): B =
    fa.resume match {
      case -\/(s) => F.foldMap(s)(foldMap(_)(f))
      case \/-(r) => f(r)
    }

  override final def foldLeft[A, B](fa: Free[F, A], z: B)(f: (B, A) => B): B =
    fa.resume match {
      case -\/(s) => F.foldLeft(s, z)((b, a) => foldLeft(a, b)(f))
      case \/-(r) => f(z, r)
    }

  override final def foldRight[A, B](fa: Free[F, A], z: => B)(f: (A, => B) => B): B =
    fa.resume match {
      case -\/(s) => F.foldRight(s, z)(foldRight(_, _)(f))
      case \/-(r) => f(r, z)
    }
}

private sealed trait FreeFoldable1[F[_]] extends Foldable1[({type λ[α] = Free[F, α]})#λ] {
  def F: Foldable1[F]
  implicit def F0: Functor[F]

  override final def foldMap1[A, B: Semigroup](fa: Free[F, A])(f: A => B): B =
    fa.resume match {
      case -\/(s) => F.foldMap1(s)(foldMap1(_)(f))
      case \/-(r) => f(r)
    }

  override final def foldMapRight1[A, B](fa: Free[F, A])(z: A => B)(f: (A, => B) => B): B =
    fa.resume match {
      case -\/(s) => F.foldMapRight1(s)(foldMapRight1(_)(z)(f))(foldRight(_, _)(f))
      case \/-(r) => z(r)
    }

  override final def foldMapLeft1[A, B](fa: Free[F, A])(z: A => B)(f: (B, A) => B): B =
    fa.resume match {
      case -\/(s) => F.foldMapLeft1(s)(foldMapLeft1(_)(z)(f))((b, a) => foldLeft(a, b)(f))
      case \/-(r) => z(r)
    }
}

private sealed trait FreeTraverse[F[_]] extends Traverse[({type λ[α] = Free[F, α]})#λ] with FreeFoldable[F]{
  implicit def F: Traverse[F]
  override final def F0 = F

  override final def map[A, B](fa: Free[F, A])(f: A => B) = fa map f

  override final def traverseImpl[G[_], A, B](fa: Free[F, A])(f: A => G[B])(implicit G: Applicative[G]): G[Free[F, B]] =
    fa.resume match {
      case -\/(s) => G.map(F.traverseImpl(s)(traverseImpl[G, A, B](_)(f)))(Suspend(_))
      case \/-(r) => G.map(f(r))(Return(_))
    }
}

private sealed abstract class FreeTraverse1[F[_]] extends Traverse1[({type λ[α] = Free[F, α]})#λ] with FreeTraverse[F] with FreeFoldable1[F]{
  implicit def F: Traverse1[F]

  override final def traverse1Impl[G[_], A, B](fa: Free[F, A])(f: A => G[B])(implicit G: Apply[G]): G[Free[F, B]] =
    fa.resume match {
      case -\/(s) => G.map(F.traverse1Impl(s)(traverse1Impl[G, A, B](_)(f)))(Suspend(_))
      case \/-(r) => G.map(f(r))(Return(_))
    }
}