probability monad(from LYAHFGG) in scala

Prob.scala

package ex01

import cats.Monad

import scala.annotation.tailrec

sealed trait Coin
case object Heads extends Coin
case object Tails extends Coin

case class Rational(n: Int, d: Int) {
  def *(r: Rational) = Rational.apply(n * r.n, d * r.d)
}

object Rational {
  implicit class PairOps[A](val t: (A, A)) extends AnyVal {
    def map[R](f: A => R): (R, R) = (f(t._1), f(t._2))
    def foldMap[R, B](f: A => R, g: (R, R) => B): B = g.tupled(t map f)
  }
  def apply(n: Int, d: Int): Rational = (n, d) foldMap (_ / gcd(n, d), new Rational(_, _))
  private def gcd(a: Int, b: Int): Int = if (b == 0) a else gcd(b, a % b)
}
import Prob._
case class Prob[+A](es: List[Event[A]]) {
  def map[B](f: A => B): Prob[B] = Prob(es.map { case (a, r) => f(a) -> r })
  def flatMap[B](f: A => Prob[B]): Prob[B] = flatten(map(f))
}
object Prob {
  type Event[+A] = (A, Rational)
  def flatten[A](ppa: Prob[Prob[A]]): Prob[A] = Prob( for {
    (pa, r1) <- ppa.es
    (a,  r2) <- pa.es
  } yield a -> r1 * r2)
}
object ProbInstances {
  implicit def probMonad: Monad[Prob] = new Monad[Prob] {
    def flatMap[A, B](pa: Prob[A])(f: A => Prob[B]): Prob[B] = pa flatMap f

    def tailRecM[A, B](a: A)(f: A => Prob[Either[A, B]]): Prob[B] = {
      val buf = List.newBuilder[(B, Rational)]
      @tailrec def go(pes: List[(Prob[Either[A, B]], Rational)]): Unit = pes match {
        case (Prob(e :: es), r0) :: tail => e match {
          case (Right(b), r) => buf += (b -> r * r0) ; go(Prob(es) -> r0 :: tail)
          case (Left(a2), r) => go(f(a2) -> r :: Prob(es) -> r :: tail)
        }
        case (Prob(Nil), _) :: tail => go(tail)
        case Nil                    => ()
      }
      go(pure(f(a)).es)
      Prob(buf.result)
    }
    override def pure[A](a: A): Prob[A] = Prob(List(a -> Rational(1, 1)))
  }
}