Fix #4442: Add support for lambda serialization

In Scala, lambdas whose SAM extend Serializable as well as lambdas whose
SAM is simply scala.Function* should be serializable but this was not
the case in Dotty so far. On the JVM, lambdas instantiated using
invokedynamic calls require some special handling to be serializable:

1. We need to use the invokedynamic bootstrap method
   `LambdaMetaFactory#altMetafactory` instead of
   `LambdaMetaFactory#metafactory`, this allows us to pass the
   FLAG_SERIALIZABLE flag. This is implemented in the backend submodule
   commit included in this commit (see lampepfl/scala#39).

2. In the enclosing class where the lambda is defined, a
   $deserializeLambda$ method needs to be generated, this is implemented
   in this commit.

Most of the logic for $deserializeLambda$ is implemented in the Scala
2.12 standard libraries class scala.runtime.LambdaDeserialize and
scala.runtime.LambdaDeserializer which can be used here as-is, the only
logic we actually need to implement here is:

1. In `collectSerializableLambdas`, we collect all serializable lambdas.
   Unlike scalac, our backend does not do any inlining currently so our
   implementation is more straightfoward than theirs.

2. In `addLambdaDeserialize`, we implement the actual
   $deserializeLambda$ method, the implementation here is directly
   copied from scalac, it's complex because it needs to work around a
   limitation of bootstrap methods (they cannot take more than 251
   arguments).

Since some of this code comes from scalac, this is:
Co-Authored-By: Jason Zaugg <[email protected]>
Co-Authored-By: Lukas Rytz <[email protected]>

Author: 3 people

compiler/src/dotty/tools/backend/jvm/DottyBackendInterface.scala

@@ -706,6 +706,7 @@ class DottyBackendInterface(outputDirectory: AbstractFile, val superCallsMap: Ma
     def isJavaEntryPoint: Boolean = CollectEntryPoints.isJavaEntryPoint(sym)
 
     def isClassConstructor: Boolean = toDenot(sym).isClassConstructor
+    def isSerializable: Boolean = toDenot(sym).isSerializable
 
     /**
      * True for module classes of modules that are top-level or owned only by objects. Module classes
@@ -855,6 +856,9 @@ class DottyBackendInterface(outputDirectory: AbstractFile, val superCallsMap: Ma
 
     def samMethod(): Symbol =
       toDenot(sym).info.abstractTermMembers.headOption.getOrElse(toDenot(sym).info.member(nme.apply)).symbol
+
+    def isFunctionClass: Boolean =
+      defn.isFunctionClass(sym)
   }
 
 

compiler/src/dotty/tools/backend/jvm/GenBCode.scala

@@ -6,6 +6,7 @@ import dotty.tools.dotc.ast.tpd
 import dotty.tools.dotc.core.Phases.Phase
 
 import scala.collection.mutable
+import scala.collection.JavaConverters._
 import scala.tools.asm.CustomAttr
 import scala.tools.nsc.backend.jvm._
 import dotty.tools.dotc.transform.SymUtils._
@@ -23,6 +24,7 @@ import java.io.DataOutputStream
 
 
 import scala.tools.asm
+import scala.tools.asm.Handle
 import scala.tools.asm.tree._
 import tpd._
 import StdNames._
@@ -304,10 +306,98 @@ class GenBCodePipeline(val entryPoints: List[Symbol], val int: DottyBackendInter
     class Worker2 {
       // lazy val localOpt = new LocalOpt(new Settings())
 
-      def localOptimizations(classNode: ClassNode): Unit = {
+      private def localOptimizations(classNode: ClassNode): Unit = {
         // BackendStats.timed(BackendStats.methodOptTimer)(localOpt.methodOptimizations(classNode))
       }
 
+
+      /* Return an array of all serializable lambdas in this class */
+      private def collectSerializableLambdas(classNode: ClassNode): Array[Handle] = {
+        val indyLambdaBodyMethods = new mutable.ArrayBuffer[Handle]
+        for (m <- classNode.methods.asScala) {
+          val iter = m.instructions.iterator
+          while (iter.hasNext) {
+            val insn = iter.next()
+            insn match {
+              case indy: InvokeDynamicInsnNode
+                  // No need to check the exact bsmArgs because we only generate
+                  // altMetafactory indy calls for serializable lambdas.
+                  if indy.bsm == BCodeBodyBuilder.lambdaMetaFactoryAltMetafactoryHandle =>
+                val implMethod = indy.bsmArgs(1).asInstanceOf[Handle]
+                indyLambdaBodyMethods += implMethod
+              case _ =>
+            }
+          }
+        }
+        indyLambdaBodyMethods.toArray
+      }
+
+      /*
+       * Add:
+       *
+       * private static Object $deserializeLambda$(SerializedLambda l) {
+       *   try return indy[scala.runtime.LambdaDeserialize.bootstrap, targetMethodGroup$0](l)
+       *   catch {
+       *     case i: IllegalArgumentException =>
+       *       try return indy[scala.runtime.LambdaDeserialize.bootstrap, targetMethodGroup$1](l)
+       *       catch {
+       *         case i: IllegalArgumentException =>
+       *           ...
+       *             return indy[scala.runtime.LambdaDeserialize.bootstrap, targetMethodGroup${NUM_GROUPS-1}](l)
+       *       }
+       *
+       * We use invokedynamic here to enable caching within the deserializer without needing to
+       * host a static field in the enclosing class. This allows us to add this method to interfaces
+       * that define lambdas in default methods.
+       *
+       * SI-10232 we can't pass arbitrary number of method handles to the final varargs parameter of the bootstrap
+       * method due to a limitation in the JVM. Instead, we emit a separate invokedynamic bytecode for each group of target
+       * methods.
+       */
+      private def addLambdaDeserialize(classNode: ClassNode, implMethodsArray: Array[Handle]): Unit = {
+        import asm.Opcodes._
+        import BCodeBodyBuilder._
+        import bTypes._
+        import coreBTypes._
+
+        val cw = classNode
+
+        // Make sure to reference the ClassBTypes of all types that are used in the code generated
+        // here (e.g. java/util/Map) are initialized. Initializing a ClassBType adds it to
+        // `classBTypeFromInternalNameMap`. When writing the classfile, the asm ClassWriter computes
+        // stack map frames and invokes the `getCommonSuperClass` method. This method expects all
+        // ClassBTypes mentioned in the source code to exist in the map.
+
+        val serlamObjDesc = MethodBType(jliSerializedLambdaRef :: Nil, ObjectReference).descriptor
+
+        val mv = cw.visitMethod(ACC_PRIVATE + ACC_STATIC + ACC_SYNTHETIC, "$deserializeLambda$", serlamObjDesc, null, null)
+        def emitLambdaDeserializeIndy(targetMethods: Seq[Handle]): Unit = {
+          mv.visitVarInsn(ALOAD, 0)
+          mv.visitInvokeDynamicInsn("lambdaDeserialize", serlamObjDesc, lambdaDeserializeBootstrapHandle, targetMethods: _*)
+        }
+
+        val targetMethodGroupLimit = 255 - 1 - 3 // JVM limit. See See MAX_MH_ARITY in CallSite.java
+        val groups: Array[Array[Handle]] = implMethodsArray.grouped(targetMethodGroupLimit).toArray
+        val numGroups = groups.length
+
+        import scala.tools.asm.Label
+        val initialLabels = Array.fill(numGroups - 1)(new Label())
+        val terminalLabel = new Label
+        def nextLabel(i: Int) = if (i == numGroups - 2) terminalLabel else initialLabels(i + 1)
+
+        for ((label, i) <- initialLabels.iterator.zipWithIndex) {
+          mv.visitTryCatchBlock(label, nextLabel(i), nextLabel(i), jlIllegalArgExceptionRef.internalName)
+        }
+        for ((label, i) <- initialLabels.iterator.zipWithIndex) {
+          mv.visitLabel(label)
+          emitLambdaDeserializeIndy(groups(i))
+          mv.visitInsn(ARETURN)
+        }
+        mv.visitLabel(terminalLabel)
+        emitLambdaDeserializeIndy(groups(numGroups - 1))
+        mv.visitInsn(ARETURN)
+      }
+
       def run(): Unit = {
         while (true) {
           val item = q2.poll
@@ -317,7 +407,11 @@ class GenBCodePipeline(val entryPoints: List[Symbol], val int: DottyBackendInter
           }
           else {
             try {
-              localOptimizations(item.plain.classNode)
+              val plainNode = item.plain.classNode
+              localOptimizations(plainNode)
+              val serializableLambdas = collectSerializableLambdas(plainNode)
+              if (serializableLambdas.nonEmpty)
+                addLambdaDeserialize(plainNode, serializableLambdas)
               addToQ3(item)
             } catch {
               case ex: Throwable =>

tests/run/inlineAddDeserializeLambda.scala

@@ -0,0 +1,20 @@
+class C { inline final def f: Int => Int = (x: Int) => x + 1 }
+
+object Test extends App {
+  import java.io._
+
+  def serialize(obj: AnyRef): Array[Byte] = {
+    val buffer = new ByteArrayOutputStream
+    val out = new ObjectOutputStream(buffer)
+    out.writeObject(obj)
+    buffer.toByteArray
+  }
+  def deserialize(a: Array[Byte]): AnyRef = {
+    val in = new ObjectInputStream(new ByteArrayInputStream(a))
+    in.readObject
+  }
+
+  def serializeDeserialize[T <: AnyRef](obj: T) = deserialize(serialize(obj)).asInstanceOf[T]
+
+  assert(serializeDeserialize((new C).f).isInstanceOf[Function1[_, _]])
+}

tests/run/lambda-serialization-gc.scala

@@ -0,0 +1,40 @@
+import java.io._
+
+import java.net.URLClassLoader
+
+class C {
+  def serializeDeserialize[T <: AnyRef](obj: T) = {
+    val buffer = new ByteArrayOutputStream
+    val out = new ObjectOutputStream(buffer)
+    out.writeObject(obj)
+    val in = new ObjectInputStream(new ByteArrayInputStream(buffer.toByteArray))
+    in.readObject.asInstanceOf[T]
+  }
+
+  serializeDeserialize((c: String) => c.length)
+}
+
+object Test {
+  def main(args: Array[String]): Unit = {
+    test()
+  }
+
+  def test(): Unit = {
+    val loader = getClass.getClassLoader.asInstanceOf[URLClassLoader]
+    val loaderCClass = classOf[C]
+    def deserializedInThrowawayClassloader = {
+      val throwawayLoader: java.net.URLClassLoader = new java.net.URLClassLoader(loader.getURLs, ClassLoader.getSystemClassLoader) {
+        val maxMemory = Runtime.getRuntime.maxMemory()
+        val junk = new Array[Long]((maxMemory / 8 / 4).toInt)
+      }
+      val clazz = throwawayLoader.loadClass("C")
+      assert(clazz != loaderCClass)
+      clazz.newInstance()
+    }
+    (1 to 5) foreach { i =>
+      // This would OOM by the fifth iteration if we leaked `throwawayLoader` during
+      // deserialization.
+      deserializedInThrowawayClassloader
+    }
+  }
+}

tests/run/lambda-serialization-others.scala

@@ -0,0 +1,67 @@
+import java.io.{ByteArrayInputStream, ByteArrayOutputStream, ObjectInputStream, ObjectOutputStream, PrintWriter, StringWriter}
+import java.lang.invoke.{MethodHandleInfo, SerializedLambda}
+
+class C1 {
+  val xxlfun = ((x1: Int,
+    x2: String,
+    x3: Int,
+    x4: Int,
+    x5: Int,
+    x6: Int,
+    x7: Int,
+    x8: Int,
+    x9: Int,
+    x10: Int,
+    x11: Int,
+    x12: Int,
+    x13: Int,
+    x14: Int,
+    x15: Int,
+    x16: Int,
+    x17: Int,
+    x18: Int,
+    x19: Int,
+    x20: Int,
+    x21: Int,
+    x22: Int,
+    x23: Int,
+    x24: Int,
+    x25: Int,
+    x26: Int) => x2 + x1)
+
+  val depfun: (x1: Int) => List[x1.type] = x1 => List(x1)
+
+  val erasedfun: erased Int => Int = erased (x1) => 0
+}
+
+class C2
+    extends Serializable /* Needed because of #5866 */ {
+  val impfun: given Int => Int = given x1 => x1
+
+  val impdepfun: given (x1: Int) => List[x1.type] = given x1 => List(x1)
+
+  val erasedimpfun: given erased Int => Int = given erased (x1) => 0
+}
+
+object Test {
+  def main(args: Array[String]): Unit = {
+    val c1 = new C1
+    serializeDeserialize(c1.xxlfun)
+    serializeDeserialize(c1.erasedfun)
+
+    val c2 = new C2
+    serializeDeserialize[given Int => Int](c2.impfun)
+    // Won't compile until #5841 is merged
+    // serializeDeserialize[given Int => Int](c2.impdepfun)
+    serializeDeserialize[given Int => Int](c2.erasedimpfun)
+  }
+
+  def serializeDeserialize[T <: AnyRef](obj: T): Unit = {
+    val buffer = new ByteArrayOutputStream
+    val out = new ObjectOutputStream(buffer)
+    out.writeObject(obj)
+    val in = new ObjectInputStream(new ByteArrayInputStream(buffer.toByteArray))
+    in.readObject.asInstanceOf[T]
+  }
+}
+

tests/run/lambda-serialization-security.scala

@@ -0,0 +1,47 @@
+import java.io.{ByteArrayInputStream, ObjectInputStream, ObjectOutputStream, ByteArrayOutputStream}
+
+trait IntToString extends java.io.Serializable { def apply(i: Int): String }
+
+object Test {
+  def main(args: Array[String]): Unit = {
+    roundTrip()
+    roundTripIndySam()
+  }
+
+  def roundTrip(): Unit = {
+    val c = new Capture("Capture")
+    val lambda = (p: Param) => ("a", p, c)
+    val reconstituted1 = serializeDeserialize(lambda).asInstanceOf[Object => Any]
+    val p = new Param
+    assert(reconstituted1.apply(p) == ("a", p, c))
+    val reconstituted2 = serializeDeserialize(lambda).asInstanceOf[Object => Any]
+    assert(reconstituted1.getClass == reconstituted2.getClass)
+
+    val reconstituted3 = serializeDeserialize(reconstituted1)
+    assert(reconstituted3.apply(p) == ("a", p, c))
+
+    val specializedLambda = (p: Int) => List(p, c).length
+    assert(serializeDeserialize(specializedLambda).apply(42) == 2)
+    assert(serializeDeserialize(serializeDeserialize(specializedLambda)).apply(42) == 2)
+  }
+
+  // lambda targeting a SAM, not a FunctionN (should behave the same way)
+  def roundTripIndySam(): Unit = {
+    val lambda: IntToString = (x: Int) => "yo!" * x
+    val reconstituted1 = serializeDeserialize(lambda).asInstanceOf[IntToString]
+    val reconstituted2 = serializeDeserialize(reconstituted1).asInstanceOf[IntToString]
+    assert(reconstituted1.apply(2) == "yo!yo!")
+    assert(reconstituted1.getClass == reconstituted2.getClass)
+  }
+
+  def serializeDeserialize[T <: AnyRef](obj: T) = {
+    val buffer = new ByteArrayOutputStream
+    val out = new ObjectOutputStream(buffer)
+    out.writeObject(obj)
+    val in = new ObjectInputStream(new ByteArrayInputStream(buffer.toByteArray))
+    in.readObject.asInstanceOf[T]
+  }
+}
+
+case class Capture(s: String) extends Serializable
+class Param

tests/run/lambda-serialization.scala

@@ -0,0 +1,66 @@
+import java.io.{ByteArrayInputStream, ByteArrayOutputStream, ObjectInputStream, ObjectOutputStream, PrintWriter, StringWriter}
+import java.lang.invoke.{MethodHandleInfo, SerializedLambda}
+
+class C extends java.io.Serializable {
+  val fs = List(
+    () => (), () => (), () => (), () => (), () => (), () => (), () => (), () => (), () => (), () => () ,() => (), () => (), () => (), () => (), () => (),
+    () => (), () => (), () => (), () => (), () => (), () => (), () => (), () => (), () => (), () => () ,() => (), () => (), () => (), () => (), () => (),
+    () => (), () => (), () => (), () => (), () => (), () => (), () => (), () => (), () => (), () => () ,() => (), () => (), () => (), () => (), () => (),
+    () => (), () => (), () => (), () => (), () => (), () => (), () => (), () => (), () => (), () => () ,() => (), () => (), () => (), () => (), () => (),
+    () => (), () => (), () => (), () => (), () => (), () => (), () => (), () => (), () => (), () => () ,() => (), () => (), () => (), () => (), () => (),
+    () => (), () => (), () => (), () => (), () => (), () => (), () => (), () => (), () => (), () => () ,() => (), () => (), () => (), () => (), () => (),
+    () => (), () => (), () => (), () => (), () => (), () => (), () => (), () => (), () => (), () => () ,() => (), () => (), () => (), () => (), () => (),
+    () => (), () => (), () => (), () => (), () => (), () => (), () => (), () => (), () => (), () => () ,() => (), () => (), () => (), () => (), () => (),
+    () => (), () => (), () => (), () => (), () => (), () => (), () => (), () => (), () => (), () => () ,() => (), () => (), () => (), () => (), () => (),
+    () => (), () => (), () => (), () => (), () => (), () => (), () => (), () => (), () => (), () => () ,() => (), () => (), () => (), () => (), () => (),
+    () => (), () => (), () => (), () => (), () => (), () => (), () => (), () => (), () => (), () => () ,() => (), () => (), () => (), () => (), () => (),
+    () => (), () => (), () => (), () => (), () => (), () => (), () => (), () => (), () => (), () => () ,() => (), () => (), () => (), () => (), () => (),
+    () => (), () => (), () => (), () => (), () => (), () => (), () => (), () => (), () => (), () => () ,() => (), () => (), () => (), () => (), () => (),
+    () => (), () => (), () => (), () => (), () => (), () => (), () => (), () => (), () => (), () => () ,() => (), () => (), () => (), () => (), () => (),
+    () => (), () => (), () => (), () => (), () => (), () => (), () => (), () => (), () => (), () => () ,() => (), () => (), () => (), () => (), () => ()
+  )
+  private def foo(): Unit = {
+    assert(false, "should not be called!!!")
+  }
+}
+
+trait FakeSam { def apply(): Unit }
+
+object Test {
+  def main(args: Array[String]): Unit = {
+    allRealLambdasRoundTrip()
+    fakeLambdaFailsToDeserialize()
+  }
+
+  def allRealLambdasRoundTrip(): Unit = {
+    new C().fs.map(x => serializeDeserialize(x).apply())
+  }
+
+  def fakeLambdaFailsToDeserialize(): Unit = {
+    val fake = new SerializedLambda(classOf[C], classOf[FakeSam].getName, "apply", "()V",
+      MethodHandleInfo.REF_invokeVirtual, classOf[C].getName, "foo", "()V", "()V", Array(new C))
+    try {
+      serializeDeserialize(fake).asInstanceOf[FakeSam].apply()
+      assert(false)
+    } catch {
+      case ex: Exception =>
+        val stackTrace = stackTraceString(ex)
+        assert(stackTrace.contains("Illegal lambda deserialization"), stackTrace)
+    }
+  }
+
+  def serializeDeserialize[T <: AnyRef](obj: T) = {
+    val buffer = new ByteArrayOutputStream
+    val out = new ObjectOutputStream(buffer)
+    out.writeObject(obj)
+    val in = new ObjectInputStream(new ByteArrayInputStream(buffer.toByteArray))
+    in.readObject.asInstanceOf[T]
+  }
+
+  def stackTraceString(ex: Throwable): String = {
+    val writer = new StringWriter
+    ex.printStackTrace(new PrintWriter(writer))
+    writer.toString
+  }
+}
+