Change retains annotation from using term arguments to using type arguments

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

@@ -2263,9 +2263,9 @@ object desugar {
               AppliedTypeTree(ref(defn.SeqType), t),
               New(ref(defn.RepeatedAnnot.typeRef), Nil :: Nil))
         else if op.name == nme.CC_REACH then
-          Apply(ref(defn.Caps_reachCapability), t :: Nil)
+          Annotated(t, New(ref(defn.ReachCapabilityAnnot.typeRef), Nil :: Nil))
         else if op.name == nme.CC_READONLY then
-          Apply(ref(defn.Caps_readOnlyCapability), t :: Nil)
+          Annotated(t, New(ref(defn.ReadOnlyCapabilityAnnot.typeRef), Nil :: Nil))
         else
           assert(ctx.mode.isExpr || ctx.reporter.errorsReported || ctx.mode.is(Mode.Interactive), ctx.mode)
           Select(t, op.name)

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

@@ -388,6 +388,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)
@@ -530,10 +532,23 @@ 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)))
-
-  def makeCapsOf(tp: RefTree)(using Context): Tree =
-    TypeApply(capsInternalDot(nme.capsOf), tp :: Nil)
+    var annot: Tree = scalaAnnotationDot(annotName)
+    if annotName == tpnme.retainsCap then
+      annot = New(annot, Nil)
+    else
+      val trefs =
+        if refs.isEmpty then
+          // The NothingType is used to represent the empty capture set.
+          ref(defn.NothingType)
+        else
+          // Treat all references as term references before typing.
+          // A dummy term symbol will be created for each capture variable,
+          // and references to them will be replaced with the corresponding
+          // type references during typing.
+          refs.map(SingletonTypeTree).reduce[Tree]((a, b) => makeOrType(a, b))
+      annot = New(AppliedTypeTree(annot, trefs :: Nil), Nil)
+      annot.putAttachment(RetainsAnnot, ())
+    Annotated(parent, annot)
 
   def makeConstructor(tparams: List[TypeDef], vparamss: List[List[ValDef]], rhs: Tree = EmptyTree)(using Context): DefDef =
     DefDef(nme.CONSTRUCTOR, joinParams(tparams, vparamss), TypeTree(), rhs)
@@ -557,6 +572,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/Capability.scala

@@ -603,6 +603,21 @@ object Capabilities:
     def assumedContainsOf(x: TypeRef)(using Context): SimpleIdentitySet[Capability] =
       CaptureSet.assumedContains.getOrElse(x, SimpleIdentitySet.empty)
 
+    /** The type representing this capability.
+     *  Note this method does not distinguish different `RootCapability` instances,
+     *  and should only be used for printing or phases not related to CC.
+     */
+    def toType(using Context): Type = this match
+      case c: RootCapability => defn.captureRoot.termRef
+      case c: CoreCapability => c
+      case c: DerivedCapability =>
+        val c1 = c.underlying.toType
+        c match
+          case _: ReadOnly => ReadOnlyCapability(c1)
+          case _: Reach => ReachCapability(c1)
+          case _: Maybe => MaybeCapability(c1)
+          case _ => c1
+
     def toText(printer: Printer): Text = printer.toTextCapability(this)
   end Capability
 

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

@@ -39,15 +39,14 @@ case class CaptureAnnotation(refs: CaptureSet, boxed: Boolean)(cls: Symbol) exte
 
   /** Reconstitute annotation tree from capture set */
   override def tree(using Context) =
-    val elems = refs.elems.toList.map {
-      case c: TermRef => ref(c)
-      case c: TermParamRef => untpd.Ident(c.paramName).withType(c)
-      case c: ThisType => This(c.cls)
-      case c: RootCapability => ref(defn.captureRoot)
-      // TODO: Will crash if the type is an annotated type, for example `cap.rd`
-    }
-    val arg = repeated(elems, TypeTree(defn.AnyType))
-    New(symbol.typeRef, arg :: Nil)
+    if symbol == defn.RetainsCapAnnot then
+      New(symbol.typeRef, Nil)
+    else
+      val elems = refs.elems.toList.map(_.toType)
+      val trefs =
+        if elems.isEmpty then defn.NothingType
+        else elems.reduce((a, b) => OrType(a, b, soft = false))
+      New(AppliedType(symbol.typeRef, trefs :: Nil), Nil)
 
   override def symbol(using Context) = cls
 

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

