Fix problems when changing between annotated can capturing types

Fix problems in denotation transformers when changing between annotated can capturing types

Author: odersky

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

@@ -0,0 +1,97 @@
+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
+
+object CaptureOps:
+  import tpd.*
+
+  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)
+      annot.tree match
+        case Apply(_, Typed(SeqLiteral(elems, _), _) :: Nil) =>
+          CaptureSet(elems.map(_.tpe.asInstanceOf[CaptureRef])*)
+            .showing(i"toCaptureSet $annot --> $result", capt)
+
+  extension (tp: AnnotatedType)
+    def toCapturingType(using Context): Type =
+      CapturingType(tp.parent, tp.annot.toCaptureSet)
+
+  extension (tp: Type)
+
+    /** 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,
+     *  that captureset counts towards the capture variables of the envirionment.
+     */
+    def boxedCaptured(using Context): CaptureSet =
+      def getBoxed(tp: Type, enabled: Boolean): CaptureSet = tp match
+        case tp: CapturingType if enabled => tp.refs
+        case tp: TypeVar => getBoxed(tp.underlying, enabled = true)
+        case tp: TypeRef if tp.symbol == defn.AnyClass && enabled => CaptureSet.universal
+        case tp: TypeProxy => getBoxed(tp.superType, enabled)
+        case tp: AndType => getBoxed(tp.tp1, enabled) ++ getBoxed(tp.tp2, enabled)
+        case tp: OrType => getBoxed(tp.tp1, enabled) ** getBoxed(tp.tp2, enabled)
+        case _ => CaptureSet.empty
+      getBoxed(tp, enabled = false)
+
+    /** If this type appears as an expected type of a term, does it imply
+     *  that the term should be boxed? ^^^ Special treat Any?
+     */
+    def needsBox(using Context): Boolean = tp match
+      case _: TypeVar => true
+      case tp: TypeRef =>
+        tp.info match
+          case TypeBounds(lo, _) => lo.needsBox
+          case _ => false
+      case tp: RefinedOrRecType => tp.parent.needsBox
+      case tp: AnnotatedType => tp.parent.needsBox
+      case tp: LazyRef => tp.ref.needsBox
+      case tp: AndType => tp.tp1.needsBox || tp.tp2.needsBox
+      case tp: OrType => tp.tp1.needsBox && tp.tp2.needsBox
+      case _ => false
+
+    def canHaveInferredCapture(using Context): Boolean = tp match
+      case tp: CapturingType =>
+        false
+      case tp: TypeRef =>
+        if tp.symbol.isClass then
+          !tp.symbol.isValueClass && tp.symbol != defn.AnyClass
+        else
+          tp.underlying.canHaveInferredCapture
+      case tp: TypeProxy =>
+        tp.underlying.canHaveInferredCapture
+      case tp: AndType =>
+        tp.tp1.canHaveInferredCapture && tp.tp2.canHaveInferredCapture
+      case tp: OrType =>
+        tp.tp1.canHaveInferredCapture || tp.tp2.canHaveInferredCapture
+      case _ =>
+        false
+
+    def addCaptureVars(using Context): Type =
+      if ctx.settings.Ycc.value && canHaveInferredCapture then
+        CapturingType(tp, CaptureSet.Var()) // ^^^ go deep
+      else
+        tp
+end CaptureOps

compiler/src/dotty/tools/dotc/Compiler.scala

@@ -86,6 +86,7 @@ class Compiler {
          new StringInterpolatorOpt) :: // Optimizes raw and s string interpolators by rewriting them to string concatenations
     List(new PreRecheck) ::          // Preparations for check captures phase, enabled under -Ycc
     List(new CheckCaptures) ::       // Check captures, enabled under -Ycc
