The_Batch_Monad.scala

/*

This code was made to investigate if it is possible
to use a batch API as if it were a single item one.

More precisely, I have this batch API:

    final case class Request(id: Int)
    final case class Response(id: Int)

    type Answers = Map[Request, Response]

    def batch(l: List[Request]): IO[Map[Request, Response]] =
      for {
        _ <- IO.log(s"Calling batch API with $l")
      } yield l.map(t => (t, Response(t.id))).toMap

But i don't want in my code to regroup request explicitly.
Instead i want to use something like:

    (l: List[Request]).traverse(sendRequest)

    def sendRequest(request: Request): IO[Response] = ???

Of course i could write `sendRequest` as

    def sendRequest(request: Request): IO[Response] =
      batch(List(request)).map(...)

But it would be silly because if my list `l` contains 100
requests, then the code `(l: List[Request]).traverse(sendRequest)`
would perform 100 calls to the batch API, which is exactly
what i want to avoid!

Instead of calling the batch API 100 times, we want to regroup
calls to `sendRequest` in order to call the `batch` function
only once.

The following code:

val requests1 = List(11,12,13,14,15,16).map(Request(_))
val requests2 = List(21,22,23,24,25,26).map(Request(_))

    val test =
      for {
        x <- requests1.traverse(Free.sendRequest).zip(
            requests2.traverse(Free.sendRequest))
        y <- x._1.traverse { case Response(n) => Free.sendRequest(Request(n + 100)) }.zip(
            x._2.traverse { case Response(n) => Free.sendRequest(Request(n + 100)) })
      } yield x._1 ++ x._2 ++ y._1 ++ y._2
    }

should only perform 2 calls to the batch API!!
Run this file to check ;)
*/
  

object TheBatchMonad {

import scala.language.higherKinds

//////////////////////////
// Usual Typeclasses
//

trait Functor[F[_]] {
  def map[A,B](fa: F[A])(f: A => B): F[B]
}

implicit class FunctorOps[F[_],A](val self: F[A])(implicit F: Functor[F]) {
  def map[B](f: A => B): F[B] = F.map(self)(f)
}

trait Applicative[F[_]] extends Functor[F] {
  def pure[A](a: A): F[A]
  def ap[A,B](f: F[A => B], a: F[A]): F[B]

  def map[A,B](fa: F[A])(f: A => B): F[B] =
    ap(pure(f), fa)

  def zip[A,B](fa: F[A], fb: F[B]): F[(A,B)] =
    ap(ap(pure((a:A) => (b:B) => (a,b)), fa), fb)
}

implicit class ApplicativeOps1[F[_],A](val self: F[A])(implicit F: Applicative[F]) {
  def zip[B](fb: F[B]): F[(A,B)] = F.zip(self, fb)
}

implicit class ApplicativeOps2[F[_],A,B](val self: F[A => B])(implicit F: Applicative[F]) {
  def ap(fa: F[A]): F[B] =
    F.ap(self, fa)
}

trait Monad[F[_]] extends Applicative[F] {
  def flatMap[A,B](fa: F[A])(f: A => F[B]): F[B]
  
