Merge pull request #428 from jvican/add-specialized-tuple-support

Add support for specialized tuples

Author: jsuereth

core/src/main/scala/scala/pickling/generator/sourcegen.scala

@@ -70,15 +70,21 @@ private[pickling] trait SourceGenerator extends Macro with tags.FastTypeTagMacro
       }
     }
 
+    val classMapper = q"$picklingPath.util.ClassMapper"
     val unrecognizedClass = q"$errorsPath.UnrecognizedClass"
 
     def genSubclassDispatch(x: SubclassDispatch): c.Tree = {
       val tpe = x.parent.tpe[c.universe.type](c.universe)
-      val clazzName = newTermName("clazz")
-        val compileTimeDispatch: List[CaseDef] = (x.subClasses map { subtpe =>
-        val tpe = subtpe.tpe[c.universe.type](c.universe)
-        CaseDef(Bind(clazzName, Ident(nme.WILDCARD)), q"clazz == classOf[$tpe]", createPickler(tpe, q"builder"))
-      })(collection.breakOut)
+      val clazz = newTermName("clazz")
+      val compileTimeDispatch: List[CaseDef] =
+        (x.subClasses map { subtpe =>
+          val tpe = subtpe.tpe[c.universe.type](c.universe)
+          CaseDef(
+            Bind(clazz, Ident(nme.WILDCARD)),
+            q"_root_.scala.pickling.util.ClassMapper.areSameClasses($clazz, classOf[$tpe])",
+            createPickler(tpe, q"builder")
+          )
+        })(collection.breakOut)
 
       val failDispatch = {
         val dispatcheeNames = x.subClasses.map(_.className).mkString(", ")
@@ -111,7 +117,8 @@ private[pickling] trait SourceGenerator extends Macro with tags.FastTypeTagMacro
         case Some(b) =>
           val parentTpe = x.parent.tpe[c.universe.type](c.universe)
           val impl = generatePickleImplFromAst(b)
-          q"""if(classOf[$parentTpe] == picklee.getClass) $impl else $subclasses"""
+          val checkIfParent = q"$classMapper.areSameClasses(picklee.getClass, classOf[$parentTpe])"
+          q"""if($checkIfParent) $impl else $subclasses"""
       }
     }
     def genPickleEntry(op: PickleEntry): c.Tree = {

core/src/main/scala/scala/pickling/internal/DefaultPicklerRegistry.scala

@@ -4,13 +4,14 @@ import scala.collection.concurrent.TrieMap
 import scala.collection.mutable
 import scala.pickling.spi.PicklerRegistry._
 import scala.reflect.runtime.universe.Mirror
-
 import scala.pickling._
+import scala.pickling.runtime.CustomRuntime
 import scala.pickling.spi.{PicklerRegistry, RuntimePicklerGenerator}
+import scala.pickling.util.ClassMapper
 
 /** Default pickle registry just uses TrieMaps and delegates behavior to a runtime pickler generator. */
 final class DefaultPicklerRegistry(generator: RuntimePicklerGenerator)
-  extends PicklerRegistry with RuntimePicklerRegistry {
+  extends PicklerRegistry with CustomRuntime {
 
   type PicklerGenerator = FastTypeTag[_] => Pickler[_]
   type UnpicklerGenerator = FastTypeTag[_] => Unpickler[_]
@@ -20,7 +21,16 @@ final class DefaultPicklerRegistry(generator: RuntimePicklerGenerator)
   private val unpicklerMap: mutable.Map[String, Unpickler[_]] = new TrieMap[String, Unpickler[_]]
   private val unpicklerGenMap: mutable.Map[String, UnpicklerGenerator] = new TrieMap[String, UnpicklerGenerator]
 
-  registerRuntimePicklersAtInit()
+  val tupleGenerators: (PicklerGen[Any], UnpicklerGen[Any]) =
+    (tuplePicklerGenerator.asInstanceOf[PicklerGen[Any]],
+      tupleUnpicklerGenerator.asInstanceOf[UnpicklerGen[Any]])
+
+  val templatesToRegister =
+    Vector("scala.Tuple2" -> tupleGenerators) ++
+      ClassMapper.specializedTupleNamesFor("scala.Tuple2")
+        .map(_ -> tupleGenerators)
+
+  registerTemplatesAtInit(templatesToRegister)
 
   override def genUnpickler(mirror: Mirror, tagKey: String)(implicit share: refs.Share): Unpickler[_] = {
     lookupUnpickler(tagKey) match {

core/src/main/scala/scala/pickling/internal/RuntimePicklerRegistry.scala

@@ -156,11 +156,11 @@ trait CustomRuntime {
     }
   }
 
-  val tuplePicklerGenerator: PicklerUnpicklerGen[(Any, Any)] = { tpe =>
+  val tuplePicklerGenerator: PicklerGen[(Any, Any)] = { tpe =>
     // TODO - Actually extract the tpe of the internal things.
     val tag = FastTypeTag.apply(tpe.toString)
     // TODO Remove this redundancy and reuse the tag above
-    Tuple2RuntimePicklerUnpickler
+    Tuple2RuntimePicklerUnpickler.asInstanceOf[Pickler[(Any, Any)]]
   }
 
   val tupleUnpicklerGenerator: UnpicklerGen[(Any,Any)] = {

core/src/main/scala/scala/pickling/runtime/CustomRuntime.scala

@@ -125,6 +125,17 @@ trait PicklerRegistry {
     */
   private[pickling] def dumpStateTo(r: PicklerRegistry): Unit
 
+  type Templates = (PicklerGen[Any], UnpicklerGen[Any])
+
+  /** Register templates (also known as generators) that know how to pickle
+    * and unpickle types at runtime without using reflection.*/
+  final def registerTemplatesAtInit(ts: Vector[(String, Templates)]): Unit = {
+    for((key, (pickler, unpickler)) <- ts) {
+      registerPicklerGenerator(key, pickler)
+      registerUnpicklerGenerator(key, unpickler)
+    }
+  }
+
 }
 
 object PicklerRegistry {

core/src/main/scala/scala/pickling/spi/PicklerRegistry.scala

@@ -0,0 +1,33 @@
+package scala.pickling.util
+
+object ClassMapper {
+
+  val specializedTuplesTemplate =
+    Vector(
+      "$mcII$sp", "$mcIJ$sp", "$mcID$sp", "$mcIZ$sp", "$mcJI$sp", "$mcJJ$sp",
+      "$mcJD$sp", "$mcJC$sp", "$mcJZ$sp", "$mcDI$sp", "$mcDJ$sp", "$mcDD$sp",
+      "$mcDC$sp", "$mcDZ$sp", "$mcCI$sp", "$mcCJ$sp", "$mcCD$sp", "$mcCC$sp",
+      "$mcCZ$sp", "$mcZI$sp", "$mcZJ$sp", "$mcZD$sp", "$mcZC$sp", "$mcZZ$sp"
+    )
+
+  /* Map specialized classes to classes. Canonical use case: tuples.
+   * We map classes instead of strings to check at runtime that they exist. */
+  val specialMappingClasses: Map[Class[_], Class[_]] =
+    mapSpecializedTuplesFor("scala.Tuple2") // add also other special cases
+
+  def specializedTupleNamesFor(tupleClassName: String): Vector[String] =
+    specializedTuplesTemplate.map(tupleClassName + _)
+
+  def mapSpecializedTuplesFor(tupleClassName: String): Map[Class[_], Class[_]] = {
+    val tupleClass = Class.forName(tupleClassName)
+    specializedTupleNamesFor(tupleClassName)
+      .map(Class.forName).map(_ -> tupleClass).toMap
+  }
+
+  @inline def isSpecializedClass(specialized: Class[_], clazz: Class[_]) =
+    specialMappingClasses.get(specialized).exists(_ == clazz)
+
+  def areSameClasses(clazz: Class[_], clazzT: Class[_]): Boolean =
+    clazz == clazzT || isSpecializedClass(clazz, clazzT)
+
+}

core/src/main/scala/scala/pickling/util/ClassMapper.scala

@@ -0,0 +1,28 @@
+package scala.pickling.generation.pickler
+
+import org.scalatest.FunSuite
+
+import scala.pickling._
+import Defaults._
+import json._
+import scala.pickling.util.ClassMapper
+import static._
+
+class Tuple2Pickler extends FunSuite {
+
+  test("pickle/unpickle (Int, String)") {
+    val t = (1, "")
+    val t2 = t.pickle.unpickle[(Int, String)]
+    assert(t === t2)
+  }
+
+  test("pickle/unpickle any specialized tuple like (Int, Int)") {
+    val t = (1, 2)
+    val clz = classOf[(Int, Int)]
+    val clzT = t.getClass
+    println(ClassMapper.areSameClasses(clzT, clz))
+    val t2 = t.pickle.unpickle[(Int, Int)]
+    assert(t === t2)
+  }
+
+}