Don't detour over tree when going from RetainingAnnotation to capture set

Author: odersky

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

@@ -54,28 +54,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.
    */
-  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 =
     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 tree: TypeTree =>
-        tree.tpe match
-          case AppliedType(_, refs :: Nil) => refs
-          case _ => NoType
       case _ => NoType
 
 extension (tp: Type)
@@ -101,8 +88,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.isStrict =>
+      ann.retainedType.retainedElementsRaw
     case tp =>
       tp.dealiasKeepAnnots match
         case tp: TypeRef if tp.symbol == defn.Caps_CapSet =>
@@ -244,10 +231,11 @@ 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 ann.isStrict && !parent.derivesFrom(defn.Caps_CapSet) =>
+        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) =>

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

@@ -1674,8 +1674,8 @@ object CaptureSet:
           else empty
         case CapturingType(parent, refs) =>
           recur(parent) ++ refs
-        case tp @ AnnotatedType(parent, ann) if ann.symbol.isRetains =>
-          recur(parent) ++ ann.tree.toCaptureSet
+        case tp @ AnnotatedType(parent, ann: RetainingAnnotation) if ann.isStrict =>
+          recur(parent) ++ ann.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

@@ -31,8 +31,13 @@ class RetainingAnnotation(tpe: Type) extends CompactAnnotation(tpe):
     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)
+
     def toCaptureSet(using Context): CaptureSet =
-      CaptureSet(argumentType(0).retainedElements*)
-*/
+      CaptureSet(retainedType.retainedElements*)
+
 end RetainingAnnotation

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

@@ -14,12 +14,7 @@ object RetainingType:
     val annotCls = if byName then defn.RetainsByNameAnnot else defn.RetainsAnnot
     AnnotatedType(tp, RetainingAnnotation(annotCls.typeRef.appliedTo(typeElems)))
 
-  def unapply(tp: AnnotatedType)(using Context): Option[(Type, Type)] =
-    val sym = tp.annot.symbol
-    if sym.isRetainsLike then
-      tp.annot match
-        case _: CaptureAnnotation => None
-        case ann => Some((tp.parent, ann.tree.retainedSet))
-    else
-      None
+  def unapply(tp: AnnotatedType)(using Context): Option[(Type, Type)] = tp.annot match
+    case ann: RetainingAnnotation => Some((tp.parent, ann.retainedType))
+    case _ => None
 end RetainingType

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)
@@ -972,15 +970,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 +1009,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)))

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

@@ -317,8 +317,8 @@ class PlainPrinter(_ctx: Context) extends Printer {
         }
       case ExprType(restp) =>
         def arrowText: Text = restp match
-          case AnnotatedType(parent, ann) if ann.symbol == defn.RetainsByNameAnnot =>
-            ann.tree.retainedSet.retainedElementsRaw match
+          case AnnotatedType(parent, ann: RetainingAnnotation) if !ann.isStrict =>
+            ann.retainedType.retainedElementsRaw match
               case ref :: Nil if ref.isCapRef => Str("=>")
               case refs => Str("->") ~ toTextRetainedElems(refs)
           case _ =>

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

@@ -705,7 +705,7 @@ class RefinedPrinter(_ctx: Context) extends PlainPrinter(_ctx) {
         toTextTemplate(tree)
       case Annotated(arg, annot) =>
         def captureSet =
-          annot.asInstanceOf[tpd.Tree].toCaptureSet
+          CaptureSet(annot.asInstanceOf[tpd.Tree].retainedSet.retainedElements*)
         def toTextAnnot =
           toTextLocal(arg) ~~ annotText(annot.symbol.enclosingClass, annot)
         def toTextRetainsAnnot =