Specialization for Scala 3

Specialization is one of the few remaining desirable features from Scala 2 that's are as yet missing in Scala 3. We could try to port the Scala 2 scheme, which would be non-trivial since the implementation is quite complex. But that scheme is problematic enough to suggest that we also look for alternatives. A possible alternative is described in this note. It is meant to complement the proposal on inline traits. That proposal also contains a more detailed critique of Scala 2 specialization.

The main problem of Scala-2 specialization is code bloat. We have to pro-actively generate up to 11 copies of functions and classes when they have a specialized type parameter, and this grows exponentially with the number of such type parameters. Miniboxing tries to reduce the number under the exponent from ~10 to 3 or 4, but it has problems dealing with arrays.

Languages like C++, Rust, Go, D, or Zig avoid the proactive generation of all possible specializations by monomorphizing the whole program. This means we only need to generate a specialized version of a function or class if it is actually used in the program. On the other hand, a global monomorphization can lead itself to code bloat and long compile times. It is also a problematic choice for binary APIs.

This note discusses a different scheme to get specialization for Scala 3, which is somewhat between Scala 2's selective specialization and full monomorphization. As in Scala 2, specialized type parameters are tagged explicitly (but not with a annotation). But as for monomorphization, specializations are only generated if a specialized type is referenced in the program. To make this work efficiently, we need a way to transport information about possible specialization types through generic code (full monomorphization does not need that since it eliminates all generic code).

We do that using a type class Specialized that is typically used as a context bound on a type parameter of some class. It indicates that we want to create specialized versions of that class where the type parameter is instantiated to the type argument. The specialized versions offer optimization opportunities compared to the generic class.

Example

As an example, consider a Vec class for vectors over a numeric type.

import scala.math.Numeric
final class Vec[T: Specialized](elems: T*)(using num: Numeric[T]):
 private val arr: Array[T] = Array(elems*)
 def length = arr.length
 def scalarProduct(other: Vec[T]): T =
 require(this.length == other.length)
 var result = num.fromInt(0)
 for i <- 0 until length do
 result = num.plus(result, num.times(arr(i), other.arr(i)))
 result
end Vec

The idea is that we want to specialize vectors on the type parameter T in order to get important efficiency gains, including the following:

• Use an array arr specialized to the actual element instead of a fully generic array that has to be accessed via reflection
• Avoid boxing for internal values like result
• Avoid boxing in the API for values like the result of scalarProduct
• Specialize on the concrete Numeric class instance for T, so that calls to num's methods have static targets and can be inlined.

Terminology and Restrictions

A class such as Vec which has some type parameter X that comes with a context parameter of type Specialized[X] is called a specialized class, and its type parameter X is called a specialized type parameter.

Restrictions:

• Specialized classes must be top-level.
• Specialized classes cannot be abstract.
• The parameters, extension clause, and body of a specialized class must be
  legal if the class is translated to an inline trait.

The types which can be substituted for a Specialized parameter are AnyRef or a primitive type Byte, Short, Char, Int, Long, Float, Double, Boolean, Unit. We call these types the specializable argument types.

Besides Specialized, there is also a SpecializedAt type class which allows to specify the specializable argument types explicitly. Given a context bound X: SpecializedAt[T_1, ..., T_n] the specializable argument types for X are T1, ..., Tn. As for Specialized, these types must be primitive types or AnyRef.

Examples:

• The specializable argument types for [X: SpecializedAt[Long | Double | AnyRef]] are Long, Double, and AnyRef.
• The specializable argument types for [X: SpecializedAt[Int]] consist of just Int.

(The precise definition of specializable argument types can be changed without changing the essence of the proposal, we keep it here for concreteness).

A Closer Look At Specialized

Specialized arguments can be inter-converted with ClassTags. For instance Specialized could be a subclass of ClassTag and an additional given could produce a Specialized value from a classtag. Specialized differs from ClassTag mainly in that it has some additional fields that help make specialization efficient. For instance, it could carry a tag field that maps every specializable type to a valid index, and maps all other types to a sentinel number.

A concrete implementation of Specialized and SpecializedAt could be as follows:

