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

allanrenucci · allanrenucci · commit b4bcd050d60a · 2018-04-11T11:44:00.000+02:00
Private fields that are only used during lzyy val initialization can be
assigned null once the lazy val is initialized. This is not just an
optimization, but is needed for correctness to prevent memory leaks.
diff --git a/compiler/src/dotty/tools/dotc/Compiler.scala b/compiler/src/dotty/tools/dotc/Compiler.scala
@@ -73,6 +73,7 @@ class Compiler {
          new LiftTry,                // Put try expressions that might execute on non-empty stacks into their own methods
          new HoistSuperArgs,         // Hoist complex arguments of supercalls to enclosing scope
          new ClassOf,                // Expand `Predef.classOf` calls.
+         new CollectNullableFields,  // Collect fields that can be null out after use in lazy initialization
          new RefChecks) ::           // Various checks mostly related to abstract members and overriding
     List(new TryCatchPatterns,       // Compile cases in try/catch
          new PatternMatcher,         // Compile pattern matches
@@ -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
diff --git a/compiler/src/dotty/tools/dotc/core/Phases.scala b/compiler/src/dotty/tools/dotc/core/Phases.scala
@@ -210,6 +210,7 @@ object Phases {
 
     private[this] var myTyperPhase: 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 = _
@@ -224,6 +225,7 @@ object Phases {
 
     final def typerPhase = myTyperPhase
     final def picklerPhase = myPicklerPhase
+    final def collectNullableFieldsPhase = myCollectNullableFieldsPhase
     final def refchecksPhase = myRefChecksPhase
     final def patmatPhase = myPatmatPhase
     final def elimRepeatedPhase = myElimRepeatedPhase
@@ -241,6 +243,7 @@ object Phases {
 
       myTyperPhase = phaseOfClass(classOf[FrontEnd])
       myPicklerPhase = phaseOfClass(classOf[Pickler])
+      myCollectNullableFieldsPhase = phaseOfClass(classOf[CollectNullableFields])
       myRefChecksPhase = phaseOfClass(classOf[RefChecks])
       myElimRepeatedPhase = phaseOfClass(classOf[ElimRepeated])
       myExtensionMethodsPhase = phaseOfClass(classOf[ExtensionMethods])
diff --git a/compiler/src/dotty/tools/dotc/transform/CollectNullableFields.scala b/compiler/src/dotty/tools/dotc/transform/CollectNullableFields.scala
@@ -0,0 +1,97 @@
+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 null 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.
+ *
+ *  {{{
+ *  class TestByNameLazy(byNameMsg: => String) {
+ *    lazy val byLazyValMsg = byNameMsg
+ *  }
+ *  }}}
+ *
+ *  Here `byNameMsg` should be null out once `byLazyValMsg` is
+ *  initialised.
+ */
+class CollectNullableFields extends MiniPhase {
+  import tpd._
+
+  override def phaseName = CollectNullableFields.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) = {
+    nullability = new IdentityHashMap
+    ctx
+  }
+
+  private def recordUse(tree: Tree)(implicit ctx: Context): Tree = {
+    val sym = tree.symbol
+
+    def isNullableType(tpe: Type) =
+      tpe.isInstanceOf[ExprType] ||
+      tpe.widenDealias.typeSymbol.isNullableClass
+    val isNullablePrivateField = sym.isField && sym.is(Private, butNot = Lazy) && isNullableType(sym.info)
+
+    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 inLazyValInitializer = from.is(Lazy, butNot = Module)
+          val info = if (inLazyValInitializer) 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): Map[Symbol, List[Symbol]] = {
+    val result = new mutable.HashMap[Symbol, mutable.ListBuffer[Symbol]]
+
+    nullability.forEach {
+      case (sym, Nullable(from)) =>
+        val bldr = result.getOrElseUpdate(from, new mutable.ListBuffer)
+        bldr += sym
+      case _ =>
+    }
+
+    result.mapValues(_.toList).toMap
+  }
+}
diff --git a/compiler/src/dotty/tools/dotc/transform/LazyVals.scala b/compiler/src/dotty/tools/dotc/transform/LazyVals.scala
@@ -39,7 +39,7 @@ class LazyVals extends MiniPhase with IdentityDenotTransformer {
 
   /** List of names of phases that should have finished processing of tree
     * before this phase starts processing same tree */
-  override def runsAfter = Set(Mixin.name)
+  override def runsAfter = Set(Mixin.name, CollectNullableFields.name)
 
   override def changesMembers = true  // the phase adds lazy val accessors
 
@@ -50,6 +50,15 @@ class LazyVals extends MiniPhase with IdentityDenotTransformer {
 
   val containerFlagsMask = Flags.Method | Flags.Lazy | Flags.Accessor | Flags.Module
 
+  /** A map of lazy values to the fields they should null after initialization. */
+  private[this] var lazyValNullables: Map[Symbol, List[Symbol]] = _
+  private def nullableFor(sym: Symbol) = lazyValNullables.getOrElse(sym, Nil)
+
+  override def prepareForUnit(tree: Tree)(implicit ctx: Context) = {
+    lazyValNullables = ctx.collectNullableFieldsPhase.asInstanceOf[CollectNullableFields].lazyValNullables
+    ctx
+  }
+
   override def transformDefDef(tree: tpd.DefDef)(implicit ctx: Context): tpd.Tree =
    transformLazyVal(tree)
 
@@ -150,7 +159,7 @@ class LazyVals extends MiniPhase with IdentityDenotTransformer {
         val initBody =
           adaptToType(
             ref(holderSymbol).select(defn.Object_synchronized).appliedTo(
-              adaptToType(mkNonThreadSafeDef(result, flag, initer), defn.ObjectType)),
+              adaptToType(mkNonThreadSafeDef(result, flag, initer, nullableFor(x.symbol)), defn.ObjectType)),
             tpe)
         val initTree = DefDef(initSymbol, initBody)
         val holderTree = ValDef(holderSymbol, New(holderImpl.typeRef, List()))
@@ -176,37 +185,50 @@ class LazyVals extends MiniPhase with IdentityDenotTransformer {
     holders:::stats
   }
 
+  private def nullOut(nullables: List[Symbol])(implicit ctx: Context): List[Tree] = {
+    val nullConst = Literal(Constants.Constant(null))
+    nullables.map { sym =>
+      val field = if (sym.isGetter) sym.field else sym
+      field.setFlag(Flags.Mutable)
+      ref(field).becomes(nullConst)
+    }
+  }
+
   /** Create non-threadsafe lazy accessor equivalent to such code
       * def methodSymbol() = {
       *   if (flag) target
       *   else {
       *     target = rhs
       *     flag = true
+      *     nullable = null
       *     target
       *     }
       *   }
       */
 
-    def mkNonThreadSafeDef(target: Tree, flag: Tree, rhs: Tree)(implicit ctx: Context) = {
+    def mkNonThreadSafeDef(target: Tree, flag: Tree, rhs: Tree, nullables: List[Symbol])(implicit ctx: Context) = {
       val setFlag = flag.becomes(Literal(Constants.Constant(true)))
-      val setTargets = if (isWildcardArg(rhs)) Nil else target.becomes(rhs) :: Nil
-      val init = Block(setFlag :: setTargets, target.ensureApplied)
+      val setNullables = nullOut(nullables)
+      val setTargetAndNullable = if (isWildcardArg(rhs)) setNullables else target.becomes(rhs) :: setNullables
+      val init = Block(setFlag :: setTargetAndNullable, target.ensureApplied)
       If(flag.ensureApplied, target.ensureApplied, init)
     }
 
     /** Create non-threadsafe lazy accessor for not-nullable types  equivalent to such code
       * def methodSymbol() = {
       *   if (target eq null) {
       *     target = rhs
+      *     nullable = null
       *     target
       *   } else target
       * }
       */
-    def mkDefNonThreadSafeNonNullable(target: Symbol, rhs: Tree)(implicit ctx: Context) = {
+    def mkDefNonThreadSafeNonNullable(target: Symbol, rhs: Tree, nullables: List[Symbol])(implicit ctx: Context) = {
       val cond = ref(target).select(nme.eq).appliedTo(Literal(Constant(null)))
       val exp = ref(target)
       val setTarget = exp.becomes(rhs)
-      val init = Block(List(setTarget), exp)
+      val setNullables = nullOut(nullables)
+      val init = Block(setTarget :: setNullables, exp)
       If(cond, init, exp)
     }
 
@@ -222,14 +244,14 @@ class LazyVals extends MiniPhase with IdentityDenotTransformer {
 
         val containerTree = ValDef(containerSymbol, defaultValue(tpe))
         if (x.tpe.isNotNull && tpe <:< defn.ObjectType) { // can use 'null' value instead of flag
-          val slowPath = DefDef(x.symbol.asTerm, mkDefNonThreadSafeNonNullable(containerSymbol, x.rhs))
+          val slowPath = DefDef(x.symbol.asTerm, mkDefNonThreadSafeNonNullable(containerSymbol, x.rhs, nullableFor(x.symbol)))
           Thicket(containerTree, slowPath)
         }
         else {
           val flagName = LazyBitMapName.fresh(x.name.asTermName)
           val flagSymbol = ctx.newSymbol(x.symbol.owner, flagName,  containerFlags | Flags.Private, defn.BooleanType).enteredAfter(this)
           val flag = ValDef(flagSymbol, Literal(Constants.Constant(false)))
-          val slowPath = DefDef(x.symbol.asTerm, mkNonThreadSafeDef(ref(containerSymbol), ref(flagSymbol), x.rhs))
+          val slowPath = DefDef(x.symbol.asTerm, mkNonThreadSafeDef(ref(containerSymbol), ref(flagSymbol), x.rhs, nullableFor(x.symbol)))
           Thicket(containerTree, flag, slowPath)
         }
     }
@@ -263,10 +285,23 @@ class LazyVals extends MiniPhase with IdentityDenotTransformer {
       *         result = $target
       *       }
       *     }
+      *   nullable = null
       *   result
       * }
       */
-    def mkThreadSafeDef(methodSymbol: TermSymbol, claz: ClassSymbol, ord: Int, target: Symbol, rhs: Tree, tp: Types.Type, offset: Tree, getFlag: Tree, stateMask: Tree, casFlag: Tree, setFlagState: Tree, waitOnLock: Tree)(implicit ctx: Context) = {
+    def mkThreadSafeDef(methodSymbol: TermSymbol,
+                        claz: ClassSymbol,
+                        ord: Int,
+                        target: Symbol,
+                        rhs: Tree,
+                        tp: Types.Type,
+                        offset: Tree,
+                        getFlag: Tree,
+                        stateMask: Tree,
+                        casFlag: Tree,
+                        setFlagState: Tree,
+                        waitOnLock: Tree,
+                        nullables: List[Symbol])(implicit ctx: Context) = {
       val initState = Literal(Constants.Constant(0))
       val computeState = Literal(Constants.Constant(1))
       val notifyState = Literal(Constants.Constant(2))
@@ -330,7 +365,8 @@ class LazyVals extends MiniPhase with IdentityDenotTransformer {
 
       val whileBody = List(ref(flagSymbol).becomes(getFlag.appliedTo(thiz, offset)), cases)
       val cycle = WhileDo(methodSymbol, whileCond, whileBody)
-      DefDef(methodSymbol, Block(resultDef :: retryDef :: flagDef :: cycle :: Nil, ref(resultSymbol)))
+      val setNullables = nullOut(nullables)
+      DefDef(methodSymbol, Block(resultDef :: retryDef :: flagDef :: cycle :: setNullables, ref(resultSymbol)))
     }
 
     def transformMemberDefVolatile(x: ValOrDefDef)(implicit ctx: Context) = {
@@ -391,7 +427,7 @@ class LazyVals extends MiniPhase with IdentityDenotTransformer {
         val state =   Select(ref(helperModule), lazyNme.RLazyVals.state)
         val cas =     Select(ref(helperModule), lazyNme.RLazyVals.cas)
 
-        val accessor = mkThreadSafeDef(x.symbol.asTerm, claz, ord, containerSymbol, x.rhs, tpe, offset, getFlag, state, cas, setFlag, wait)
+        val accessor = mkThreadSafeDef(x.symbol.asTerm, claz, ord, containerSymbol, x.rhs, tpe, offset, getFlag, state, cas, setFlag, wait, nullableFor(x.symbol))
         if (flag eq EmptyTree)
           Thicket(containerTree, accessor)
         else Thicket(containerTree, flag, accessor)
