Optimization: Introduce CompactAnnotation

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

@@ -1035,14 +1035,13 @@ object Capabilities:
       else t match
         case t @ CapturingType(_, _) =>
           mapOver(t)
+        case t @ AnnotatedType(parent, ann: RetainingAnnotation)
+        if ann.isStrict && ann.toCaptureSet.containsCap =>
+          // Applying `this` can cause infinite recursion in some cases during printing.
+          // scalac -Xprint:all tests/pos/i23885/S_1.scala tests/pos/i23885/S_2.scala
+          mapOver(CapturingType(this(parent), ann.toCaptureSet))
         case t @ AnnotatedType(parent, ann) =>
-          val parent1 = this(parent)
-          if ann.symbol.isRetains && ann.tree.toCaptureSet.containsCap then
-            // Applying `this` can cause infinite recursion in some cases during printing.
-            // scalac -Xprint:all tests/pos/i23885/S_1.scala tests/pos/i23885/S_2.scala
-            mapOver(CapturingType(parent1, ann.tree.toCaptureSet))
-          else
-            t.derivedAnnotatedType(parent1, ann)
+          t.derivedAnnotatedType(this(parent), ann)
         case defn.RefinedFunctionOf(_) =>
           t  // stop at dependent function types
         case _ =>

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

@@ -14,14 +14,11 @@ import Annotations.Annotation
 import CaptureSet.VarState
 import Capabilities.*
 import Mutability.isStatefulType
-import StdNames.nme
+import StdNames.{nme, tpnme}
 import config.Feature
 import NameKinds.TryOwnerName
 import typer.ProtoTypes.WildcardSelectionProto
 
-/** Attachment key for capturing type trees */
-private val Captures: Key[CaptureSet] = Key()
-
 /** Are we at checkCaptures phase? */
 def isCaptureChecking(using Context): Boolean =
   ctx.phaseId == Phases.checkCapturesPhaseId
@@ -54,24 +51,15 @@ def ccState(using Context): CCState =
 
 extension (tree: Tree)
 
-  /** Convert a @retains or @retainsByName annotation tree to the capture set it represents.
-   *  For efficience, the result is cached as an Attachment on the tree.
+  /** The type representing the capture set of @retains, @retainsCap or @retainsByName
+   *  annotation tree (represented as an Apply node).
    */
-  def toCaptureSet(using Context): CaptureSet =
-    tree.getAttachment(Captures) match
-      case Some(refs) => refs
-      case None =>
-        val refs = CaptureSet(tree.retainedSet.retainedElements*)
-        tree.putAttachment(Captures, refs)
-        refs
-
-  /** The type representing the capture set of @retains, @retainsCap or @retainsByName annotation. */
-  def retainedSet(using Context): Type =
-    tree match
-      case Apply(TypeApply(_, refs :: Nil), _) => refs.tpe
-      case _ =>
-        if tree.symbol.maybeOwner == defn.RetainsCapAnnot
-        then defn.captureRoot.termRef else NoType
+  def retainedSet(using Context): Type = tree match
+    case Apply(TypeApply(_, refs :: Nil), _) => refs.tpe
+    case _ =>
+      if tree.symbol.maybeOwner == defn.RetainsCapAnnot
+      then defn.captureRoot.termRef
+      else NoType
 
 extension (tp: Type)
 
@@ -96,8 +84,8 @@ extension (tp: Type)
   def retainedElementsRaw(using Context): List[Type] = tp match
     case OrType(tp1, tp2) =>
       tp1.retainedElementsRaw ++ tp2.retainedElementsRaw
-    case AnnotatedType(tp1, ann) if tp1.derivesFrom(defn.Caps_CapSet) && ann.symbol.isRetains =>
-      ann.tree.retainedSet.retainedElementsRaw
+    case AnnotatedType(tp1, ann: RetainingAnnotation) if tp1.derivesFrom(defn.Caps_CapSet) =>
+      ann.retainedType.retainedElementsRaw
     case tp =>
       tp.dealiasKeepAnnots match
         case tp: TypeRef if tp.symbol == defn.Caps_CapSet =>
@@ -239,10 +227,12 @@ extension (tp: Type)
       case tp @ CapturingType(parent, refs) =>
         if tp.isBoxed || parent.derivesFrom(defn.Caps_CapSet) then tp
         else tp.boxed
+      case tp @ AnnotatedType(parent, ann: RetainingAnnotation)
+      if !parent.derivesFrom(defn.Caps_CapSet) =>
+        assert(ann.isStrict)
+        CapturingType(parent, ann.toCaptureSet, boxed = true)
       case tp @ AnnotatedType(parent, ann) =>