class Specialized[T](val ct: ClassTag[T]):
 import Specialized.Tag
 lazy val tag: Tag ct match
 case ClassTag.Byte => Tag.Byte
 case ClassTag.Short => Tag.Short
 case ClassTag.Char => Tag.Char
 case ClassTag.Int => Tag.Int
 case ClassTag.Long => Tag.Long
 case ClassTag.Float => Tag.Float
 case ClassTag.Double => Tag.Double
 case ClassTag.Boolean => Tag.Boolean
 case ClassTag.Unit => Tag.Unit
 case ClassTag.AnyRef => Tag.AnyRef
 case _ => Tag.Other
type SpecializedAt[S] = [T] =>> Specialized[T]
object Specialized:
 enum Tag:
 case Byte, Short, Char, Int, Long, Float, Double, Boolean, Unit, AnyRef, Other
 given fromClassTag[T](using ct: ClassTag[T]): Specialized[T] = Specialized(ct)
object ClassTag:
 ...
 given fromSpecialized[T](using sp: Specialized[T]): ClassTag[T] = sp.ct

Note that the first S parameter of SpecializedAt is a phantom type that is not used on the right hand side. In fact the type argument for S is only used to determine the specialized argument types for a type variable with the SpecializedAt annotation. When it comes to parameter passing, all SpecializedAt dictionaries are equivalent to Specialized.

Expansion of Specialized Classes

A general form of a specialized class modelled after Vec is:

 class C[X](a: A)(using s: Specialized[X], b: B[X]) extends ...
 ...

We will use this class as an example to explain how specialized classes are compiled. In general, we can have

• one or more type parameters like X,
• zero or more class arguments like a,
• one or more Specialized context parameters like s, each over a separate type parameter,
• one or more typeclass parameters like b that take specialized type parameters as arguments.

A specialized class expands to several artifacts.

• First, there is an inline trait with the parents and body as given in the class.
  For instance, here is the inline trait for class C above:

   inline trait C[X](a: A)(using s: Specialized[X], b: B[X]) extends ...
   ...

• Second, there are zero or more specialized versions of the class.The name of the specialized version is of the form C$sp$N where C is the specialized class and N is a tag of a specialized argument type. We propose for now to use the same naming schemes as for Scala 2 specialization, but this aspect could be modified.

  The definition of the specialized version depends on whether the specialized argument T is final or not. If T is final, we determine a canonical implementations of each typeclass argument over T using

   class C$sp$T[X <: T](a: A)(using s: Specialized[X], b: B[X])
   extends C[X](a)(using s, summon[B[T]]):
   assert(b.getClass == summon[B[T]].getClass)

  In this case, we assume that any specialized typeclass arguments like b will be the same whereever C[T] gets instantiated. We test that this assumption is correct at least at the toplevel by comparing the class of any specialized instantiation argument with the class of the corresponding argument with which C$sp$T was first instantiated. Fixing type class arguments in this way creates important potential for optimizations when inlining trait C at T.

  If T is not final, we expand instead to:

  class C$sp$T[X <: T](a1: A)(using s: Specialized[X], b: B[X]) extends C[X](a)(using s, b)

  If T is not final, we can still specialize the class to an element type upper-bounded by T. This could be important in some cases. For instance, if class C is specialized at AnyRef and there are generic arrays of type Array[X] in C, those arrays would be specialized to an Array[Object] representation. This is much more efficient than an unconstrained Array[X] which needs reflection for element access.

  On the other hand, we cannot fix the type class argument b since different subtypes of T might have different instances, even if we assume global coherence.

• Third, there is a generic implementation of the class which is used for all type arguments that are not specializable. Here is the one for C:

   class C$generic[X](a: A)(using s: Specialized[X], b: B[X]) extends C[X](a)(using s, b)

Creating Specialized Versions

Scala erasure maps the type C[T] to C$sp$S or C$sp$S[T] if T conforms to a specializable argument S for C. This step is important for performance. Without it, any interaction with class C would go through the inefficient generic API of C that relies on boxing.

