Scala: Is a for comprehension a loop? If not, is it any better than a loop?

When we write Scala code that needs to be repeated we often put it into a for comprehension. Many folks who code Scala call this is loop.

for( i <- 1 until 5){
  println(i)
}

But that is not 100% accurate. A for comprehension does not have any mutable loop variables which allows for some interesting optimizations. Scala does have true loops though. We can build one using while.

var i=0
while(i  < 5){
  print(i)
  i += 1
}

Just like Java this does use a mutable loop variable. So let’s consider what the difference between these two approaches is. The for comprehension above is turned int the following by the compiler:

Range(1, 5).foreach(i => println(i))

So when using a for comprehension, we generate a sequence of numbers and call the loop body as a function with them. This allows the JVM to move these calls onto different cores when using for comprehensions. But when using loops, such as with while above, the reading and updating of i implies synchronization on i, which is a mutex lock. Therefore the loop bodies cannot be computed independently. This coordination slows things down.

So we should always use for comprehensions right? Well, it’s not always possible. In some cases we really do have loop variables. Consider a case where we compute the sum of each i. We need to maintain a loop variable in this case for sum.

var i=0
var sum=0
while(i  < 5){
  print(i)
  i += 1
  sum += i
}
println(s"sum=$sum")

We can get around this problem using recursion. We make the loop body a function and we pass the loop variables in as arguments, Then we return the accumulator(s) from the function like this:

@tailrec
private def loop(i: Int, upperBound: Int, curSum: Int): Int ={
  print(i)
  if(i >= upperBound)
    curSum
  else
    loop(i+1, upperBound, curSum + i)
}
println(s"sum=${loop(0, 5, 0)}")

Notice how we had our cake and ate it too: We have loop variables but no mutables. This enables the compiler to perform its optimizations that are limited to immutable data structures.

The code above uses one extra trick that will help the compiler further optimize the code. If the recursion is the last statement of a function and the function is private, we can annotate it with @tailerec, which asks the compiler to flatten out the recursion such that the underlying byte code would not use recursive calls. Besides performance benefits, this will also save us from running out of stack space when using these techniques for long iterations.

Don’t use Java style while loops with var in Scala. Scala’s better than that. Use for comprehensions when you can. Use recursion when you have accumulators with @tailrec. Your code will be faster and easier to maintain.