Fixes for cleanup retains scheme

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

@@ -778,6 +778,7 @@ object Trees {
     override def isEmpty: Boolean = !hasType
     override def toString: String =
       s"TypeTree${if (hasType) s"[$typeOpt]" else ""}"
+    def isInferred = false
   }
 
   /** Tree that replaces a level 1 splices in pickled (level 0) quotes.
@@ -800,6 +801,7 @@ object Trees {
    */
   class InferredTypeTree[+T <: Untyped](implicit @constructorOnly src: SourceFile) extends TypeTree[T]:
     type ThisTree[+T <: Untyped] <: InferredTypeTree[T]
+    override def isInferred = true
 
   /** ref.type */
   case class SingletonTypeTree[+T <: Untyped] private[ast] (ref: Tree[T])(implicit @constructorOnly src: SourceFile)

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

@@ -806,10 +806,10 @@ class TreePickler(pickler: TastyPickler, attributes: Attributes) {
           report.error(ex.toMessage, tree.srcPos.focus)
           pickleErrorType()
         case ex: AssertionError =>
-          println(i"error when pickling tree $tree")
+          println(i"error when pickling tree $tree of class ${tree.getClass}")
           throw ex
         case ex: MatchError =>
-          println(i"error when pickling tree $tree")
+          println(i"error when pickling tree $tree of class ${tree.getClass}")
           throw ex
       }
   }

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

