The power of monads and for comprehensions

If you have ever worked with Scala and its Future type, you may have seen that using a Future in combination with other collections in a for comprehensions isn't very pretty. If you haven't heard of futures before, check out what Daniel Westheide says about them in his guide. This blog post will show you how you can use monads in combination with for comprehensions to clean up your code. Let's get right into it. For comprehensions are a way to combine the filtration and the transformation of lists. However, when working with collections inside futures it can get quite messy. Say that we for some reason need to sum two collections that are wrapped in separate futures, but only if they meet a requirement. In the following example, we only sum the elements which are greater than 2.

for {
    a1 <- Future(Set(1, 2, 3))
    b1 <- Future(Set(1, 2, 3))
} yield {
    for {
        x <- a1
        y <- b1 if y > 2
    } yield {
        x + y
    }
}
//Result: Future[Set[4,5,6]]

The reason that we need two separate for comprehensions is that we need all types in one for comprehension to be the same type. The first line of the for comprehension sets the type, and all succeeding lines with <- needs to be the same type. Let's investigate how for comprehensions work under the hood. For comprehensions are nothing more than syntactic sugar for chains of map, flatMap, and withFilter. We can decompose the for comprehension into a nested expression of map, flatMap, and withFilter. To do this I used the integrated function from JetBrains Intellij called Desugar Scala Code...

def example() = {
    val a1 = Set(1, 2, 3)
    val b2 = Set(1, 2, 3)
    for {
        x <- a1
        y <- b2 if y > 2
    } yield {
        x + y
    }
}

When you de-sugar the for comprehension above, it becomes:

def example() = {
    val a1 = Set(1, 2, 3)
    val b2 = Set(1, 2, 3)
    a1
    .flatMap(x =>
        b2
        .withFilter(y => y > 2)
        .map(y => x + y)
    )
}

As we see, there are only three methods being used in a for comprehension, map, flatMap and withFilter. The withFilter method works as filter except that filter creates a new collection, while withFilter does not. In the for comprehension, it is aliased using the if condition. Knowing what a for comprehension consists of, let's look at how we can clean up the previous example using for comprehension in combination with a nifty concept called Monads.

Monads

Monads are a design pattern in functional programming. A Monad works as a wrapper around types, and can be wrapped around for example List, Set, and Option. Option is in fact a monad itself, but we can also wrap monads with monads. If you want a better understanding of monads, you can check out the blogpost Exploring monads in Scala Collections. To clean up our example, we need a monad for Future[Set[A]] which we can call SetF[A]. So how can we implement the SetF monads for our for comprehension? All we need is to implement the operations map, flatMap, and withFilter to make it usable in a for comprehension.

final case class SetF[A](_future: Future[Iterable[A]]) extends AnyVal {

    def future: Future[Set[A]] = _future.map(_.toSet)

    def flatMap[B](f: A => SetF[B]): SetF[B] = {
        val newFuture: Future[Set[B]] = for {
          list: Set[A]        <- future
          result: Set[Set[B]] <- Future.sequence(list.map(f(_).future))
        } yield {
          result.flatten
        }
        SetF(newFuture)
    }

    def withFilter(f: A => Boolean): SetF[A] = SetF(future.map(_.filter(f)))

    def map[B](f: A => B): SetF[B] = SetF(future.map(option => option.map(f)))
}

Let's revisit our previous example, but now using SetF.

def addAllCombinations() = {
    val a = Future(Set(1, 2, 3))
    val b = Future(Set(1, 2, 3))
    for {
        x <- SetF(a)
        y <- SetF(b) if y > 2
    } yield {
        x + y
    }
}

As you see, it's quite the improvement. There is not that much code to implement SetF, but it can contribute greatly to simplifying the code you write everyday.

To make the monad more flexible, we can create apply methods for it, so it can wrap around any Iterable you want. Additionally, we'll implement a method called lift. It makes it possible to use futures of scalar values in the same comprehension, e.g. Future[Int].

object SetF {
def lift[A](f: Future[A]): SetF[A] = SetF(f.map(Set(_)))
def apply[A](o: Iterable[A]): SetF[A] = 
    SetF(Future.successful(o match {
      case o: Set[_] => o.asInstanceOf[Set[A]]
      case o         => o.toSet
      }))
}

With our newly implemented methods we can apply our SetF to all of Future[A], Set[A], Future[Set[A]]. If you only want to multiply the Future[Set[Int]] with a Future[Double] you can use lift to get the value out from the Future. Let's look at an example with lift.

def example2() = {
    val a = Future(0.5)
    val b = Future(Set(1, 2, 3))
    for {
        x <- SetF.lift(a)
        y <- SetF(b)
    } yield {
        x * y
    }
}

Monads are incredibly versatile, and can be used in many ways. I hope this post helped you get a better grasp of how they can be used.