+    List(new PostCapture) ::         // Reset denotations transformed by CheckCaptures back to annotated types
     List(new PruneErasedDefs,        // Drop erased definitions from scopes and simplify erased expressions
          new UninitializedDefs,      // Replaces `compiletime.uninitialized` by `_`
          new InlinePatterns,         // Remove placeholders of inlined patterns

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

@@ -351,12 +351,4 @@ object CaptureSet:
         empty
     recur(tp)
       .showing(i"capture set of $tp = $result", capt)
-
-  def fromAnnotation(annot: Annotation)(using Context): CaptureSet =
-    import ast.Trees.*
-    assert(annot.symbol == defn.RetainsAnnot)
-    annot.tree match
-      case Apply(_, Typed(SeqLiteral(elems, _), _) :: Nil) =>
-        CaptureSet(elems.map(_.tpe.asInstanceOf[CaptureRef])*)
-
 end CaptureSet

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

@@ -0,0 +1,41 @@
+package dotty.tools.dotc
+package transform
+
+import core.*
+import Phases.Phase
+import DenotTransformers.InfoTransformer
+import Contexts.{Context, ctx}
+import Types.*, Symbols.*
+import Denotations.SingleDenotation
+import SymDenotations.SymDenotation
+import cc.CaptureOps.toAnnotation
+
+/** A denotation transformer that comes after the CheckCaptures. It resets
+ *  all CapturingTypes in the info of derived SingletonDenotations to AnnotatedTypes.
+ *  SymDenotations are left unchaged. For them, the info is already reset in
+ *  Rechecl#updateInfo.
+ */
+class PostCapture extends Phase, InfoTransformer:
+  thisPhase =>
+
+  def phaseName: String = "postCapture"
+
+  override def isEnabled(using Context) = prev.isEnabled
+  override def changesBaseTypes: Boolean = true
+
+  override def transformInfo(tp: Type, sym: Symbol)(using Context): Type =
+    val mapType = new TypeMap:
+      def apply(t: Type) = t match
+        case CapturingType(parent, cs) =>
+          AnnotatedType(this(parent), toAnnotation(cs))
+        case _ =>
+          mapOver(t)
+    mapType(tp)
+
+  override def transform(ref: SingleDenotation)(using Context): SingleDenotation =
+    if ref.isInstanceOf[SymDenotation] then ref
+    else super.transform(ref)
+
+  def run(using Context): Unit = ()
+
+end PostCapture

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

@@ -54,28 +54,16 @@ abstract class Recheck extends Phase, IdentityDenotTransformer:
           ).installAfter(preRecheckPhase)
 
     /** Hooks to be overridden */
-    protected def reinfer(tp: Type)(using Context): Type = tp
-    protected def transformType(tp: Type)(using Context): Type = tp
-
-    def reinferResult(info: Type)(using Context): Type = info match
-      case info: MethodOrPoly =>
-        info.derivedLambdaType(resType = reinferResult(info.resType))
-      case info: ExprType =>
-        info.derivedExprType(resType = reinferResult(info.resType))
-      case _ =>
-        reinfer(info)
+    protected def transformType(tp: Type, inferred: Boolean)(using Context): Type = tp
 
     def enterDef(stat: Tree)(using Context): Unit =
       val sym = stat.symbol
-      val transInfo = transformType(sym.info)
-      val newInfo = stat match
-        case stat: ValOrDefDef if stat.tpt.isInstanceOf[InferredTypeTree] =>
-          reinferResult(transInfo)
-        case stat: Bind =>
-          reinferResult(transInfo)
-        case _ =>
-          transInfo
-      //println(i"update info $sym: ${sym.info} --> $transInfo --> $newInfo")
+      val isInferred = stat match
+        case stat: ValOrDefDef => stat.tpt.isInstanceOf[InferredTypeTree]
+        case stat: Bind => true
+        case _ => false
+      val newInfo = transformType(sym.info, isInferred)
+      //println(i"update info $sym: ${sym.info} --> $newInfo")
       sym.updateInfo(newInfo)
 
     def constFold(tree: Tree, tp: Type)(using Context): Type =
@@ -234,9 +222,8 @@ abstract class Recheck extends Phase, IdentityDenotTransformer:
       val elemTypes = tree.elems.map(recheck(_, elemProto))
       seqLitType(tree, TypeComparer.lub(declaredElemType :: elemTypes))
 
