Merge pull request #162 from tresata/feat-sortedmap

support serialization of scala sortedmap

chill-scala/src/main/scala/com/twitter/chill/ScalaKryoInstantiator.scala

@@ -22,6 +22,7 @@ import scala.collection.immutable.{
   NumericRange,
   Range,
   SortedSet,
+  SortedMap,
   ListMap,
   HashMap,
   Queue
@@ -140,6 +141,7 @@ class ScalaCollectionsRegistrar extends IKryoRegistrar {
                            classOf[NumericRange.Inclusive[_]],
                            classOf[NumericRange.Exclusive[_]]))
       // Add some maps
+      .forSubclass[SortedMap[Any, Any]](new SortedMapSerializer)
       .forTraversableSubclass(ListMap.empty[Any,Any])
       .forTraversableSubclass(HashMap.empty[Any,Any])
       // The above ListMap/HashMap must appear before this:

chill-scala/src/main/scala/com/twitter/chill/SortedMapSerializer.scala

@@ -0,0 +1,54 @@
+/*
+Copyright 2012 Twitter, Inc.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package com.twitter.chill
+
+import scala.collection.immutable.SortedMap
+
+class SortedMapSerializer[A, B] extends KSerializer[SortedMap[A, B]] {
+  type M = SortedMap[A, B]
+
+  def write(kser: Kryo, out: Output, map: M) {
+    //Write the size
+    out.writeInt(map.size, true)
+
+    // Write the ordering
+    kser.writeClassAndObject(out, map.ordering.asInstanceOf[AnyRef])
+    map.foreach { t  =>
+      val tRef = t.asInstanceOf[AnyRef]
+      kser.writeClassAndObject(out, tRef)
+      // After each intermediate object, flush
+      out.flush()
+    }
+  }
+
+  def read(kser: Kryo, in: Input, cls: Class[M]): M = {
+    val size = in.readInt(true)
+    val ordering = kser.readClassAndObject(in).asInstanceOf[Ordering[A]]
+
+    // Go ahead and be faster, and not as functional cool, and be mutable in here
+    var idx = 0
+    val builder = SortedMap.canBuildFrom[A, B](ordering)()
+    builder.sizeHint(size)
+
+    while (idx < size) {
+      val item = kser.readClassAndObject(in).asInstanceOf[(A, B)]
+      builder += item
+      idx += 1
+    }
+    builder.result()
+  }
+}

chill-scala/src/test/scala/com/twitter/chill/KryoSpec.scala

@@ -18,7 +18,7 @@ package com.twitter.chill
 
 import org.specs._
 
-import scala.collection.immutable.{SortedSet, BitSet, ListSet, ListMap, HashMap}
+import scala.collection.immutable.{SortedSet, BitSet, ListSet, SortedMap, ListMap, HashMap}
 import scala.collection.mutable.{ArrayBuffer => MArrayBuffer, HashMap => MHashMap}
 import _root_.java.util.PriorityQueue
 import _root_.java.util.Locale
@@ -76,6 +76,8 @@ class KryoSpec extends Specification with BaseProperties {
                       SortedSet(1L, 2L, 3L, 4L),
                       BitSet(),
                       BitSet((0 until 1000).map{ x : Int => x*x } : _*),
+                      SortedMap[Long, String](),
+                      SortedMap("b" -> 2, "a" -> 1),
                       ListMap("good" -> 0.5, "bad" -> -1.0),
                       HashMap("good" -> 0.5, "bad" -> -1.0),
                       TestCaseClassForSerialization("case classes are: ", 10),
@@ -98,6 +100,8 @@ class KryoSpec extends Specification with BaseProperties {
       val rtTest = test map { serialize(_) } map { deserialize[AnyRef](_) }
       rtTest.zip(test).foreach { case (serdeser, orig) =>
         serdeser must be_==(orig)
+        //orig.getClass.asInstanceOf[Class[Any]] must be_==(serdeser.getClass.asInstanceOf[Class[Any]])
+        //println("orig " + orig.getClass.asInstanceOf[Class[Any]] + " => serde " + serdeser.getClass.asInstanceOf[Class[Any]])
       }
     }
     "round trip a SortedSet" in {