@@ -47,57 +47,59 @@ def ccState(using Context): CCState =
 
 extension (tree: Tree)
 
-  /** Map tree with a Capability type to the corresponding capability,
-   *  map CapSet^{refs} to the `refs` references,
-   *  throw IllegalCaptureRef otherwise
-   */
-  def toCapabilities(using Context): List[Capability] = tree match
-    case ReachCapabilityApply(arg) =>
-      arg.toCapabilities.map(_.reach)
-    case ReadOnlyCapabilityApply(arg) =>
-      arg.toCapabilities.map(_.readOnly)
-    case CapsOfApply(arg) =>
-      arg.toCapabilities
-    case _ => tree.tpe.dealiasKeepAnnots match
-      case ref: TermRef if ref.isCapRef =>
-        GlobalCap :: Nil
-      case ref: Capability if ref.isTrackableRef =>
-        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
-
   /** 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.
    */
   def toCaptureSet(using Context): CaptureSet =
     tree.getAttachment(Captures) match
       case Some(refs) => refs
       case None =>
-        val refs = CaptureSet(tree.retainedElems.flatMap(_.toCapabilities)*)
-          //.showing(i"toCaptureSet $tree --> $result", capt)
+        val refs = CaptureSet(tree.retainedSet.retainedElements*)
         tree.putAttachment(Captures, refs)
         refs
 
-  /** 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(defn.captureRoot) :: Nil
-      else Nil
+  /** The type representing the capture set of @retains, @retainsCap or @retainsByName annotation. */
+  def retainedSet(using Context): Type =
+    tree match
+      case Apply(TypeApply(_, refs :: Nil), _) => refs.tpe
+      case _ =>
+        if tree.symbol.maybeOwner == defn.RetainsCapAnnot
+        then defn.captureRoot.termRef else NoType
 
 extension (tp: Type)
 
+  def toCapability(using Context): Capability = tp match
+    case ReachCapability(tp1) =>
+      tp1.toCapability.reach
+    case ReadOnlyCapability(tp1) =>
+      tp1.toCapability.readOnly
+    case ref: TermRef if ref.isCapRef =>
+      GlobalCap
+    case ref: Capability if ref.isTrackableRef =>
+      ref
+    case _ =>
+      // if this was compiled from cc syntax, problem should have been reported at Typer
+      throw IllegalCaptureRef(tp)
+
+  /** A list of raw elements of a retained set.
+   *  This will not crash even if it contains a non-wellformed Capability.
+   */
+  def retainedElementsRaw(using Context): List[Type] = tp match
+    case OrType(tp1, tp2) =>
+      tp1.retainedElementsRaw ++ tp2.retainedElementsRaw
+    case tp =>
+      // Nothing is a special type to represent the empty set
+      if tp.isNothingType then Nil
+      else tp :: Nil // should be checked by wellformedness
+
+  /** A list of capabilities of a retained set. */
+  def retainedElements(using Context): List[Capability] =
+    retainedElementsRaw.map(_.toCapability)
+
   /** Is this type a Capability that can be tracked?
    *  This is true for
    *    - all ThisTypes and all TermParamRef,
    *    - stable TermRefs with NoPrefix or ThisTypes as prefixes,
-   *    - the root capability `caps.cap`
-   *    - abstract or parameter TypeRefs that derive from caps.CapSet
    *    - annotated types that represent reach or maybe capabilities
    */
   final def isTrackableRef(using Context): Boolean = tp match
@@ -408,7 +410,7 @@ extension (cls: ClassSymbol)
       || bc.is(CaptureChecked)
           && bc.givenSelfType.dealiasKeepAnnots.match
             case CapturingType(_, refs) => refs.isAlwaysEmpty
-            case RetainingType(_, refs) => refs.isEmpty
+            case RetainingType(_, refs) => refs.retainedElements.isEmpty
             case selfType =>
               isCaptureChecking  // At Setup we have not processed self types yet, so
                                  // unless a self type is explicitly given, we can't tell
@@ -523,32 +525,36 @@ class CleanupRetains(using Context) extends TypeMap:
   def apply(tp: Type): Type =
     tp match
       case AnnotatedType(tp, annot) if annot.symbol == defn.RetainsAnnot || annot.symbol == defn.RetainsByNameAnnot =>
