Optimization: Introduce CompactAnnotation

This one takes a type argument instead of as a tree argument. For now
it's reserved for retains-like annotations.

CompactAnnotations don't need tree maps and annotated types containing them
hash properly.

compiler/src/dotty/tools/dotc/cc/CaptureOps.scala

@@ -67,11 +67,16 @@ extension (tree: Tree)
 
   /** The type representing the capture set of @retains, @retainsCap or @retainsByName annotation. */
   def retainedSet(using Context): Type =
-    tree match
+    val rcap = defn.RetainsCapAnnot
+    if tree.symbol == rcap || tree.symbol.maybeOwner == rcap then
+      defn.captureRoot.termRef
+    else tree match
       case Apply(TypeApply(_, refs :: Nil), _) => refs.tpe
-      case _ =>
-        if tree.symbol.maybeOwner == defn.RetainsCapAnnot
-        then defn.captureRoot.termRef else NoType
+      case tree: TypeTree =>
+        tree.tpe match
+          case AppliedType(_, refs :: Nil) => refs
+          case _ => NoType
+      case _ => NoType
 
 extension (tp: Type)
 

compiler/src/dotty/tools/dotc/cc/SepCheck.scala

@@ -919,7 +919,7 @@ class SepCheck(checker: CheckCaptures.CheckerAPI) extends tpd.TreeTraverser:
     assert(mtps.hasSameLengthAs(argss), i"diff for $fn: ${fn.symbol} /// $mtps /// $argss")
     val mtpsWithArgs = mtps.zip(argss)
     val argMap = mtpsWithArgs.toMap
-    val deps = mutable.HashMap[Tree, List[Tree]]().withDefaultValue(Nil)
+    val deps = mutable.LinkedHashMap[Tree, List[Tree]]().withDefaultValue(Nil)
 
     def argOfDep(dep: Capability): Option[Tree] =
       dep.stripReach match

compiler/src/dotty/tools/dotc/cc/ccConfig.scala

@@ -59,7 +59,7 @@ object ccConfig:
 
   /** Not used currently. Handy for trying out new features */
   def newScheme(using ctx: Context): Boolean =
-    Feature.sourceVersion.stable.isAtLeast(SourceVersion.`3.7`)
+    Feature.sourceVersion.stable.isAtLeast(SourceVersion.`3.9`)
 
   def allowUse(using Context): Boolean =
     Feature.sourceVersion.stable.isAtMost(SourceVersion.`3.7`)

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

