Fix #4291: Copy sticky attachments in TreeCopier

compiler/src/dotty/tools/dotc/ast/Trees.scala

@@ -31,7 +31,7 @@ object Trees {
   @sharable var ntrees = 0
 
   /** Property key for trees with documentation strings attached */
-  val DocComment = new Property.Key[Comment]
+  val DocComment = new Property.StickyKey[Comment]
 
   @sharable private[this] var nextId = 0 // for debugging
 
@@ -913,10 +913,10 @@ object Trees {
       def postProcess(tree: Tree, copied: untpd.MemberDef): copied.ThisTree[T]
 
       def finalize(tree: Tree, copied: untpd.Tree): copied.ThisTree[T] =
-        postProcess(tree, copied withPos tree.pos)
+        postProcess(tree, copied.withPos(tree.pos).withAttachmentsFrom(tree))
 
       def finalize(tree: Tree, copied: untpd.MemberDef): copied.ThisTree[T] =
-        postProcess(tree, copied withPos tree.pos)
+        postProcess(tree, copied.withPos(tree.pos).withAttachmentsFrom(tree))
 
       def Ident(tree: Tree)(name: Name): Ident = tree match {
         case tree: BackquotedIdent =>

compiler/src/dotty/tools/dotc/util/Attachment.scala

@@ -2,9 +2,12 @@ package dotty.tools.dotc.util
 
 /** A class inheriting from Attachment.Container supports
  *  adding, removing and lookup of attachments. Attachments are typed key/value pairs.
+ *
+ *  Attachments whose key is an instance of `StickyKey` will be kept when the attachments
+ *  are copied using `withAttachmentsFrom`.
  */
 object Attachment {
-  import Property.Key
+  import Property.{Key, StickyKey}
 
   /** An implementation trait for attachments.
    *  Clients should inherit from Container instead.
@@ -88,6 +91,16 @@ object Attachment {
   trait Container extends LinkSource {
     private[Attachment] var next: Link[_] = null
 
+    /** Copy the sticky attachments from `container` to this container. */
+    final def withAttachmentsFrom(container: Container): this.type = {
+      var current: Link[_] = container.next
+      while (current != null) {
+        if (current.key.isInstanceOf[StickyKey[_]]) pushAttachment(current.key, current.value)
+        current = current.next
+      }
+      this
+    }
+
     final def pushAttachment[V](key: Key[V], value: V): Unit = {
       assert(!getAttachment(key).isDefined, s"duplicate attachment for key $key")
       next = new Link(key, value, next)

compiler/src/dotty/tools/dotc/util/Property.scala

@@ -7,4 +7,12 @@ object Property {
 
   /** The class of keys for properties of type V */
   class Key[+V]
-}
+
+  /**
+   * The class of keys for sticky properties of type V
+   *
+   * Sticky properties are properties that should be copied over when their container
+   * is copied.
+   */
+  class StickyKey[+V] extends Key[V]
+}

compiler/test/dotty/tools/dotc/ast/AttachmentsTest.scala

@@ -0,0 +1,63 @@
+package dotty.tools.dotc.ast
+
+import dotty.tools.DottyTest
+import dotty.tools.dotc.ast.Trees._
+import dotty.tools.dotc.util.Property
+import dotty.tools.dotc.transform.PostTyper
+
+import org.junit.Test
+import org.junit.Assert.{assertEquals, assertTrue, fail}
+
+class AttachmentsTests extends DottyTest {
+
+  private val TestKey = new Property.Key[String]
+  private val StickyTestKey = new Property.StickyKey[String]
+  private val StickyTestKey2 = new Property.StickyKey[String]
+
+  @Test
+  def attachmentsAreNotCopiedOver: Unit = {
+    checkCompile("frontend", "class A") {
+      case (PackageDef(_, (clazz: tpd.TypeDef) :: Nil), context) =>
+        assertTrue("Attachment shouldn't be present", clazz.getAttachment(TestKey).isEmpty)
+
+        val msg = "hello"
+        clazz.putAttachment(TestKey, msg)
+        assertEquals(Some(msg), clazz.getAttachment(TestKey))
+
+        val copy = tpd.cpy.TypeDef(clazz)(rhs = tpd.EmptyTree)
+        assertTrue("A copy should have been returned", clazz ne copy)
+        assertTrue("Attachment shouldn't be present", copy.getAttachment(TestKey).isEmpty)
+
+      case _ =>
+        fail
+    }
+  }
+
+  @Test
+  def stickyAttachmentsAreCopiedOver: Unit = {
+    checkCompile("frontend", "class A") {
+      case (PackageDef(_, (clazz: tpd.TypeDef) :: Nil), context) =>
+        assertTrue("Attachment shouldn't be present", clazz.getAttachment(StickyTestKey).isEmpty)
+        assertTrue("Attachment shouldn't be present", clazz.getAttachment(StickyTestKey2).isEmpty)
+        assertTrue("Attachment shouldn't be present", clazz.getAttachment(TestKey).isEmpty)
+
+        val msg = "hello"
+        clazz.putAttachment(StickyTestKey, msg)
+        clazz.putAttachment(TestKey, msg)
+        clazz.putAttachment(StickyTestKey2, msg)
+        assertEquals(Some(msg), clazz.getAttachment(StickyTestKey))
+        assertEquals(Some(msg), clazz.getAttachment(TestKey))
+        assertEquals(Some(msg), clazz.getAttachment(StickyTestKey))
+
+        val copy = tpd.cpy.TypeDef(clazz)(rhs = tpd.EmptyTree)
+        assertTrue("A copy should have been returned", clazz ne copy)
+        assertTrue("Attachment should be present", copy.getAttachment(StickyTestKey).isDefined)
+        assertTrue("Attachment shouldn't be present", copy.getAttachment(TestKey).isEmpty)
+        assertTrue("Attachment should be present", copy.getAttachment(StickyTestKey2).isDefined)
+
+      case _ =>
+        fail
+    }
+  }
+
+}