Clean up forcing logic for lazy tree fields

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/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,27 +859,17 @@ object Trees {
    *                            if this is of class untpd.DerivingTemplate.
    *                            Typed templates only have parents.
    */
-  case class Template[+T <: Untyped] private[ast] (constr: DefDef[T], private var myParentsOrDerived: LazyTreeList[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
-
-    // The post-condition of forceIfLazy is that all lazy fields are trees, so
-    // we need to force parentsAndDerived as well as body.
-    override def forceIfLazy(using Context) =
-      parentsOrDerived
-      super.forceIfLazy
 
-    def body(using Context): List[Tree[T]] = forceIfLazy
+    def forceFields()(using Context): Unit =
+      preParentsOrDerived = force(preParentsOrDerived)
+      preBody = force(preBody)
 
-    def parentsOrDerived(using Context): List[Tree[T]] = myParentsOrDerived match
-      case latePs: Lazy[List[Tree[T]]] @unchecked =>
-        myParentsOrDerived = latePs.complete
-        parentsOrDerived
-      case ps: List[Tree[T]] @unchecked => ps
+    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
@@ -1027,17 +1021,11 @@ object Trees {
    *  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
-    }
-  }
+  trait WithLazyFields:
+    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
+    def forceFields()(using Context): Unit
 
   /** A base trait for lazy tree fields.
    *  These can be instantiated with Lazy instances which

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/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/TreeChecker.scala

@@ -303,9 +303,11 @@ object TreeChecker {
       // case tree: untpd.Select =>
       // case tree: untpd.Bind =>
       case vd : ValDef =>
-        assertIdentNotJavaClass(vd.forceIfLazy)
+        vd.forceFields()
+        assertIdentNotJavaClass(vd)
       case dd : DefDef =>
-        assertIdentNotJavaClass(dd.forceIfLazy)
+        dd.forceFields()
+        assertIdentNotJavaClass(dd)
       // case tree: untpd.TypeDef =>
       case Apply(fun, args) =>
         assertIdentNotJavaClass(fun)