-        RetainingType(tp, Nil, byName = annot.symbol == defn.RetainsByNameAnnot)
+        RetainingType(tp, defn.NothingType, byName = annot.symbol == defn.RetainsByNameAnnot)
       case _ => mapOver(tp)
 
-/** An extractor for `caps.reachCapability(ref)`, which is used to express a reach
- *  capability as a tree in a @retains annotation.
+/** A base class for extractors that match annotated types with a specific
+ *  Capability annotation.
  */
-object ReachCapabilityApply:
-  def unapply(tree: Apply)(using Context): Option[Tree] = tree match
-    case Apply(reach, arg :: Nil) if reach.symbol == defn.Caps_reachCapability => Some(arg)
+abstract class AnnotatedCapability(annotCls: Context ?=> ClassSymbol):
+  def apply(tp: Type)(using Context): AnnotatedType =
+    AnnotatedType(tp, Annotation(annotCls, util.Spans.NoSpan))
+
+  def unapply(tree: AnnotatedType)(using Context): Option[Type] = tree match
+    case AnnotatedType(parent: Type, ann) if ann.hasSymbol(annotCls) => Some(parent)
     case _ => None
 
-/** An extractor for `caps.readOnlyCapability(ref)`, which is used to express a read-only
- *  capability as a tree in a @retains annotation.
+end AnnotatedCapability
+
+/** An extractor for `ref @readOnlyCapability`, which is used to express
+ *  the read-only capability `ref.rd` as a type.
  */
-object ReadOnlyCapabilityApply:
-  def unapply(tree: Apply)(using Context): Option[Tree] = tree match
-    case Apply(ro, arg :: Nil) if ro.symbol == defn.Caps_readOnlyCapability => Some(arg)
-    case _ => None
+object ReadOnlyCapability extends AnnotatedCapability(defn.ReadOnlyCapabilityAnnot)
 
-/** An extractor for `caps.capsOf[X]`, which is used to express a generic capture set
- *  as a tree in a @retains annotation.
+/** An extractor for `ref @reachCapability`, which is used to express
+ *  the reach capability `ref*` as a type.
  */
-object CapsOfApply:
-  def unapply(tree: TypeApply)(using Context): Option[Tree] = tree match
-    case TypeApply(capsOf, arg :: Nil) if capsOf.symbol == defn.Caps_capsOf => Some(arg)
-    case _ => None
+object ReachCapability extends AnnotatedCapability(defn.ReachCapabilityAnnot)
+
+/** An extractor for `ref @amaybeCapability`, which is used to express
+ *  the maybe capability `ref?` as a type.
+ */
+object MaybeCapability extends AnnotatedCapability(defn.MaybeCapabilityAnnot)
 
 /** An extractor for all kinds of function types as well as method and poly types.
  *  It includes aliases of function types such as `=>`. TODO: Can we do without?

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

@@ -100,25 +100,24 @@ object CheckCaptures:
    *  This check is performed at Typer.
    */
   def checkWellformed(parent: Tree, ann: Tree)(using Context): Unit =
-    def check(elem: Tree, pos: SrcPos): Unit = elem.tpe match
+    def check(elem: Type, pos: SrcPos): Unit = elem match
       case ref: Capability =>
         if !ref.isTrackableRef && !ref.isCapRef then
           report.error(em"$elem cannot be tracked since it is not a parameter or local value", pos)
       case tpe =>
         report.error(em"$elem: $tpe is not a legal element of a capture set", pos)
-    for elem <- ann.retainedElems do
+    for elem <- ann.retainedSet.retainedElementsRaw do
       elem match
-        case CapsOfApply(arg) =>
-          def isLegalCapsOfArg =
-            arg.symbol.isType && arg.symbol.info.derivesFrom(defn.Caps_CapSet)
-          if !isLegalCapsOfArg then
-            report.error(
-              em"""$arg is not a legal prefix for `^` here,
-                  |is must be a type parameter or abstract type with a caps.CapSet upper bound.""",
-              elem.srcPos)
-        case ReachCapabilityApply(arg) => check(arg, elem.srcPos)
-        case ReadOnlyCapabilityApply(arg) => check(arg, elem.srcPos)
-        case _ => check(elem, elem.srcPos)
+        case ref: TypeRef =>
+          val refSym = ref.symbol
+          if refSym.isType && !refSym.info.derivesFrom(defn.Caps_CapSet) then
+            report.error(em"$elem is not a legal element of a capture set", ann.srcPos)
+        case ReachCapability(ref) =>
+          check(ref, ann.srcPos)
+        case ReadOnlyCapability(ref) =>
+          check(ref, ann.srcPos)
+        case _ =>
+          check(elem, ann.srcPos)
 
   /** Under the sealed policy, report an error if some part of `tp` contains the
    *  root capability in its capture set or if it refers to a type parameter that

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

@@ -12,24 +12,20 @@ import Decorators.i
  */
 object RetainingType:
 