-        if ann.symbol.isRetains && !parent.derivesFrom(defn.Caps_CapSet)
-        then CapturingType(parent, ann.tree.toCaptureSet, boxed = true)
-        else tp.derivedAnnotatedType(parent.boxDeeply, ann)
+        tp.derivedAnnotatedType(parent.boxDeeply, ann)
       case tp: (Capability & SingletonType) if tp.isTrackableRef && !tp.isAlwaysPure =>
         recur(CapturingType(tp, CaptureSet(tp)))
       case tp1 @ AppliedType(tycon, args) if defn.isNonRefinedFunction(tp1) =>
@@ -502,8 +492,7 @@ extension (cls: ClassSymbol)
       defn.pureBaseClasses.contains(bc)
       || bc.is(CaptureChecked)
           && bc.givenSelfType.dealiasKeepAnnots.match
-            case CapturingType(_, refs) => refs.isAlwaysEmpty
-            case RetainingType(_, refs) => refs.retainedElements.isEmpty
+            case CapturingOrRetainsType(_, refs) => refs.isAlwaysEmpty
             case selfType =>
               isCaptureChecking  // At Setup we have not processed self types yet, so
                                  // unless a self type is explicitly given, we can't tell
@@ -540,13 +529,23 @@ extension (cls: ClassSymbol)
 
 extension (sym: Symbol)
 
-  /** This symbol is one of `retains` or `retainsCap` */
+  private def inScalaAnnotation(using Context): Boolean =
+    sym.maybeOwner.name == tpnme.annotation
+    && sym.owner.owner == defn.ScalaPackageClass
+
+  /** Is this symbol one of `retains` or `retainsCap`?
+   *  Try to avoid cycles by not forcing definition symbols except scala package.
+   */
   def isRetains(using Context): Boolean =
-    sym == defn.RetainsAnnot || sym == defn.RetainsCapAnnot
+    (sym.name == tpnme.retains || sym.name == tpnme.retainsCap)
+    && inScalaAnnotation
 
-  /** This symbol is one of `retains`, `retainsCap`, or`retainsByName` */
+  /** Is this symbol one of `retains`, `retainsCap`, or`retainsByName`?
+   *  Try to avoid cycles by not forcing definition symbols except scala package.
+   */
   def isRetainsLike(using Context): Boolean =
-    isRetains || sym == defn.RetainsByNameAnnot
+    (sym.name == tpnme.retains || sym.name == tpnme.retainsCap || sym.name == tpnme.retainsByName)
+    && inScalaAnnotation
 
   /** A class is pure if:
    *   - one its base types has an explicitly declared self type with an empty capture set
@@ -653,16 +652,16 @@ class PathSelectionProto(val select: Select, val pt: Type) extends typer.ProtoTy
   def selector(using Context): Symbol = select.symbol
 
 /** Drop retains annotations in the inferred type if CC is not enabled
- *  or transform them into RetainingTypes if CC is enabled.
+ *  or transform them into retains annotations with Nothing (i.e. empty set) as
+ *   argument if CC is enabled (we need to do that to keep by-name status).
  */
 class CleanupRetains(using Context) extends TypeMap:
   def apply(tp: Type): Type = tp match
-    case AnnotatedType(parent, annot) if annot.symbol.isRetainsLike =>
+    case tp @ AnnotatedType(parent, annot: RetainingAnnotation) =>
       if Feature.ccEnabled then
-        if annot.symbol == defn.RetainsAnnot || annot.symbol == defn.RetainsByNameAnnot then
-          RetainingType(parent, defn.NothingType, byName = annot.symbol == defn.RetainsByNameAnnot)
-        else mapOver(tp)
-      else apply(parent)
+        if annot.symbol == defn.RetainsCapAnnot then tp
+        else AnnotatedType(this(parent), RetainingAnnotation(annot.symbol.asClass, defn.NothingType))
+      else this(parent)
     case _ => mapOver(tp)
 
 /** A base class for extractors that match annotated types with a specific

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

@@ -536,10 +536,6 @@ sealed abstract class CaptureSet extends Showable:
   /** More info enabled by -Y flags */
   def optionalInfo(using Context): String = ""
 
-  /** A regular @retains or @retainsByName annotation with the elements of this set as arguments. */
-  def toRegularAnnotation(cls: Symbol)(using Context): Annotation =
-    Annotation(CaptureAnnotation(this, boxed = false)(cls).tree)
-
   override def toText(printer: Printer): Text =
     printer.toTextCaptureSet(this) ~~ description
 
