Add extendible way to pickle specialized classes

* Reuse code in `RuntimePicklerRegistry` and `sourcegen`. * Create `ClassMapper` to delegate the tasks of checking class equality.
scala · May 14, 2016 · 9c9af01 · 9c9af01
1 parent 00a198d
commit 9c9af01
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 37 deletions.
diff --git a/core/src/main/scala/scala/pickling/generator/sourcegen.scala b/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 = {

diff --git a/core/src/main/scala/scala/pickling/internal/RuntimePicklerRegistry.scala b/core/src/main/scala/scala/pickling/internal/RuntimePicklerRegistry.scala
@@ -3,6 +3,7 @@ package internal
 
 import scala.pickling.runtime.CustomRuntime
 import scala.pickling.spi.PicklerRegistry
+import scala.pickling.util.ClassMapper
 
 /** Helper class to register picklers and unpicklers that
   * handle special/crazy classes for runtime generation.
@@ -19,37 +20,10 @@ trait RuntimePicklerRegistry extends CustomRuntime {
    * this list because this is meant to be used
    * internally. If a user wants to do the same,
    * she just needs to use the [[PicklerRegistry]]. */
-  val picklersToRegister = List(
-    ("scala.Tuple2", tupleGenerators),
-    ("scala.Tuple2$mcII$sp", tupleGenerators),
-    ("scala.Tuple2$mcIJ$sp", tupleGenerators),
-    ("scala.Tuple2$mcID$sp", tupleGenerators),
-    ("scala.Tuple2$mcIZ$sp", tupleGenerators),
-
-    ("scala.Tuple2$mcJI$sp", tupleGenerators),
-    ("scala.Tuple2$mcJJ$sp", tupleGenerators),
-    ("scala.Tuple2$mcJD$sp", tupleGenerators),
-    ("scala.Tuple2$mcJC$sp", tupleGenerators),
-    ("scala.Tuple2$mcJZ$sp", tupleGenerators),
-
-    ("scala.Tuple2$mcDI$sp", tupleGenerators),
-    ("scala.Tuple2$mcDJ$sp", tupleGenerators),
-    ("scala.Tuple2$mcDD$sp", tupleGenerators),
-    ("scala.Tuple2$mcDC$sp", tupleGenerators),
-    ("scala.Tuple2$mcDZ$sp", tupleGenerators),
-
-    ("scala.Tuple2$mcCI$sp", tupleGenerators),
-    ("scala.Tuple2$mcCJ$sp", tupleGenerators),
-    ("scala.Tuple2$mcCD$sp", tupleGenerators),
-    ("scala.Tuple2$mcCC$sp", tupleGenerators),
-    ("scala.Tuple2$mcCZ$sp", tupleGenerators),
-
-    ("scala.Tuple2$mcZI$sp", tupleGenerators),
-    ("scala.Tuple2$mcZJ$sp", tupleGenerators),
-    ("scala.Tuple2$mcZD$sp", tupleGenerators),
-    ("scala.Tuple2$mcZC$sp", tupleGenerators),
-    ("scala.Tuple2$mcZZ$sp", tupleGenerators)
-  )
+  val picklersToRegister =
+    List("scala.Tuple2" -> tupleGenerators) ++
+    ClassMapper.specializedTupleNamesFor("scala.Tuple2")
+      .map(_ -> tupleGenerators)
 
   final def registerRuntimePicklersAtInit(): Unit = {
     for((key, (pickler, unpickler)) <- picklersToRegister) {

diff --git a/core/src/main/scala/scala/pickling/util/ClassMapper.scala b/core/src/main/scala/scala/pickling/util/ClassMapper.scala
@@ -0,0 +1,33 @@
+package scala.pickling.util
+
+object ClassMapper {
+
+  val specializedTuplesTemplate =
+    List(
+      "$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): List[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).contains(clazz)
+
+  def areSameClasses(clazz: Class[_], clazzT: Class[_]): Boolean =
+    clazz == clazzT || isSpecializedClass(clazz, clazzT)
+
+}