How can one provide manually specialized implementations with Scala specialization?

Question

Specialization promises to provide high-efficiency implmentations for primitive types with minimal extra boilerplate. But specialization seems to be too eager for its own g

Answer 1

I know it's quite old, but I bumped at it looking for something else and maybe it'll prove useful. I had a similar motivation, and answered it in how to check I'm inside a specialized function or class

I used a reverse lookup table - SpecializedKey is a specialized class which equals all other instances with the same specialization, so I can perform a check like this

def onlyBytes[@specialized E](arg :E) :Option[E] =
    if (specializationFor[E]==specializationFor[Byte]) Some(arg)
    else None

Of course, there's no performance benefit when working with individual primitive values, but with collections, especially iterators, it becomes useful.

final val AllButUnit = new Specializable.Group((Byte, Short, Int, Long, Char, Float, Double, Boolean, AnyRef))

def specializationFor[@specialized(AllButUnit) E] :ResolvedSpecialization[E] =
   Specializations(new SpecializedKey[E]).asInstanceOf[ResolvedSpecialization[E]]


private val Specializations = Seq(
    resolve[Byte],
    resolve[Short],
    resolve[Int],
    resolve[Long],
    resolve[Char],
    resolve[Float],
    resolve[Double],
    resolve[Boolean],
    resolve[Unit],
    resolve[AnyRef]
).map(
    spec => spec.key -> spec :(SpecializedKey[_], ResolvedSpecialization[_])
).toMap.withDefaultValue(resolve[AnyRef])

private def resolve[@specialized(AllButUnit) E :ClassTag] :ResolvedSpecialization[E] =
    new ResolvedSpecialization[E](new SpecializedKey[E], new Array[E](0))


class ResolvedSpecialization[@specialized(AllButUnit) E] private[SpecializedCompanion]
    (val array :Array[E], val elementType :Class[E], val classTag :ClassTag[E], private[SpecializedCompanion] val key :SpecializedKey[E]) {

    private[SpecializedCompanion] def this(key :SpecializedKey[E], array :Array[E]) =
    this(array, array.getClass.getComponentType.asInstanceOf[Class[E]], ClassTag(array.getClass.getComponentType.asInstanceOf[Class[E]]), key)

    override def toString = s"@specialized($elementType)"

    override def equals(that :Any) = that match {
        case r :ResolvedSpecialization[_] => r.elementType==elementType
        case _ => false
    }

    override def hashCode = elementType.hashCode
}

private class SpecializedKey[@specialized(AllButUnit) E] {
    override def equals(that :Any) = that.getClass==getClass
    override def hashCode = getClass.hashCode

    def className = getClass.getName
    override def toString = className.substring(className.indexOf("$")+1)
}

Answer 2

This is an answer from the scala internals mailing list:

With miniboxing specialization, you can use the reflection feature:

import MbReflection._
import MbReflection.SimpleType._
import MbReflection.SimpleConv._

object Test {
  def bippy[@miniboxed A, @miniboxed B](a: A, b: B): B =
    (reifiedType[A], reifiedType[B]) match {
      case (`int`, `int`) => (a.as[Int] + b.as[Int]).as[B]
      case (  _  , `int`) => (b.as[Int] + 1).as[B]
      case (`int`,   _  ) =>  b
      case (  _  ,   _  ) =>  b
    }

  def main(args: Array[String]): Unit = {
    def x = 1.0
    assert(bippy(3,4) == 7)
    assert(bippy(x,4) == 5)
    assert(bippy(3,x) == x)
    assert(bippy(x,x) == x)
  }
}

This way, you can choose the exact behavior of the bippy method based on the type arguments without defining any implicit classes.

Answer 3

This is my best attempt so far. It works but the implementation isn't pretty (even if the results are). Improvements are welcome!

There is a macro-free way to do this, both at the class and method level, and it does involve type classes--quite a lot of them! And the answer is not exactly the same for classes and methods. So bear with me.

Manually Specialized Classes

