Revive or replace @specialized?

Specialization was rather messy in Scala 2. It had a complicated implementation and hard to understand restrictions and inefficiencies for the user. Nevertheless, when used right, specialization can gain significant performance improvements, which are still unavailable in Scala 3. I am opening this issue to start a discussion what one could do to achieve the essential benefits of specialization in Scala 3, while hopefully avoiding some of its shortcomings.

Pain Points of Scala 2 Specialization

Here are some point points of Scala 2 specialization as far as I can remember them. It might be slightly off or incomplete. Please fill in details that are missing.

1. Irregular class inheritance. When extending a specialized class, we really need multiple inheritance. For instance, say class A is specialized, class B extends A, then we have the following extension relationship between A, B and specialized versions A_Int, B_Int:

        B
      /   \
    A     B_Int
      \   /
      A_Int
   

   Since multiple inheritance of classes is not supported by the JVM, specialization will in this case drop the edge from A_Int to B_Int, meaning that some natural specializations are lost.
2. Duplicate class fields. If A[T] contains a field x: T, then A_Int needs to contain a field x_Int: Int. Since A_Int is a subclass of A, we end up with two sets of fields, one of them unused.
3. Possible code explosion. Since specialized variants of classes and methods have to be compiled ahead of time, we can get many variants to compile: 9 for a single specialized parameter, 81 for two parameters, and so on. This quickly gets out of hand, leading to code bloat and long compile times. As a mitigation, we can restrict the set of specialized types in the @specialized annotation. There’s also the experimental miniboxing work where instead of 9 specialization targets we only have 3 (Long, Double, and Object). The main problem with miniboxing is that without further complications is gets inefficient for defining arrays.
4. Complicated implementation, to a large degree forced by the complications described above.

Elements of a Redesign

Given the problems with 1. and 2. above I think we should avoid specialized classes. We could instead experiment with some combination of the following ideas.

Inline Traits

An inline trait is defined like a normal trait with an inline modifier.

inline trait MatrixLib[A]:
  opaque type Matrix = ...
  extension (x: Matrix)
     def * (y: Matrix): Matrix = ...
     ...

An inline trait is expanded when it is inherited by an object or class (not by another trait). Example:

object FloatMatrices extends MatrixLib[Float]   // MatrixLib inlined and specialized to `Float` here.

An inline trait itself translates to a pure interface. All its contents are inlined each time the trait is extended by a class or object. Inline traits avoid all of the pain points of Scala 2 specialization. They can do more than primitive specialization since they also specialize on value parameters and reference types. This helps avoid megamorphic dispatch. Inline traits also profit from all the optimizations available for inline methods, including inline matches, summonFrom, embedded splices. Indeed the analogy of inline traits and inline methods is strong: inline calls correspond to supercalls in extending classes and objects, inline parameters are inline parameters of the trait, inline traits can have nested inline matches, etc. Inline traits should not be too hard to implement, since they can probably draw on much of the implementation for method inlining.

Inline trait expansions are only generated on demand when a class or object extends an inline trait. This avoids the up-front cost of creating specialized copies which might never be needed. On the other hand, one needs to be more careful about de-duplication. If an inline trait is expanded with the same arguments twice, we get two expansions.

I see inline traits primarily as building blocks for large generic, specializable libraries. Typically, an application would instantiate such a library only once or a small number of times.

Compared to specialization, inline traits have one shortcoming, namely that interfaces are not specialized. If some code is parameterized by MatrixLib[T] for type parameter T that code will use unspecialized interfaces (which are implemented via bridge methods using standard erasure). Unspecialized interfaces might impose a boxing/unboxing overhead. This leads us to consider also the following extensions.

Specialized Traits

We could do Scala 2 style specialization for traits only. This avoids pain points (1) and (2) and could simplify the implementation, in particular since we could probably make use of the inlining improvements. We can also combine specialized and inline for traits. An inline trait with specialized parameters translates to the specialization of a pure interface. This means a group of interfaces with bridge methods, which is relatively lightweight. So making specialized traits inline can reduce the up-front code generation overhead.

Specialized Methods

Say we have an inline trait like MatrixLib. Consider a factory method for MatrixLib like this:

object MatrixLib:
  def apply[T](...): MatrixLib[T] = new MatrixLib[T](...) {}

Unfortunately, that would instantiate MatrixLib at generic T, without any gain in customization (and without any additional code overhead either). We can get better customization by making the apply an inline method:

object MatrixLib:
  inline def apply[T](...): MatrixLib[T] = new MatrixLib[T](...) {}

Now, we get a new, specialized copy of MatrixLib at each call of apply. Whether this is good or bad depends on how many call sites there are. But if we had specialized methods, we could also do this:

object MatrixLib:
  def apply[@specialized(Int, Float, Double) T](...): MatrixLib[T] = new MatrixLib[T](...) {}

This will produce three copies of apply for matrix libraries over Ints, Floats or Doubles.

Summary

Inline traits cover some of the use cases of Scala 2 specialization with completely different tradeoffs, and potentially more opportunities for optimization and customization. They can be complemented with specialized traits and methods which each addresses one issue of inline traits:

• specialized traits mean that efficiency is not lost at interfaces,
• specialized methods mean that we can tailor the set of expansions of inline traits to support a specific set of primitive types, without risk of duplicate expansions, by using a specialized factory method.

[odersky]

To summarize: The new design has three elements

• Inline traits, which allow to customize large parts of object constructions,
• a Specialized type class trait, which transports knowledge about primitive vs reference types,
• a Specialized(...) application, which splits a generic code block according to a specialized type parameter.

inline traits and Specialized interact in that we know we can create specialized subtraits T_X for inline traits and have erasure cast from T[X] to T_X.

[nicolasstucki]

More detailed version

inline trait Numeric[T: Specialized]:
  inline def plus(x: T, y: T): T
  def times(x: T, y: T): T

// generated
trait Numeric_Int extends Numeric[Int]/*removed*/:
  inline def plus(x: Int, y: Int): Int
  def times(x: Int, y: Int): Int

object NumericInt extends Numeric[Int]/*erases to Numeric_Int*/:
  inline def plus(x: Int, y: Int): Int = x + y
  def times(x: Int, y: Int): Int = x * y

inline trait A[T]:
  transparent inline given n: Numeric[T] = summonInline
  def f(x: T): T = n.times(n.plus(x, x), x)

class B extends A[Int]/*removed*/::
  type T = Int                            // generated
  transparent inline given n: Numeric[T] = NumericInt  // generated
  override def f(x: T): T = 
     n.times(n.plus(x, x), x)       // generated
     // then inlined to
     // NumericInt.times(x + x, x)