Change capture elements to union type

Author: noti0na1

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

@@ -386,6 +386,8 @@ object untpd extends Trees.Instance[Untyped] with UntypedTreeInfo {
   /** Property key for contextual Apply trees of the form `fn given arg` */
   val KindOfApply: Property.StickyKey[ApplyKind] = Property.StickyKey()
 
+  val RetainsAnnot: Property.StickyKey[Unit] = Property.StickyKey()
+
   // ------ Creation methods for untyped only -----------------
 
   def Ident(name: Name)(implicit src: SourceFile): Ident = new Ident(name)
@@ -528,7 +530,16 @@ object untpd extends Trees.Instance[Untyped] with UntypedTreeInfo {
     Select(scalaDot(nme.caps), nme.CAPTURE_ROOT)
 
   def makeRetaining(parent: Tree, refs: List[Tree], annotName: TypeName)(using Context): Annotated =
-    Annotated(parent, New(scalaAnnotationDot(annotName), List(refs)))
+    var annot: Tree = scalaAnnotationDot(annotName)
+    if annotName == tpnme.retainsCap then
+      annot = New(annot, Nil)
+    else
+      val trefs =
+        if refs.isEmpty then ref(defn.NothingType)
+        else refs.map(SingletonTypeTree).reduce[Tree]((a, b) => makeOrType(a, b))
+      annot = New(AppliedTypeTree(annot, trefs :: Nil), Nil)
+      annot.putAttachment(RetainsAnnot, ())
+    Annotated(parent, annot)
 
   def makeCapsOf(tp: RefTree)(using Context): Tree =
     TypeApply(capsInternalDot(nme.capsOf), tp :: Nil)
@@ -555,6 +566,9 @@ object untpd extends Trees.Instance[Untyped] with UntypedTreeInfo {
   def makeAndType(left: Tree, right: Tree)(using Context): AppliedTypeTree =
     AppliedTypeTree(ref(defn.andType.typeRef), left :: right :: Nil)
 
+  def makeOrType(left: Tree, right: Tree)(using Context): AppliedTypeTree =
+    AppliedTypeTree(ref(defn.orType.typeRef), left :: right :: Nil)
+
   def makeParameter(pname: TermName, tpe: Tree, mods: Modifiers, isBackquoted: Boolean = false)(using Context): ValDef = {
     val vdef = ValDef(pname, tpe, EmptyTree)
     if (isBackquoted) vdef.pushAttachment(Backquoted, ())

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

@@ -75,11 +75,32 @@ extension (tree: Tree)
     tree.getAttachment(Captures) match
       case Some(refs) => refs
       case None =>
-        val refs = CaptureSet(tree.retainedElems.flatMap(_.toCaptureRefs)*)
-          //.showing(i"toCaptureSet $tree --> $result", capt)
+        val refs =
+          tree match
+            case Apply(_: TypeApply, _) =>
+              CaptureSet(tree.retainedElemsFromType*)
+            case _ =>
+              CaptureSet(tree.retainedElems.flatMap(_.toCaptureRefs)*)
+        // println(s"toCaptureSet: $tree -> $refs")
         tree.putAttachment(Captures, refs)
         refs
 
+  def retainedElemsFromType(using Context): List[CaptureRef] =
+    def collectRefs(tp: Type): List[CaptureRef] = tp match
+      case tp: CaptureRef if tp.isTrackableRef =>
+        tp :: Nil
+      case tp: TypeRef if tp.symbol.isType && tp.derivesFrom(defn.Caps_CapSet) =>
+        tp :: Nil
+      case OrType(tp1, tp2) =>
+        collectRefs(tp1) ++ collectRefs(tp2)
+      case _ =>
+        Nil
+    tree match
+      case Apply(TypeApply(_, refs :: Nil), _) =>
+        collectRefs(refs.tpe)
+      case _ =>
+        Nil
+
   /** The arguments of a @retains, @retainsCap or @retainsByName annotation */
   def retainedElems(using Context): List[Tree] = tree match
     case Apply(_, Typed(SeqLiteral(elems, _), _) :: Nil) =>

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

@@ -1117,9 +1117,7 @@ trait Applications extends Compatibility {
             val fun2 = Applications.retypeSignaturePolymorphicFn(fun1, methType)
             simpleApply(fun2, proto)
           case funRef: TermRef =>
-            // println(i"typedApply: $funRef, ${tree.args}, ${funRef.symbol.maybeOwner.isRetains}")
-            val applyCtx = if funRef.symbol.maybeOwner.isRetains then ctx.addMode(Mode.InCaptureSet) else ctx
-            val app = ApplyTo(tree, fun1, funRef, proto, pt)(using applyCtx)
+            val app = ApplyTo(tree, fun1, funRef, proto, pt)
             convertNewGenericArray(
               widenEnumCase(
                 postProcessByNameArgs(funRef, app).computeNullable(),

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

@@ -718,7 +718,8 @@ class Typer(@constructorOnly nestingLevel: Int = 0) extends Namer
     else if name.isTermName && ctx.mode.is(Mode.InCaptureSet) then
       // If we are in a capture set and the identifier is not a term name,
       // try to type it with the same name but as a type
-      typed(untpd.makeCapsOf(untpd.cpy.Ident(tree)(name.toTypeName)), pt)
+      // typed(untpd.makeCapsOf(untpd.cpy.Ident(tree)(name.toTypeName)), pt)
+      typed(untpd.cpy.Ident(tree)(name.toTypeName), pt)
     else
       errorTree(tree, MissingIdent(tree, kind, name, pt))
   end typedIdent
@@ -2540,8 +2541,12 @@ class Typer(@constructorOnly nestingLevel: Int = 0) extends Namer
       errorTree(tree,
         em"Illegal context bound: ${tycon.tpe} does not take type parameters$selfNote.")
 
-  def typedSingletonTypeTree(tree: untpd.SingletonTypeTree)(using Context): SingletonTypeTree = {
+  def typedSingletonTypeTree(tree: untpd.SingletonTypeTree)(using Context): Tree = {
     val ref1 = typedExpr(tree.ref, SingletonTypeProto)
+    // println(i"typed singleton type $ref1 : ${ref1.tpe}, ${ctx.mode.is(Mode.InCaptureSet)}")
+    ref1.tpe match
+      case _: TypeRef if ctx.mode.is(Mode.InCaptureSet) => return ref1
+      case _ =>
     checkStable(ref1.tpe, tree.srcPos, "singleton type")
     assignType(cpy.SingletonTypeTree(tree)(ref1), ref1)
   }
@@ -3404,7 +3409,10 @@ class Typer(@constructorOnly nestingLevel: Int = 0) extends Namer
   end typedPackageDef
 
   def typedAnnotated(tree: untpd.Annotated, pt: Type)(using Context): Tree = {
-    val annot0 = withMode(Mode.InAnnotation)(typedExpr(tree.annot))
+    var annotCtx = ctx.addMode(Mode.InAnnotation)
+    if tree.annot.hasAttachment(untpd.RetainsAnnot) then
+      annotCtx = annotCtx.addMode(Mode.InCaptureSet)
+    val annot0 = typedExpr(tree.annot)(using annotCtx)
     val annot1 = checkAnnotClass(annot0)
     val annotCls = Annotations.annotClass(annot1)
     if annotCls == defn.NowarnAnnot then

library/src/scala/annotation/retains.scala

@@ -12,7 +12,7 @@ package scala.annotation
  *  non-standard capturing type syntax.
  */
 @experimental
-class retains(xs: (Any@retainsArg)*) extends annotation.StaticAnnotation
+class retains[Args](xs: (Any @retainsArg)*) extends annotation.StaticAnnotation
 
 /** Equivalent in meaning to `@retains(cap)`, but consumes less bytecode.
  */

library/src/scala/annotation/retainsByName.scala

@@ -2,5 +2,5 @@ package scala.annotation
 
 /** An annotation that indicates capture of an enclosing by-name type
  */
-@experimental class retainsByName(xs: (Any@retainsArg)*) extends annotation.StaticAnnotation
+@experimental class retainsByName[Args](xs: (Any@retainsArg)*) extends annotation.StaticAnnotation