Make some tree fields lazy

Lazy fields are
 - the rhs field of a ValDef or DefDef
 - the body field of a Template

These can be instantiated with Lazy instances.
The scheme is such that lazy fields are completely
transparent for users of the Trees API.

The only downside is that the parameter used to initialize
a potentially lazy field has a weak type (now it's Any, with
Dotty it would be a union type of the form `T | Lazy[T]`.
Therefore, the parameter cannot be recovered through pattern
matching.

6 files changed, 71 insertions, 37 deletions

diff --git a/src/dotty/tools/dotc/ast/Desugar.scala b/src/dotty/tools/dotc/ast/Desugar.scala
index cd198818a..60ca45f28 100644
--- a/src/dotty/tools/dotc/ast/Desugar.scala
+++ b/src/dotty/tools/dotc/ast/Desugar.scala
@@ -218,7 +218,7 @@ object desugar {
 
   /** The expansion of a class definition. See inline comments for what is involved */
   def classDef(cdef: TypeDef)(implicit ctx: Context): Tree = {
-    val TypeDef(name, impl @ Template(constr0, parents, self, body)) = cdef
+    val TypeDef(name, impl @ Template(constr0, parents, self, _)) = cdef
     val mods = cdef.mods
 
     val (constr1, defaultGetters) = defDef(constr0, isPrimaryConstructor = true) match {
@@ -242,7 +242,7 @@ object desugar {
     val constr = cpy.DefDef(constr1)(tparams = constrTparams, vparamss = constrVparamss)
 
     // Add constructor type parameters to auxiliary constructors
-    val normalizedBody = body map {
+    val normalizedBody = impl.body map {
       case ddef: DefDef if ddef.name.isConstructorName =>
         cpy.DefDef(ddef)(tparams = constrTparams)
       case stat =>
@@ -425,10 +425,10 @@ object desugar {
       val modul = ValDef(name, clsRef, New(clsRef, Nil))
         .withMods(mods | ModuleCreationFlags)
         .withPos(mdef.pos)
-      val ValDef(selfName, selfTpt, selfRhs) = tmpl.self
+      val ValDef(selfName, selfTpt, _) = tmpl.self
       val selfMods = tmpl.self.mods
       if (!selfTpt.isEmpty) ctx.error("object definition may not have a self type", tmpl.self.pos)
-      val clsSelf = ValDef(selfName, SingletonTypeTree(Ident(name)), selfRhs)
+      val clsSelf = ValDef(selfName, SingletonTypeTree(Ident(name)), tmpl.self.rhs)
         .withMods(selfMods)
         .withPos(tmpl.self.pos orElse tmpl.pos.startPos)
       val clsTmpl = cpy.Template(tmpl)(self = clsSelf, body = tmpl.body)
diff --git a/src/dotty/tools/dotc/ast/TreeInfo.scala b/src/dotty/tools/dotc/ast/TreeInfo.scala
index 3d633c58d..6ebf29ffa 100644
--- a/src/dotty/tools/dotc/ast/TreeInfo.scala
+++ b/src/dotty/tools/dotc/ast/TreeInfo.scala
@@ -28,8 +28,7 @@ trait TreeInfo[T >: Untyped <: Type] { self: Trees.Instance[T] =>
    */
   def isPureInterfaceMember(tree: Tree): Boolean = unsplice(tree) match {
     case EmptyTree | Import(_, _) | TypeDef(_, _) => true
-    case DefDef(_, _, _, _, rhs) => rhs.isEmpty
-    case ValDef(_, _, rhs) => rhs.isEmpty
+    case defn: ValOrDefDef => defn.rhs.isEmpty
     case _ => false
   }
 
@@ -91,7 +90,7 @@ trait TreeInfo[T >: Untyped <: Type] { self: Trees.Instance[T] =>
 
   /** If tree is a closure, it's body, otherwise tree itself */
   def closureBody(tree: tpd.Tree): tpd.Tree = tree match {
-    case Block(DefDef(nme.ANON_FUN, _, _, _, rhs) :: Nil, Closure(_, _, _)) => rhs
+    case Block((meth @ DefDef(nme.ANON_FUN, _, _, _, _)) :: Nil, Closure(_, _, _)) => meth.rhs
     case _ => tree
   }
 
@@ -287,8 +286,8 @@ trait TypedTreeInfo extends TreeInfo[Type] { self: Trees.Instance[Type] =>
        | Import(_, _)
        | DefDef(_, _, _, _, _) =>
       Pure
-    case vdef @ ValDef(_, _, rhs) =>
-      if (vdef.mods is Mutable) Impure else exprPurity(rhs)
+    case vdef @ ValDef(_, _, _) =>
+      if (vdef.mods is Mutable) Impure else exprPurity(vdef.rhs)
     case _ =>
       Impure
   }
@@ -478,7 +477,7 @@ trait TypedTreeInfo extends TreeInfo[Type] { self: Trees.Instance[Type] =>
           if (stats exists (definedSym(_) == sym)) stats else Nil
         encl match {
           case Block(stats, _) => verify(stats)
-          case Template(_, _, _, stats) => verify(stats)
+          case encl: Template => verify(encl.body)
           case PackageDef(_, stats) => verify(stats)
           case _ => Nil
         }
diff --git a/src/dotty/tools/dotc/ast/TreeTypeMap.scala b/src/dotty/tools/dotc/ast/TreeTypeMap.scala
index 98027cd39..d2ec3ea10 100644
--- a/src/dotty/tools/dotc/ast/TreeTypeMap.scala
+++ b/src/dotty/tools/dotc/ast/TreeTypeMap.scala
@@ -78,22 +78,22 @@ final class TreeTypeMap(
     }
 
   override def transform(tree: tpd.Tree)(implicit ctx: Context): tpd.Tree = treeMap(tree) match {
-    case impl @ Template(constr, parents, self, body) =>
+    case impl @ Template(constr, parents, self, _) =>
       val tmap = withMappedSyms(localSyms(impl :: self :: Nil))
       cpy.Template(impl)(
           constr = tmap.transformSub(constr),
           parents = parents mapconserve transform,
           self = tmap.transformSub(self),
-          body = body mapconserve tmap.transform
+          body = impl.body mapconserve tmap.transform
         ).withType(tmap.mapType(impl.tpe))
     case tree1 =>
       tree1.withType(mapType(tree1.tpe)) match {
         case id: Ident if tpd.needsSelect(id.tpe) =>
           ref(id.tpe.asInstanceOf[TermRef]).withPos(id.pos)
-        case ddef @ DefDef(name, tparams, vparamss, tpt, rhs) =>
+        case ddef @ DefDef(name, tparams, vparamss, tpt, _) =>
           val (tmap1, tparams1) = transformDefs(ddef.tparams)
           val (tmap2, vparamss1) = tmap1.transformVParamss(vparamss)
-          cpy.DefDef(ddef)(name, tparams1, vparamss1, tmap2.transform(tpt), tmap2.transform(rhs))
+          cpy.DefDef(ddef)(name, tparams1, vparamss1, tmap2.transform(tpt), tmap2.transform(ddef.rhs))
         case blk @ Block(stats, expr) =>
           val (tmap1, stats1) = transformDefs(stats)
           val expr1 = tmap1.transform(expr)
diff --git a/src/dotty/tools/dotc/ast/Trees.scala b/src/dotty/tools/dotc/ast/Trees.scala
index ff44a48c5..57cbb7f8f 100644
--- a/src/dotty/tools/dotc/ast/Trees.scala
+++ b/src/dotty/tools/dotc/ast/Trees.scala
@@ -334,9 +334,9 @@ object Trees {
   }
 
   /** A ValDef or DefDef tree */
-  trait ValOrDefDef[-T >: Untyped] extends MemberDef[T] {
+  trait ValOrDefDef[-T >: Untyped] extends MemberDef[T] with WithLazyField[Tree[T]] {
     def tpt: Tree[T]
-    def rhs: Tree[T]
+    def rhs: Tree[T] = forceIfLazy
   }
 
   // ----------- Tree case classes ------------------------------------
@@ -619,17 +619,22 @@ object Trees {
   }
 
   /** mods val name: tpt = rhs */
-  case class ValDef[-T >: Untyped] private[ast] (name: TermName, tpt: Tree[T], rhs: Tree[T])
+  case class ValDef[-T >: Untyped] private[ast] (name: TermName, tpt: Tree[T], private var preRhs: Any /*Tree[T] | Lazy[Tree[T]]*/)
     extends ValOrDefDef[T] {
     type ThisTree[-T >: Untyped] = ValDef[T]
     assert(isEmpty || tpt != genericEmptyTree)
+    protected def maybeLazy = preRhs
+    protected def maybeLazy_=(x: Any) = preRhs = x
   }
 
   /** mods def name[tparams](vparams_1)...(vparams_n): tpt = rhs */
-  case class DefDef[-T >: Untyped] private[ast] (name: TermName, tparams: List[TypeDef[T]], vparamss: List[List[ValDef[T]]], tpt: Tree[T], rhs: Tree[T])
+  case class DefDef[-T >: Untyped] private[ast] (name: TermName, tparams: List[TypeDef[T]], 
+      vparamss: List[List[ValDef[T]]], tpt: Tree[T], private var preRhs: Any /*Tree[T] | Lazy[Tree[T]]*/)
     extends ValOrDefDef[T] {
     type ThisTree[-T >: Untyped] = DefDef[T]
     assert(tpt != genericEmptyTree)
+    protected def maybeLazy = preRhs
+    protected def maybeLazy_=(x: Any) = preRhs = x
   }
 
   /** mods class name template     or
@@ -652,9 +657,12 @@ object Trees {
   }
 
   /** extends parents { self => body } */
-  case class Template[-T >: Untyped] private[ast] (constr: DefDef[T], parents: List[Tree[T]], self: ValDef[T], body: List[Tree[T]])
-    extends DefTree[T] {
+  case class Template[-T >: Untyped] private[ast] (constr: DefDef[T], parents: List[Tree[T]], self: ValDef[T], private var preBody: Any /*List[Tree[T]] | Lazy[List[Tree]]]*/)
+    extends DefTree[T] with WithLazyField[List[Tree[T]]] {
     type ThisTree[-T >: Untyped] = Template[T]
+    protected def maybeLazy = preBody
+    protected def maybeLazy_=(x: Any) = preBody = x
+    def body: List[Tree[T]] = forceIfLazy
   }
 
   /** import expr.selectors
@@ -746,6 +754,33 @@ object Trees {
     if (buf != null) buf.toList else trees
   }
 
+  // ----- Lazy trees and tree sequences
+
+  /** A tree that can have a lazy field 
+   *  The field is represented by some private `var` which is
+   *  proxied by the `maybeLazy` accessors. Forcing the field will
+   *  set the `var` to the underlying value. 
+   */
+  trait WithLazyField[+T] {
+    protected def maybeLazy: Any
+    protected def maybeLazy_=(x: Any): Unit
+    def forceIfLazy: T = maybeLazy match {
+      case lzy: Lazy[T] => 
+        val x = lzy.complete
+        maybeLazy = x
+        x
+      case x: T @ unchecked => x
+    }
+  }
+
+  /** A base trait for lazy tree fields.
+   *  These can be instantiated with Lazy instances which
+   *  can delay tree construction until the field is first demanded.
+   */
+  trait Lazy[T] {
+    def complete: T
+  }
+
   // ----- Generic Tree Instances, inherited from  `tpt` and `untpd`.
 
   abstract class Instance[T >: Untyped <: Type] extends DotClass { inst =>
@@ -1114,16 +1149,16 @@ object Trees {
           cpy.UnApply(tree)(transform(fun), transform(implicits), transform(patterns))
         case EmptyValDef =>
           tree
-        case tree @ ValDef(name, tpt, rhs) =>
+        case tree @ ValDef(name, tpt, _) =>
           val tpt1 = transform(tpt)
-          val rhs1 = transform(rhs)
-          cpy.ValDef(tree)(name, tpt1, rhs1)
-        case tree @ DefDef(name, tparams, vparamss, tpt, rhs) =>
-          cpy.DefDef(tree)(name, transformSub(tparams), vparamss mapConserve (transformSub(_)), transform(tpt), transform(rhs))
+          val rhs1 = transform(tree.rhs)
+          cpy.ValDef(tree)(name, transform(tpt1), transform(rhs1))
+        case tree @ DefDef(name, tparams, vparamss, tpt, _) =>
+          cpy.DefDef(tree)(name, transformSub(tparams), vparamss mapConserve (transformSub(_)), transform(tpt), transform(tree.rhs))
         case tree @ TypeDef(name, rhs) =>
           cpy.TypeDef(tree)(name, transform(rhs), tree.tparams)
-        case Template(constr, parents, self, body) =>
-          cpy.Template(tree)(transformSub(constr), transform(parents), transformSub(self), transformStats(body))
+        case tree @ Template(constr, parents, self, _) =>
+          cpy.Template(tree)(transformSub(constr), transform(parents), transformSub(self), transformStats(tree.body))
         case Import(expr, selectors) =>
           cpy.Import(tree)(transform(expr), selectors)
         case PackageDef(pid, stats) =>
@@ -1213,14 +1248,14 @@ object Trees {
           this(x, trees)
         case UnApply(fun, implicits, patterns) =>
           this(this(this(x, fun), implicits), patterns)
-        case ValDef(name, tpt, rhs) =>
-          this(this(x, tpt), rhs)
-        case DefDef(name, tparams, vparamss, tpt, rhs) =>
-          this(this((this(x, tparams) /: vparamss)(apply), tpt), rhs)
+        case tree @ ValDef(name, tpt, _) =>
+          this(this(x, tpt), tree.rhs)
+        case tree @ DefDef(name, tparams, vparamss, tpt, _) =>
+          this(this((this(x, tparams) /: vparamss)(apply), tpt), tree.rhs)
         case TypeDef(name, rhs) =>
           this(x, rhs)
-        case Template(constr, parents, self, body) =>
-          this(this(this(this(x, constr), parents), self), body)
+        case tree @ Template(constr, parents, self, _) =>
+          this(this(this(this(x, constr), parents), self), tree.body)
         case Import(expr, selectors) =>
           this(x, expr)
         case PackageDef(pid, stats) =>
diff --git a/src/dotty/tools/dotc/ast/tpd.scala b/src/dotty/tools/dotc/ast/tpd.scala
index 52a617ea2..49358a1c0 100644
--- a/src/dotty/tools/dotc/ast/tpd.scala
+++ b/src/dotty/tools/dotc/ast/tpd.scala
@@ -161,7 +161,7 @@ object tpd extends Trees.Instance[Type] with TypedTreeInfo {
   def UnApply(fun: Tree, implicits: List[Tree], patterns: List[Tree], proto: Type)(implicit ctx: Context): UnApply =
     ta.assignType(untpd.UnApply(fun, implicits, patterns), proto)
 
-  def ValDef(sym: TermSymbol, rhs: Tree = EmptyTree)(implicit ctx: Context): ValDef =
+  def ValDef(sym: TermSymbol, rhs: Any /*Tree | Lazy[Tree]*/ = EmptyTree)(implicit ctx: Context): ValDef =
     ta.assignType(untpd.ValDef(sym.name, TypeTree(sym.info), rhs), sym)
 
   def SyntheticValDef(name: TermName, rhs: Tree)(implicit ctx: Context): ValDef =
diff --git a/src/dotty/tools/dotc/ast/untpd.scala b/src/dotty/tools/dotc/ast/untpd.scala
index dfff17514..d7791bfa5 100644
--- a/src/dotty/tools/dotc/ast/untpd.scala
+++ b/src/dotty/tools/dotc/ast/untpd.scala
@@ -142,10 +142,10 @@ object untpd extends Trees.Instance[Untyped] with UntypedTreeInfo {
   def Bind(name: Name, body: Tree): Bind = new Bind(name, body)
   def Alternative(trees: List[Tree]): Alternative = new Alternative(trees)
   def UnApply(fun: Tree, implicits: List[Tree], patterns: List[Tree]): UnApply = new UnApply(fun, implicits, patterns)
-  def ValDef(name: TermName, tpt: Tree, rhs: Tree): ValDef = new ValDef(name, tpt, rhs)
-  def DefDef(name: TermName, tparams: List[TypeDef], vparamss: List[List[ValDef]], tpt: Tree, rhs: Tree): DefDef = new DefDef(name, tparams, vparamss, tpt, rhs)
+  def ValDef(name: TermName, tpt: Tree, rhs: Any /*Tree | Lazy[Tree]*/): ValDef = new ValDef(name, tpt, rhs)
+  def DefDef(name: TermName, tparams: List[TypeDef], vparamss: List[List[ValDef]], tpt: Tree, rhs: Any /*Tree | Lazy[Tree]*/): DefDef = new DefDef(name, tparams, vparamss, tpt, rhs)
   def TypeDef(name: TypeName, rhs: Tree): TypeDef = new TypeDef(name, rhs)
-  def Template(constr: DefDef, parents: List[Tree], self: ValDef, body: List[Tree]): Template = new Template(constr, parents, self, body)
+  def Template(constr: DefDef, parents: List[Tree], self: ValDef, body:Any /* List[Tree] | lazy[List[Tree]]*/): Template = new Template(constr, parents, self, body)
   def Import(expr: Tree, selectors: List[untpd.Tree]): Import = new Import(expr, selectors)
   def PackageDef(pid: RefTree, stats: List[Tree]): PackageDef = new PackageDef(pid, stats)
   def Annotated(annot: Tree, arg: Tree): Annotated = new Annotated(annot, arg)