Fix posttyper

Author: noti0na1

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

@@ -59,22 +59,14 @@ extension (tree: Tree)
         tree.putAttachment(Captures, refs)
         refs
 
+  /** The type representing the capture set of retains annotation.
+    */
   def retainedSet(using Context): Type =
     tree match
       case Apply(TypeApply(_, refs :: Nil), _) => refs.tpe
       case _ =>
         if tree.symbol.maybeOwner == defn.RetainsCapAnnot
-        then root.cap
-        else NoType
-
-  /** The arguments of a @retains, @retainsCap or @retainsByName annotation */
-  def retainedElems(using Context): List[Tree] = tree match
-    case Apply(_, Typed(SeqLiteral(elems, _), _) :: Nil) =>
-      elems
-    case _ =>
-      if tree.symbol.maybeOwner == defn.RetainsCapAnnot
-      then ref(root.cap) :: Nil
-      else Nil
+        then root.cap else NoType
 
 extension (tp: Type)
 

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

@@ -311,9 +311,9 @@ 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.retainedElems match
-              case ref :: Nil if ref.symbol == defn.captureRoot => Str("=>")
-              case refs => Str("->") ~ toTextRetainedElems(refs)
+            ann.tree.retainedSet.retainedElements match
+              case ref :: Nil if ref.isCap => Str("=>")
+              case refs => Str("->") ~ toTextRetainedElems(refs.map(r => ast.tpd.TypeTree(r)))
           case _ =>
             if Feature.pureFunsEnabled then "->" else "=>"
         changePrec(GlobalPrec)(arrowText ~ " " ~ toText(restp))

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

@@ -196,7 +196,10 @@ class PostTyper extends MacroTransform with InfoTransformer { thisPhase =>
       val saved = inJavaAnnot
       inJavaAnnot = annot.symbol.is(JavaDefined)
       if (inJavaAnnot) checkValidJavaAnnotation(annot)
-      try transform(annot)
+      try
+        val annotCtx = if annot.hasAttachment(untpd.RetainsAnnot)
+          then ctx.addMode(Mode.InCaptureSet) else ctx
+        transform(annot)(using annotCtx)
       finally inJavaAnnot = saved
     }
 

compiler/src/dotty/tools/dotc/typer/Typer.scala

@@ -2532,7 +2532,6 @@ class Typer(@constructorOnly nestingLevel: Int = 0) extends Namer
   def typedSingletonTypeTree(tree: untpd.SingletonTypeTree)(using Context): Tree = {
     val ref1 = typedExpr(tree.ref, SingletonTypeProto)
     if ctx.mode.is(Mode.InCaptureSet) && ref1.symbol.is(Flags.CaptureParam) then
-      // println(s"typedSingletonTypeTree: $ref1 -> ${ref1.tpe.widen}")
       return Ident(ref1.tpe.widen.asInstanceOf[TypeRef]).withSpan(tree.span)
     checkStable(ref1.tpe, tree.srcPos, "singleton type")
     assignType(cpy.SingletonTypeTree(tree)(ref1), ref1)