Add capture annotations

Author: odersky

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

@@ -0,0 +1,56 @@
+package dotty.tools
+package dotc
+package cc
+
+import core.*
+import Types.*, Symbols.*, Contexts.*, Annotations.*
+import ast.Trees.*
+import ast.{tpd, untpd}
+import Decorators.*
+import config.Printers.capt
+import printing.Printer
+import printing.Texts.Text
+
+
+case class CaptureAnnotation(refs: CaptureSet) extends Annotation:
+  import CaptureAnnotation.*
+  import tpd.*
+
+  override def tree(using Context) =
+    val elems = refs.elems.toList.map {
+      case cr: TermRef => ref(cr)
+      case cr: TermParamRef => untpd.Ident(cr.paramName).withType(cr)
+      case cr: ThisType => This(cr.cls)
+    }
+    val arg = repeated(elems, TypeTree(defn.AnyType))
+    New(symbol.typeRef, arg :: Nil)
+
+  override def symbol(using Context) = defn.RetainsAnnot
+
+  override def derivedAnnotation(tree: Tree)(using Context): Annotation =
+    unsupported("derivedAnnotation(Tree)")
+
+  def derivedAnnotation(refs: CaptureSet)(using Context): Annotation =
+    if this.refs eq refs then this else CaptureAnnotation(refs)
+
+  override def sameAnnotation(that: Annotation)(using Context): Boolean = that match
+    case CaptureAnnotation(refs2) => refs == refs2
+    case _ => false
+
+  override def mapWith(tp: TypeMap)(using Context) =
+    val elems = refs.elems.toList
+    val elems1 = elems.mapConserve(tp)
+    if elems1 eq elems then this
+    else if elems1.forall(_.isInstanceOf[CaptureRef])
+    then CaptureAnnotation(CaptureSet(elems1.asInstanceOf[List[CaptureRef]]*))
+    else EmptyAnnotation
+
+  override def refersToParamOf(tl: TermLambda)(using Context): Boolean =
+    refs.elems.exists {
+      case TermParamRef(tl1, _) => tl eq tl1
+      case _ => false
+    }
+
+  override def toText(printer: Printer): Text = refs.toText(printer)
+
+end CaptureAnnotation

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

@@ -8,43 +8,37 @@ import ast.Trees.*
 import ast.{tpd, untpd}
 import Decorators.*
 import config.Printers.capt
+import util.Property.Key
 
 object CaptureOps:
   import tpd.*
 
+  private val Captures: Key[CaptureSet] = Key()
+
   def retainedElems(tree: Tree)(using Context): List[Tree] = tree match
     case Apply(_, Typed(SeqLiteral(elems, _), _) :: Nil) => elems
     case _ => Nil
 
   extension (tree: Tree)
+
     def toCaptureRef(using Context): CaptureRef = tree.tpe.asInstanceOf[CaptureRef]
 
-  extension (cs: CaptureSet)
-    def toAnnotation(using Context): Annotation =
-      val refs = cs.elems.toList.map {
-        case cr: TermRef => ref(cr)
-        case cr: TermParamRef => untpd.Ident(cr.paramName).withType(cr)
-        case cr: ThisType => This(cr.cls)
-      }
-      val arg = repeated(refs, TypeTree(defn.AnyType))
-      Annotation(defn.RetainsAnnot.typeRef, arg)
-
-  extension (tp: CapturingType)
-    def toAnnotatedType(using Context): AnnotatedType =
-      AnnotatedType(tp.parent, tp.refs.toAnnotation)
-
-  extension (annot: Annotation)
     def toCaptureSet(using Context): CaptureSet =
-      assert(annot.symbol == defn.RetainsAnnot)
-      CaptureSet(retainedElems(annot.tree).map(_.toCaptureRef)*)
-        .showing(i"toCaptureSet $annot --> $result", capt)
-
-  extension (tp: AnnotatedType)
-    def toCapturingType(using Context): Type =
-      CapturingType(tp.parent, tp.annot.toCaptureSet)
+      tree.getAttachment(Captures) match
+        case Some(refs) => refs
+        case None =>
+          val refs = CaptureSet(retainedElems(tree).map(_.toCaptureRef)*)
+            .showing(i"toCaptureSet $tree --> $result", capt)
+          tree.putAttachment(Captures, refs)
+          refs
 
   extension (tp: Type)
 
