LazyAnnotation: avoid unnecessary memory leaks

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

@@ -35,6 +35,10 @@ object Annotations {
     def argumentConstant(i: Int)(implicit ctx: Context): Option[Constant] =
       for (ConstantType(c) <- argument(i) map (_.tpe)) yield c
 
+    /** The tree evaluaton is in progress. */
+    def isEvaluating: Boolean = false
+
+    /** The tree evaluation has finished. */
     def isEvaluated: Boolean = true
 
     def ensureCompleted(implicit ctx: Context): Unit = tree
@@ -48,16 +52,32 @@ object Annotations {
   }
 
   abstract class LazyAnnotation extends Annotation {
-    override def symbol(implicit ctx: Context): Symbol
-    def complete(implicit ctx: Context): Tree
-
-    private var myTree: Tree = null
-    def tree(implicit ctx: Context): Tree = {
-      if (myTree == null) myTree = complete(ctx)
-      myTree
-    }
+    protected var mySym: Symbol | (Context => Symbol)
+    override def symbol(using ctx: Context): Symbol =
+      assert(mySym != null)
+      mySym match {
+        case symFn: (Context => Symbol) @unchecked =>
+          mySym = null
+          mySym = symFn(ctx)
+        case sym: Symbol if sym.defRunId != ctx.runId =>
+          mySym = sym.denot.current.symbol
+        case _ =>
+      }
+      mySym.asInstanceOf[Symbol]
+
+    protected var myTree: Tree | (Context => Tree)
+    def tree(using ctx: Context): Tree =
+      assert(myTree != null)
+      myTree match {
+        case treeFn: (Context => Tree) @unchecked =>
+          myTree = null
+          myTree = treeFn(ctx)
+        case _ =>
+      }
+      myTree.asInstanceOf[Tree]
 
-    override def isEvaluated: Boolean = myTree != null
+    override def isEvaluating: Boolean = myTree == null
+    override def isEvaluated: Boolean = myTree.isInstanceOf[Tree]
   }
 
   /** An annotation indicating the body of a right-hand side,
@@ -72,24 +92,31 @@ object Annotations {
     override def ensureCompleted(implicit ctx: Context): Unit = ()
   }
 
-  case class ConcreteBodyAnnotation(body: Tree) extends BodyAnnotation {
+  class ConcreteBodyAnnotation(body: Tree) extends BodyAnnotation {
     def tree(implicit ctx: Context): Tree = body
   }
 
-  case class LazyBodyAnnotation(private var bodyExpr: Context => Tree) extends BodyAnnotation {
-    // TODO: Make `bodyExpr` an IFT once #6865 os in bootstrap
-    private var evaluated = false
-    private var myBody: Tree = _
-    def tree(implicit ctx: Context): Tree = {
-      if (evaluated) assert(myBody != null)
-      else {
-        evaluated = true
-        myBody = bodyExpr(ctx)
-        bodyExpr = null
+  abstract class LazyBodyAnnotation extends BodyAnnotation {
+    // Copy-pasted from LazyAnnotation to avoid having to turn it into a trait
+    protected var myTree: Tree | (Context => Tree)
+    def tree(using ctx: Context): Tree =
+      assert(myTree != null)
+      myTree match {
+        case treeFn: (Context => Tree) @unchecked =>
+          myTree = null
+          myTree = treeFn(ctx)
+        case _ =>
       }
-      myBody
-    }
-    override def isEvaluated: Boolean = evaluated
+      myTree.asInstanceOf[Tree]
+
+    override def isEvaluating: Boolean = myTree == null
+    override def isEvaluated: Boolean = myTree.isInstanceOf[Tree]
+  }
+
+  object LazyBodyAnnotation {
+    def apply(bodyFn: Context ?=> Tree): LazyBodyAnnotation =
+      new LazyBodyAnnotation:
+        protected var myTree: Tree | (Context => Tree) = ctx => bodyFn(using ctx)
   }
 
   object Annotation {
@@ -120,23 +147,15 @@ object Annotations {
     /** Create an annotation where the tree is computed lazily. */
     def deferred(sym: Symbol)(treeFn: Context ?=> Tree)(implicit ctx: Context): Annotation =
       new LazyAnnotation {
-        override def symbol(implicit ctx: Context): Symbol = sym
-        def complete(implicit ctx: Context) = treeFn(using ctx)
+        protected var myTree: Tree | (Context => Tree) = ctx => treeFn(using ctx)
+        protected var mySym: Symbol | (Context => Symbol) = sym
       }
 
     /** Create an annotation where the symbol and the tree are computed lazily. */
-    def deferredSymAndTree(symf: Context ?=> Symbol)(treeFn: Context ?=> Tree)(implicit ctx: Context): Annotation =
+    def deferredSymAndTree(symFn: Context ?=> Symbol)(treeFn: Context ?=> Tree)(implicit ctx: Context): Annotation =
       new LazyAnnotation {
-        private var mySym: Symbol = _
-
-        override def symbol(implicit ctx: Context): Symbol = {
-          if (mySym == null || mySym.defRunId != ctx.runId) {
-            mySym = symf(using ctx)
-            assert(mySym != null)
-          }
-          mySym
-        }
-        def complete(implicit ctx: Context) = treeFn(using ctx)
+        protected var mySym: Symbol | (Context => Symbol) = ctx => symFn(using ctx)
+        protected var myTree: Tree | (Context => Tree) = ctx => treeFn(using ctx)
       }
 
     def deferred(atp: Type, args: List[Tree])(implicit ctx: Context): Annotation =

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

@@ -581,9 +581,9 @@ class TreeUnpickler(reader: TastyReader,
         forkAt(templateStart).indexTemplateParams()(localContext(sym))
       }
       else if (sym.isInlineMethod)
-        sym.addAnnotation(LazyBodyAnnotation { ctx0 =>
+        sym.addAnnotation(LazyBodyAnnotation { (using ctx0: Context) =>
           val ctx1 = localContext(sym)(ctx0).addMode(Mode.ReadPositions)
-          implicit val ctx: Context = sourceChangeContext(Addr(0))(ctx1)
+          given Context = sourceChangeContext(Addr(0))(ctx1)
             // avoids space leaks by not capturing the current context
           forkAt(rhsStart).readTerm()
         })

compiler/src/dotty/tools/dotc/typer/PrepareInlineable.scala

@@ -216,12 +216,12 @@ object PrepareInlineable {
       inlined: Symbol, treeExpr: Context => Tree)(implicit ctx: Context): Unit =
     inlined.unforcedAnnotation(defn.BodyAnnot) match {
       case Some(ann: ConcreteBodyAnnotation) =>
-      case Some(ann: LazyBodyAnnotation) if ann.isEvaluated =>
+      case Some(ann: LazyBodyAnnotation) if ann.isEvaluated || ann.isEvaluating =>
       case _ =>
         if (!ctx.isAfterTyper) {
           val inlineCtx = ctx
-          inlined.updateAnnotation(LazyBodyAnnotation { _ =>
-            implicit val ctx = inlineCtx
+          inlined.updateAnnotation(LazyBodyAnnotation {
+            given ctx as Context = inlineCtx
             val initialErrorCount = ctx.reporter.errorCount
             var inlinedBody = treeExpr(ctx)
             if (ctx.reporter.errorCount == initialErrorCount) {