Unpickle arguments of parent constructors in Templates lazily

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))