-    def recheckTypeTree(tree: TypeTree)(using Context): Type = tree match
-      case tree: InferredTypeTree => reinfer(tree.tpe)
-      case _ => transformType(tree.tpe)
+    def recheckTypeTree(tree: TypeTree)(using Context): Type =
+      transformType(tree.tpe, tree.isInstanceOf[InferredTypeTree])
 
     def recheckAnnotated(tree: Annotated)(using Context): Type =
       tree.tpe match

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

@@ -28,6 +28,7 @@ import reporting._
 import ProtoTypes._
 import dotty.tools.backend.jvm.DottyBackendInterface.symExtensions
 import CaptureSet.CompareResult
+import cc.CaptureOps.toCaptureSet
 
 object CheckCaptures:
   case class Env(owner: Symbol, captured: CaptureSet, isBoxed: Boolean, outer: Env):
@@ -85,17 +86,45 @@ class CheckCaptures extends Recheck:
   class CaptureChecker(ictx: Context) extends Rechecker(ictx):
     import ast.tpd.*
 
-    override def reinfer(tp: Type)(using Context): Type =
-      CapturingType(tp, CaptureSet.Var()) // ^^^ go deep
-
-    override def transformType(tp: Type)(using Context): Type =
+    override def transformType(tp: Type, inferred: Boolean)(using Context): Type =
       val mapType = new TypeMap:
-        def apply(t: Type) = mapOver(t) match
+        def apply(t: Type) = t match
           case AnnotatedType(parent, annot) if annot.symbol == defn.RetainsAnnot =>
-            CapturingType(parent, CaptureSet.fromAnnotation(annot))
-          case t1 =>
-            t1
-      mapType(tp)
+            CapturingType(this(parent), annot.toCaptureSet)
+          case t @ RefinedType(core, nme.apply, appInfo) if defn.isFunctionType(t) =>
+            val core1 = this(core)
+            val appInfo1 = this(appInfo)
+            if appInfo1 ne appInfo then appInfo1.toFunctionType(isJava = false)
+            else t.derivedRefinedType(core1, nme.apply, appInfo)
+          case _ =>
+            mapOver(t)
+
+      def reinfer(tp: Type): Type =
+        val cleanType = new TypeMap:
+          def apply(t: Type) = t match
+            case AnnotatedType(parent, annot) if annot.symbol == defn.RetainsAnnot =>
+              parent
+            case _ =>
+              mapOver(t)
+        def addVars(tp: Type): Type =
+          val tp1 = tp match
+            case tp @ AppliedType(tycon, args) =>
+              tp.derivedAppliedType(tycon, args.map(addVars))
+            case _ => // ^^^ go deeper in other types as well
+              tp
+          CapturingType(tp1, CaptureSet.Var())
+        tp match
+          case tp: MethodOrPoly =>
+            tp.derivedLambdaType(resType = reinfer(tp.resType))
+          case tp: ExprType =>
+            tp.derivedExprType(resType = reinfer(tp.resType))
+          case _ =>
+            val tp1 = cleanType(tp)
+            addVars(tp1)
+
+      val tp1 = if inferred then reinfer(tp) else tp
+      mapType(tp1)
+    end transformType
 
     private var curEnv: Env = Env(NoSymbol, CaptureSet.empty, false, null)
 

tests/pos-custom-args/captures/capt-depfun.scala

@@ -1,11 +1,12 @@
 class C
-type Cap = C retains *
-type Top = Any retains *
+type Cap = C @retains(*)
 
-type T = (x: Cap) => String retains x.type
+type T = (x: Cap) => String @retains(x)
 
-def f(y: Cap): String retains * =
-  val a: T = (x: Cap) => ""
+val aa: ((x: Cap) => String @retains(x)) = (x: Cap) => ""
+
+def f(y: Cap): String @retains(*) =
+  val a: ((x: Cap) => String @retains(x)) = (x: Cap) => ""
   val b = a(y)
-  val c: String retains y.type = b
+  val c: String @retains(y) = b
   c