From 4cacc0cee9499360af016cc200e53fd007c1c24f Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@apache.org>
Date: Sat, 28 Jun 2014 01:32:31 -0700
Subject: [PATCH 1/2] Support serializing asJavaIterable.

---
 .../chill/JavaIterableWrapperSerializer.scala | 48 +++++++++++++++++++
 .../twitter/chill/ScalaKryoInstantiator.scala |  2 +
 .../scala/com/twitter/chill/KryoSpec.scala    |  4 ++
 3 files changed, 54 insertions(+)
 create mode 100644 chill-scala/src/main/scala/com/twitter/chill/JavaIterableWrapperSerializer.scala

diff --git a/chill-scala/src/main/scala/com/twitter/chill/JavaIterableWrapperSerializer.scala b/chill-scala/src/main/scala/com/twitter/chill/JavaIterableWrapperSerializer.scala
new file mode 100644
index 00000000..9bc588c6
--- /dev/null
+++ b/chill-scala/src/main/scala/com/twitter/chill/JavaIterableWrapperSerializer.scala
@@ -0,0 +1,48 @@
+package com.twitter.chill
+
+import _root_.java.lang.{ Iterable => JIterable }
+
+/**
+ * A Kryo serializer for serializing results returned by asJavaIterable.
+ *
+ * The underlying object is scala.collection.convert.Wrappers$IterableWrapper.
+ * Kryo deserializes this into an AbstractCollection, which unfortunately doesn't work.
+ *
+ * Ported from Apache Spark's KryoSerializer.scala.
+ */
+private class JavaIterableWrapperSerializer extends KSerializer[JIterable[_]] {
+
+  import JavaIterableWrapperSerializer._
+
+  override def write(kryo: Kryo, out: Output, obj: JIterable[_]): Unit = {
+    // If the object is the wrapper, simply serialize the underlying Scala Iterable object.
+    // Otherwise, serialize the object itself.
+    if (obj.getClass == wrapperClass && underlyingMethodOpt.isDefined) {
+      kryo.writeClassAndObject(out, underlyingMethodOpt.get.invoke(obj))
+    } else {
+      kryo.writeClassAndObject(out, obj)
+    }
+  }
+
+  override def read(kryo: Kryo, in: Input, clz: Class[JIterable[_]]): JIterable[_] = {
+    kryo.readClassAndObject(in) match {
+      case scalaIterable: Iterable[_] =>
+        scala.collection.JavaConversions.asJavaIterable(scalaIterable)
+      case javaIterable: JIterable[_] =>
+        javaIterable
+    }
+  }
+}
+
+private object JavaIterableWrapperSerializer {
+  // The class returned by asJavaIterable (scala.collection.convert.Wrappers$IterableWrapper).
+  val wrapperClass = scala.collection.JavaConversions.asJavaIterable(Seq(1)).getClass
+
+  // Get the underlying method so we can use it to get the Scala collection for serialization.
+  private val underlyingMethodOpt = {
+    try Some(wrapperClass.getDeclaredMethod("underlying")) catch {
+      case e: Exception =>
+        None
+    }
+  }
+}
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 0cd429dd..6f2ba3d6 100644
--- a/chill-scala/src/main/scala/com/twitter/chill/ScalaKryoInstantiator.scala
+++ b/chill-scala/src/main/scala/com/twitter/chill/ScalaKryoInstantiator.scala
@@ -160,6 +160,8 @@ class ScalaCollectionsRegistrar extends IKryoRegistrar {
       .forTraversableSubclass(MMap.empty[Any, Any], isImmutable = false)
       .forTraversableSubclass(MSet.empty[Any], isImmutable = false)
       .forTraversableSubclass(ListBuffer.empty[Any], isImmutable = false)
+
+    newK.register(JavaIterableWrapperSerializer.wrapperClass, new JavaIterableWrapperSerializer)
   }
 }
 
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 b3e7dd32..4c941c85 100644
--- a/chill-scala/src/test/scala/com/twitter/chill/KryoSpec.scala
+++ b/chill-scala/src/test/scala/com/twitter/chill/KryoSpec.scala
@@ -172,6 +172,10 @@ class KryoSpec extends Specification with BaseProperties {
     "handle scala enums" in {
       WeekDay.values.foreach { _ must roundtrip }
     }
+    "handle asJavaIterable" in {
+      val col = scala.collection.JavaConversions.asJavaIterable(Seq(12345))
+      col must roundtrip
+    }
     "use java serialization" in {
       val kinst = { () => getKryo.javaForClass[TestCaseClassForSerialization] }
       rtEquiv(kinst, TestCaseClassForSerialization("hey", 42)) must beTrue

From 9cd7ba4ad1ab4a5283f5e26b08e0c8d716d34443 Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@apache.org>
Date: Mon, 30 Jun 2014 14:26:30 -0700
Subject: [PATCH 2/2] Separated the registration of
 JavaIterableWrapperSerializer into a new class.

---
 .../scala/com/twitter/chill/ScalaKryoInstantiator.scala   | 8 ++++++++
 1 file changed, 8 insertions(+)

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 6f2ba3d6..20cf4943 100644
--- a/chill-scala/src/main/scala/com/twitter/chill/ScalaKryoInstantiator.scala
+++ b/chill-scala/src/main/scala/com/twitter/chill/ScalaKryoInstantiator.scala
@@ -160,7 +160,11 @@ class ScalaCollectionsRegistrar extends IKryoRegistrar {
       .forTraversableSubclass(MMap.empty[Any, Any], isImmutable = false)
       .forTraversableSubclass(MSet.empty[Any], isImmutable = false)
       .forTraversableSubclass(ListBuffer.empty[Any], isImmutable = false)
+  }
+}
 
+class JavaWrapperCollectionRegistrar extends IKryoRegistrar {
+  def apply(newK: Kryo) {
     newK.register(JavaIterableWrapperSerializer.wrapperClass, new JavaIterableWrapperSerializer)
   }
 }
@@ -170,6 +174,10 @@ class AllScalaRegistrar extends IKryoRegistrar {
   def apply(k: Kryo) {
     val col = new ScalaCollectionsRegistrar
     col(k)
+
+    val jcol = new JavaWrapperCollectionRegistrar
+    jcol(k)
+
     // Register all 22 tuple serializers and specialized serializers
     ScalaTupleSerialization.register(k)
     k.forClass[Symbol](new KSerializer[Symbol] {