Merge pull request #411 from jvican/lazy-init-generated-picklers-unpi…

…cklers

Add lazy initialization and isLookupEnabled

core/src/main/scala/scala/pickling/Pickler.scala

@@ -102,27 +102,11 @@ object PicklerUnpickler {
   }
 }
 
-object AutoRegister {
-  import scala.pickling.internal.currentRuntime
-
-  private def isLookupEnabled: Boolean =
-    !currentRuntime.isInstanceOf[NoReflectionRuntime]
-
-  private[pickling] def existsPicklerFor(key: String) =
-    isLookupEnabled && currentRuntime.picklers.lookupExistingPickler(key).isEmpty
-
-  private[pickling] def existsUnpicklerFor(key: String) =
-    isLookupEnabled && currentRuntime.picklers.lookupExistingUnpickler(key).isEmpty
-
-}
-
 trait AutoRegisterPickler[T] {
   this: Pickler[T] =>
 
   locally {
-    if (AutoRegister.existsPicklerFor(tag.key)) {
-      currentRuntime.picklers.registerPickler(this)
-    }
+    currentRuntime.picklers.registerPickler(this)
   }
 
 }
@@ -131,9 +115,7 @@ trait AutoRegisterUnpickler[T] {
   this: Unpickler[T] =>
 
   locally {
-    if (AutoRegister.existsUnpicklerFor(tag.key)) {
-      currentRuntime.picklers.registerUnpickler(this)
-    }
+    currentRuntime.picklers.registerUnpickler(this)
   }
 
 }

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

@@ -384,63 +384,114 @@ private[pickling]  trait SourceGenerator extends Macro with tags.FastTypeTagMacr
   }
 
   val picklingPath = q"_root_.scala.pickling"