@@ -303,20 +303,19 @@ class PostTyper extends MacroTransform with InfoTransformer { thisPhase =>
       if !tree.symbol.is(Package) then tree
       else errorTree(tree, em"${tree.symbol} cannot be used as a type")
 
-    // Cleans up retains annotations in inferred type trees. This is needed because
-    // during the typer, it is infeasible to correctly infer the capture sets in most
-    // cases, resulting ill-formed capture sets that could crash the pickler later on.
-    // See #20035.
-    private def cleanupRetainsAnnot(symbol: Symbol, tpt: Tree)(using Context): Tree =
+    /** Make result types of ValDefs and DefDefs that override some other definitions
+     *  declared types rather than InferredTypes. This is necessary since we otherwise
+     *  clean retains annotations from such types. But for an overriding symbol the
+     *  retains annotations come from the explicitly declared parent types, so should
+     *  be kept.
+     */
+    private def makeOverrideTypeDeclared(symbol: Symbol, tpt: Tree)(using Context): Tree =
       tpt match
         case tpt: InferredTypeTree
-        if !symbol.allOverriddenSymbols.hasNext =>
-          // if there are overridden symbols, the annotation comes from an explicit type of the overridden symbol
-          // and should be retained.
-          val tm = new CleanupRetains
-          val tpe1 = tm(tpt.tpe)
-          tpt.withType(tpe1)
-        case _ => tpt
+        if symbol.allOverriddenSymbols.hasNext =>
+          TypeTree(tpt.tpe, inferred = false).withSpan(tpt.span).withAttachmentsFrom(tpt)
+        case _ =>
+          tpt
 
     override def transform(tree: Tree)(using Context): Tree =
       try tree match {
@@ -432,7 +431,7 @@ class PostTyper extends MacroTransform with InfoTransformer { thisPhase =>
           registerIfHasMacroAnnotations(tree)
           checkErasedDef(tree)
           Checking.checkPolyFunctionType(tree.tpt)
-          val tree1 = cpy.ValDef(tree)(tpt = cleanupRetainsAnnot(tree.symbol, tree.tpt), rhs = normalizeErasedRhs(tree.rhs, tree.symbol))
+          val tree1 = cpy.ValDef(tree)(tpt = makeOverrideTypeDeclared(tree.symbol, tree.tpt), rhs = normalizeErasedRhs(tree.rhs, tree.symbol))
           if tree1.removeAttachment(desugar.UntupledParam).isDefined then
             checkStableSelection(tree.rhs)
           processValOrDefDef(super.transform(tree1))
@@ -441,7 +440,7 @@ class PostTyper extends MacroTransform with InfoTransformer { thisPhase =>
           checkErasedDef(tree)
           Checking.checkPolyFunctionType(tree.tpt)
           annotateContextResults(tree)
-          val tree1 = cpy.DefDef(tree)(tpt = cleanupRetainsAnnot(tree.symbol, tree.tpt), rhs = normalizeErasedRhs(tree.rhs, tree.symbol))
+          val tree1 = cpy.DefDef(tree)(tpt = makeOverrideTypeDeclared(tree.symbol, tree.tpt), rhs = normalizeErasedRhs(tree.rhs, tree.symbol))
           processValOrDefDef(superAcc.wrapDefDef(tree1)(super.transform(tree1).asInstanceOf[DefDef]))
         case tree: TypeDef =>
           registerIfHasMacroAnnotations(tree)
@@ -524,12 +523,12 @@ class PostTyper extends MacroTransform with InfoTransformer { thisPhase =>
               report.error(em"type ${alias.tpe} outside bounds $bounds", tree.srcPos)
           super.transform(tree)
         case tree: TypeTree =>
-          tree.withType(
-            tree.tpe match {
-              case AnnotatedType(tpe, annot) => AnnotatedType(tpe, transformAnnot(annot))
-              case tpe => tpe
-            }
-          )
+          val tpe = if tree.isInferred then CleanupRetains()(tree.tpe) else tree.tpe
+          tree.withType:
+            tpe match
+              case AnnotatedType(parent, annot) =>
+                AnnotatedType(parent, transformAnnot(annot)) // TODO: Also map annotations embedded in type?
+              case _ => tpe
         case Typed(Ident(nme.WILDCARD), _) =>
           withMode(Mode.Pattern)(super.transform(tree))
             // The added mode signals that bounds in a pattern need not

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

@@ -883,7 +883,7 @@ trait Applications extends Compatibility {
     def makeVarArg(n: Int, elemFormal: Type): Unit = {
       val args = typedArgBuf.takeRight(n).toList
       typedArgBuf.dropRightInPlace(n)
-      val elemtpt = TypeTree(elemFormal)
+      val elemtpt = TypeTree(elemFormal, inferred = true)
       typedArgBuf += seqToRepeated(SeqLiteral(args, elemtpt))
     }
 

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

@@ -51,7 +51,7 @@ trait TypeAssigner {
     else sym.info
 
   private def toRepeated(tree: Tree, from: ClassSymbol)(using Context): Tree =
-    Typed(tree, TypeTree(tree.tpe.widen.translateToRepeated(from)))
+    Typed(tree, TypeTree(tree.tpe.widen.translateToRepeated(from), inferred = true))
 
   def seqToRepeated(tree: Tree)(using Context): Tree = toRepeated(tree, defn.SeqClass)
 

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

@@ -1457,7 +1457,7 @@ class Typer(@constructorOnly nestingLevel: Int = 0) extends Namer
         cpy.Block(block)(stats, expr1) withType expr1.tpe // no assignType here because avoid is redundant
       case _ =>
         val target = pt.simplified
-        val targetTpt = InferredTypeTree().withType(target)
+        val targetTpt = TypeTree(target, inferred = true)
         if tree.tpe <:< target then Typed(tree, targetTpt)
         else
           // This case should not normally arise. It currently does arise in test cases
@@ -2092,7 +2092,7 @@ class Typer(@constructorOnly nestingLevel: Int = 0) extends Namer
                 // TODO: move the check above to patternMatcher phase
                 val uncheckedTpe = AnnotatedType(sel.tpe.widen, Annotation(defn.UncheckedAnnot, tree.selector.span))
                 tpd.cpy.Match(result)(
-                  selector = tpd.Typed(sel, new tpd.InferredTypeTree().withType(uncheckedTpe)),
+                  selector = tpd.Typed(sel, tpd.TypeTree(uncheckedTpe, inferred = true)),
                   cases = result.cases
                 )
               case _ =>

library/src/scala/caps.scala

@@ -29,6 +29,9 @@ import annotation.{experimental, compileTimeOnly, retainsCap}
    */
   given containsImpl[C <: CapSet @retainsCap, R <: Singleton]: Contains[C, R]()
 
+  /** A wrapper indicating a type variable in a capture argument list of a
+   *  @retains annotation. E.g. `^{x, Y^}` is represented as `@retains(x, capsOf[Y])`.
+   */
   @compileTimeOnly("Should be be used only internally by the Scala compiler")
   def capsOf[CS]: Any = ???
 

tests/pos-custom-args/captures/cc-poly-varargs.scala

@@ -0,0 +1,20 @@
+trait Cancellable
+
+abstract class Source[+T, Cap^]
+
+extension[T, Cap^](src: Source[T, Cap]^)
+  def transformValuesWith[U](f: (T -> U)^{Cap^}): Source[U, Cap]^{src, f} = ???
+
+def race[T, Cap^](sources: Source[T, Cap]^{Cap^}*): Source[T, Cap]^{Cap^} = ???
+
+def either[T1, T2, Cap^](src1: Source[T1, Cap]^{Cap^}, src2: Source[T2, Cap]^{Cap^}): Source[Either[T1, T2], Cap]^{Cap^} =
+  val left = src1.transformValuesWith(Left(_))
+  val right = src2.transformValuesWith(Right(_))
+  race(left, right)
+
+
+
+
+
+
+