Specialization of generic functions in Scala (or Java)

Question

Is it possible to specialize generic functions (or class) in Scala? For example, I want to write a generic function that writes data into a ByteBuffer:

def writeData[T](buffer: ByteBuffer, data: T) = buffer.put(data) 

But as the put method takes only a byte and put it into the buffer, I need to specialize it for Ints and Longs as follows:

def writeData[Int](buffer: ByteBuffer, data: Int) = buffer.putInt(data) def writeData[Long](buffer: ByteBuffer, data: Long) = buffer.putLong(data) 

and it won't compile. Of course, I could instead write 3 different functions writeByte, writeInt and writeLong respectively, but let's say there is another function for an array:

def writeArray[T](buffer: ByteBuffer, array: Array[T]) {   for (elem <- array) writeData(buffer, elem) } 

and this wouldn't work without the specialized writeData functions: I'll have to deploy another set of functions writeByteArray, writeIntArray, writeLongArray. Having to deal with the situation this way whenever I need to use type-dependent write functions is not cool. I did some research and one possible workaround is to test the type of the parameter:

def writeArray[T](buffer: ByteBuffer, array: Array[T]) {   if (array.isInstanceOf[Array[Byte]])     for (elem <- array) writeByte(buffer, elem)   else if (array.isInstanceOf[Array[Int]])     for (elem <- array) writeInt(buffer, elem)   ... } 

This might work but it's less efficient because type-checking is done in runtime unlike the specialized function version.

So my question is, what is the most desirable and preferred way to solve this kind of problem in Scala or Java? I appreciate your help in advance!

Answer 1

Wouldn't it be nice if you could have both a compact and efficient solution? It turns out that you can, given Scala's @specialized feature. First a warning: the feature is somewhat buggy, and may break if you try to use it for something too complicated. But for this case, it's almost perfect.

The @specialized annotation creates separate classes and/or methods for each primitive type, and then calls that instead of the generic version whenever the compiler knows for sure what the primitive type is. The only drawback is that it does all of this completely automatically--you don't get to fill in your own method. That's kind of a shame, but you can overcome the problem using type classes.

Let's look at some code:

import java.nio.ByteBuffer trait BufferWriter[@specialized(Byte,Int) A]{   def write(b: ByteBuffer, a: A): Unit } class ByteWriter extends BufferWriter[Byte] {   def write(b: ByteBuffer, a: Byte) { b.put(a) } } class IntWriter extends BufferWriter[Int] {   def write(b: ByteBuffer, a: Int) { b.putInt(a) } } object BufferWriters {   implicit val byteWriter = new ByteWriter   implicit val intWriter = new IntWriter } 

This gives us a BufferWriter trait which is generic, but we override each of the specific primitive types that we want (in this case Byte and Int) with an appropriate implementation. Specialization is smart enough to link up this explicit version with the hidden one it normally uses for specialization. So you've got your custom code, but how do you use it? This is where the implicit vals come in (I've done it this way for speed and clarity):

import BufferWriters._ def write[@specialized(Byte,Int) A: BufferWriter](b: ByteBuffer, ar: Array[A]) {   val writer = implicitly[BufferWriter[A]]   var i = 0   while (i < ar.length) {     writer.write(b, ar(i))     i += 1   } } 

The A: BufferWriter notation means that in order to call this write method, you need to have an implicit BufferWriter[A] handy. We've supplied them with the vals in BufferWriters, so we should be set. Let's see if this works.

val b = ByteBuffer.allocate(6) write(b, Array[Byte](1,2)) write(b, Array[Int](0x03040506)) scala> b.array res3: Array[Byte] = Array(1, 2, 3, 4, 5, 6) 

If you put these things in a file and start poking around the classes with javap -c -private you'll see that the appropriate primitive methods are being used.

(Note that if you didn't use specialization, this strategy would still work, but it would have to box values inside the loop to copy the array out.)