-  def apply(tp: Type, refs: List[Tree], byName: Boolean = false)(using Context): Type =
+  def apply(tp: Type, typeElems: Type, byName: Boolean = false)(using Context): Type =
     val annotCls = if byName then defn.RetainsByNameAnnot else defn.RetainsAnnot
-    val annotTree =
-      New(annotCls.typeRef,
-        Typed(
-          SeqLiteral(refs, TypeTree(defn.AnyType)),
-          TypeTree(defn.RepeatedParamClass.typeRef.appliedTo(defn.AnyType))) :: Nil)
+    val annotTree = New(AppliedType(annotCls.typeRef, typeElems :: Nil), Nil)
     AnnotatedType(tp, Annotation(annotTree))
 
-  def unapply(tp: AnnotatedType)(using Context): Option[(Type, List[Tree])] =
+  def unapply(tp: AnnotatedType)(using Context): Option[(Type, Type)] =
     val sym = tp.annot.symbol
     if sym.isRetainsLike then
       tp.annot match
         case _: CaptureAnnotation =>
           assert(ctx.mode.is(Mode.IgnoreCaptures), s"bad retains $tp at ${ctx.phase}")
           None
         case ann =>
-          Some((tp.parent, ann.tree.retainedElems))
+          Some((tp.parent, ann.tree.retainedSet))
     else
       None
 end RetainingType

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

@@ -806,7 +806,7 @@ class Setup extends PreRecheck, SymTransformer, SetupAPI:
       case CapturingType(_, refs) =>
         !refs.isAlwaysEmpty
       case RetainingType(parent, refs) =>
-        !refs.isEmpty
+        !refs.retainedElements.isEmpty
       case tp: (TypeRef | AppliedType) =>
         val sym = tp.typeSymbol
         if sym.isClass
@@ -852,7 +852,7 @@ class Setup extends PreRecheck, SymTransformer, SetupAPI:
         && !refs.isUniversal  // if refs is {cap}, an added variable would not change anything
       case RetainingType(parent, refs) =>
         needsVariable(parent)
-        && !refs.tpes.exists:
+        && !refs.retainedElements.exists:
             case ref: TermRef => ref.isCapRef
             case _ => false
       case AnnotatedType(parent, _) =>
@@ -947,19 +947,13 @@ class Setup extends PreRecheck, SymTransformer, SetupAPI:
    *  @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.retainedElems} in $tpt")
-    var retained = ann.retainedElems.toArray
-    for i <- 0 until retained.length do
-      val refTree = retained(i)
-      val refs =
-        try refTree.toCapabilities
-        catch case ex: IllegalCaptureRef =>
-          report.error(em"Illegal capture reference: ${ex.getMessage}", refTree.srcPos)
-          Nil
-      for ref <- refs do
+    capt.println(i"checkWF post $parent ${ann.retainedSet} in $tpt")
+    try
+      var retained = ann.retainedSet.retainedElements.toArray
+      for i <- 0 until retained.length do
+        val ref = retained(i)
         def pos =
-          if refTree.span.exists then refTree.srcPos
-          else if ann.span.exists then ann.srcPos
+          if ann.span.exists then ann.srcPos
           else tpt.srcPos
 
         def check(others: CaptureSet, dom: Type | CaptureSet): Unit =
@@ -976,13 +970,14 @@ class Setup extends PreRecheck, SymTransformer, SetupAPI:
         val others =
           for
             j <- 0 until retained.length if j != i
-            r <- retained(j).toCapabilities
+            r = retained(j)
             if !r.isTerminalCapability
           yield r
         val remaining = CaptureSet(others*)
         check(remaining, remaining)
       end for
-    end for
+    catch case ex: IllegalCaptureRef =>
+      report.error(em"Illegal capture reference: ${ex.getMessage}", tpt.srcPos)
   end checkWellformed
 
   /** Check well formed at post check time. We need to wait until after