Unpickle arguments of parent constructors in Templates lazily (#16688)

Avoid reading the arguments of parent constructors unless someone forces
them. We don't need them to determine the parent types of the class to
which the template belongs. This makes TreeUnpickler as lazy as Namer in
this respect and therefore avoids CyclicReferences during unpickling.

Fixes #16673

Author: odersky

compiler/src/dotty/tools/dotc/ast/NavigateAST.scala

@@ -4,7 +4,7 @@ package ast
 import core.Contexts._
 import core.Decorators._
 import util.Spans._
-import Trees.{MemberDef, DefTree, WithLazyField}
+import Trees.{MemberDef, DefTree, WithLazyFields}
 import dotty.tools.dotc.core.Types.AnnotatedType
 import dotty.tools.dotc.core.Types.ImportType
 import dotty.tools.dotc.core.Types.Type
@@ -106,16 +106,14 @@ object NavigateAST {
         // FIXME: We shouldn't be manually forcing trees here, we should replace
         // our usage of `productIterator` by something in `Positioned` that takes
         // care of low-level details like this for us.
-        p match {
-          case p: WithLazyField[?] =>
-            p.forceIfLazy
+        p match
+          case p: WithLazyFields => p.forceFields()
           case _ =>
-        }
         val iterator = p match
           case defdef: DefTree[?] =>
             p.productIterator ++ defdef.mods.productIterator
           case _ =>
-            p.productIterator 
+            p.productIterator
         childPath(iterator, p :: path)
       }
       else {

compiler/src/dotty/tools/dotc/ast/TreeInfo.scala

@@ -313,7 +313,7 @@ trait TreeInfo[T <: Untyped] { self: Trees.Instance[T] =>
    */
   def parentsKind(parents: List[Tree])(using Context): FlagSet = parents match {
     case Nil => NoInitsInterface
-    case Apply(_, _ :: _) :: _ => EmptyFlags
+    case Apply(_, _ :: _) :: _ | Block(_, _) :: _ => EmptyFlags
     case _ :: parents1 => parentsKind(parents1)
   }
 

compiler/src/dotty/tools/dotc/ast/TreeMapWithImplicits.scala

@@ -55,10 +55,10 @@ class TreeMapWithImplicits extends tpd.TreeMapWithPreciseStatContexts {
             transform(tree.tpt),
             transform(tree.rhs)(using nestedScopeCtx(tree.paramss.flatten)))
         }
-      case impl @ Template(constr, parents, self, _) =>
+      case impl @ Template(constr, _, self, _) =>
         cpy.Template(tree)(
           transformSub(constr),
-          transform(parents)(using ctx.superCallContext),
+          transform(impl.parents)(using ctx.superCallContext),
           Nil,
           transformSelf(self),
           transformStats(impl.body, tree.symbol))

compiler/src/dotty/tools/dotc/ast/TreeTypeMap.scala

@@ -92,11 +92,11 @@ class TreeTypeMap(
     cpy.Inlined(tree)(call, bindings1, expanded1)
 
   override def transform(tree: tpd.Tree)(using Context): tpd.Tree = treeMap(tree) match {
-    case impl @ Template(constr, parents, self, _) =>
+    case impl @ Template(constr, _, self, _) =>
       val tmap = withMappedSyms(localSyms(impl :: self :: Nil))
       cpy.Template(impl)(
           constr = tmap.transformSub(constr),
-          parents = parents.mapconserve(transform),
+          parents = impl.parents.mapconserve(transform),
           self = tmap.transformSub(self),
           body = impl.body mapconserve
             (tmap.transform(_)(using ctx.withOwner(mapOwner(impl.symbol.owner))))

compiler/src/dotty/tools/dotc/ast/Trees.scala

@@ -407,12 +407,12 @@ object Trees {
   }
 
   /** A ValDef or DefDef tree */
-  abstract class ValOrDefDef[+T <: Untyped](implicit @constructorOnly src: SourceFile) extends MemberDef[T] with WithLazyField[Tree[T]] {
+  abstract class ValOrDefDef[+T <: Untyped](implicit @constructorOnly src: SourceFile) extends MemberDef[T], WithLazyFields {
     type ThisTree[+T <: Untyped] <: ValOrDefDef[T]
     def name: TermName
     def tpt: Tree[T]
-    def unforcedRhs: LazyTree[T] = unforced
-    def rhs(using Context): Tree[T] = forceIfLazy
+    def unforcedRhs: LazyTree[T]
+    def rhs(using Context): Tree[T]
   }
 
   trait ValOrTypeDef[+T <: Untyped] extends MemberDef[T]:
@@ -808,8 +808,10 @@ object Trees {
     extends ValOrDefDef[T], ValOrTypeDef[T] {
     type ThisTree[+T <: Untyped] = ValDef[T]
     assert(isEmpty || (tpt ne genericEmptyTree))
-    def unforced: LazyTree[T] = preRhs
-    protected def force(x: Tree[T @uncheckedVariance]): Unit = preRhs = x
+
+    def unforcedRhs: LazyTree[T] = preRhs
+    def forceFields()(using Context): Unit = preRhs = force(preRhs)
+    def rhs(using Context): Tree[T] = { forceFields(); preRhs.asInstanceOf[Tree[T]] }
   }
 
   /** mods def name[tparams](vparams_1)...(vparams_n): tpt = rhs */
@@ -818,8 +820,10 @@ object Trees {
     extends ValOrDefDef[T] {
     type ThisTree[+T <: Untyped] = DefDef[T]
     assert(tpt ne genericEmptyTree)
-    def unforced: LazyTree[T] = preRhs
-    protected def force(x: Tree[T @uncheckedVariance]): Unit = preRhs = x
+
+    def unforcedRhs: LazyTree[T] = preRhs
+    def forceFields()(using Context): Unit = preRhs = force(preRhs)
+    def rhs(using Context): Tree[T] = { forceFields(); preRhs.asInstanceOf[Tree[T]] }
 
     def leadingTypeParams(using Context): List[TypeDef[T]] = paramss match
       case (tparams @ (tparam: TypeDef[_]) :: _) :: _ => tparams.asInstanceOf[List[TypeDef[T]]]
@@ -855,16 +859,20 @@ object Trees {
    *                            if this is of class untpd.DerivingTemplate.
    *                            Typed templates only have parents.
    */
-  case class Template[+T <: Untyped] private[ast] (constr: DefDef[T], parentsOrDerived: List[Tree[T]], self: ValDef[T], private var preBody: LazyTreeList[T])(implicit @constructorOnly src: SourceFile)
-    extends DefTree[T] with WithLazyField[List[Tree[T]]] {
+  case class Template[+T <: Untyped] private[ast] (constr: DefDef[T], private var preParentsOrDerived: LazyTreeList[T], self: ValDef[T], private var preBody: LazyTreeList[T])(implicit @constructorOnly src: SourceFile)
+    extends DefTree[T] with WithLazyFields {
     type ThisTree[+T <: Untyped] = Template[T]
-    def unforcedBody: LazyTreeList[T] = unforced
-    def unforced: LazyTreeList[T] = preBody
-    protected def force(x: List[Tree[T @uncheckedVariance]]): Unit = preBody = x
-    def body(using Context): List[Tree[T]] = forceIfLazy
 
-    def parents: List[Tree[T]] = parentsOrDerived // overridden by DerivingTemplate
-    def derived: List[untpd.Tree] = Nil           // overridden by DerivingTemplate
+    def forceFields()(using Context): Unit =
+      preParentsOrDerived = force(preParentsOrDerived)
+      preBody = force(preBody)
+
+    def unforcedBody: LazyTreeList[T] = preBody
+    def body(using Context): List[Tree[T]] = { forceFields(); preBody.asInstanceOf[List[Tree[T]]] }
+    def parentsOrDerived(using Context): List[Tree[T]] = { forceFields(); preParentsOrDerived.asInstanceOf[List[Tree[T]]] }
+
+    def parents(using Context): List[Tree[T]] = parentsOrDerived // overridden by DerivingTemplate
+    def derived: List[untpd.Tree] = Nil // overridden by DerivingTemplate
   }
 
 
@@ -1008,30 +1016,27 @@ object Trees {
 
   // ----- Lazy trees and tree sequences
 
-  /** A tree that can have a lazy field
-   *  The field is represented by some private `var` which is
-   *  accessed by `unforced` and `force`. Forcing the field will
-   *  set the `var` to the underlying value.
-   */
-  trait WithLazyField[+T <: AnyRef] {
-    def unforced: T | Lazy[T]
-    protected def force(x: T @uncheckedVariance): Unit
-    def forceIfLazy(using Context): T = unforced match {
-      case lzy: Lazy[T @unchecked] =>
-        val x = lzy.complete
-        force(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 <: AnyRef] {
+  trait Lazy[+T <: AnyRef]:
     def complete(using Context): T
-  }
+
+  /** A tree that can have a lazy fields.
+   *  Such fields are variables of type `T | Lazy[T]`, for some tyope `T`.
+   */
+  trait WithLazyFields:
+
+    /** If `x` is lazy, computes the underlying value */
+    protected def force[T <: AnyRef](x: T | Lazy[T])(using Context): T = x match
+      case x: Lazy[T] @unchecked => x.complete
+      case x: T @unchecked => x
+
+    /** Assigns all lazy fields their underlying non-lazy value. */
+    def forceFields()(using Context): Unit
+
+  end WithLazyFields
 
   // ----- Generic Tree Instances, inherited from `tpt` and `untpd`.
 
@@ -1355,7 +1360,7 @@ object Trees {
         DefDef(tree: Tree)(name, paramss, tpt, rhs)
       def TypeDef(tree: TypeDef)(name: TypeName = tree.name, rhs: Tree = tree.rhs)(using Context): TypeDef =
         TypeDef(tree: Tree)(name, rhs)
-      def Template(tree: Template)(constr: DefDef = tree.constr, parents: List[Tree] = tree.parents, derived: List[untpd.Tree] = tree.derived, self: ValDef = tree.self, body: LazyTreeList = tree.unforcedBody)(using Context): Template =
+      def Template(tree: Template)(using Context)(constr: DefDef = tree.constr, parents: List[Tree] = tree.parents, derived: List[untpd.Tree] = tree.derived, self: ValDef = tree.self, body: LazyTreeList = tree.unforcedBody): Template =
         Template(tree: Tree)(constr, parents, derived, self, body)
       def Hole(tree: Hole)(isTerm: Boolean = tree.isTerm, idx: Int = tree.idx, args: List[Tree] = tree.args, content: Tree = tree.content, tpt: Tree = tree.tpt)(using Context): Hole =
         Hole(tree: Tree)(isTerm, idx, args, content, tpt)
@@ -1618,8 +1623,8 @@ object Trees {
               inContext(localCtx(tree)) {
                 this(x, rhs)
               }
-            case tree @ Template(constr, parents, self, _) if tree.derived.isEmpty =>
-              this(this(this(this(x, constr), parents), self), tree.body)
+            case tree @ Template(constr, _, self, _) if tree.derived.isEmpty =>
+              this(this(this(this(x, constr), tree.parents), self), tree.body)
             case Import(expr, _) =>
               this(x, expr)
             case Export(expr, _) =>

compiler/src/dotty/tools/dotc/ast/untpd.scala

@@ -54,7 +54,8 @@ object untpd extends Trees.Instance[Untyped] with UntypedTreeInfo {
    */
   class DerivingTemplate(constr: DefDef, parentsOrDerived: List[Tree], self: ValDef, preBody: LazyTreeList, derivedCount: Int)(implicit @constructorOnly src: SourceFile)
   extends Template(constr, parentsOrDerived, self, preBody) {
-    override val parents = parentsOrDerived.dropRight(derivedCount)
+    private val myParents = parentsOrDerived.dropRight(derivedCount)
+    override def parents(using Context) = myParents
     override val derived = parentsOrDerived.takeRight(derivedCount)
   }
 
@@ -415,6 +416,8 @@ object untpd extends Trees.Instance[Untyped] with UntypedTreeInfo {
   def Template(constr: DefDef, parents: List[Tree], derived: List[Tree], self: ValDef, body: LazyTreeList)(implicit src: SourceFile): Template =
     if (derived.isEmpty) new Template(constr, parents, self, body)
     else new DerivingTemplate(constr, parents ++ derived, self, body, derived.length)
+  def Template(constr: DefDef, parents: LazyTreeList, self: ValDef, body: LazyTreeList)(implicit src: SourceFile): Template =
+    new Template(constr, parents, self, body)
   def Import(expr: Tree, selectors: List[ImportSelector])(implicit src: SourceFile): Import = new Import(expr, selectors)
   def Export(expr: Tree, selectors: List[ImportSelector])(implicit src: SourceFile): Export = new Export(expr, selectors)
   def PackageDef(pid: RefTree, stats: List[Tree])(implicit src: SourceFile): PackageDef = new PackageDef(pid, stats)

compiler/src/dotty/tools/dotc/core/tasty/PositionPickler.scala

@@ -8,7 +8,7 @@ import dotty.tools.tasty.TastyBuffer
 import TastyBuffer._
 
 import ast._
-import Trees.WithLazyField
+import Trees.WithLazyFields
 import util.{SourceFile, NoSource}
 import core._
 import Annotations._, Decorators._
@@ -80,7 +80,7 @@ class PositionPickler(
     def alwaysNeedsPos(x: Positioned) = x match {
       case
           // initialSpan is inaccurate for trees with lazy field
-          _: WithLazyField[?]
+          _: WithLazyFields
 
           // A symbol is created before the corresponding tree is unpickled,
           // and its position cannot be changed afterwards.

compiler/src/dotty/tools/dotc/core/tasty/TreeUnpickler.scala

@@ -959,6 +959,51 @@ class TreeUnpickler(reader: TastyReader,
       tree.setDefTree
     }
 
+    /** Read enough of parent to determine its type, without reading arguments
+     *  of applications. This is necessary to make TreeUnpickler as lazy as Namer
+     *  in this regard. See i16673 for a test case.
+     */
+    private def readParentType()(using Context): Type =
+      readByte() match
+        case TYPEAPPLY =>
+          val end = readEnd()
+          val tycon = readParentType()
+          if tycon.typeParams.isEmpty then
+            goto(end)
+            tycon
+          else
+            val args = until(end)(readTpt())
+            val cls = tycon.classSymbol
+            assert(cls.typeParams.hasSameLengthAs(args))
+            cls.typeRef.appliedTo(args.tpes)
+        case APPLY | BLOCK =>
+          val end = readEnd()
+          try readParentType()
+          finally goto(end)
+        case SELECTin =>
+          val end = readEnd()
+          readName()
+          readTerm() match
+            case nu: New =>
+              try nu.tpe
+              finally goto(end)
+        case SHAREDterm =>
+          forkAt(readAddr()).readParentType()
+
+    /** Read template parents
+     *  @param  withArgs if false, only read enough of parent trees to determine their type
+     *                   but skip constructor arguments. Return any trees that were partially
+     *                   parsed in this way as InferredTypeTrees.
+     */
+    def readParents(withArgs: Boolean)(using Context): List[Tree] =
+      collectWhile(nextByte != SELFDEF && nextByte != DEFDEF) {
+        nextUnsharedTag match
+          case APPLY | TYPEAPPLY | BLOCK =>
+            if withArgs then readTerm()
+            else InferredTypeTree().withType(readParentType())
+          case _ => readTpt()
+      }
+
     private def readTemplate(using Context): Template = {
       val start = currentAddr
       assert(sourcePathAt(start).isEmpty)
@@ -981,12 +1026,8 @@ class TreeUnpickler(reader: TastyReader,
         while (bodyIndexer.reader.nextByte != DEFDEF) bodyIndexer.skipTree()
         bodyIndexer.indexStats(end)
       }
-      val parents = collectWhile(nextByte != SELFDEF && nextByte != DEFDEF) {
-        nextUnsharedTag match {
-          case APPLY | TYPEAPPLY | BLOCK => readTerm()(using parentCtx)
-          case _ => readTpt()(using parentCtx)
-        }
-      }
+      val parentReader = fork
+      val parents = readParents(withArgs = false)(using parentCtx)
       val parentTypes = parents.map(_.tpe.dealias)
       val self =
         if (nextByte == SELFDEF) {
@@ -1000,7 +1041,13 @@ class TreeUnpickler(reader: TastyReader,
         selfInfo = if (self.isEmpty) NoType else self.tpt.tpe)
         .integrateOpaqueMembers
       val constr = readIndexedDef().asInstanceOf[DefDef]
-      val mappedParents = parents.map(_.changeOwner(localDummy, constr.symbol))
+      val mappedParents: LazyTreeList =
+        if parents.exists(_.isInstanceOf[InferredTypeTree]) then
+          // parents were not read fully, will need to be read again later on demand
+          new LazyReader(parentReader, localDummy, ctx.mode, ctx.source,
+            _.readParents(withArgs = true)
+             .map(_.changeOwner(localDummy, constr.symbol)))
+        else parents
 
       val lazyStats = readLater(end, rdr => {
         val stats = rdr.readIndexedStats(localDummy, end)
@@ -1009,7 +1056,7 @@ class TreeUnpickler(reader: TastyReader,
       defn.patchStdLibClass(cls)
       NamerOps.addConstructorProxies(cls)
       setSpan(start,
-        untpd.Template(constr, mappedParents, Nil, self, lazyStats)
+        untpd.Template(constr, mappedParents, self, lazyStats)
           .withType(localDummy.termRef))
     }
 

compiler/src/dotty/tools/dotc/interactive/Interactive.scala

@@ -313,8 +313,8 @@ object Interactive {
             case _ =>
           }
           localCtx
-        case tree @ Template(constr, parents, self, _) =>
-          if ((constr :: self :: parents).contains(nested)) outer
+        case tree @ Template(constr, _, self, _) =>
+          if ((constr :: self :: tree.parentsOrDerived).contains(nested)) outer
           else contextOfStat(tree.body, nested, tree.symbol, outer.inClassContext(self.symbol))
         case _ =>
           outer

compiler/src/dotty/tools/dotc/sbt/ExtractAPI.scala

@@ -737,8 +737,7 @@ private class ExtractAPICollector(using Context) extends ThunkHolder {
       var h = initHash
 
       p match
-        case p: WithLazyField[?] =>
-          p.forceIfLazy
+        case p: WithLazyFields => p.forceFields()
         case _ =>
 
       if inlineOrigin.exists then

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

@@ -38,10 +38,10 @@ abstract class MacroTransform extends Phase {
             tree
           case _: PackageDef | _: MemberDef =>
             super.transform(tree)(using localCtx(tree))
-          case impl @ Template(constr, parents, self, _) =>
+          case impl @ Template(constr, _, self, _) =>
             cpy.Template(tree)(
               transformSub(constr),
-              transform(parents)(using ctx.superCallContext),
+              transform(impl.parents)(using ctx.superCallContext),
               Nil,
               transformSelf(self),
               transformStats(impl.body, tree.symbol))