Can scala compiler do the micro optimization for simple pattern matching？

hepin1989 · January 20, 2024, 9:27pm

I resently submited a change to pekko, which changes from pattern matching to eq for performance reason, which generates less bytecode, but very annoying.

Can Scala compiler do this for us and then no one need to write it. And you can find this kind of trick in Scala stdlib too:

github.com

scala/scala/blob/2.13.x/src/library/scala/collection/IterableOnce.scala#L1178


      
           *  @example    `Seq("a", 1, 5L).collectFirst({ case x: Int => x*10 }) = Some(10)`
           */
          def collectFirst[B](pf: PartialFunction[A, B]): Option[B] = {
            // Presumably the fastest way to get in and out of a partial function is for a sentinel function to return itself
            // (Tested to be lower-overhead than runWith.  Would be better yet to not need to (formally) allocate it)
            val sentinel: scala.Function1[A, Any] = new AbstractFunction1[A, Any] {
              def apply(a: A): AbstractFunction1[A, Any] = this
            }
            val it = iterator
            while (it.hasNext) {
              val x = pf.applyOrElse(it.next(), sentinel)
              if (x.asInstanceOf[AnyRef] ne sentinel) return Some(x.asInstanceOf[B])
            }
            None
          }
          
          /** Aggregates the results of applying an operator to subsequent elements.
           *
           *  Since this method degenerates to `foldLeft` for sequential (non-parallel) collections,
           *  where the combining operation is ignored, it is advisable to prefer `foldLeft` for that case.
           *

[1] chore: Micro optimization for `collect*` operator. by He-Pin · Pull Request #983 · apache/incubator-pekko · GitHub

lrytz · January 22, 2024, 10:22am

The spec requries calling equals for case X patterns (Pattern Matching | Scala 3.4).

A static analyis could in principle identify cases where X is private and does not override equals. But the type system is not enough for that, X could have static type Object but be initialized with a class that does have an override.

So I don’t see a good way of implementing this optimization.

morgen-peschke · January 22, 2024, 4:08pm

Could this be feasible as a method on PartialFunction, or would allocating an Option erase the gains?

trait PartialFunction[-A, +B] extends (A => B) { self =>
  def applyOrNone[A1 <: A](a: A1): Option[B] = {
    val r = applyOrElse(a, PartialFunction.ApplyOrNoneSentinel)
    if (r.asInstanceOf[AnyRef] ne ApplyOrNoneSentinel) Some(r.asInstanceOf[B]) 
    else None
  }
}
object PartialFunction {
  private val ApplyOrNoneSentinel: scala.Function1[Any, Any] = (a: Any) => this
}

som-snytt · January 22, 2024, 10:20pm

The idiom to use reference eq is to match the singleton type:

scala> def f[A](xs: List[A]) = xs match { case _: Nil.type => 0 case _ => 1 }
def f[A](xs: List[A]): Int

scala> :javap #f
  public <A extends java.lang.Object> int f(scala.collection.immutable.List<A>);
    descriptor: (Lscala/collection/immutable/List;)I
    flags: (0x0001) ACC_PUBLIC
    Code:
      stack=2, locals=2, args_size=2
         0: aload_1
         1: getstatic     #22                 // Field scala/collection/immutable/Nil$.MODULE$:Lscala/collection/immutable/Nil$;
         4: if_acmpne     9
         7: iconst_0
         8: ireturn
         9: iconst_1
        10: ireturn
      StackMapTable: number_of_entries = 1
        frame_type = 9 /* same */
      LineNumberTable:
        line 1: 0
      LocalVariableTable:
        Start  Length  Slot  Name   Signature
            0      11     0  this   L$line8/$read$$iw;
            0      11     1    xs   Lscala/collection/immutable/List;
    Signature: #15                          // <A:Ljava/lang/Object;>(Lscala/collection/immutable/List<TA;>;)I
    MethodParameters:
      Name                           Flags
      xs                             final

scala>

in the linked case,

       pf.applyOrElse(grab(in), NotApplied) match {
          case _: NotApplied.type             => completeStage()

That is per spec, and not an implementation detail.

hepin1989 · January 23, 2024, 5:22am

Thanks, a follow up in ↓
[1] chore: Make use of pattern matching on singleton type instead. by He-Pin · Pull Request #1026 · apache/incubator-pekko · GitHub