@@ -1676,10 +1672,8 @@ object CaptureSet:
         case tp: (TypeRef | TypeParamRef) =>
           if tp.derivesFrom(defn.Caps_CapSet) then tp.captureSet
           else empty
-        case CapturingType(parent, refs) =>
+        case CapturingOrRetainsType(parent, refs) =>
           recur(parent) ++ refs
-        case tp @ AnnotatedType(parent, ann) if ann.symbol.isRetains =>
-          recur(parent) ++ ann.tree.toCaptureSet
         case tpd @ defn.RefinedFunctionOf(rinfo: MethodOrPoly) if followResult =>
           ofType(tpd.parent, followResult = false)            // pick up capture set from parent type
           ++ recur(rinfo.resType).freeInResult(rinfo)         // add capture set of result

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

@@ -60,7 +60,7 @@ object CapturingType:
     case AnnotatedType(parent, ann: CaptureAnnotation)
     if isCaptureCheckingOrSetup =>
       Some((parent, ann.refs))
-    case AnnotatedType(parent, ann) if ann.symbol.isRetains && alsoRetains =>
+    case AnnotatedType(parent, ann: RetainingAnnotation) if ann.isStrict && alsoRetains =>
       // There are some circumstances where we cannot map annotated types
       // with retains annotations to capturing types, so this second recognizer
       // path still has to exist. One example is when checking capture sets
@@ -75,7 +75,7 @@ object CapturingType:
       //
       // TODO In other situations we expect that the type is already transformed to a
       // CapturingType and we should crash if this not the case.
-      try Some((parent, ann.tree.toCaptureSet))
+      try Some((parent, ann.toCaptureSet))
       catch case ex: IllegalCaptureRef => None
     case _ =>
       None

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

@@ -0,0 +1,49 @@
+package dotty.tools
+package dotc
+package cc
+
+import core.*
+import Types.*, Symbols.*, Contexts.*
+import Annotations.{Annotation, CompactAnnotation, EmptyAnnotation}
+import config.Feature
+
+/** A class for annotations @retains, @retainsByName and @retainsCap */
+class RetainingAnnotation(tpe: Type) extends CompactAnnotation(tpe):
+
+  def this(cls: ClassSymbol, args: Type*)(using Context) = this(cls.typeRef.appliedTo(args.toList))
+
+  /** 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.
+   */
+  override protected def sanitize(tp: Type)(using Context): Type = tp match
+    case SkolemType(_) =>
+      SkolemType(defn.AnyType)
+    case tp @ AnnotatedType(parent, ann: RetainingAnnotation)
+    if parent.typeSymbol != defn.Caps_CapSet && ann.symbol != defn.RetainsCapAnnot =>
+      AnnotatedType(parent, RetainingAnnotation(defn.RetainsCapAnnot))
+    case tp @ OrType(tp1, tp2) =>
+      tp.derivedOrType(sanitize(tp1), sanitize(tp2))
+    case _ =>
+      tp
+
+  override def mapWith(tm: TypeMap)(using Context): Annotation =
+    if Feature.ccEnabledSomewhere then mapWithCtd(tm) else EmptyAnnotation
+
+  def isStrict(using Context): Boolean = symbol.isRetains
+
+  def retainedType(using Context): Type =
+    if symbol == defn.RetainsCapAnnot then defn.captureRoot.termRef
+    else argumentType(0)
+
+  private var myCaptureSet: CaptureSet | Null = null
+
+  def toCaptureSet(using Context): CaptureSet =
+    if myCaptureSet == null then
+      myCaptureSet = CaptureSet(retainedType.retainedElements*)
+    myCaptureSet.nn
+
+end RetainingAnnotation

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/Setup.scala

@@ -433,23 +433,21 @@ class Setup extends PreRecheck, SymTransformer, SetupAPI:
           case t @ CapturingType(parent, refs) =>
             checkRetainsOK:
               t.derivedCapturingType(stripImpliedCaptureSet(this(parent)), refs)
-          case t @ AnnotatedType(parent, ann) =>
-            val parent1 = this(parent)
-            if ann.symbol.isRetains then
-              val parent2 = stripImpliedCaptureSet(parent1)
+          case t @ AnnotatedType(parent, ann: RetainingAnnotation) if ann.isStrict =>
+            val parent1 = stripImpliedCaptureSet(this(parent))
+            if !tptToCheck.isEmpty then
+              checkWellformedLater(parent1, ann, tptToCheck)
+            try
+              checkRetainsOK:
+                CapturingType(parent1, ann.toCaptureSet)
+            catch case ex: IllegalCaptureRef =>
               if !tptToCheck.isEmpty then
