Emit more efficient code for simple for-comprehensions

[danlehmann]

I've been using Scala for years to write performance-critical code. One the pain points has always been the sub-par performance or seemingly benign for-comprehensions.

While it is clear that complex for-comprehensions emit inefficient code, a very common use-case is using them as a replacement for simple for-loops.

As an example, the two loops below do the same thing. The for-comprehension is cleaner (i is scoped more narrowly, no need for var). But one look at the decompiled code shows you just how much more inefficient the for-comprehension is.

Minimized code

class LoopTest {
  final val count = 1000

  def testWhile(): Unit = {
    var i = 0
    while (i < count) {
      Console.println("Number: " + i)
      i += 1
    }
  }

  def testFor(): Unit = {
    for (i <- 0 until count) {
      Console.println("Number: " + i)
    }
  }
}

Output

I compiled using dotty-0.24.0-RC1. This is the decompiled output:

public void testWhile() {
  for (int i = 0; i < 1000; ++i) {
    Console$.MODULE$.println((Object)("Number: " + i));
  }
}
    
public void testFor() {
  RichInt$.MODULE$.until$extension(
    Predef$.MODULE$.intWrapper(0), 1000).foreach(
      (Function1)(i -> Console$.MODULE$.println((Object)("Number: " + i))));
}

Expectation

Is there some way to allow using for-comprehensions with a guarantee that performant code is being produced? Maybe it would be possible to introduce something like @scala.annotation.for (analogous to @scala.annotation.switch)?

For me, such an improvement would have a big impact: In my project, I spent hours translating for-comprehensions into simple while loops to reap (measurable) performance benefits. But it is unfortunate that in the process I had to make my code less readable and more error prone. Scala 2.13's inliner helped somewhat, but the results weren't reliable, so I mostly ended up avoiding for-comprehensions.

[nicolasstucki]

A simple way to remove the foreach is using inline

case class MyRange(start: Int, end: Int) {
  inline def foreach(inline f: Int => Unit) =
    var i = start
    val e = end
    while (i < e) {
      f(i)
      i += 1
    }
}

object App {
  val count: Int = 4
  for (i <- MyRange(0, count)) {
    Console.println("Number: " + i)
  }
}

Which generates the code

val count: Int = 4
    {
      val MyRange_this: MyRange = MyRange.apply(0, App.count)
      {
        var i: Int = MyRange_this.start
        val e: Int = MyRange_this.end
        while(i < e) {
          Console.println("Number: ".+(i))
          i = i.+(1)
        }
      }:Unit
    }

We could also try to use opaque types to avoid the creation of the object in some cases.

[smarter]

inline def foreach(inline f: Int => Unit) =

We cannot do this in practice currently because Range#foreach overrides a non-inline foreach so its parameter cannot be marked inline, so we would need to allow inline overloads which only differ in whether their parameter is inline or not (also the version that takes an inline parameter should not be called when the argument we pass to it has side effects like foreach({println("a"); x => ()), that would break semantics). Alternatively, maybe the inline typer could automatically inline pure arguments like lambdas ?

[LPTK]