You manually specialize classes the same way that you manually provide any kind of different implementation for classes: your superclass is abstract (or is a trait), and the subclasses provide the implementation details.

abstract class Bippy[@specialized(Int) B] {
  def b: B
  def next: Bippy[B]
}

class BippyInt(initial: Int) extends Bippy[Int] {
  private var myB: Int = initial
  def b: Int = myB
  def next = { myB += 1; this }
}

class BippyObject(initial: Object) extends Bippy[Object] {
  private var myB: Object = initial
  def b: B = myB
  def next = { myB = myB.toString; this }
}

Now, if only we had a specialized method to pick out the right implementations, we'd be done:

object Bippy{
  def apply[@specialized(Int) B](initial: B) = ???  // Now what?
}

So we've converted our problem of providing custom specialized classes and methods into just needing to provide custom specialized methods.

Manually Specialized Methods

Manually specializing a method requires a way to write one implementation that can nonetheless select which implementation you want (at compile time). Type classes are great at this. Suppose we already had type classes that implemented all of our functionality, and that the compiler would select the right one. Then we could just write

def foo[@specialized(Int) A: SpecializedFooImpl](a: A): String =
  implicitly[SpecializedFooImpl[A]](a)

...or we could if implicitly was guaranteed to preserve specialization and if we only ever wanted a single type parameter. In general these things are not true, so we'll write our type class out as an implicit parameter rather than relying on the A: TC syntactic sugar.

def foo[@specialized(Int) A](a: A)(implicit impl: SpecializedFooImpl[A]): String =
  impl(a)