+    def derivedCapturingType(parent: Type, refs: CaptureSet)(using Context): Type = tp match
+      case CapturingType(p, r) =>
+        if (parent eq p) && (refs eq r) then tp
+        else CapturingType(parent, refs)
+
     /** If this is  type variable instantiated or upper bounded with a capturing type,
      *  the capture set associated with that type. Extended to and-or types and
      *  type proxies in the obvious way. If a term has a type with a boxed captureset,
@@ -91,4 +85,24 @@ object CaptureOps:
       case _ =>
         false
 
-end CaptureOps
+  object CapturingAnnotType:
+
+    def apply(parent: Type, refs: CaptureSet)(using Context): Type =
+      AnnotatedType(parent, CaptureAnnotation(refs))
+
+    def unapply(tp: AnnotatedType)(using Context) = tp.annot match
+      case ann: CaptureAnnotation =>
+        Some((tp.parent, ann.refs))
+      case ann =>
+        if ann.symbol == defn.RetainsAnnot
+        then Some((tp.parent, ann.tree.toCaptureSet))
+        else None
+  end CapturingAnnotType
+
+  object PreCapturingType:
+    def unapply(tp: AnnotatedType)(using Context) = tp.annot match
+      case ann: CaptureAnnotation => Some((tp.parent, ann.refs))
+      case _ =>
+        None
+
+end CaptureOps

compiler/src/dotty/tools/dotc/core/Annotations.scala

@@ -82,7 +82,7 @@ object Annotations {
   }
 
   case class ConcreteAnnotation(t: Tree) extends Annotation {
-  
+
     def tree(using Context): Tree = t
 
     override def refersToParamOf(tl: TermLambda)(using Context): Boolean =
@@ -92,7 +92,7 @@ object Annotations {
           case _ => false
         case _ => false
       }
-      
+
     override def argsText(printer: Printer): Text =
       def toTextArg(arg: Tree): Text = arg match
         case Typed(SeqLiteral(elems, _), _) => printer.toTextGlobal(elems, ", ")

compiler/src/dotty/tools/dotc/core/CaptureSet.scala

@@ -79,7 +79,7 @@ sealed abstract class CaptureSet extends Showable:
     reporting.trace(i"$this accountsFor $x, ${x.captureSetOfInfo}?", show = true) {
       elems.contains(x)
       || !x.isRootCapability && (x.captureSetOfInfo frozen_<:< this) == CompareResult.OK
-   }
+    }
 
   /** The subcapturing test */
   def <:< (that: CaptureSet)(using Context): CompareResult =
@@ -89,6 +89,10 @@ sealed abstract class CaptureSet extends Showable:
   def frozen_<:<(that: CaptureSet)(using Context): CompareResult =
     subcaptures(that)(using ctx, FrozenState)
 
+  def =:= (that: CaptureSet)(using Context): Boolean =
+       (this frozen_<:< that) == CompareResult.OK
+    && (that frozen_<:< this) == CompareResult.OK
+
   private def subcaptures(that: CaptureSet)(using Context, VarState): CompareResult =
     val result = that.tryInclude(elems)
     if result == CompareResult.OK then addSuper(that) else varState.abort()

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

@@ -8,7 +8,7 @@ import Contexts.{Context, ctx}
 import Types.*, Symbols.*
 import Denotations.SingleDenotation
 import SymDenotations.SymDenotation
-import cc.CaptureOps.toAnnotation
+import cc.CaptureAnnotation
 
 /** A denotation transformer that comes after the CheckCaptures. It resets
  *  all CapturingTypes in the info of derived SingletonDenotations to AnnotatedTypes.
@@ -27,7 +27,7 @@ class PostCapture extends Phase, InfoTransformer:
     val mapType = new TypeMap:
       def apply(t: Type) = t match
         case CapturingType(parent, cs) =>
-          AnnotatedType(this(parent), toAnnotation(cs))
+          AnnotatedType(this(parent), CaptureAnnotation(cs))
         case _ =>
           mapOver(t)
     mapType(tp)

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

@@ -114,7 +114,7 @@ class CheckCaptures extends Recheck:
       val mapType = new TypeMap:
         def apply(t: Type) = t match
           case AnnotatedType(parent, annot) if annot.symbol == defn.RetainsAnnot =>
-            CapturingType(this(parent), annot.toCaptureSet)
+            CapturingType(this(parent), annot.tree.toCaptureSet)
           case t @ RefinedType(core, nme.apply, appInfo) =>
             mapRefined(t, this(core), this(appInfo))
           case _ =>