Fix #5340: Improve runtime implementation of local lazy vals

Reduce the amount of code generated for local lazy val by moving some of
the generated code to the runtime library.

This is a port of scala/scala#5374

Author: allanrenucci

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

@@ -897,10 +897,25 @@ class Definitions {
   // ----- Symbol sets ---------------------------------------------------
 
   lazy val AbstractFunctionType: Array[TypeRef] = mkArityArray("scala.runtime.AbstractFunction", MaxImplementedFunctionArity, 0)
-  val AbstractFunctionClassPerRun: PerRun[Array[Symbol]] = new PerRun[Array[Symbol]](implicit ctx => AbstractFunctionType.map(_.symbol.asClass))
+  val AbstractFunctionClassPerRun: PerRun[Array[Symbol]] = new PerRun(implicit ctx => AbstractFunctionType.map(_.symbol.asClass))
   def AbstractFunctionClass(n: Int)(implicit ctx: Context): Symbol = AbstractFunctionClassPerRun()(ctx)(n)
   private lazy val ImplementedFunctionType = mkArityArray("scala.Function", MaxImplementedFunctionArity, 0)
-  def FunctionClassPerRun: PerRun[Array[Symbol]] = new PerRun[Array[Symbol]](implicit ctx => ImplementedFunctionType.map(_.symbol.asClass))
+  def FunctionClassPerRun: PerRun[Array[Symbol]] = new PerRun(implicit ctx => ImplementedFunctionType.map(_.symbol.asClass))
+
+  val LazyHolder: PerRun[Map[Symbol, Symbol]] = new PerRun({ implicit ctx =>
+    def holderImpl(holderType: String) = ctx.requiredClass("scala.runtime." + holderType)
+    Map[Symbol, Symbol](
+      IntClass     -> holderImpl("LazyInt"),
+      LongClass    -> holderImpl("LazyLong"),
+      BooleanClass -> holderImpl("LazyBoolean"),
+      FloatClass   -> holderImpl("LazyFloat"),
+      DoubleClass  -> holderImpl("LazyDouble"),
+      ByteClass    -> holderImpl("LazyByte"),
+      CharClass    -> holderImpl("LazyChar"),
+      ShortClass   -> holderImpl("LazyShort")
+    )
+    .withDefaultValue(holderImpl("LazyRef"))
+  })
 
   lazy val TupleType: Array[TypeRef] = mkArityArray("scala.Tuple", MaxTupleArity, 1)
 

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

@@ -129,50 +129,58 @@ class LazyVals extends MiniPhase with IdentityDenotTransformer {
     Thicket(field, getter)
   }
 
-  /** Replace a local lazy val inside a method,
-    * with a LazyHolder from
-    * dotty.runtime(eg dotty.runtime.LazyInt)
-    */
+  /** Desugar a local `lazy val x: Int = <RHS>` into:
+   *
+   *  ```
+   *  val x$lzy = new scala.runtime.LazyInt()
+   *
+   *  def x$lzycompute(): Int = x$lzy.synchronized {
+   *    if (x$lzy.initialized()) x$lzy.value()
+   *    else x$lzy.initialize(<RHS>)
+   *      // for a Unit-typed lazy val, this becomes `{ rhs ; x$lzy.initialize() }`
+   *      // to avoid passing around BoxedUnit
+   *  }
+   *
+   *  def x(): Int = if (x$lzy.initialized()) x$lzy.value() else x$lzycompute()
+   *  ```
+   *
+   *  TODO: Implement Unit-typed lazy val optimization described above
+   */
   def transformLocalDef(x: ValOrDefDef)(implicit ctx: Context): Thicket = {
-      val valueInitter = x.rhs
-      val xname = x.name.asTermName
-      val holderName = LazyLocalName.fresh(xname)
-      val initName = LazyLocalInitName.fresh(xname)
-      val tpe = x.tpe.widen.resultType.widen
-
-      val holderType =
-        if (tpe isRef defn.IntClass) "LazyInt"
-        else if (tpe isRef defn.LongClass) "LazyLong"
-        else if (tpe isRef defn.BooleanClass) "LazyBoolean"
-        else if (tpe isRef defn.FloatClass) "LazyFloat"
-        else if (tpe isRef defn.DoubleClass) "LazyDouble"
-        else if (tpe isRef defn.ByteClass) "LazyByte"
-        else if (tpe isRef defn.CharClass) "LazyChar"
-        else if (tpe isRef defn.ShortClass) "LazyShort"
-        else "LazyRef"
-
-
-      val holderImpl = ctx.requiredClass("dotty.runtime." + holderType)
-
-      val holderSymbol = ctx.newSymbol(x.symbol.owner, holderName, containerFlags, holderImpl.typeRef, coord = x.pos)
-      val initSymbol = ctx.newSymbol(x.symbol.owner, initName, initFlags, MethodType(Nil, tpe), coord = x.pos)
-      val result = ref(holderSymbol).select(lazyNme.value).withPos(x.pos)
-      val flag = ref(holderSymbol).select(lazyNme.initialized)
-      val initer = valueInitter.changeOwnerAfter(x.symbol, initSymbol, this)
-      val initBody =
-        adaptToType(
-          ref(holderSymbol).select(defn.Object_synchronized).appliedTo(
-            adaptToType(mkNonThreadSafeDef(result, flag, initer, nullables = Nil), defn.ObjectType)),
-          tpe)
-      val initTree = DefDef(initSymbol, initBody)
-      val holderTree = ValDef(holderSymbol, New(holderImpl.typeRef, List()))
-      val methodBody = tpd.If(flag.ensureApplied,
-        result.ensureApplied,
-        ref(initSymbol).ensureApplied).ensureConforms(tpe)
-
-      val methodTree = DefDef(x.symbol.asTerm, methodBody)
-      ctx.debuglog(s"found a lazy val ${x.show},\nrewrote with ${holderTree.show}")
-      Thicket(holderTree, initTree, methodTree)
+    val xname = x.name.asTermName
+    val holderName = LazyLocalName.fresh(xname)
+    val initName = LazyLocalInitName.fresh(xname)
+    val tpe = x.tpe.widen.resultType.widen
+
+    // val x$lzy = new scala.runtime.LazyInt()
+    val holderImpl = defn.LazyHolder()(ctx)(tpe.typeSymbol)
+    val holderSymbol = ctx.newSymbol(x.symbol.owner, holderName, containerFlags, holderImpl.typeRef, coord = x.pos)
+    val holderTree = ValDef(holderSymbol, New(holderImpl.typeRef, List()))
+
+    val holderRef = ref(holderSymbol)
+    val getValue = holderRef.select(lazyNme.value).ensureApplied.withPos(x.pos)
+    val initialized = holderRef.select(lazyNme.initialized).ensureApplied
+
+    // def x$lzycompute(): Int = x$lzy.synchronized {
+    //   if (x$lzy.initialized()) x$lzy.value()
+    //   else x$lzy.initialize(<RHS>)
+    // }
+    val initSymbol = ctx.newSymbol(x.symbol.owner, initName, initFlags, MethodType(Nil, tpe), coord = x.pos)
+    val rhs = x.rhs.changeOwnerAfter(x.symbol, initSymbol, this)
+    val initialize = holderRef.select(lazyNme.initialize).appliedTo(rhs)
+    val initBody =
+      adaptToType(
+        holderRef.select(defn.Object_synchronized).appliedTo(
+          adaptToType(If(initialized, getValue, initialize), defn.ObjectType)),
+        tpe)
+    val initTree = DefDef(initSymbol, initBody)
+
+    // def x(): Int = if (x$lzy.initialized()) x$lzy.value() else x$lzycompute()
+    val accessorBody = If(initialized, getValue, ref(initSymbol).ensureApplied).ensureConforms(tpe)
+    val accessor = DefDef(x.symbol.asTerm, accessorBody)
+
+    ctx.debuglog(s"found a lazy val ${x.show},\nrewrote with ${holderTree.show}")
+    Thicket(holderTree, initTree, accessor)
   }
 
 
@@ -458,6 +466,7 @@ object LazyVals {
     val result: TermName      = "result".toTermName
     val value: TermName       = "value".toTermName
     val initialized: TermName = "initialized".toTermName
+    val initialize: TermName  = "initialize".toTermName
     val retry: TermName       = "retry".toTermName
   }
 }

library/src/dotty/runtime/LazyHolders.scala

@@ -0,0 +1,8 @@
+Int
+Long
+Float
+Double
+Byte
+Char
+Short
+String

tests/run/i5340.check

@@ -0,0 +1,23 @@
+object Test {
+  def main(args: Array[String]): Unit = {
+    lazy val a: Int = { println("Int"); 1 }
+    lazy val b: Long = { println("Long"); 1L }
+    lazy val c: Float = { println("Float"); 1F }
+    lazy val d: Double = { println("Double"); 1.0 }
+    lazy val e: Byte = { println("Byte"); 1 }
+    lazy val f: Char = { println("Char"); '1' }
+    lazy val g: Short = { println("Short"); 1 }
+    lazy val h: String = { println("String"); "1" }
+    lazy val i: Unit = { println("Unit"); () }
+
+    a
+    b
+    c
+    d
+    e
+    f
+    g
+    h
+    // i // FIXME: See #5350
+  }
+}