Use BodyAnnot to indicate rhs of inline method

Since fundamental operations such as treeCopy have to know
about inline bodies, it seems better to represent this
information directly in an annotation.

7 files changed, 60 insertions, 25 deletions

diff --git a/src/dotty/annotation/internal/Body.scala b/src/dotty/annotation/internal/Body.scala
new file mode 100644
index 000000000..7e26b02f2
--- /dev/null
+++ b/src/dotty/annotation/internal/Body.scala
@@ -0,0 +1,8 @@
+package dotty.annotation.internal
+
+import scala.annotation.Annotation
+
+/** The class associated with a `BodyAnnotation`, which indicates
+ *  an inline method's right hand side
+ */
+final class Body() extends Annotation
diff --git a/src/dotty/tools/dotc/ast/TreeTypeMap.scala b/src/dotty/tools/dotc/ast/TreeTypeMap.scala
index 6d0c7d8e3..efa46debb 100644
--- a/src/dotty/tools/dotc/ast/TreeTypeMap.scala
+++ b/src/dotty/tools/dotc/ast/TreeTypeMap.scala
@@ -138,10 +138,7 @@ final class TreeTypeMap(
 
   def apply[ThisTree <: tpd.Tree](tree: ThisTree): ThisTree = transform(tree).asInstanceOf[ThisTree]
 
-  def apply(annot: Annotation): Annotation = {
-    val tree1 = apply(annot.tree)
-    if (tree1 eq annot.tree) annot else ConcreteAnnotation(tree1)
-  }
+  def apply(annot: Annotation): Annotation = annot.derivedAnnotation(apply(annot.tree))
 
   /** The current tree map composed with a substitution [from -> to] */
   def withSubstitution(from: List[Symbol], to: List[Symbol]): TreeTypeMap =
diff --git a/src/dotty/tools/dotc/core/Annotations.scala b/src/dotty/tools/dotc/core/Annotations.scala
index 5f96a60e6..fde41ef76 100644
--- a/src/dotty/tools/dotc/core/Annotations.scala
+++ b/src/dotty/tools/dotc/core/Annotations.scala
@@ -42,6 +42,34 @@ object Annotations {
     override def symbol(implicit ctx: Context): Symbol = sym
   }
 
+  /** An annotation indicating the body of a right-hand side,
+   *  typically of an inline method. Treated specially in
+   *  pickling/unpickling and treecopies
+   */
+  abstract class BodyAnnotation extends Annotation {
+    override def symbol(implicit ctx: Context) = defn.BodyAnnot
+    override def derivedAnnotation(tree: Tree)(implicit ctx: Context) =
+      if (tree eq this.tree) this else ConcreteBodyAnnotation(tree)
+    override def arguments(implicit ctx: Context) = Nil
+  }
+
+  case class ConcreteBodyAnnotation(body: Tree) extends BodyAnnotation {
+    def tree(implicit ctx: Context) = body
+  }
+
+  case class LazyBodyAnnotation(bodyExpr: Context => Tree) extends BodyAnnotation {
+    private var evaluated = false
+    private var myBody: Tree = _
+    def tree(implicit ctx: Context) = {
+      if (evaluated) assert(myBody != null)
+      else {
+        evaluated = true
+        myBody = bodyExpr(ctx)
+      }
+      myBody
+    }
+  }
+
   object Annotation {
 
     def apply(tree: Tree) = ConcreteAnnotation(tree)
diff --git a/src/dotty/tools/dotc/core/Definitions.scala b/src/dotty/tools/dotc/core/Definitions.scala
index 12677edb6..75b75d3d5 100644
--- a/src/dotty/tools/dotc/core/Definitions.scala
+++ b/src/dotty/tools/dotc/core/Definitions.scala
@@ -456,6 +456,8 @@ class Definitions {
   def AliasAnnot(implicit ctx: Context) = AliasAnnotType.symbol.asClass
   lazy val AnnotationDefaultAnnotType = ctx.requiredClassRef("dotty.annotation.internal.AnnotationDefault")
   def AnnotationDefaultAnnot(implicit ctx: Context) = AnnotationDefaultAnnotType.symbol.asClass
+  lazy val BodyAnnotType = ctx.requiredClassRef("dotty.annotation.internal.Body")
+  def BodyAnnot(implicit ctx: Context) = BodyAnnotType.symbol.asClass
   lazy val ChildAnnotType = ctx.requiredClassRef("dotty.annotation.internal.Child")
   def ChildAnnot(implicit ctx: Context) = ChildAnnotType.symbol.asClass
   lazy val CovariantBetweenAnnotType = ctx.requiredClassRef("dotty.annotation.internal.CovariantBetween")
diff --git a/src/dotty/tools/dotc/core/tasty/TreePickler.scala b/src/dotty/tools/dotc/core/tasty/TreePickler.scala
index 99a3ff85c..b6f52c0ec 100644
--- a/src/dotty/tools/dotc/core/tasty/TreePickler.scala
+++ b/src/dotty/tools/dotc/core/tasty/TreePickler.scala
@@ -565,10 +565,11 @@ class TreePickler(pickler: TastyPickler) {
     sym.annotations.foreach(pickleAnnotation)
   }
 
-  def pickleAnnotation(ann: Annotation)(implicit ctx: Context) = {
-    writeByte(ANNOTATION)
-    withLength { pickleType(ann.symbol.typeRef); pickleTree(ann.tree) }
-  }
+  def pickleAnnotation(ann: Annotation)(implicit ctx: Context) =
+    if (ann.symbol != defn.BodyAnnot) { // inline bodies are reconstituted automatically when unpickling
+      writeByte(ANNOTATION)
+      withLength { pickleType(ann.symbol.typeRef); pickleTree(ann.tree) }
+    }
 
   def pickle(trees: List[Tree])(implicit ctx: Context) = {
     trees.foreach(tree => if (!tree.isEmpty) pickleTree(tree))
diff --git a/src/dotty/tools/dotc/core/tasty/TreeUnpickler.scala b/src/dotty/tools/dotc/core/tasty/TreeUnpickler.scala
index 1b8434129..4a94e5d71 100644
--- a/src/dotty/tools/dotc/core/tasty/TreeUnpickler.scala
+++ b/src/dotty/tools/dotc/core/tasty/TreeUnpickler.scala
@@ -488,13 +488,12 @@ class TreeUnpickler(reader: TastyReader, tastyName: TastyName.Table, posUnpickle
         sym.completer.withDecls(newScope)
         forkAt(templateStart).indexTemplateParams()(localContext(sym))
       }
-      else if (annots.exists(_.symbol == defn.InlineAnnot)) {
-        Inliner.registerInlineInfo(
-            sym,
-            implicit ctx => forkAt(rhsStart).readTerm(),
-            implicit ctx => localContext(sym).addMode(Mode.ReadPositions)) 
-              // Previous line avoids space leaks because it does not capture the current context.
-      }
+      else if (annots.exists(_.symbol == defn.InlineAnnot))
+        sym.addAnnotation(LazyBodyAnnotation { ctx0 =>
+          implicit val ctx: Context = localContext(sym)(ctx0).addMode(Mode.ReadPositions)
+            // avoids space leaks by not capturing the current context
+          forkAt(rhsStart).readTerm()
+        })
       goto(start)
       sym
     }
diff --git a/src/dotty/tools/dotc/typer/Inliner.scala b/src/dotty/tools/dotc/typer/Inliner.scala
index ef67384ff..079154f26 100644
--- a/src/dotty/tools/dotc/typer/Inliner.scala
+++ b/src/dotty/tools/dotc/typer/Inliner.scala
@@ -16,7 +16,7 @@ import Contexts.Context
 import Names.{Name, TermName}
 import NameOps._
 import SymDenotations.SymDenotation
-import Annotations.Annotation
+import Annotations._
 import transform.ExplicitOuter
 import Inferencing.fullyDefinedType
 import config.Printers.inlining
@@ -223,9 +223,6 @@ object Inliner {
   /** A key to be used in an attachment for `@inline` annotations */
   private val InlineInfo = new Property.Key[InlineInfo]
 
-  /** A key to be used in an attachment for `@inline` annotations */
-  private val InlineBody = new Property.Key[Tree]
-
   /** A key to be used in a context property that tracks enclosing inlined calls */
   private val InlinedCalls = new Property.Key[List[Tree]] // to be used in context
 
@@ -240,20 +237,22 @@ object Inliner {
    */
   def registerInlineInfo(
       sym: SymDenotation, treeExpr: Context => Tree, inlineCtxFn: Context => Context)(implicit ctx: Context): Unit = {
-    val inlineAnnotTree = sym.unforcedAnnotation(defn.InlineAnnot).get.tree
-    if (inlineAnnotTree.getAttachment(InlineBody).isEmpty)
+    if (sym.unforcedAnnotation(defn.BodyAnnot).isEmpty) {
+      val inlineAnnotTree = sym.unforcedAnnotation(defn.InlineAnnot).get.tree
       inlineAnnotTree.getAttachment(InlineInfo) match {
         case Some(inlineInfo) if inlineInfo.isEvaluated => // keep existing attachment
         case _ =>
           if (!ctx.isAfterTyper)
           inlineAnnotTree.putAttachment(InlineInfo, new InlineInfo(treeExpr, inlineCtxFn))
       }
+    }
   }
 
   /** Register an evaluated inline body for `sym` */
   def updateInlineBody(sym: SymDenotation, body: Tree)(implicit ctx: Context): Unit = {
-    val inlineAnnot = sym.unforcedAnnotation(defn.InlineAnnot).get
-    assert(inlineAnnot.tree.putAttachment(InlineBody, body).isDefined)
+    assert(sym.unforcedAnnotation(defn.BodyAnnot).isDefined)
+    sym.removeAnnotation(defn.BodyAnnot)
+    sym.addAnnotation(ConcreteBodyAnnotation(body))
   }
 
   /** Optionally, the inline info attached to the `@inline` annotation of `sym`. */
@@ -262,7 +261,7 @@ object Inliner {
 
   /** Optionally, the inline body attached to the `@inline` annotation of `sym`. */
   private def inlineBody(sym: SymDenotation)(implicit ctx: Context): Option[Tree] =
-    sym.getAnnotation(defn.InlineAnnot).get.tree.getAttachment(InlineBody)
+    sym.getAnnotation(defn.BodyAnnot).map(_.tree)
 
     /** Definition is an inline method with a known body to inline (note: definitions coming
    *  from Scala2x class files might be `@inline`, but still lack that body.
@@ -282,7 +281,8 @@ object Inliner {
   def removeInlineAccessors(sym: SymDenotation)(implicit ctx: Context): List[MemberDef] = {
     val inlineAnnotTree = sym.getAnnotation(defn.InlineAnnot).get.tree
     val inlineInfo = inlineAnnotTree.removeAttachment(InlineInfo).get
-    inlineAnnotTree.putAttachment(InlineBody, inlineInfo.body)
+    sym.addAnnotation(ConcreteBodyAnnotation(inlineInfo.body))
+    assert(sym.getAnnotation(defn.BodyAnnot).isDefined)
     inlineInfo.accessors
   }