(Actually, that's less boilerplate anyway.)

So we've converted our problem of providing custom specialized methods into just needing to write specialized typeclasses and getting the compiler to fill in the correct ones.

Manually Specialized Type Classes

Type classes are just classes, and now we have to write specialized classes again, but there's a critical difference. The user isn't the one asking for arbitrary instances. This gives us just enough extra flexibility for it to work.

For foo, we need an Int version and a fully generic version.

trait SpecFooImpl[@specialized (Int), A] {
  def apply(param: A): String
}

final class SpecFooImplAny[A] extends SpecFooImpl[A] {
  def apply(param: A) = param.toString
}

final class SpecFooImplInt extends SpecFooImpl[Int] {
  def apply(param: Int) = "!" * math.max(0, param)
}

Now we could create implicits to supply those type classes like so

implicit def specFooAsAny[A] = new SpecFooImplAny[A]

implicit val specFooAsInt = new SpecFooImplInt

except we have a problem: if we actually try to call foo: Int, both implicits will apply. So if we just had a way to prioritize which type class we chose, we'd be all set.

Selection of type classes (and implicits in general)

One of the secret ingredients the compiler uses to determine the right implicit to use is inheritance. If implicits come from A via B extends A, but B declares its own that also could apply, those in B win if all else is equal. So we put the ones we want to win deeper in the inheritance hierarchy.

Also, since you're free to define implicits in traits, you can mix them in anywhere.

So the last piece of our puzzle is to pop our type class implicits into a chain of traits that extend each other, with the more generic ones appearing earlier.

trait LowPriorityFooSpecializers {
  implicit def specializeFooAsAny[A] = new SpecializedFooImplAny[A]
}

trait FooSpecializers extends LowPriorityFooSpecializers {
  implicit val specializeFooAsInt = new SpecializedFooImplInt
}

Mix in the highest-priority trait to wherever the implicits are needed, and the type classes will be picked as desired.

Note that the type classes will be as specialized as you make them even if the specialized annotation is not used. So you can do without specialized at all, as long as you know the type precisely enough, unless you want to use specialized functions or interoperate with other specialized classes. (And you probably do.)

A complete example

Let's suppose we want to make a two-parameter specialized bippy function that will do apply the following transformation:

bippy(a, b) -> b
bippy(a, b: Int) -> b+1
bippy(a: Int, b) -> b
bippy(a: Int, b: Int) -> a+b

We should be able to achieve this with three type classes and a single specialized method. Let's try, first the method:

def bippy[@specialized(Int) A, @specialized(Int) B](a: A, b: B)(implicit impl: SpecBippy[A, B]) =
  impl(a, b)

Then the type classes:

trait SpecBippy[@specialized(Int) A, @specialized(Int) B] {
  def apply(a: A, b: B): B
}

final class SpecBippyAny[A, B] extends SpecBippy[A, B] {
  def apply(a: A, b: B) = b
}

final class SpecBippyAnyInt[A] extends SpecBippy[A, Int] {
  def apply(a: A, b: Int) = b + 1
}

final class SpecBippyIntInt extends SpecBippy[Int, Int] {
  def apply(a: Int, b: Int) = a + b
}

Then the implicits in chained traits:

trait LowerPriorityBippySpeccer {
  // Trick to avoid allocation since generic case is erased anyway!
  private val mySpecBippyAny = new SpecBippyAny[AnyRef, AnyRef]
  implicit def specBippyAny[A, B] = mySpecBippyAny.asInstanceOf[SpecBippyAny[A, B]]
}

trait LowPriorityBippySpeccer extends LowerPriorityBippySpeccer {
  private val mySpecBippyAnyInt = new SpecBippyAnyInt[AnyRef]
  implicit def specBippyAnyInt[A] = mySpecBippyAnyInt.asInstanceOf[SpecBippyAnyInt[A]]
}

// Make this last one an object so we can import the contents
object BippySpeccer extends LowPriorityBippySpeccer {
  implicit val specBippyIntInt = new SpecBippyIntInt
}

and finally we'll try it out (after pasting everything in together in :paste in the REPL):

scala> import Speccer._
import Speccer._

scala> bippy(Some(true), "cod")
res0: String = cod

scala> bippy(1, "salmon")
res1: String = salmon

scala> bippy(None, 3)
res2: Int = 4

scala> bippy(4, 5)
res3: Int = 9

It works--our custom implementations are enabled. Just to check that we can use any type, but we don't leak into the wrong implementation:

scala> bippy(4, 5: Short)
res4: Short = 5

scala> bippy(4, 5: Double)
res5: Double = 5.0

scala> bippy(3: Byte, 2)
res6: Int = 3

And finally, to verify that we have actually avoided boxing, we'll time bippy at summing a bunch of integers:

scala> val th = new ichi.bench.Thyme
th: ichi.bench.Thyme = ichi.bench.Thyme@1130520d

scala> val adder = (i: Int, j: Int) => i + j
adder: (Int, Int) => Int = <function2>

scala> var a = Array.fill(1024)(util.Random.nextInt)
a: Array[Int] = Array(-698116967, 2090538085, -266092213, ...

scala> th.pbenchOff(){
  var i, s = 0
  while (i < 1024) { s = adder(a(i), s); i += 1 }
  s 
}{ 
  var i, s = 0
  while (i < 1024) { s = bippy(a(i), s); i += 1 }
  s
}

Benchmark comparison (in 1.026 s)
Not significantly different (p ~= 0.2795)
  Time ratio:    0.99424   95% CI 0.98375 - 1.00473   (n=30)
    First     330.7 ns   95% CI 328.2 ns - 333.1 ns
    Second    328.8 ns   95% CI 326.3 ns - 331.2 ns

So we can see that our specialized bippy-adder achieves the same kind of performance as specialized Function2 does (about 3 adds per ns, which is about right for a modern machine).

Summary

To write custom specialized code using the @specialized annotation,

1. Make the specialized class abstract and manually supply concrete implementations
2. Make specialized methods (including generators for a specialized class) take typeclasses that do the real work
3. Make the base typeclass trait @specialized and provide concrete implementations
4. Provide implicit vals or defs in an inheritance-hierarchy of traits so the correct one is selected

It's a lot of boilerplate, but at the end of it all you get a seamless custom-specialized experience.