  def ap[A,B](ff: F[A => B], fa: F[A]): F[B] =
    flatMap(ff) { (f:A => B) =>
      flatMap(fa) { (a:A) =>
        pure(f(a))
      }
    }
}

implicit class MonadOps[F[_],A](val self: F[A])(implicit F: Monad[F]) {
  def flatMap[B](f: A => F[B]): F[B] =
    F.flatMap(self)(f)
}


implicit class ListOps[A](val self: List[A]) {
  def traverse[F[_],B](f: A => F[B])(implicit F: Applicative[F]): F[List[B]] =
    self match {
      case Nil =>
        F.pure(Nil)
      case hd :: tl =>
        val fhd: F[B]       = f(hd)
        val ftl: F[List[B]] = tl.traverse(f)
        F.pure((b:B) => (l:List[B]) => b :: l).ap(fhd).ap(ftl)
    }
}

//////////////////////////
// Naive implementation of IO that will
// be more than enough here
//

final case class IO[A](run: () => A)

implicit object IO extends Monad[IO] {
  def pure[A](a: A): IO[A] =
    IO { () => a }

  def flatMap[A,B](fa: IO[A])(f: A => IO[B]): IO[B] =
    IO { () => f(fa.run()).run() }

  def log(message: String): IO[Unit] =
    IO{ () => println(message) }
}

//////////////////////////
// The Batch API 
//

final case class Request(id: Int)
final case class Response(id: Int)

type Answers = Map[Request, Response]

def batch(l: List[Request]): IO[Answers] =
  for {
    _ <- IO.log(s"Calling batch API with $l")
  } yield l.map(t => (t, Response(t.id))).toMap


//////////////////////////
// The Answer to the problem!!
// Don't read beyond this point if
// you want to solve it yourself.
// Or go beyond it if you want to
// read one solution of the problem.
//


/** The free monad, adapted to fit our needs */
sealed abstract class Free[A](val requests: Set[Request]) {
  def run: IO[A] = Free.run(Map.empty)(this)
}
final case class Pure[A](value: A) extends Free[A](Set.empty)
final case class Lift[A](io: IO[A]) extends Free[A](Set.empty)
final case class Ap[A,B](fun: Free[A => B], arg: Free[A]) extends Free[B](fun.requests ++ arg.requests)
final case class Flat[A](value: Free[Free[A]]) extends Free[A](value.requests)
final case class SendRequest(request: Request) extends Free[Response](Set(request))


/** Now we can write our single request sending abstraction.
  * As easy as it can be ;)
  */
def sendRequest(request: Request): Free[Response] =
  SendRequest(request)

/** Obviously, the free monad is a monad */
implicit object Free extends Monad[Free] {
  def pure[A](a: A): Free[A] = Pure(a)
    
  override def ap[A,B](fun: Free[A => B], arg: Free[A]): Free[B] =
    (fun, arg) match {
      case (Pure(f), Pure(a)) => Pure(f(a))
      case (_,_) =>
        Ap(fun, arg)
    }

  def flat[A](value: Free[Free[A]]): Free[A] =
    value match {
      case Pure(m) => m
      case _ => Flat(value)
    }
  
  def flatMap[A,B](fa: Free[A])(f: A => Free[B]): Free[B] =
    fa match {
      case Pure(a) => f(a)
      case _ => flat(map(fa)(f))
    }

  def replace[A](answers: Answers)(fa: Free[A]): Free[A] =
    if (fa.requests.intersect(answers.keySet).isEmpty)
      fa
    else
      replace(answers) {
        fa match {
          case SendRequest(t)  =>
            answers.get(t) match {
              case Some(a) => Pure(a)
              case _       => fa
            }
          case Ap(fun, arg) => Free.ap(replace(answers)(fun), replace(answers)(arg))
          case Flat(ffa)    => Free.flat(replace(answers)(ffa))
          case _            => fa
        }
      }

  /** Unsurprisingly, this is only useful if we can run it! */
  def run[A](answers: Answers)(fa: Free[A]): IO[A] =
    if (fa.requests.nonEmpty) {       
      val normal  = replace(answers)(fa)
      val newReqs = normal.requests -- answers.keySet

      if (newReqs.isEmpty)
        run(answers)(normal)
      else
        batch(newReqs.toList).flatMap { (newAnswers: Answers) =>
          val totalAnswers = answers ++ newAnswers
          run(totalAnswers)(replace(totalAnswers)(normal))
        }
    } else
      // SO fa.tarifs IS EMPTY
      fa match {
        case Pure(a)        => IO.pure(a)
        case Lift(io)       => io
        case SendRequest(t) => throw new Exception("Impossible because fa.requests is empty")
        case Ap(fun, arg)   =>
          val ioFun = run(answers)(fun)
          val ioArg = run(answers)(arg)
          IO.ap(ioFun, ioArg)
        case Flat(ffa) =>
          run(answers)(ffa).flatMap { (fa: Free[A]) =>
            run(answers)(fa)
          }      
      }
}
}
import TheBatchMonad._

val requests1 = List(11,12,13,14,15,16).map(Request(_))
val requests2 = List(21,22,23,24,25,26).map(Request(_))

val test: Free[List[Response]] =
  for {
    x <- requests1.traverse(sendRequest).zip(
         requests2.traverse(sendRequest))
    y <- x._1.traverse { case Response(n) => sendRequest(Request(n + 100)) }.zip(
         x._2.traverse { case Response(n) => sendRequest(Request(n + 100)) })
  } yield x._1 ++ x._2 ++ y._1 ++ y._2

val testIO: IO[List[Response]] =
  test.run

println(testIO.run())