Allow for embedded CapSet^{refs} entries in capture sets

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

@@ -64,7 +64,7 @@ def depFun(args: List[Type], resultType: Type, isContextual: Boolean, paramNames
   mt.toFunctionType(alwaysDependent = true)
 
 /** An exception thrown if a @retains argument is not syntactically a CaptureRef */
-class IllegalCaptureRef(tpe: Type) extends Exception(tpe.toString)
+class IllegalCaptureRef(tpe: Type)(using Context) extends Exception(tpe.show)
 
 /** Capture checking state, which is known to other capture checking components */
 class CCState:
@@ -127,15 +127,21 @@ class NoCommonRoot(rs: Symbol*)(using Context) extends Exception(
 
 extension (tree: Tree)
 
-  /** Map tree with CaptureRef type to its type, throw IllegalCaptureRef otherwise */
-  def toCaptureRef(using Context): CaptureRef = tree match
+  /** Map tree with CaptureRef type to its type,
+   *  map CapSet^{refs} to the `refs` references,
+   *  throw IllegalCaptureRef otherwise
+   */
+  def toCaptureRefs(using Context): List[CaptureRef] = tree match
     case ReachCapabilityApply(arg) =>
-      arg.toCaptureRef.reach
+      arg.toCaptureRefs.map(_.reach)
     case CapsOfApply(arg) =>
-      arg.toCaptureRef
-    case _ => tree.tpe match
+      arg.toCaptureRefs
+    case _ => tree.tpe.dealiasKeepAnnots match
       case ref: CaptureRef if ref.isTrackableRef =>
-        ref
+        ref :: Nil
+      case AnnotatedType(parent, ann)
+      if ann.symbol.isRetains && parent.derivesFrom(defn.Caps_CapSet) =>
+        ann.tree.toCaptureSet.elems.toList
       case tpe =>
         throw IllegalCaptureRef(tpe) // if this was compiled from cc syntax, problem should have been reported at Typer
 
@@ -146,7 +152,7 @@ extension (tree: Tree)
     tree.getAttachment(Captures) match
       case Some(refs) => refs
       case None =>
-        val refs = CaptureSet(tree.retainedElems.map(_.toCaptureRef)*)
+        val refs = CaptureSet(tree.retainedElems.flatMap(_.toCaptureRefs)*)
           //.showing(i"toCaptureSet $tree --> $result", capt)
         tree.putAttachment(Captures, refs)
         refs

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

@@ -729,25 +729,28 @@ class Setup extends PreRecheck, SymTransformer, SetupAPI:
     var retained = ann.retainedElems.toArray
     for i <- 0 until retained.length do
       val refTree = retained(i)
-      val ref = refTree.toCaptureRef
-
-      def pos =
-        if refTree.span.exists then refTree.srcPos
-        else if ann.span.exists then ann.srcPos
-        else tpt.srcPos
-
-      def check(others: CaptureSet, dom: Type | CaptureSet): Unit =
-        if others.accountsFor(ref) then
-          report.warning(em"redundant capture: $dom already accounts for $ref", pos)
-
-      if ref.captureSetOfInfo.elems.isEmpty && !ref.derivesFrom(defn.Caps_Capability) then
-        report.error(em"$ref cannot be tracked since its capture set is empty", pos)
-      check(parent.captureSet, parent)
-
-      val others =
-        for j <- 0 until retained.length if j != i yield retained(j).toCaptureRef
-      val remaining = CaptureSet(others*)
-      check(remaining, remaining)
+      for ref <- refTree.toCaptureRefs do
+        def pos =
+          if refTree.span.exists then refTree.srcPos
+          else if ann.span.exists then ann.srcPos
+          else tpt.srcPos
+
+        def check(others: CaptureSet, dom: Type | CaptureSet): Unit =
+          if others.accountsFor(ref) then
+            report.warning(em"redundant capture: $dom already accounts for $ref", pos)
+
+        if ref.captureSetOfInfo.elems.isEmpty && !ref.derivesFrom(defn.Caps_Capability) then
+          report.error(em"$ref cannot be tracked since its capture set is empty", pos)
+        check(parent.captureSet, parent)
+
+        val others =
+          for
+            j <- 0 until retained.length if j != i
+            r <- retained(j).toCaptureRefs
+          yield r
+        val remaining = CaptureSet(others*)
+        check(remaining, remaining)
+      end for
     end for
   end checkWellformedPost
 

tests/pos/cc-poly-source-capability.scala

@@ -4,7 +4,9 @@ import caps.{CapSet, Capability}
 
 @experimental object Test:
 
-  class Label extends Capability
+  class Async extends Capability
+
+  def listener(async: Async): Listener^{async} = ???
 
   class Listener
 
@@ -15,22 +17,16 @@ import caps.{CapSet, Capability}
 
     def allListeners: Set[Listener^{X^}] = listeners
 
-  def test1(lbl1: Label, lbl2: Label) =
-    val src = Source[CapSet^{lbl1, lbl2}]
-    def l1: Listener^{lbl1} = ???
-    val l2: Listener^{lbl2} = ???
-    src.register{l1}
-    src.register{l2}
-    val ls = src.allListeners
-    val _: Set[Listener^{lbl1, lbl2}] = ls
-
-  def test2(lbls: List[Label]) =
-    def makeListener(lbl: Label): Listener^{lbl} = ???
-    val listeners = lbls.map(makeListener)
-    val src = Source[CapSet^{lbls*}]
-    for l <- listeners do
-      src.register(l)
+  def test1(async1: Async, others: List[Async]) =
+    val src = Source[CapSet^{async1, others*}]
+    val lst1 = listener(async1)
+    val lsts = others.map(listener)
+    val _: List[Listener^{others*}] = lsts
+    src.register{lst1}
+    src.register(listener(async1))
+    lsts.foreach(src.register)
+    others.map(listener).foreach(src.register)
     val ls = src.allListeners
-    val _: Set[Listener^{lbls*}] = ls
+    val _: Set[Listener^{async1, others*}] = ls