+  val picklersRegistry = q"$picklingPath.internal.currentRuntime.picklers"
+  val generated = tq"$picklingPath.Generated"
+
+  def lazyInit(lookup: c.Tree, key: c.Tree, tpe: c.Tree,
+               notInitialized: c.TermName): c.Tree = {
+    q"""
+      if($picklersRegistry.isLookupEnabled) {
+        $lookup($key) match {
+          case Some(p) => p.asInstanceOf[$tpe]
+          case None => $notInitialized: $tpe
+        }
+      } else $notInitialized: $tpe
+    """
+  }
+
+  def picklerUnpicklerLazyInit(lookup1: c.Tree, lookup2: c.Tree, key: c.Tree,
+                               tpe: c.Tree, notInitialized: c.TermName) = {
+    q"""
+      if($picklersRegistry.isLookupEnabled) {
+        ($lookup1($key), $lookup2($key)) match {
+          case (Some(p), Some(u)) if p == u => p.asInstanceOf[$tpe]
+          case _ => $notInitialized: $tpe
+        }
+      } else $notInitialized: $tpe
+    """
+  }
 
   /** generates the tree which will construct + return a new instance of a Pickler class, capable of
     * pickling an instance of type T, using the behavior outlined by the PicklerAst.
     */
   def generatePicklerClass[T: c.WeakTypeTag](picklerAst: PicklerAst): c.Tree = {
+
     val tpe = computeType[T]
-    val picklerName = c.fresh((syntheticBaseName(tpe) + "Pickler"): TermName)
+    val picklerName = c.fresh(newTermName(syntheticBaseName(tpe) + "Pickler"))
     val createTagTree = super[FastTypeTagMacros].impl[T]
+    val picklerType = tq"$picklingPath.Pickler[$tpe]"
+    val key = q"$createTagTree.key"
+    val lookup = q"$picklersRegistry.lookupExistingPickler"
+
     q"""
       _root_.scala.Predef.locally {
-        implicit object $picklerName extends $picklingPath.Pickler[$tpe]
-          with $picklingPath.Generated with $picklingPath.AutoRegisterPickler[$tpe]{
+        implicit object $picklerName extends $picklerType with $picklingPath.Generated
+            with $picklingPath.AutoRegisterPickler[$tpe]{
+
           lazy val tag: $picklingPath.FastTypeTag[$tpe] = $createTagTree
           def pickle(picklee: $tpe, builder: $picklingPath.PBuilder): _root_.scala.Unit = ${genPicklerLogic[T](picklerAst)}
+
         }
-        $picklerName
+        ${lazyInit(lookup, key, picklerType, picklerName)}
       }
     """
+
   }
 
   def generateUnpicklerClass[T: c.WeakTypeTag](unpicklerAst: UnpicklerAst): c.Tree = {
+
     val tpe = computeType[T]
-    val unpicklerName = c.fresh((syntheticBaseName(tpe) + "Unpickler"): TermName)
+    val unpicklerName = c.fresh(newTermName(syntheticBaseName(tpe) + "Unpickler"))
     val createTagTree = super[FastTypeTagMacros].impl[T]
     val unpickleLogic = genUnpicklerLogic[T](unpicklerAst)
+    val unpicklerType = tq"$picklingPath.Unpickler[$tpe]"
+    val genUnpicklerType = tq"$unpicklerType with $generated"
+    val key = q"$createTagTree.key"
+    val lookup = q"$picklersRegistry.lookupExistingUnpickler"
+
     q"""
        _root_.scala.Predef.locally {
-          implicit object $unpicklerName extends $picklingPath.Unpickler[$tpe]
-            with $picklingPath.Generated with $picklingPath.AutoRegisterUnpickler[$tpe]{
+          implicit object $unpicklerName extends $unpicklerType with $generated
+              with $picklingPath.AutoRegisterUnpickler[$tpe] {
+
             lazy val tag: $picklingPath.FastTypeTag[$tpe] = $createTagTree
             def unpickle(tagKey: _root_.java.lang.String, reader: $picklingPath.PReader): _root_.scala.Any = $unpickleLogic
+
           }
-          $unpicklerName : $picklingPath.Unpickler[$tpe] with $picklingPath.Generated
+          ${lazyInit(lookup, key, genUnpicklerType, unpicklerName)}
        }
      """
+
   }
 
   def generatePicklerUnpicklerClass[T: c.WeakTypeTag](impl: PickleUnpickleImplementation): c.Tree = {
+
     val tpe = computeType[T]
-    val name = c.fresh((syntheticBaseName(tpe) + "PicklerUnpickler"): TermName)
+    val name = c.fresh(newTermName(syntheticBaseName(tpe) + "PicklerUnpickler"))
     val createTagTree = super[FastTypeTagMacros].impl[T]
     val unpickleLogic = genUnpicklerLogic[T](impl.unpickle)
     val pickleLogic = genPicklerLogic[T](impl.pickle)
+    val key = q"$createTagTree.key"
+    val picklerUnpicklerType = tq"$picklingPath.AbstractPicklerUnpickler[$tpe]"
+    val genPicklerUnpicklerType = tq"$picklerUnpicklerType with $generated"
+    val lookup1 = q"$picklersRegistry.lookupExistingPickler"
+    val lookup2 = q"$picklersRegistry.lookupExistingUnpickler"
+
     q"""
        _root_.scala.Predef.locally {
-          implicit object $name extends $picklingPath.AbstractPicklerUnpickler[$tpe]
-            with $picklingPath.Generated with $picklingPath.AutoRegister[$tpe] {
-            //import _root_.scala.language.existentials
+          implicit object $name extends $picklerUnpicklerType with $generated
+              with $picklingPath.AutoRegister[$tpe] {
+
             override lazy val tag: $picklingPath.FastTypeTag[$tpe] = $createTagTree
             override def pickle(picklee: $tpe, builder: $picklingPath.PBuilder): _root_.scala.Unit = $pickleLogic
             override def unpickle(tagKey: _root_.java.lang.String, reader: $picklingPath.PReader): _root_.scala.Any = $unpickleLogic
+
           }
-          $name : $picklingPath.AbstractPicklerUnpickler[$tpe] with $picklingPath.Generated
+          ${picklerUnpicklerLazyInit(lookup1, lookup2, key, genPicklerUnpicklerType, name)}
        }
      """
-  }
 
