C-style for loops in Scala 3

C-style for loops in Scala 3 :: Sep 20, 2021

With Scala 3, we can easily implement fast C-style loops.

/** standard C-style for loop */
inline def loop[A](
  inline start: A,
  inline condition: A => Boolean,
  inline advance: A => A
)(inline loopBody: A => Any): Unit =
  var a = start
  while condition(a) do
    loopBody(a)
    a = advance(a)

The usage looks like this.

loop(0, _ < 9, _ + 1) { i =>
  println("number: " + i)
}

The generated byte-code viewed in the CFR decompiler looks perfect; a simple Java for loop.

for (int a = 0; a < 9; ++a) {
    Predef$.MODULE$.println((Object)("number: " + a));
}

Lets also define some ergonomic variants:

/** Java-style for-of loop (ex, for(x : myList) {...}) */
inline def loop[A](inline iterable: Iterable[A])(inline loopBody: A => Any): Unit =
  val it = iterable.iterator
  while it.hasNext do
    loopBody(it.next())

import java.lang.Iterable as JIterable

/** Java-style for-of loop (ex, for(x : myList) {...}) */
inline def loop[A](inline iterable: JIterable[A])(inline loopBody: A => Any): Unit =
  val it = iterable.iterator
  while it.hasNext do
    loopBody(it.next())

/** for-of loop, for Arrays */
inline def loop[A](arr: Array[A])(inline loopBody: A => Any): Unit =
  var i = 0
  while i < arr.length do
    loopBody(arr(i))
    i += 1

/** simple loop for ranges */
inline def loop(
  start: Int,
  until: Int,
  inline advance: Int => Int = _ + 1
)(inline loopBody: Int => Any): Unit =
  var i = start
  while i < until do
    loopBody(i)
    i = advance(i)

But why not use for-do or while? Consider an example: given a two dimensional array, return true if the array contains some integer x.

It’s easy enough with Scala’s for-do:

def arrContains(arr: Array[Array[Int]], x: Int): Boolean =
  for row <- arr do
    for i <- row do
      if i == x then return true
  false

However if you decompile the output you’ll see a monster:

    public boolean arrContains(int[][] arr, int x) {
        boolean bl;
        Object object = new Object();
        try {
            Object object2 = Predef$.MODULE$.refArrayOps((Object[])arr);
            ArrayOps$.MODULE$.foreach$extension(object2, (Function1)(JProcedure1 & Serializable)row -> {
                Object object = Predef$.MODULE$.intArrayOps(row);
                ArrayOps$.MODULE$.foreach$extension(object, (Function1)(JFunction1.mcVI.sp & Serializable)i -> {
                    if (i == x) {
                        throw new NonLocalReturnControl(object, (Object)BoxesRunTime.boxToBoolean((boolean)true));
                    }
                });
            });
            bl = false;
        }
        catch (NonLocalReturnControl ex) {
            if (ex.key() == object) {
                bl = BoxesRunTime.unboxToBoolean((Object)ex.value());
            }
            throw ex;
        }
        return bl;
    }

Numbers are being boxed, Early Return is implemented by throwing Exceptions, and because x is in an outer scope, the second foreach lambda is capturing, which generates an anonymous class.

The performance is good with while; the byte code looks exactly like the source. But it’s hard to read.

  def arrContains(arr: Array[Array[Int]], x: Int): Boolean =
    var row = 0
    while row < arr.length do
      val rowArr = arr(row)
      var col = 0
      while col < rowArr.size do
        if rowArr(col) == x then return true
        col += 1
      row += 1
    false

And finally, using loop():

def arrContains(arr: Array[Array[Int]], x: Int): Boolean =
  loop(arr) { rowArray =>
    loop(rowArray) { i =>
      if i == x then return true
    }
  }
  false

Nearly as ergonomic as for-do, and the byte code looks much cleaner:

    public boolean arrContains(int[][] arr, int x) {
        for (int i = 0; i < arr.length; ++i) {
            int[] rowArray$proxy1 = arr[i];
            for (int i2 = 0; i2 < rowArray$proxy1.length; ++i2) {
                int i$proxy1 = rowArray$proxy1[i2];
                if (i$proxy1 != x) continue;
                return true;
            }
        }
        return false;
    }

I wonder if for-do could be updated to behave like the loop() overloads, falling back to foreach(..) when required. I’m sure it would make some codebases a lot faster.