diff --git a/chill-scala/src/main/scala/com/twitter/chill/ScalaKryoInstantiator.scala b/chill-scala/src/main/scala/com/twitter/chill/ScalaKryoInstantiator.scala index d9c0c773..80a04f29 100644 --- a/chill-scala/src/main/scala/com/twitter/chill/ScalaKryoInstantiator.scala +++ b/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: diff --git a/chill-scala/src/main/scala/com/twitter/chill/SortedMapSerializer.scala b/chill-scala/src/main/scala/com/twitter/chill/SortedMapSerializer.scala new file mode 100644 index 00000000..bea288ca --- /dev/null +++ b/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() + } +} diff --git a/chill-scala/src/test/scala/com/twitter/chill/KryoSpec.scala b/chill-scala/src/test/scala/com/twitter/chill/KryoSpec.scala index d7f66f9d..9ed53324 100644 --- a/chill-scala/src/test/scala/com/twitter/chill/KryoSpec.scala +++ b/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 {