-                checkWellformedLater(parent2, ann.tree, tptToCheck)
-              try
-                checkRetainsOK:
-                  CapturingType(parent2, ann.tree.toCaptureSet)
-              catch case ex: IllegalCaptureRef =>
-                if !tptToCheck.isEmpty then
-                  report.error(em"Illegal capture reference: ${ex.getMessage}", tptToCheck.srcPos)
-                parent2
-            else if ann.symbol == defn.UncheckedCapturesAnnot then
-              makeUnchecked(apply(parent))
-            else
-              t.derivedAnnotatedType(parent1, ann)
+                report.error(em"Illegal capture reference: ${ex.getMessage}", tptToCheck.srcPos)
+              parent1
+          case t @ AnnotatedType(parent, ann) =>
+            if ann.symbol == defn.UncheckedCapturesAnnot
+            then makeUnchecked(this(parent))
+            else t.derivedAnnotatedType(this(parent), ann)
           case throwsAlias(res, exc) =>
             this(expandThrowsAlias(res, exc, Nil))
           case t @ AppliedType(tycon, args)
@@ -815,10 +813,8 @@ class Setup extends PreRecheck, SymTransformer, SetupAPI:
    */
   private def instanceCanBeImpure(tp: Type)(using Context): Boolean = {
     tp.dealiasKeepAnnots match
-      case CapturingType(_, refs) =>
+      case CapturingOrRetainsType(_, refs) =>
         !refs.isAlwaysEmpty
-      case RetainingType(parent, refs) =>
-        !refs.retainedElements.isEmpty
       case tp: (TypeRef | AppliedType) =>
         val sym = tp.typeSymbol
         if sym.isClass
@@ -858,15 +854,10 @@ class Setup extends PreRecheck, SymTransformer, SetupAPI:
         needsVariable(tp.tp1) && needsVariable(tp.tp2)
       case tp: OrType =>
         needsVariable(tp.tp1) || needsVariable(tp.tp2)
-      case CapturingType(parent, refs) =>
+      case CapturingOrRetainsType(parent, refs) =>
         needsVariable(parent)
         && refs.isConst       // if refs is a variable, no need to add another
         && !refs.isUniversal  // if refs is {cap}, an added variable would not change anything
-      case RetainingType(parent, refs) =>
-        needsVariable(parent)
-        && !refs.retainedElements.exists:
-            case ref: TermRef => ref.isCapRef
-            case _ => false
       case AnnotatedType(parent, _) =>
         needsVariable(parent)
       case _ =>
@@ -972,15 +963,13 @@ class Setup extends PreRecheck, SymTransformer, SetupAPI:
    *  @param ann      the original retains annotation
    *  @param tpt      the tree for which an error or warning should be reported
    */
-  private def checkWellformed(parent: Type, ann: Tree, tpt: Tree)(using Context): Unit =
-    capt.println(i"checkWF post $parent ${ann.retainedSet} in $tpt")
+  private def checkWellformed(parent: Type, ann: RetainingAnnotation, tpt: Tree)(using Context): Unit =
+    capt.println(i"checkWF post $parent ${ann.retainedType} in $tpt")
     try
-      var retained = ann.retainedSet.retainedElements.toArray
+      var retained = ann.retainedType.retainedElements.toArray
       for i <- 0 until retained.length do
         val ref = retained(i)
-        def pos =
-          if ann.span.exists then ann.srcPos
-          else tpt.srcPos
+        def pos = tpt.srcPos
 
         def check(others: CaptureSet, dom: Type | CaptureSet): Unit =
           if others.accountsFor(ref) then
@@ -1013,7 +1002,7 @@ class Setup extends PreRecheck, SymTransformer, SetupAPI:
    *  recheck because we find out only then whether capture sets are empty or
    *  capabilities are redundant.
    */
-  private def checkWellformedLater(parent: Type, ann: Tree, tpt: Tree)(using Context): Unit =
+  private def checkWellformedLater(parent: Type, ann: RetainingAnnotation, tpt: Tree)(using Context): Unit =
     if !tpt.span.isZeroExtent && enclosingInlineds.isEmpty then
       todoAtPostCheck += (ctx1 =>
         checkWellformed(parent, ann, tpt)(using ctx1.withOwner(ctx.owner)))