Merge pull request #4289 from dotty-staging/fix-1692

Fix #1692: Null out fields after use in lazy initialization

Author: allanrenucci

compiler/src/dotty/tools/dotc/Compiler.scala

@@ -87,6 +87,7 @@ class Compiler {
          new InterceptedMethods,     // Special handling of `==`, `|=`, `getClass` methods
          new Getters,                // Replace non-private vals and vars with getter defs (fields are added later)
          new ElimByName,             // Expand by-name parameter references
+         new CollectNullableFields,  // Collect fields that can be nulled out after use in lazy initialization
          new ElimOuterSelect,        // Expand outer selections
          new AugmentScala2Traits,    // Expand traits defined in Scala 2.x to simulate old-style rewritings
          new ResolveSuper,           // Implement super accessors and add forwarders to trait methods
@@ -97,7 +98,7 @@ class Compiler {
     List(new Erasure) ::             // Rewrite types to JVM model, erasing all type parameters, abstract types and refinements.
     List(new ElimErasedValueType,    // Expand erased value types to their underlying implmementation types
          new VCElideAllocations,     // Peep-hole optimization to eliminate unnecessary value class allocations
-         new Mixin,                   // Expand trait fields and trait initializers
+         new Mixin,                  // Expand trait fields and trait initializers
          new LazyVals,               // Expand lazy vals
          new Memoize,                // Add private fields to getters and setters
          new NonLocalReturns,        // Expand non-local returns

compiler/src/dotty/tools/dotc/core/Phases.scala

@@ -211,6 +211,7 @@ object Phases {
     private[this] var myTyperPhase: Phase = _
     private[this] var mySbtExtractDependenciesPhase: Phase = _
     private[this] var myPicklerPhase: Phase = _
+    private[this] var myCollectNullableFieldsPhase: Phase = _
     private[this] var myRefChecksPhase: Phase = _
     private[this] var myPatmatPhase: Phase = _
     private[this] var myElimRepeatedPhase: Phase = _
@@ -226,6 +227,7 @@ object Phases {
     final def typerPhase = myTyperPhase
     final def sbtExtractDependenciesPhase = mySbtExtractDependenciesPhase
     final def picklerPhase = myPicklerPhase
+    final def collectNullableFieldsPhase = myCollectNullableFieldsPhase
     final def refchecksPhase = myRefChecksPhase
     final def patmatPhase = myPatmatPhase
     final def elimRepeatedPhase = myElimRepeatedPhase
@@ -244,6 +246,7 @@ object Phases {
       myTyperPhase = phaseOfClass(classOf[FrontEnd])
       mySbtExtractDependenciesPhase = phaseOfClass(classOf[sbt.ExtractDependencies])
       myPicklerPhase = phaseOfClass(classOf[Pickler])
+      myCollectNullableFieldsPhase = phaseOfClass(classOf[CollectNullableFields])
       myRefChecksPhase = phaseOfClass(classOf[RefChecks])
       myElimRepeatedPhase = phaseOfClass(classOf[ElimRepeated])
       myExtensionMethodsPhase = phaseOfClass(classOf[ExtensionMethods])

compiler/src/dotty/tools/dotc/transform/CollectNullableFields.scala

@@ -0,0 +1,112 @@
+package dotty.tools.dotc.transform
+
+import dotty.tools.dotc.ast.tpd
+import dotty.tools.dotc.core.Contexts.Context
+import dotty.tools.dotc.core.Flags._
+import dotty.tools.dotc.core.Symbols.Symbol
+import dotty.tools.dotc.core.Types.{Type, ExprType}
+import dotty.tools.dotc.transform.MegaPhase.MiniPhase
+import dotty.tools.dotc.transform.SymUtils._
+
+import scala.collection.JavaConverters._
+import scala.collection.mutable
+
+import java.util.IdentityHashMap
+
+object CollectNullableFields {
+  val name = "collectNullableFields"
+}
+
+/** Collect fields that can be nulled out after use in lazy initialization.
+ *
+ *  This information is used during lazy val transformation to assign null to private
+ *  fields that are only used within a lazy val initializer. This is not just an optimization,
+ *  but is needed for correctness to prevent memory leaks. E.g.
+ *
+ *  ```scala
+ *  class TestByNameLazy(byNameMsg: => String) {
+ *    lazy val byLazyValMsg = byNameMsg
+ *  }
+ *  ```
+ *
+ *  Here `byNameMsg` should be null out once `byLazyValMsg` is
+ *  initialised.
+ *
+ *  A field is nullable if all the conditions below hold:
+ *    - belongs to a non trait-class
+ *    - is private[this]
+ *    - is not lazy
+ *    - its type is nullable
+ *    - is only used in a lazy val initializer
+ *    - defined in the same class as the lazy val
+ */
+class CollectNullableFields extends MiniPhase {
+  import tpd._
+
+  override def phaseName = CollectNullableFields.name
+
+  /** Running after `ElimByName` to see by names as nullable types. */
+  override def runsAfter = Set(ElimByName.name)
+
+  private[this] sealed trait FieldInfo
+  private[this] case object NotNullable extends FieldInfo
+  private[this] case class Nullable(by: Symbol) extends FieldInfo
+
+  /** Whether or not a field is nullable */
+  private[this] var nullability: IdentityHashMap[Symbol, FieldInfo] = _
+
+  override def prepareForUnit(tree: Tree)(implicit ctx: Context) = {
+    if (nullability == null) nullability = new IdentityHashMap
+    ctx
+  }
+
+  private def recordUse(tree: Tree)(implicit ctx: Context): Tree = {
+    val sym = tree.symbol
+    val isNullablePrivateField =
+      sym.isField &&
+      !sym.is(Lazy) &&
+      !sym.owner.is(Trait) &&
+      sym.initial.is(PrivateLocal) &&
+      sym.info.widenDealias.typeSymbol.isNullableClass
+
+    if (isNullablePrivateField)
+      nullability.get(sym) match {
+        case Nullable(from) if from != ctx.owner => // used in multiple lazy val initializers
+          nullability.put(sym, NotNullable)
+        case null => // not in the map
+          val from = ctx.owner
+          val isNullable =
+            from.is(Lazy, butNot = Module) && // is lazy val
+            from.owner.isClass &&             // is field
+            from.owner.eq(sym.owner)          // is lazy val and field defined in the same class
+          val info = if (isNullable) Nullable(from) else NotNullable
+          nullability.put(sym, info)
+        case _ =>
+          // Do nothing for:
+          //  - NotNullable
+          //  - Nullable(ctx.owner)
+      }
+
+    tree
+  }
+
+  override def transformIdent(tree: Ident)(implicit ctx: Context) =
+    recordUse(tree)
+
+  override def transformSelect(tree: Select)(implicit ctx: Context) =
+    recordUse(tree)
+
+  /** Map lazy values to the fields they should null after initialization. */
+  def lazyValNullables(implicit ctx: Context): IdentityHashMap[Symbol, mutable.ListBuffer[Symbol]] = {
+    val result = new IdentityHashMap[Symbol, mutable.ListBuffer[Symbol]]
+
+    nullability.forEach {
+      case (sym, Nullable(from)) =>
+        val bldr = result.computeIfAbsent(from, _ => new mutable.ListBuffer)
+        bldr += sym
+      case _ =>
+    }
+
+    result
+  }
+}