Make LazyBodyAnnotation more like LazyAnnotation

And address the TODO comment by adding an apply method that takes a
`Context ?=> Tree`.

Also add `Annotation#notEvaluated`, previously
`LazyBodyAnnotation#isEvaluated` was also returning true when evaluation
was started but not finished which was confusing.

Author: smarter

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
 
+    /** Evaluation has not started */
+    def notEvaluated: Boolean = false
+
+    /** Evaluation has finished */
     def isEvaluated: Boolean = true
 
     def ensureCompleted(implicit ctx: Context): Unit = tree
@@ -75,6 +79,7 @@ object Annotations {
       myTree.asInstanceOf[Tree]
 
     override def isEvaluated: Boolean = myTree.isInstanceOf[Tree]
+    override def notEvaluated: Boolean = !isEvaluated && myTree != null
   }
 
   /** An annotation indicating the body of a right-hand side,
@@ -89,24 +94,32 @@ 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 =>
+          val fn = treeFn
+          myTree = null
+          myTree = fn(ctx)
+        case _ =>
       }
-      myBody
-    }
-    override def isEvaluated: Boolean = evaluated
+      myTree.asInstanceOf[Tree]
+
+    override def isEvaluated: Boolean = myTree.isInstanceOf[Tree]
+    override def notEvaluated: Boolean = !isEvaluated && myTree != null
+  }
+
+  object LazyBodyAnnotation {
+    def apply(bodyFn: Context ?=> Tree): LazyBodyAnnotation =
+      new LazyBodyAnnotation:
+        protected var myTree: Tree | (Context => Tree) = ctx => bodyFn(using ctx)
   }
 
   object 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.notEvaluated =>
       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) {