+  }
 
   /** Given a tree which unpickles a value, we regsiter that value with its OID hint so that
     * we can effectively deserialized later references.
@@ -449,16 +500,13 @@ private[pickling]  trait SourceGenerator extends Macro with tags.FastTypeTagMacr
     // TODO - We may not want to ALWAYS do this, some kind of enabling flag...
     val instance = c.fresh(newTermName("instance"))
     q"""
-              val oid = _root_.scala.pickling.internal.`package`.currentRuntime.refRegistry.unpickle.preregisterUnpicklee()
-              val $instance = $instantiationLogic
-              _root_.scala.pickling.internal.`package`.currentRuntime.refRegistry.unpickle.registerUnpicklee($instance, oid)
-              $instance
-            """
+      val oid = _root_.scala.pickling.internal.`package`.currentRuntime.refRegistry.unpickle.preregisterUnpicklee()
+      val $instance = $instantiationLogic
+      _root_.scala.pickling.internal.`package`.currentRuntime.refRegistry.unpickle.registerUnpicklee($instance, oid)
+      $instance
+    """
   }
 
-
-
-
   def computeType[T: c.WeakTypeTag]: Type = {
     val originalTpe = weakTypeOf[T]
     // Note: this makes it so modules work, things like foo.type.
@@ -480,6 +528,7 @@ private[pickling]  trait SourceGenerator extends Macro with tags.FastTypeTagMacr
         )
      """
   }
+
   def genExternalizablUnPickle(reader: TermName, pe: UnpickleExternalizable): c.Tree = {
     val tpe = pe.tpe.tpe[c.universe.type](c.universe)
     val readerName = c.fresh(newTermName("readerName"))
@@ -499,7 +548,6 @@ private[pickling]  trait SourceGenerator extends Macro with tags.FastTypeTagMacr
         """
   }
 
-
   // ---- Reflective Helper Methods ----
 
   def reflectivelyGet(target: TermName, value: IrMember)(body: c.Tree => c.Tree): List[c.Tree] = {

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

@@ -52,6 +52,14 @@ final class DefaultPicklerRegistry(generator: RuntimePicklerGenerator) extends P
   override def registerUnpickler[T](key: String, p: Unpickler[T]): Unit =
     unpicklerMap.put(key, p)
 
+  override private[pickling] def clearRegisteredPicklerUnpicklerFor[T: FastTypeTag]: Unit = {
+    val tag = implicitly[FastTypeTag[T]]
+    picklerMap -= tag.key
+    unpicklerMap -= tag.key
+  }
+
+  override val isLookupEnabled = true
+
   /** Checks the existence of a pickler ignoring the registered generators. */
   override def lookupExistingPickler(key: String): Option[Pickler[_]] =
     picklerMap.get(key)

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