To make this erasure work, the compiler needs to create a specialized version of C for every specializable argument type S where C[S] is referred to in a program. A compiler can scan its classpath to see whether a required specialized version already exists. If none exists, the specialized version of class C at S will be created on the fly. This typically involves reading the inline trait for C from Tasty and expanding it in the freshly created class C$sp$S.

Now consider a construction of a specializable class, say new C[T](a)(using s, b). If T is statically known to conform to a specializable type S , we can rewrite this to new C$sp$S(a) or new C$sp$[T](a)(using s, b) depending on whether S is final or not. If T is a monomorphic type that does not conform to any specializable type S, we rewrite to new C$generic[T](a)(s, b) instead.

That leaves type variables T where the actual runtime instance is unknown at compile time. In this case we must still create the appropriate value of class C$sp$T in the case where at runtime the Specialized value s indicates T is a specializable
argument for C. This is necessary, since otherwise it would not be sound to erase all occurrences of type C[T] to their specialized versions.

We do this by implementing new expressions as calls to a special factory method in the companion object of the specializable class. For instance, the new expression above would be rewritten to C.newInstance(s)[T](a, s, b), where the companion object of C contains a factory method newInstance defined like this:

object C:
 def newInstance(s: Specialized[?]): [X] => (A, Specialized[X], B[X]) => C[X] = ...

The implementation of newInstance depends on whether we assume static linking or not.

With static linking, the linker could generate code that maps every Specialized[T] entry where T is a specializable argument for C mentioned in the code to the function

 [X] => (a: A, s: Specialized[X], b: B[X]) => C$sp$T(a)(using s, b)

Specialized entries could contain consecutive tags for specializable types and a separate tag for other types, so this mapping could be a simple table lookup on the tag.

If we don't assume static linking, we can get equivalent functionality using reflection. For instance, on the JVM, we can implement newInstance by calling Class.forName with the name of the specialized version corresponding to a type and the class loader of
trait C itself. If such a class exists, we can then extract the calling function by getting a method handle on the constructor.
If the specialized version does not exist we would instead return the creation method for the generic version of C (this could happen if classes are created only from generic code, but never mentioned explicitly at the specialized type in the source). We can use hashes to make sure that any specialized version found matches the version of the inline trait it implements. We would cache these computations in a lambda factory so all calls after the first one are simple lambda invocations.

AOT Specialization

Our scheme does specialization on demand. A type C[T] gets specialized if T is a specializable argument for C and the type is mentioned in the program (either explicitly or inferred).

A typical use case is where a library defines specializable generic classes and then applications using that library would use these classes at specific argument types. Then the classes would be specialized on demand in the binary code of such applications.

We might also want to define some common specializations that are generated with the library and thus can be shared by all applications using them. To do this it is enough to mention the specialized type in the library at any point whatsoever.

As a convenience feature, we can define a method

 inline def specialize[T]: Unit = ... // macro checking that `T` is specializable

Then libraries can force specializations by calling specialized like this:

 specialize[Vec[Int]]
 specialize[Matrix[Float]]

Discussion

Upsides of this proposal relative to Scala 2 specialization:

• No exponential code explosion.
• Hence, potential to specialize many more classes.
• It's also possible to inline calls to typeclasses such as Numeric which could be very expensive if compiled as is.

Downsides:

• Some overhead on creation of specialized class instances since a classtag needs to be passed and the right lambda needs to be looked up
• The scheme needs either reflection support for Class.forName or a linking step.
• No support for specialized polymorphic methods.
• No support for specialized traits (but traits can be inlined).

The characteristics of the new scheme mean that it is ideally suited for larger classes such as

• collection classes
• vectors, matrices or other tensors,
• a whole library containing abstractions as inner classes.

The larger the specialized code body is, the more likely it is that we will not cross abstraction boundaries requiring boxing, and the fewer lookups are needed to find the right constructor. On the other hand, for Scala 2 specialization, exponential code blow up makes all these usage scenarios unrealistic. So Scala 2 specialization is best used for small classes and traits such as functions and tuples.

Scala 3 already supports specialized functions in the Scala 2 style. It should not be too hard to support specialized tuples in that style as well. For most other things, the scheme described here looks like a better option.