Improve mapping and pickling of annotated types

`Annotation.mapWith` maps an `Annotation` with a type map `tm`. As an optimization, this function first checks if `tm` would result in any change (by traversing the annotation’s argument trees with a `TreeAccumulator`) before applying `tm` to the whole annotation tree. This optimization had two problems: 1. it didn’t include type parameters, and  2. it used `frozen_=:=`  to compare types, which didn’t work as expected with `NoType`. This commit fixes these issues.

Additionally, positions of trees that appear only inside `AnnotatedType` were not pickled. This commit also fixes this.

Author: mbovel

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

@@ -141,9 +141,17 @@ trait TreeInfo[T <: Untyped] { self: Trees.Instance[T] =>
     loop(tree, Nil)
 
   /** All term arguments of an application in a single flattened list */
+  def allTermArguments(tree: Tree): List[Tree] = unsplice(tree) match {
+    case Apply(fn, args) => allArguments(fn) ::: args
+    case TypeApply(fn, args) => allArguments(fn)
+    case Block(_, expr) => allArguments(expr)
+    case _ => Nil
+  }
+
+  /** All type and term arguments of an application in a single flattened list */
   def allArguments(tree: Tree): List[Tree] = unsplice(tree) match {
     case Apply(fn, args) => allArguments(fn) ::: args
-    case TypeApply(fn, _) => allArguments(fn)
+    case TypeApply(fn, args) => allArguments(fn) ::: args
     case Block(_, expr) => allArguments(expr)
     case _ => Nil
   }

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

@@ -30,8 +30,8 @@ object Annotations {
     def derivedAnnotation(tree: Tree)(using Context): Annotation =
       if (tree eq this.tree) this else Annotation(tree)
 
-    /** All arguments to this annotation in a single flat list */
-    def arguments(using Context): List[Tree] = tpd.allArguments(tree)
+    /** All term arguments of this annotation in a single flat list */
+    def arguments(using Context): List[Tree] = tpd.allTermArguments(tree)
 
     def argument(i: Int)(using Context): Option[Tree] = {
       val args = arguments
@@ -54,15 +54,18 @@ object Annotations {
      *  type, since ranges cannot be types of trees.
      */
     def mapWith(tm: TypeMap)(using Context) =
-      val args = arguments
+      val args = tpd.allArguments(tree)
       if args.isEmpty then this
       else
+        // Checks if `tm` would result in any change by applying it to types
+        // inside the annotations' arguments and checking if the resulting types
+        // are different.
         val findDiff = new TreeAccumulator[Type]:
           def apply(x: Type, tree: Tree)(using Context): Type =
             if tm.isRange(x) then x
             else
               val tp1 = tm(tree.tpe)
-              foldOver(if tp1 frozen_=:= tree.tpe then x else tp1, tree)
+              foldOver(if !tp1.exists || (tp1 frozen_=:= tree.tpe) then x else tp1, tree)
         val diff = findDiff(NoType, args)
         if tm.isRange(diff) then EmptyAnnotation
         else if diff.exists then derivedAnnotation(tm.mapOver(tree))

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

@@ -33,6 +33,7 @@ object PositionPickler:
       pickler: TastyPickler,
       addrOfTree: TreeToAddr,
       treeAnnots: untpd.MemberDef => List[tpd.Tree],
+      typeAnnots: List[tpd.Tree],
       relativePathReference: String,
       source: SourceFile,
       roots: List[Tree],
@@ -136,6 +137,9 @@ object PositionPickler:
     }
     for (root <- roots)
       traverse(root, NoSource)
+
+    for annotTree <- typeAnnots do
+      traverse(annotTree, NoSource)
   end picklePositions
 end PositionPickler
 

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

@@ -41,6 +41,10 @@ class TreePickler(pickler: TastyPickler, attributes: Attributes) {
    */
   private val annotTrees = util.EqHashMap[untpd.MemberDef, mutable.ListBuffer[Tree]]()
 
+  /** A set of annotation trees appearing in annotated types.
+   */
+  private val annotatedTypeTrees = mutable.ListBuffer[Tree]()
+
   /** A map from member definitions to their doc comments, so that later
    *  parallel comment pickling does not need to access symbols of trees (which
    *  would involve accessing symbols of named types and possibly changing phases
@@ -57,6 +61,8 @@ class TreePickler(pickler: TastyPickler, attributes: Attributes) {
     val ts = annotTrees.lookup(tree)
     if ts == null then Nil else ts.toList
 
+  def typeAnnots: List[Tree] = annotatedTypeTrees.toList
+
   def docString(tree: untpd.MemberDef): Option[Comment] =
     Option(docStrings.lookup(tree))
 
@@ -278,6 +284,7 @@ class TreePickler(pickler: TastyPickler, attributes: Attributes) {
     case tpe: AnnotatedType =>
       writeByte(ANNOTATEDtype)
       withLength { pickleType(tpe.parent, richTypes); pickleTree(tpe.annot.tree) }
+      annotatedTypeTrees += tpe.annot.tree
     case tpe: AndType =>
       writeByte(ANDtype)
       withLength { pickleType(tpe.tp1, richTypes); pickleType(tpe.tp2, richTypes) }

compiler/src/dotty/tools/dotc/quoted/PickledQuotes.scala

@@ -224,7 +224,7 @@ object PickledQuotes {
     if tree.span.exists then
       val positionWarnings = new mutable.ListBuffer[Message]()
       val reference = ctx.settings.sourceroot.value
-      PositionPickler.picklePositions(pickler, treePkl.buf.addrOfTree, treePkl.treeAnnots, reference,
+      PositionPickler.picklePositions(pickler, treePkl.buf.addrOfTree, treePkl.treeAnnots, treePkl.typeAnnots, reference,
         ctx.compilationUnit.source, tree :: Nil, positionWarnings)
       positionWarnings.foreach(report.warning(_))
 

compiler/src/dotty/tools/dotc/transform/Pickler.scala

@@ -322,7 +322,7 @@ class Pickler extends Phase {
           if tree.span.exists then
             val reference = ctx.settings.sourceroot.value
             PositionPickler.picklePositions(
-                pickler, treePkl.buf.addrOfTree, treePkl.treeAnnots, reference,
+                pickler, treePkl.buf.addrOfTree, treePkl.treeAnnots, treePkl.typeAnnots, reference,
                 unit.source, tree :: Nil, positionWarnings,
                 scratch.positionBuffer, scratch.pickledIndices)
 

tests/pos/annot-17939b.scala

@@ -0,0 +1,10 @@
+import scala.annotation.Annotation
+class myRefined(f: ? => Boolean) extends Annotation
+
+def test(axes: Int) = true
+
+trait Tensor:
+  def mean(axes: Int): Int @myRefined(_ => test(axes))
+
+class TensorImpl() extends Tensor:
+  def mean(axes: Int) = ???

tests/pos/annot-18064.scala

@@ -0,0 +1,9 @@
+//> using options "-Xprint:typer"
+
+class myAnnot[T]() extends annotation.Annotation
+
+trait Tensor[T]:
+  def add: Tensor[T] @myAnnot[T]()
+
+class TensorImpl[A]() extends Tensor[A]:
+  def add /* : Tensor[A] @myAnnot[A] */ = this

tests/pos/annot-5789.scala

@@ -0,0 +1,10 @@
+class Annot[T] extends scala.annotation.Annotation
+
+class D[T](val f: Int@Annot[T])
+
+object A{
+  def main(a:Array[String]) = {
+    val c = new D[Int](1)
+    c.f
+  }
+}

tests/printing/annot-18064.check

@@ -0,0 +1,16 @@
+[[syntax trees at end of                     typer]] // tests/printing/annot-18064.scala
+package <empty> {
+  class myAnnot[T >: Nothing <: Any]() extends annotation.Annotation() {
+    T
+  }
+  trait Tensor[T >: Nothing <: Any]() extends Object {
+    T
+    def add: Tensor[Tensor.this.T] @myAnnot[T]
+  }
+  class TensorImpl[A >: Nothing <: Any]() extends Object(), Tensor[
+    TensorImpl.this.A] {
+    A
+    def add: Tensor[A] @myAnnot[A] = this
+  }
+}
+

tests/printing/annot-18064.scala

@@ -0,0 +1,9 @@
+//> using options "-Xprint:typer"
+
+class myAnnot[T]() extends annotation.Annotation
+
+trait Tensor[T]:
+  def add: Tensor[T] @myAnnot[T]()
+
+class TensorImpl[A]() extends Tensor[A]:
+  def add /* : Tensor[A] @myAnnot[A] */ = this

tests/printing/annot-19846b.check

@@ -0,0 +1,33 @@
+[[syntax trees at end of                     typer]] // tests/printing/annot-19846b.scala
+package <empty> {
+  class lambdaAnnot(g: () => Int) extends scala.annotation.Annotation(),
+    annotation.StaticAnnotation {
+    private[this] val g: () => Int
+  }
+  final lazy module val Test: Test = new Test()
+  final module class Test() extends Object() { this: Test.type =>
+    val y: Int = ???
+    val z:
+      Int @lambdaAnnot(
+        {
+          def $anonfun(): Int = Test.y
+          closure($anonfun)
+        }
+      )
+     = f(Test.y)
+  }
+  final lazy module val annot-19846b$package: annot-19846b$package =
+    new annot-19846b$package()
+  final module class annot-19846b$package() extends Object() {
+    this: annot-19846b$package.type =>
+    def f(x: Int):
+      Int @lambdaAnnot(
+        {
+          def $anonfun(): Int = x
+          closure($anonfun)
+        }
+      )
+     = x
+  }
+}
+

tests/printing/annot-19846b.scala

@@ -0,0 +1,7 @@
+class lambdaAnnot(g: () => Int) extends annotation.StaticAnnotation
+
+def f(x: Int): Int @lambdaAnnot(() => x) = x
+
+object Test:
+  val y: Int = ???
+  val z /* : Int @lambdaAnnot(() => y) */ = f(y)