Puzzler 8

Map Comprehension

Scala's for comprehensions provide an elegant syntax for invocations of powerful functional constructs based on map and flatMap. In Scala, for comprehensions are so widespread that understanding how they are desugared to map, flatMap, withFilter, and foreach calls is a common exercise when learning the language. This is especially useful because it may sometimes be preferable to desugar a for comprehension "by hand," if only for debugging purposes.

In this puzzler, we compare a for comprehension that also uses a pattern match with a desugared version of the same for expression obtained by transforming:

  for (i <- expr) yield fun(i)

to:

  expr map { i => fun(i) }

What is the result of executing the following code in the REPL?

  val xs = Seq(Seq("a", "b", "c"), Seq("d", "e", "f"),
               Seq("g", "h"), Seq("i", "j", "k"))
  
  val ys = for (Seq(x, y, z) <- xs) yield x + y + z
  
  val zs = xs map { case Seq(x, y, z) => x + y + z }

Possibilities

1. Evaluating both ys and zs throws a MatchError.
2. Both ys and zs evaluate to:

     Seq(abc, def, ijk)
   

3. Evaluating ys throws a MatchError, and zs evaluates to:

     Seq(abc, def, ijk)
   

4. ys evaluates to:

     Seq(abc, def, ijk)
   

   and evaluating zs throws a MatchError.

Explanation

Looking at the code sample in detail, you may note that Seq("g", "h") in xs doesn't actually match the Seq(x, y, z) pattern used in the for comprehension and the desugared map. So if you suspect that ys and zs don't transparently skip this value, you would assume that both would throw a MatchError. But surely the for comprehension and the desugared map call will behave in the same way?

That assumption would be wrong. The correct answer is number 4:

  scala> val ys = for (Seq(x, y, z) <- xs) yield x + y + z
  ys: Seq[String] = List(abc, def, ijk)
  
  scala> val zs = xs map { case Seq(x, y, z) => x + y + z }
  scala.MatchError: List(g, h) (of class 
      scala.collection.immutable.$colon$colon)
    at $anonfun$1.apply(<console>:8)
    at $anonfun$1.apply(<console>:8)
    at scala.collection.TraversableLike$$anonfun$map$1.
      apply(TraversableLike.scala:245)
    ...

While the for comprehension skips the problematic value, the direct map invocation fails. To explain what is going on here, start by looking at this very simple for comprehension:

  for (i <- 0 to 1) yield i + 1

This is desugared using the same scheme applied in the example:

  0 to 1 map { i => i + 1 }

You can observe this desugaring in the REPL by starting the session with the -Xprint:parser setting:[1]

  scala> for (i <- 0 to 1) yield i + 1
  [[syntax trees at end of        parser]] 
// <console>
      ...
      val res2 = 0.to(1).map(((i) => i.$plus(1)))

In Scala, the i <- 0 to 1 syntax is called a generator. With generators that perform a simple variable assignment, it's easy to forget that the generator's left-hand side is not a simple variable, it's a pattern, as demonstrated by Seq(x, y, z) <- xs in the code sample. The Scala compiler desugars generators with non-trivial patterns (i.e., patterns that constistute more than a simple variable assignment) differently. This expression:

  for (pattern <- expr) yield fun

ends up being rewritten as:

  expr map { case pattern => fun }

For example:

  scala> for (i@j <- 0 to 1) yield i + j
  [[syntax trees at end of        parser]] 
// <console>
      ...
      val res0 = 0.to(1).map(((x$1) => x$1: 
            @scala.unchecked match {
        case (i @ (j @ _)) => i.$plus(j)
      }))

So far, this is identical to our own desugaring in the example. But what the language specification also stipulates for a non-trivial pattern is the addition of a withFilter invocation.[2] Thus the following expression:

  for (pattern <- expr) yield fun

actually becomes:

  expr withFilter {
    case pattern => true
    case _ => false
  } map { case pattern => fun }

It is this withFilter invocation that transparently "strips out" the non-matching value that causes the MatchError in our attempted desugaring.

Discussion

You can treat pattern matching generators in for comprehensions as including an "if matches" guard—guards are desugared to withFilter invocations. Adding this into the desugared version produces the expected result:

  scala> xs withFilter { case Seq(x, y, z) => true; case _ =>
           false } map { case Seq(x, y, z) => x + y + z }
  res1: Seq[String] = List(abc, def, ijk)

Because for comprehensions are very common in Scala code, familiarizing yourself with how for expressions are desugared is time well spent.

---

Footnotes for Chapter 8:

[1] The compiler argument -Xprint:<phase> prints the program code after specific compiler phases. scala -Xshow-phases displays a list of the phases, the first of which is parser.

[2] Unlike filter, which creates a new collection and so incurs the overhead of an entire run through the source collection, withFilter is simply a "view." It lazily restricts the items passed on to subsequent map, flatMap, foreach, and withFilter calls and is specifically designed for efficient chaining of these operations.