Tips for forcing the scala compiler to use specialized methods

Question

The scala standard library has some classes/methods specialized for primitives to avoid boxing, e.g.

def at3(f: Int => Boolean): Boolean = {
  f(3)
}

compiles to byte code:

  public boolean at3(scala.Function1);
    Code:
       0: aload_1
       1: iconst_3
       2: invokeinterface #35,  2           // InterfaceMethod scala/Function1.apply$mcZI$sp:(I)Z
       7: ireturn

It's using Function1.apply$mcZI$sp(int): boolean instead of regular old Function1.apply(Object): boolean because that particular combo of types is specialized on Function1. That means no boxing.

However in my codebase, I have some "strong types" which are primitives with a trait mixed in that disappears at runtime. Here's a simplified recreation using Natural as an example:

sealed trait NaturalTag
type Natural = Int & NaturalTag
                                                                                                 
def int2NatUnsafe(i: Int): Natural =
  if (i >= 0) i.asInstanceOf[Natural]
  else throw new Exception(s"$i is negative, not a Natural")

// Example                                                                                                 
def at3Nat(f: Natural => Boolean): Boolean = {
  f(int2NatUnsafe(3)) // <--- f.apply(...)
}

and the example code decompiles to:

  public boolean at3Nat(scala.Function1);
    Code:
       0: aload_1
       1: aload_0
       2: iconst_3
       3: invokevirtual #72                 // Method int2NatUnsafe:(I)I
       6: invokestatic  #78                 // Method scala/runtime/BoxesRunTime.boxToInteger:(I)Ljava/lang/Integer;
       9: invokeinterface #82,  2           // InterfaceMethod scala/Function1.apply:(Ljava/lang/Object;)Ljava/lang/Object;
      14: invokestatic  #86                 // Method scala/runtime/BoxesRunTime.unboxToBoolean:(Ljava/lang/Object;)Z
      17: ireturn

Unfortunately the compiler uses Function1.apply(Object): Object which causes boxing of the int and unboxing of the bool. I understand why it's doing it, but I'm wondering if there's a way to trick it into using the specialized method.

The compiler won't let me directly use the specialized method:

def at3Direct(f: Natural => Boolean): Boolean = {
  f.apply$mcZI$sp(3)
}

-- [E008] Not Found Error: Demo.scala:17:6 --------------------------------------------------------------------------------------------------------------------------------------------------------------------------
17 |    f.apply$mcZI$sp(3)
   |    ^^^^^^^^^^^^^^^
   |    value apply$mcZI$sp is not a member of Demo.Natural => Boolean
1 error found

For some background, I'm using scala 3.5.0, and I'm aware that specialization is a bit vaguely defined in scala 3.

This kind of pattern was happening in a lot of really hot code paths that I'm trying to optimise in my app. I already have workarounds for it that avoid boxing related to inlining or casting the Function1, so I'm more just curious for general background information or tips.

I've played around a bit with opaque types as an alternative way to encode strong types, but they have their own issues. At least with opaque types the scala compiler understands very clearly that at runtime the opaque type will be the primitive it is hiding and it will specialize when applicable.

Thanks in advance for any tips or general background knowledge!

Tips for forcing the scala compiler to use specialized methods

Answers (1)

Related Questions