For-Comprehensions in Scala

This post is going to talk about what for-comprehensions are and how to use them in your code.

Background - Map and FlatMap

If you aren’t familiar with the map and flatMap functions, I recommend you read about them before proceeding.

When working with a series of flatMaps in a row, the amount of nesting necessary can damage readability. For example:

val six = future1.flatMap { one =>
  future2.flatMap { two =>
    future3.map { three =>
      one + two + three
    }
  }
}

Understanding this code takes a significant amount of scanning back and forth and manually keeping track of the different layers involved. Luckily, Scala provides a better syntax for this: for-comprehensions.

Using For-Comprehensions

For-comprehensions are nothing more than “syntactic sugar” for a series of flatMaps followed by a map. The example we looked at above, rewritten to use a for-comprehension, looks like:

val six = for {
  one <- future1
  two <- future2
  three <- future3
} yield one + two + three

When this code is compiled, it will be equivalent to the example we looked at above. However, this example is much quicker to read. It gives us everything important right in a row.

For-comprehensions can be thought of as each <- representing a call to flatMap. The exception is the final <- in the for-comprehension which instead represents a call to map.

Using For-Comprehensions on Custom Types

Let’s say that I have a custom type, MyFuture, which looks like:

trait MyFuture[A] {
}

In order to be able to use MyFuture in a for-comprehension, all I need to do is implement the map and flatMap functions for it. The type signature for these functions must match the following:

trait MyFuture[A] {
  def map[B](f: A => B): MyFuture[B]
  def flatMap[B](f: A => MyFuture[B]): MyFuture[B]
}

Once I’ve implemented these functions, I’ll be able to use MyFuture in a for-comprehension.

better-monadic-for

better-monadic-for is a compiler plugin that changes the way the Scala compiler desugars for-comprehensions. Below I will talk about a few of the features it provides and what they mean.

Final Map Optimization

In cases where you have a for-comprehension such as:

for {
  one <- future1
  two <- getFuture2(one)
} yield two

The standard Scala compiler would desugar this to:

future1.flatMap(one => getFuture2(one).map(two => two))

As you can see, the desugared code contains an unnecessary map at the end of it. better-monadic-for will remove this extra map for you so it instead looks like:

future1.flatMap(one => getFuture2(one))

This is a small change, but it is a nice optimization nonetheless and it can lead to some important implications regarding stack safety^[1] when dealing with certain effect types.

Improved Desugaring

In the native Scala compiler, you can’t do the following:

for {
  (a, b) <- IO((1, 2))
} yield a + b

Because of the way Scala desugars this, putting (a, b) on the left hand side of a for-comprehension will not compile for many types. This is because the native for-comprehension desugaring relies on a function called withFilter that can't always be implemented. Adding better-monadic-for is an easy way to unlock this syntax for types that normally wouldn’t allow it. better-monadic-for is able to do this because it does not rely on withFilter and instead uses a pattern match directly inside of its desugaring.

Conclusion

For-comprehensions are great tools to use to make your Scala code more readable. They allow you to write code utilizing flatMaps without worrying about high levels of nested blocks.

Keep in mind that you can use for-comprehensions on any type that has both the map and flatMap functions implemented for it.

I recommend that you take a deeper look at better-monadic-for and consider using it in all of your Scala projects.

Additional Resources