@@ -20,6 +20,7 @@ final class NoReflectionRuntime() extends PicklingRuntime {
       throw new UnsupportedOperationException(s"Runtime pickler generation is disabled.  Cannot create unpickler for $tagKey")
     override def genPickler(classLoader: ClassLoader, clazz: Class[_], tag: FastTypeTag[_])(implicit share: Share): Pickler[_] =
       throw new UnsupportedOperationException(s"Runtime pickler generation is disabled.  Cannot create pickler for $tag")
+    override val isLookupEnabled: Boolean = false
     override def lookupUnpickler(key: String): Option[Unpickler[_]] = None
     override def lookupPickler(key: String): Option[Pickler[_]] = None
     override def lookupExistingPickler(key: String): Option[Pickler[_]] = None
@@ -30,6 +31,7 @@ final class NoReflectionRuntime() extends PicklingRuntime {
     override def registerPicklerGenerator[T](typeConstructorKey: String, generator: (FastTypeTag[_]) => Pickler[T]): Unit = ()
     override def registerUnpickler[T](key: String, p: Unpickler[T]): Unit = ()
     override def registerPicklerUnpickler[T](key: String, p: Pickler[T] with Unpickler[T]): Unit = ()
+    override private[pickling] def clearRegisteredPicklerUnpicklerFor[T: FastTypeTag]: Unit = ()
   }
   override val refRegistry: RefRegistry = new DefaultRefRegistry()
   override val GRL: ReentrantLock = new ReentrantLock()

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

@@ -29,12 +29,15 @@ trait PicklerRegistry {
     */
   def genPickler(classLoader: ClassLoader, clazz: Class[_], tag: FastTypeTag[_])(implicit share: refs.Share): Pickler[_]
 
+  /** Checks if lookup is enabled for this registry */
+  def isLookupEnabled: Boolean
 
   /** Checks the existince of an unpickler.
     *
     * This will also check any registered generator functions.
     */
   def lookupUnpickler(key: String): Option[Unpickler[_]]
+
   /** Looks for a pickler with the given FastTypeTag string.
     *
     * This will also check any registered generator functions.
@@ -95,5 +98,12 @@ trait PicklerRegistry {
     *                   this type.
     */
   def registerPicklerUnpicklerGenerator[T](typeConstructorKey: String, generator: FastTypeTag[_] => (Pickler[T] with Unpickler[T])): Unit
-  // TODO - Some kind of clean or inspect what we have?
+
+  /** Clear the registered pickler/unpickler for a given type.
+    *
+    * Useful for avoiding conflict between picklers registered with different
+    * sharing strategies and are cached when they're initialised.
+    */
+  private[pickling] def clearRegisteredPicklerUnpicklerFor[T: FastTypeTag]: Unit
+
 }

core/src/test/scala/pickling/run/share-binary-any.scala

@@ -6,6 +6,9 @@ import scala.pickling._, scala.pickling.Defaults._, binary._
 class C(val name: String, val desc: String, var c: C, val arr: Array[Int])
 
 class ShareBinaryAnyTest extends FunSuite {
+
+  import scala.pickling.internal.currentRuntime
+
   val c1 = new C("c1", "desc", null, Array(1))
   val c2 = new C("c2", "desc", c1, Array(1))
   val c3 = new C("c3", "desc", c2, Array(1))
@@ -31,6 +34,10 @@ class ShareBinaryAnyTest extends FunSuite {
   }*/
 
   test("loop-share-nothing") {
+
+    currentRuntime.picklers.clearRegisteredPicklerUnpicklerFor[Any]
+    currentRuntime.picklers.clearRegisteredPicklerUnpicklerFor[C]
+
     intercept[StackOverflowError] {
       import shareNothing._
       c1.c = c3
@@ -59,6 +66,9 @@ class ShareBinaryAnyTest extends FunSuite {
   }*/
 
   test("noloop-share-non-primitives") {
+
+    currentRuntime.picklers.clearRegisteredPicklerUnpicklerFor[C]
+
     import shareNothing._
     c1.c = null
     val pickle = (c3: Any).pickle

core/src/test/scala/pickling/run/share-binary.scala

@@ -15,11 +15,17 @@ final class Simple(x: Int) {
 
 
 class ShareBinaryTest extends FunSuite {
+
+  import scala.pickling.internal.currentRuntime
+
   val c1 = new C("c1", "desc", null, Array(1))
   val c2 = new C("c2", "desc", c1, Array(1))
   val c3 = new C("c3", "desc", c2, Array(1))
 
   test("loop-share-nonprimitives") {
+
+    currentRuntime.picklers.clearRegisteredPicklerUnpicklerFor[C]
+
     c1.c = c3
     val pickle = c1.pickle
     //val expected = "BinaryPickle([0,0,0,29,115,99,97,108,97,46,112,105,99,107,108,105,110,103,46,115,104,97,114,101,46,98,105,110,97,114,121,46,67,0,0,0,2,99,49,0,0,0,4,100,101,115,99,-5,0,0,0,2,99,51,0,0,0,4,100,101,115,99,-5,0,0,0,2,99,50,0,0,0,4,100,101,115,99,-3,0,0,0,0,0,0,0,1,1,0,0,0,0,0,0,1,1,0,0,0,0,0,0,1,1,0,0,0])"
@@ -42,6 +48,9 @@ class ShareBinaryTest extends FunSuite {
   }
 
   test("loop-share-nothing") {
+
+    currentRuntime.picklers.clearRegisteredPicklerUnpicklerFor[C]
+
     intercept[StackOverflowError] {
       import shareNothing._
       c1.c = c3
@@ -50,6 +59,9 @@ class ShareBinaryTest extends FunSuite {
   }
 
   test("loop-share-everything") {
+
+    currentRuntime.picklers.clearRegisteredPicklerUnpicklerFor[C]
+
     import shareEverything._
     c1.c = c3
     val pickle = c1.pickle
@@ -71,6 +83,9 @@ class ShareBinaryTest extends FunSuite {
   }
 
   test("noloop-share-non-primitives") {
+
+    currentRuntime.picklers.clearRegisteredPicklerUnpicklerFor[C]
+
     import shareNothing._
     c1.c = null
     val pickle = c3.pickle