Boilerplate-free Scala ArrayBuilder specialization

Question

I'm looking for good practices to avoid rewriting the same code over and over again to achieve unboxedness. Say I have something like this:

def speedyArrayMaker[@specialized(Long) A: ClassTag](...): Array[A] = {
  val builder = Array.newBuilder[A]
  // do stuff with builder
  builder.result
}

This will result in unboxed storage underlying my builder when possible, but as I understand it, no unboxed method calls to it because I'm going through the unspecialized ArrayBuilder trait.

In a monomorphic world specialized to Long, I'd write val builder = new ArrayBuilder.ofLong() and avoid boxing at all, but short of telling ArrayBuilder/Builder to be specialized on all primitive types, I can't think of a nice way to avoid duplicating effort here. One approach I've thought of might be, within speedyArrayMaker:

val (add, builder): (A => Unit, ArrayBuilder[A]) = implicitly[ClassTag[A]].runtimeClass match {
  case java.lang.Long.TYPE =>
    val builder = new ArrayBuilder.ofLong()
    ((x: Long) => builder += x, builder).asInstanceOf
  case _ =>
    val builder = Array.newBuilder[A]
    ((x: A) => builder += x, builder)
}

Since it's only the += method we really care to get specialized, and then we get a Function1 for add which is specialized on Long. Checking with javap, indeed I get

90:  invokestatic    #118; //Method scala/runtime/BoxesRunTime.boxToLong:(J)Ljava/lang/Long;
93:  invokeinterface #127,  2; //InterfaceMethod scala/collection/mutable/Builder.$plus$eq:(Ljava/lang/Object;)Lscala/collection/mutable/Builder;

for the Array.newBuilder[A] version (even in the specialized output) and:

252: invokeinterface #204,  3; //InterfaceMethod scala/Function1.apply$mcVJ$sp:(J)V

for the convoluted version. I can box up this pattern into a "specialized builder helper" function, but it feels ugly especially when it's still dispatching at runtime based on something known at compile time during specialization. Ultimately I'd say my proposal here is to piggyback on the fact that Function1 is already specialized, and I don't particularly like it.

Are there clever tricks I can use to make this more pleasant? I realize this is a really low-level detail and will rarely be performance-critical, but given the amount of effort/code duplication that went into all the ArrayBuilder.of* specialized classes, it seems a pity to throw away some of their advantages in exchange for being polymorphic.

Edit I thought of something ugly but that I was hoping would work:

def builderOf(x: Array[Int]): ArrayBuilder.ofInt = new ArrayBuilder.ofInt()
def builderOf(x: Array[Long]): ArrayBuilder.ofLong = new ArrayBuilder.ofLong()
def builderOf[A: ClassTag](x: Array[A]): ArrayBuilder[A] = ArrayBuilder.make[A]

and then inside my specialized function:

val witness: Array[A] = null
val builder = builderOf(witness)

but it seems to call the generic builderOf even in the specialized version (even though enough type information is available to call the Array[Long] version). Anyone know why this doesn't work? The approach seems fairly clean, compared to the other one I was proposing. I guess I was hoping for a more "macro-like" approach to specialization, but I guess there's no guarantee that it'll be type-correct for all instantiations unless it picks the same method for each specialization :(

Answer 1

You could try something along the lines of (excuse the barbaric names),

import scala.collection.mutable.ArrayBuilder
import scala.reflect.ClassTag

trait SpecializedArrayBuilder[@specialized(Long) A] {
  def +=(a: A)
  def result: Array[A]
}

trait LowPrioritySpecializedArrayBuilder {
  implicit def defaultBuilder[A: ClassTag] = new SpecializedArrayBuilder[A] {
    val builder = ArrayBuilder.make[A]
    def +=(a: A) = builder += a
    def result = builder.result
  }
}

object SpecializedArrayBuilder extends LowPrioritySpecializedArrayBuilder {
  implicit val longBuilder = new SpecializedArrayBuilder[Long] {
    val builder = new ArrayBuilder.ofLong
    def +=(a: Long) = builder += a
    def result = builder.result
  }
}

object Specialized {
  def speedyArrayMaker[@specialized(Long) A](a: A)
    (implicit builder: SpecializedArrayBuilder[A]): Array[A] = {
    builder += a
    builder.result
  }

  def main(arg: Array[String]) {
    val arr = speedyArrayMaker(1L)
    println(arr)
  }
}