Branching in Free

FreeBranch.scala

// Eda and I are using Free. We have a need to handle situations where the right hand side of the for comprehension produces
// an Either and Left means that we need to 'branch' to do something to handle the failure, but handling the failure
// will not result in something that allows the rest of the main program to continue - we still want it to short circuit.
// The branch is itself a Free program, so whatever mechanism is uses needs to put the result of the branch onto the
// Left so that the main program still short circuits.
//
// Question: Is this completely the wrong way to think about this?
//
// If not, I found this: https://gist.github.com/EECOLOR/c312bdf54039a42a3058
//
// So we translated it to Cats (removing the implicit views, because they're EVIL):

import cats.Monad
import cats.free.Free
import cats.syntax.either._

object FreeBranching {
  case class Branch[F[_], L, R](value: F[Either[L, R]]) {
    /*
     * This flatMap allows you to change the L type
     *
     * L1 >: L allows you to change the L type from Nothing to Something
     * Z => L1 allows you to flatMap with Branch[F, Nothing, R1] without changing
     *         the L type
     */
    def flatMap[L1 >: L, R1, Z](f: R => Branch[F, Z, R1])(implicit ev: Z => L1, F: Monad[F]): Branch[F, L1, R1] = {
      def doLeft(l: L): F[Either[L1, R1]] = F.pure(Left(l))
      def doRight(r: R): F[Either[L1, R1]] = F.map(f(r).value)(_.left.map(ev))

      Branch(F.flatMap(value)(_.fold(doLeft, doRight)))
    }

    def map[R1](f: R => R1)(implicit F: Monad[F]): Branch[F, L, R1] =
      flatMap[L, R1, L](r => Branch(F.pure(Right(f(r)))))

    def merge[A](implicit MF: Monad[F], ev1: L => A, ev2: R => A): F[A] =
      MF.map(value)(_.left.map(ev1).right.map(ev2).merge)
  }

  type Partial[F[_]] = {
    type FreeA[A] = Free[F, A]
  }

  type FreeBranch[F[_], L, R] = Branch[Partial[F]#FreeA, L, R]

  object FreeBranch {
    // Needs type annotation because Branch expects F[_] and Free is F[_, _]
    def apply[F[_], L, R](fa: Free[F, Either[L, R]]): FreeBranch[F, L, R] = Branch[Partial[F]#FreeA, L, R](fa)

    def liftFR[F[_], L, R](fr: Free[F, R]): FreeBranch[F, L, R] = FreeBranch(fr.map(_.asRight[L]))
  }

  implicit class FreeBranchLiftSyntax[F[_], R](fr: Free[F, R]) {
    def liftToFreeBranch[L](): FreeBranch[F, L, R] = FreeBranch.liftFR(fr)
    def liftEitherToFreeBranch[A, B](implicit ev: R <:< Either[A, B]): FreeBranch[F, A, B] = FreeBranch.apply(fr.map(ev(_)))
  }

  implicit class FreeBranchOptionSyntax[R, F[_]](right: Free[F, Option[R]]) {
    def ifEmpty[L](left: ⇒ Free[F, L]): FreeBranch[F, L, R] =
      FreeBranch(right.flatMap {
        case Some(r) ⇒ freeMonad[F].pure(r.asRight[L])
        case None ⇒ left.map(_.asLeft[R])
      })
  }

  implicit class FreeBranchEitherSyntax[L, R, F[_]](right: Free[F, Either[L, R]]) {
    def ifLeft[L2](left: L ⇒ Free[F, L2]): FreeBranch[F, L2, R] = {
      FreeBranch(right.flatMap {
        case Right(r) ⇒ freeMonad[F].pure(r.asRight[L2])
        case Left(l) ⇒ left(l).map(_.asLeft[R])
      })
    }
  }

  def freeMonad[F[_]] = new Monad[Partial[F]#FreeA] {
    override def pure[A](x: A): Free[F, A] = Free.pure(x)

    override def flatMap[A, B](fa: Free[F, A])(f: (A) ⇒ Free[F, B]): Free[F, B] = fa.flatMap(f)

    override def tailRecM[A, B](a: A)(f: (A) ⇒ Free[F, Either[A, B]]): Free[F, B] = f(a).flatMap {
      case Left(a1) ⇒ tailRecM(a1)(f)
      case Right(b) ⇒ pure(b)
    }
  }
}

// We also introduced the FreeBranchEitherSyntax to handle our problem. However, this doesn't actually work for us, because
// we're using a monad stack of EitherT[IO, Failure, A]:

class MonadStacks {
  type Failure = String
  type MonadStack[A] = EitherT[IO, Failure, A]

  implicit class LiftSyntax[A](value: Either[Failure, A]) {
    def lift(): MonadStack[A] = EitherT(IO(value))
  }

  sealed trait Foo[A]
  
  // The 'ms' in Bar below is just a cheap way for us to get an instance of MonadStack into the interpreter below. In production
  // code, Bar would have other parameters and call some service or something that would return an Either[Failure, Int] and
  // the interpreter would need to lift that into the MonadStack
  case class Bar(ms: MonadStack[Int]) extends Foo[Either[Failure, Int]]

  val interpreter: Foo ~> MonadStack =
    new (Foo ~> MonadStack) {
      def apply[A](fa: Foo[A]): MonadStack[A] =
        fa match {
          case Bar(ms) ⇒ ms // Doesn't compile: A is an Either[Failure, Int], but ms is MonadStack[Int], not MonadStack[Either[Failure, Int]]
        }
    }

  "ifLeft" must "keep a value on the Right" in {
    val program = for {
      i ← Free.liftF(Bar(1.asRight[Failure].lift)).ifLeft(_ ⇒ Free.pure(""))
    } yield i

    program.value.foldMap(interpreter) must be(1.asRight)
  }
}

// So, while Eda goes to figure out the right way to solve this, I thought I'd see if I can produce an implementation of 
// FreeBranching that uses EitherT instead of Either:

object FreeBranchingEitherT {
  case class Branch[F[_], L, R, M[_]](value: F[EitherT[M, L, R]]) {
    def flatMap[L1 >: L, R1, Z](f: R => Branch[F, Z, R1, M])(implicit ev: Z => L1, monadF: Monad[F], monadM: Monad[M], traverseM: Traverse[M]): Branch[F, L1, R1, M] = {
      def doLeft(l: L): F[EitherT[M, L1, R1]] = monadF.pure(EitherT { monadM.pure(Left(l)) })
      def doRight(r: R): F[EitherT[M, L1, R1]] = monadF.map(f(r).value)((x: EitherT[M, Z, R1]) ⇒ x.leftMap(ev))

      val fmfeT: F[M[F[EitherT[M, L1, R1]]]] = monadF.map(value) { _.fold(doLeft, doRight) }
      val fmeT: F[M[EitherT[M, L1, R1]]] = monadF.flatMap(fmfeT) { m ⇒ monadF.sequence(m) }
      val fme: F[M[Either[L1, R1]]] = monadF.map(fmeT) { me ⇒ monadM.flatMap(me) { _.value } }
      val feT: F[EitherT[M, L1, R1]] = monadF.map(fme) { m ⇒ EitherT[M, L1, R1](m) }

      Branch(feT)
    }

    def map[R1](f: R => R1)(implicit monadF: Monad[F]): Branch[F, L, R1, M] = Branch[F, L, R1, M](monadF.map(value)(_.map(f)))
  }
  
  ...
}