also the version that takes an inline parameter should not be called when the argument we pass to it has side effects like foreach({println("a"); x => ()), that would break semantics

Ouch, that sounds like a pretty big semantic landmine... almost at the level of surprise factor as the old (deprecated) Map#mapValues. It's not uncommon to see code like xs.foreach(foo(bar)) where foo is a curried function. One wouldn't expect foo(bar) to be called again for each element of the collection (that could be terrible for performance, even if foo was pure).

Are there legitimate uses of inline parameters which do not create these semantic landmines?

[danlehmann]

We could also try to use opaque types

I gave that a shot but pretty quickly ran into compiler issues. It seems that inline and opaque types are mutually exclusive (see #6802 )

[neko-kai]

@dhoepelman What if you worked around it by defining extension methods outside of the companion object in a trait, then implement the trait as given in companion object?

object X {
  given ExtensionOps
}
trait ExtensionOps {
  inline final def map(...)...
}

[danlehmann]

With extension methods, I got pretty close.

object Ranges {
  opaque type ZeroBasedRange = Int

  object ZeroBasedRange {
    given ExtensionOps

    def apply(end: Int): ZeroBasedRange = end
  }

  trait ExtensionOps {
    inline final def (range: ZeroBasedRange) foreach(inline f: Int => Unit) = {
      var i = 0
      val e = range.asInstanceOf[Int]
      while (i < e) {
        f(i)
        i += 1
      }
    }
  }
}

object App {
  final val count: Int = 4
  for (i <- Ranges.ZeroBasedRange(count)) {
    Console.println("Number: " + i)
  }
}

This compiles into this:

    private App$() {
        MODULE$ = this;
        this.count = 4;
        final Ranges$ module$ = Ranges$.MODULE$;
        final int range = Ranges.ZeroBasedRange$.MODULE$.apply(this.count());
        for (int i = 0, e = range; i < e; ++i) {
            Console$.MODULE$.println((Object)("Number: " + i));
        }
        final BoxedUnit unit = BoxedUnit.UNIT;
    }

Apart from that one apply function call (which I'd assume hotspot would optimize out), this seems to be optimal code. Of course, it only works for ranges starting with 0.

[danlehmann]

I wanted to see if there's a way to generalize this to Ranges that don't start at 0. The only way I could think of is by packing two Ints into a Long like this:

object Ranges {
  opaque type AnyRange = Long

  object AnyRange {
    given AnyRangeExtensionOps

    def apply(start: Int, end: Int): AnyRange = start.toLong << 32 | end.toLong
  }

  trait AnyRangeExtensionOps {
    inline final def (range: AnyRange) foreach(inline f: Int => Unit) = {
      var i = (range.asInstanceOf[Long] >> 32).toInt
      var e = range.asInstanceOf[Long].toInt
      while (i < e) {
        f(i)
        i += 1
      }
    }
  }
}

object App {
  final val count: Int = 4
  for (i <- Ranges.AnyRange(0, count)) {
    Console.println("Number: " + i)
  }
}

This produces relatively efficient code:

    private App$() {
        MODULE$ = this;
        this.count = 4;
        final Ranges$ module$ = Ranges$.MODULE$;
        final long range = Ranges.AnyRange$.MODULE$.apply(0, this.count());
        for (int i = (int)(range >> 32), e = (int)range; i < e; ++i) {
            Console$.MODULE$.println((Object)("Number: " + i));
        }
        final BoxedUnit unit = BoxedUnit.UNIT;
    }

    // AnyRange$.apply:
    public long apply(final int start, final int end) {
        return (long)start << 32 | (long)end;
    }
    

While there's some bit-shifting, we at least don't pay any cost for object creation.

[smarter]

Note that using bitwise tricks like that can potentially confuse the JIT and prevent it from eliding bounds checks.

[neko-kai]

You can use inline matches to create and then pattern-match on a phantom compile-time only case class - this should allow you to omit the structure completely when it's known at compile-time and fallback to accessing the structure fields otherwise.

[danlehmann]

Note that using bitwise tricks like that can potentially confuse the JIT and prevent it from eliding bounds checks.

Any idea how to accomplish this with opaque types? Also, I'd assume that while it isn't ideal, the current Scala/Dotty solution of creating a Range instance would have the same issue?

[danlehmann]

You can use inline matches to create and then pattern-match on a phantom compile-time only case class - this should allow you to omit the structure completely when it's known at compile-time and fallback to accessing the structure fields otherwise.

Would you mind giving me a pointer to how this could be done? I looked up phantom classes, but it seems like they were replaced by unused, which was then renamed to ghost. And I can't find any current document about that now :-)

[smarter]

Any idea how to accomplish this with opaque types?

No, but I don't think that's a good solution anyway compared to getting inlining working right.

Also, I'd assume that while it isn't ideal, the current Scala/Dotty solution of creating a Range instance would have the same issue?

Depends what code the JIT manages to inline.