@@ -33,9 +33,16 @@ object Annotations {
     def derivedAnnotation(tree: Tree)(using Context): Annotation =
       if (tree eq this.tree) this else Annotation(tree)
 
+    def derivedClassAnnotation(cls: ClassSymbol)(using Context) =
+      Annotation(cls, tree.span)
+
     /** All term arguments of this annotation in a single flat list */
     def arguments(using Context): List[Tree] = tpd.allTermArguments(tree)
 
+    /** All type arguments of this annotation in a single flat list */
+    def argumentTypes(using Context): List[Type] =
+      tpd.allArguments(tree).filterConserve(_.isType).tpes
+
     def argument(i: Int)(using Context): Option[Tree] = {
       val args = arguments
       if (i < args.length) Some(args(i)) else None
@@ -66,23 +73,8 @@ object Annotations {
           //     the original tree. TODO Try to use this scheme for other annotations that
           //     take only type arguments as well. We should wait until after 3.9 LTS to
           //     do this, though.
-          //  2. Map all skolems (?n: T) to (?n: Any), and map all recursive captures of
-          //     that are not on CapSet to `^`. Skolems and capturing types on types
-          //     other than CapSet are not allowed in a retains annotation anyway,
-          //     so the underlying type does not matter. This simplification prevents
-          //     exponential blowup in some cases. See i24556.scala and i24556a.scala.
-          //  3. Drop the annotation entirely if CC is not enabled somehwere.
-
-          def sanitize(tp: Type): Type = tp match
-            case SkolemType(_) =>
-              SkolemType(defn.AnyType)
-            case tp @ AnnotatedType(parent, ann)
-            if ann.symbol.isRetainsLike && parent.typeSymbol != defn.Caps_CapSet =>
-              tp.derivedAnnotatedType(parent, Annotation(defn.RetainsCapAnnot, ann.tree.span))
-            case tp @ OrType(tp1, tp2) =>
-              tp.derivedOrType(sanitize(tp1), sanitize(tp2))
-            case _ =>
-              tp
+          //  2. Sanitize the arguments to prevent compilation time blowup.
+          //  3. Drop the annotation entirely if CC is not enabled somewhere.
 
           def rebuild(tree: Tree, mappedType: Type): Tree = tree match
             case Apply(fn, Nil) => cpy.Apply(tree)(rebuild(fn, mappedType), Nil)
@@ -93,8 +85,12 @@ object Annotations {
             EmptyAnnotation // strip retains-like annotations unless capture checking is enabled
           else
             val mappedType = sanitize(tm(arg.tpe))
-            if mappedType `eql` arg.tpe then this
-            else derivedAnnotation(rebuild(tree, mappedType))
+            if mappedType `eql` arg.tpe then
+              this
+            else if cc.ccConfig.newScheme then
+              CompactAnnotation(symbol.typeRef.appliedTo(mappedType))
+            else
+              derivedAnnotation(rebuild(tree, mappedType))
 
         case args =>
           // Checks if `tm` would result in any change by applying it to types
@@ -114,17 +110,12 @@ object Annotations {
 
     /** Does this annotation refer to a parameter of `tl`? */
     def refersToParamOf(tl: TermLambda)(using Context): Boolean =
-      def isLambdaParam(t: Type) = t match
-        case TermParamRef(tl1, _) => tl eq tl1
-        case _ => false
-
       val acc = new TreeAccumulator[Boolean]:
         def apply(x: Boolean, t: Tree)(using Context) =
           if x then true
-          else if t.isType then
-            t.tpe.existsPart(isLambdaParam, stopAt = StopAt.Static)
+          else if t.isType then refersToLambdaParam(t.tpe, tl)
           else t match
-            case id: (Ident | This) => isLambdaParam(id.tpe.stripped)
+            case id: (Ident | This) => isLambdaParam(id.tpe.stripped, tl)
             case _ => foldOver(x, t)
 
       tpd.allArguments(tree).exists(acc(false, _))
@@ -163,6 +154,82 @@ object Annotations {
   case class ConcreteAnnotation(t: Tree) extends Annotation:
     def tree(using Context): Tree = t
 
+  case class CompactAnnotation(tp: Type) extends Annotation:
+    assert(tp.isInstanceOf[AppliedType | TypeRef], tp)
+
+    def tree(using Context) = TypeTree(tp)
+
+    override def symbol(using Context) = tp.typeSymbol
+
+    override def derivedAnnotation(tree: Tree)(using Context): Annotation =
+      derivedAnnotation(tree.tpe)
+
+    override def derivedClassAnnotation(cls: ClassSymbol)(using Context) =
+      derivedAnnotation(cls.typeRef)
+
+    def derivedAnnotation(tp: Type)(using Context): Annotation =
+      if tp eq this.tp then this else CompactAnnotation(tp)
+
+    override def arguments(using Context): List[Tree] =
+      argumentTypes.map(TypeTree(_))
+
+    override def argumentTypes(using Context): List[Type] = tp.argTypes
+
+    def argumentType(i: Int)(using Context): Type =
+      val args = argumentTypes
+      if i < args.length then args(i) else NoType
+
+    override def argumentConstant(i: Int)(using Context): Option[Constant] =
+      argumentType(i).normalized match
+        case ConstantType(c) => Some(c)
+        case _ => None
+
+    override def mapWith(tm: TypeMap)(using Context): Annotation =
+      def derived(tp: Type) =
+        if tm.isRange(tp) then EmptyAnnotation else derivedAnnotation(tp)
+      def sanitizeArg(tp: Type) = tp match
+        case tp @ AppliedType(tycon, args) =>
+          tp.derivedAppliedType(tycon, args.mapConserve(sanitize))
+        case _ =>
+          tp
+      if !symbol.isRetainsLike then derived(tm(tp))
+      else if Feature.ccEnabledSomewhere then derived(sanitizeArg(tm(tp)))
+      else EmptyAnnotation // strip retains-like annotations unless capture checking is enabled
+
+    override def refersToParamOf(tl: TermLambda)(using Context): Boolean =
+      refersToLambdaParam(tp, tl)
+
+    override def hash: Int = tp.hash
+    override def eql(that: Annotation) = that match
+      case that: CompactAnnotation => this.tp `eql` that.tp
+      case _ => false
+  end CompactAnnotation
+
+  /** Sanitize @retains arguments to approximate illegal types that could cause a compilation
+   *  time blowup before they are dropped ot detected. This means mapping all all skolems
+   *  (?n: T) to (?n: Any), and mapping all recursive captures that are not on CapSet to `^`.
+   *  Skolems and capturing types on types other than CapSet are not allowed in a
+   *  @retains annotation anyway, so the underlying type does not matter as long as it is also
+   *  illegal. See i24556.scala and i24556a.scala.
+   */
+  private def sanitize(tp: Type)(using Context): Type = tp match
+    case SkolemType(_) =>
+      SkolemType(defn.AnyType)
+    case tp @ AnnotatedType(parent, ann)
+    if ann.symbol.isRetainsLike && parent.typeSymbol != defn.Caps_CapSet =>
+      tp.derivedAnnotatedType(parent, ann.derivedClassAnnotation(defn.RetainsCapAnnot))
+    case tp @ OrType(tp1, tp2) =>
+      tp.derivedOrType(sanitize(tp1), sanitize(tp2))
+    case _ =>
+      tp
+
+  private def isLambdaParam(t: Type, tl: TermLambda): Boolean = t match
+    case TermParamRef(tl1, _) => tl eq tl1
+    case _ => false
+
+  private def refersToLambdaParam(tp: Type, tl: TermLambda)(using Context): Boolean =
+    tp.existsPart(isLambdaParam(_, tl), stopAt = StopAt.Static)
+
   abstract class LazyAnnotation extends Annotation {
     protected var mySym: Symbol | (Context ?=> Symbol) | Null
     override def symbol(using parentCtx: Context): Symbol =
@@ -244,7 +311,9 @@ object Annotations {
 
   object Annotation {
 
-    def apply(tree: Tree): ConcreteAnnotation = ConcreteAnnotation(tree)
+    def apply(tree: Tree): Annotation = tree match
+      case tree: TypeTree => CompactAnnotation(tree.tpe)
+      case _ => ConcreteAnnotation(tree)
 
     def apply(cls: ClassSymbol, span: Span)(using Context): Annotation =
       apply(cls, Nil, span)

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

@@ -283,8 +283,13 @@ class TreePickler(pickler: TastyPickler, attributes: Attributes) {
       }
     case tpe: AnnotatedType =>
       writeByte(ANNOTATEDtype)
-      withLength { pickleType(tpe.parent, richTypes); pickleTree(tpe.annot.tree) }
-      annotatedTypeTrees += tpe.annot.tree
+      withLength:
+        pickleType(tpe.parent, richTypes)
+        tpe.annot match
+          case CompactAnnotation(tp) => pickleType(tp)
+          case _ =>
+            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/core/tasty/TreeUnpickler.scala

@@ -354,6 +354,11 @@ class TreeUnpickler(reader: TastyReader,
         op
       }
 
+      /** Can `tag` start a type argument of a CompactAnnotation? */
+      def isCompactAnnotTypeTag(tag: Int): Boolean = tag match
+        case APPLIEDtype | SHAREDtype | TYPEREF | TYPEREFdirect | TYPEREFsymbol | TYPEREFin => true
+        case _ => false
+
       def readLengthType(): Type = {
         val end = readEnd()
 
@@ -420,7 +425,12 @@ class TreeUnpickler(reader: TastyReader,
                 val hi = readVariances(readType())
                 createNullableTypeBounds(lo, hi)
             case ANNOTATEDtype =>
-              AnnotatedType(readType(), Annotation(readTree()))
+              val parent = readType()
+              val ann =
+                if isCompactAnnotTypeTag(nextByte)
+                then CompactAnnotation(readType())
+                else Annotation(readTree())
+              AnnotatedType(parent, ann)
             case ANDtype =>
               AndType(readType(), readType())
             case ORtype =>

compiler/src/dotty/tools/dotc/printing/RefinedPrinter.scala

@@ -1165,6 +1165,7 @@ class RefinedPrinter(_ctx: Context) extends PlainPrinter(_ctx) {
       case Select(qual, nme.CONSTRUCTOR) => recur(qual)
       case id @ Ident(tpnme.BOUNDTYPE_ANNOT) => "@" ~ toText(id.symbol.name)
       case New(tpt) => recur(tpt)
+      case t: tpd.TypeTree if t.tpe.isInstanceOf[AppliedType] => "@" ~ toText(t.tpe)
       case _ =>
         val annotSym = sym.orElse(tree.symbol.enclosingClass)
         if annotSym.exists then annotText(annotSym) else s"@${t.show}"