Scala version of Prob Monad (from LYAHFGG)

ex01_Prob.scala

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)))
  }
}