Sunday, December 4, 2011

Generic Rotation in Scala

My son is preparing a talk on cryptography. I told him I would write some code for demo purposes, and we decided to start with a simple substitution cipher, one that - according to history - was used by Julius Ceasar. It's a really simple mechanism: you just rotate an alphabet and replace every character in your message with the character in the rotated alphabet.

I figured it would be really easy to do in Scala. I mean, rotation on a collection should be fairly easy to implement, right? Well, it is, but not if you want to write it in such a way that the output collection is of the same type as the input collection. I mean, you could define the input parameter as a Seq[T], but then - if you would pass in a String - you would not get a String as a result. So what do you do?

After some poking around, I eventually managed to write a rotation function that behaves as you would expect.

def rotate[A,C](coll: C, number: Int)(
  implicit c2i: C => Iterable[A], cbf: collection.generic.CanBuildFrom[C,A,C]
): Option[C] = 
  if (coll.hasDefiniteSize) {
    val positions = number % coll.size match {
      case i if i < 0 => i + coll.size
      case i => i
    val builder = cbf()
    builder ++= coll.drop(positions)
    builder ++= coll.take(positions)
  } else None

Since rotation only works for collections of a definite size, and since not all collections have a definite size, it returns a None in case the collection does not have a definitive size. By default, it rotates to the left. By passing a negative number, you rotate right.

scala> rotate(List(1, 2, 3, 4), 2)
res15: Option[List[Int]] = Some(List(3, 4, 1, 2))

scala> rotate("abcdefghi", 2)    
res16: Option[java.lang.String] = Some(cdefghiab)

scala> rotate("abcdefghi", -2) // Rotate right
res17: Option[java.lang.String] = Some(hiabcdefg)


Andrew said...

Nice! Two main questions: how easy do you feel this code is to read? How easy do you feel it was to write, or perhaps: how much poking around to find the solution was required?

And a curiosity question: why the declaration of an implicit conversion from C to Iterable[A] rather than a type restriction on C?