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

Use a typeclass pattern. It has the advantage over the instanceOf checking (or pattern matching) of being typesafe.

import java.nio.ByteBuffer

trait BufferWriter[A] {
  def write(buffer: ByteBuffer, a: A)
}

class BuffPimp(buffer: ByteBuffer) {
  def writeData[A: BufferWriter](data: A) = { 
    implicitly[BufferWriter[A]].write(buffer, data)
  }
}

object BuffPimp {
  implicit def intWriter = new BufferWriter[Int] {
    def write(buffer: ByteBuffer, a: Int) = buffer.putInt(a)
  }
  implicit def doubleWriter = new BufferWriter[Double] {
    def write(buffer: ByteBuffer, a: Double) = buffer.putDouble(a)
  }
  implicit def longWriter = new BufferWriter[Long] {
    def write(buffer: ByteBuffer, a: Long) = buffer.putLong(a)
  }
  implicit def wrap(buffer: ByteBuffer) = new BuffPimp(buffer)
}

object Test {
  import BuffPimp._
  val someByteBuffer: ByteBuffer
  someByteBuffer.writeData(1)
  someByteBuffer.writeData(1.0)
  someByteBuffer.writeData(1L)
}

So this code isn't the best demonstration of typeclasses. I am still very new to them. This video gives a really solid overview of their benefits and how you can use them: http://www.youtube.com/watch?v=sVMES4RZF-8

Answer 2

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.)

Answer 3

1. The declarations

   def writeData[Int](buffer: ByteBuffer, data: Int) 
   def writeData[Long](buffer: ByteBuffer, data: Long)
   

do not compile because they are equivalent, as Int and Long are formal type parameters and not the standard Scala types. To define functions with standard Scala types just write:

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

This way you declare different functions with the same name.

1. Since they are different functions, you cannot apply them to elements of a List of statically unknown type. You have first to determine the type of the List. Note it can happen the type of the List is AnyRef, then you have dynamically determine the type of each element. The determination can be done with isInstanceOf as in your original code, or with pattern matching, as rolve suggested. I think this would produce the same bytecode.

2. In sum, you have to choose:

   • fast code with multiple functions like writeByteArray, writeIntArray etc. They all can have the same name writeArray but can be statically distinguished by their actual parameters. The variant suggested by Dominic Bou-Sa is of this kind.

   • concise but slow code with run-time type determination

Unfortunately, you cannot have both fast and concise code.

Answer 4

How about this:

def writeData(buffer: ByteBuffer, data: AnyVal) {
  data match {
    case d: Byte => buffer put d
    case d: Int  => buffer putInt d
    case d: Long => buffer putLong d
    ...
  }
}

Here, you make the case distinction in the writeData method, which makes all further methods very simple:

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

---

Advantages: Simple, short, easy to understand.

Disadvantages: Not completely type-safe if you don't handle all AnyVal types: Someone may call writeData(buffer, ()) (the second argument being of type Unit), which may result in an error at runtime. But you can also make the handling of () a no-op, which solves the problem. The complete method would look like this:

def writeData(buffer: ByteBuffer, data: AnyVal) {
  data match {
    case d: Byte   => buffer put d
    case d: Short  => buffer putShort d
    case d: Int    => buffer putInt d
    case d: Long   => buffer putLong d
    case d: Float  => buffer putFloat d
    case d: Double => buffer putDouble d
    case d: Char   => buffer putChar d
    case true      => buffer put 1.asInstanceOf[Byte]
    case false     => buffer put 0.asInstanceOf[Byte]
    case ()        =>
  }
}

By the way, this only works so easily because of Scala's strict object-oriented nature. In Java, where primitive types are not objects, this would be much more cumbersome. There, you would actually have to create a separate method for each primitive type, unless you want to do some ugly boxing and unboxing.