Refactor type maps to split out capability handling

Note i15923 does not signal a leak anymore. I moved it and some variants to pending.
Note: There seems to be something more fundamentally wrong with this test: We get an
infinite recursion for variant i15923b.

Author: odersky

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

@@ -63,9 +63,11 @@ case class CaptureAnnotation(refs: CaptureSet, boxed: Boolean)(cls: Symbol) exte
 
   override def mapWith(tm: TypeMap)(using Context) =
     val elems = refs.elems.toList
-    val elems1 = elems.mapConserve(tm)
+    val elems1 = elems.mapConserve(tm.mapCapability(_))
     if elems1 eq elems then this
-    else if elems1.forall(_.isTrackableRef)
+    else if elems1.forall:
+      case elem1: CaptureRef => elem1.isTrackableRef
+      case _ => false
     then derivedAnnotation(CaptureSet(elems1.asInstanceOf[List[CaptureRef]]*), boxed)
     else EmptyAnnotation
 

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

@@ -114,6 +114,7 @@ extension (tp: Type)
       !tp.underlying.exists // might happen during construction of lambdas
       || tp.derivesFrom(defn.Caps_CapSet)
     case root.Result(_) => true
+    case root.Fresh(_) => true
     case AnnotatedType(parent, annot) =>
       defn.capabilityWrapperAnnots.contains(annot.symbol) && parent.isTrackableRef
     case _ =>
@@ -143,9 +144,9 @@ extension (tp: Type)
     if dcs.isAlwaysEmpty then tp.captureSet
     else tp match
       case tp @ ReachCapability(_) =>
-        tp.singletonCaptureSet
+        assert(false); tp.singletonCaptureSet
       case ReadOnlyCapability(ref) =>
-        ref.deepCaptureSet(includeTypevars).readOnly
+        assert(false); ref.deepCaptureSet(includeTypevars).readOnly
       case tp: SingletonCaptureRef if tp.isTrackableRef =>
         tp.reach.singletonCaptureSet
       case _ =>
@@ -195,9 +196,12 @@ extension (tp: Type)
    *  are of the form this.C but their pathroot is still this.C, not this.
    */
   final def pathRoot(using Context): Type = tp.dealias match
-    case tp1: NamedType
-    if tp1.symbol.maybeOwner.isClass && tp1.symbol != defn.captureRoot && !tp1.symbol.is(TypeParam) =>
-      tp1.prefix.pathRoot
+    case tp1: NamedType =>
+      if tp1.symbol.maybeOwner.isClass && tp1.symbol != defn.captureRoot && !tp1.symbol.is(TypeParam) then
+        tp1.prefix match
+          case pre: CaptureRef => pre.pathRoot
+          case _ => tp1
+      else tp1
     case tp1 => tp1
 
   /** If this part starts with `C.this`, the class `C`.
@@ -214,7 +218,8 @@ extension (tp: Type)
       tp1.prefix match
         case _: ThisType | NoPrefix =>
           tp1.symbol.is(Param) || tp1.symbol.is(ParamAccessor)
-        case prefix => prefix.isParamPath
+        case prefix: CaptureRef => prefix.isParamPath
+        case _ => false
     case _: ParamRef => true
     case _ => false
 

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

@@ -319,7 +319,7 @@ sealed abstract class CaptureSet extends Showable:
   def map(tm: TypeMap)(using Context): CaptureSet =
     tm match
       case tm: BiTypeMap =>
-        val mappedElems = elems.map(tm.forward)
+        val mappedElems = elems.map(tm.mapCapability(_))
         if isConst then
           if mappedElems == elems then this
           else Const(mappedElems)
@@ -487,7 +487,7 @@ object CaptureSet:
     override def toString = elems.toString
   end Const
 
-  case class EmptyWithProvenance(ref: CaptureRef, mapped: Type) extends Const(SimpleIdentitySet.empty):
+  case class EmptyWithProvenance(ref: CaptureRef, mapped: CaptureSet) extends Const(SimpleIdentitySet.empty):
     override def optionalInfo(using Context): String =
       if ctx.settings.YccDebug.value
       then i" under-approximating the result of mapping $ref to $mapped"
@@ -587,8 +587,7 @@ object CaptureSet:
      */
     private def checkSkippedMaps(elem: CaptureRef)(using Context): Unit =
       for tm <- skippedMaps do
-        val elem1 = tm(elem)
-        for elem1 <- tm(elem).captureSet.elems do
+        for elem1 <- extrapolateCaptureRef(elem, tm, variance = 1).elems do
           assert(elem.subsumes(elem1),
             i"Skipped map ${tm.getClass} maps newly added $elem to $elem1 in $this")
 
@@ -817,14 +816,14 @@ object CaptureSet:
 
     override def tryInclude(elem: CaptureRef, origin: CaptureSet)(using Context, VarState): CompareResult =
       if origin eq source then
-        val mappedElem = bimap.forward(elem)
+        val mappedElem = bimap.mapCapability(elem)
         if accountsFor(mappedElem) then CompareResult.OK
         else addNewElem(mappedElem)
       else if accountsFor(elem) then
         CompareResult.OK
       else
         try
-          source.tryInclude(bimap.backward(elem), this)
+          source.tryInclude(bimap.inverse.mapCapability(elem), this)
             .showing(i"propagating new elem $elem backward from $this to $source = $result", captDebug)
             .andAlso(addNewElem(elem))
         catch case ex: AssertionError =>
@@ -1031,15 +1030,12 @@ object CaptureSet:
    *        - Otherwise assertion failure
    */
   def extrapolateCaptureRef(r: CaptureRef, tm: TypeMap, variance: Int)(using Context): CaptureSet =
-    val r1 = tm(r)
-    val upper = r1.captureSet
-    def isExact =
-      upper.isAlwaysEmpty
-      || upper.isConst && upper.elems.size == 1 && upper.elems.contains(r1)
-      || r.derivesFrom(defn.Caps_CapSet)
-    if variance > 0 || isExact then upper
-    else if variance < 0 then CaptureSet.EmptyWithProvenance(r, r1)
-    else upper.maybe
+    tm.mapCapability(r) match
+      case c: CaptureRef => c.captureSet
+      case (cs: CaptureSet, exact) =>
+        if cs.isAlwaysEmpty || exact || variance > 0 then cs
+        else if variance < 0 then CaptureSet.EmptyWithProvenance(r, cs)
+        else cs.maybe
 
   /** Apply `f` to each element in `xs`, and join result sets with `++` */
   def mapRefs(xs: Refs, f: CaptureRef => CaptureSet)(using Context): CaptureSet =
@@ -1289,9 +1285,11 @@ object CaptureSet:
 
     def mapRef(ref: CaptureRef): CaptureRef
 
-    def apply(t: Type) = t match
-      case t: CaptureRef if t.isTrackableRef => mapRef(t)
-      case _ => mapOver(t)
+    def apply(t: Type) = mapOver(t)
+
+    override def mapCapability(c: CaptureRef, deep: Boolean): CaptureRef = c match
+      case c: CaptureRef if c.isTrackableRef => mapRef(c)
+      case _ => super.mapCapability(c, deep)
 
     override def fuse(next: BiTypeMap)(using Context) = next match
       case next: Inverse if next.inverse.getClass == getClass => assert(false); Some(IdentityTypeMap)

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

@@ -441,7 +441,7 @@ class CheckCaptures extends Recheck, SymTransformer:
       markFree(sym, sym.termRef, tree)
 
     def markFree(sym: Symbol, ref: CaptureRef, tree: Tree)(using Context): Unit =
-      if sym.exists && ref.isTracked then markFree(ref.captureSet, tree)
+      if sym.exists && ref.isTracked then markFree(ref.singletonCaptureSet, tree)
 
     /** Make sure the (projected) `cs` is a subset of the capture sets of all enclosing
      *  environments. At each stage, only include references from `cs` that are outside

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

@@ -209,6 +209,9 @@ class Setup extends PreRecheck, SymTransformer, SetupAPI:
           innerApply(tp)
       finally isTopLevel = saved
 
+    override def mapArg(arg: Type, tparam: ParamInfo): Type =
+      super.mapArg(Recheck.mapExprType(arg), tparam)
+
     /** Map parametric functions with results that have a capture set somewhere
      *  to dependent functions.
      */
@@ -504,12 +507,15 @@ class Setup extends PreRecheck, SymTransformer, SetupAPI:
     def add = new TypeTraverser:
       var reach = false
       def traverse(t: Type): Unit = t match
-        case root.Fresh(hidden) =>
-          if reach then hidden.elems += ref.reach
-          else if ref.isTracked then hidden.elems += ref
-        case t @ CapturingType(_, _) if t.isBoxed && !reach =>
-          reach = true
-          try traverseChildren(t) finally reach = false
+        case t @ CapturingType(parent, refs) =>
+          val saved = reach
+          reach |= t.isBoxed
+          try
+            traverse(parent)
+            for case root.Fresh(hidden) <- refs.elems.iterator do
+              if reach then hidden.elems += ref.reach
+              else if ref.isTracked then hidden.elems += ref
+          finally reach = saved
         case _ =>
           traverseChildren(t)
     if ref.isTrackableRef then add.traverse(tp)
@@ -660,9 +666,13 @@ class Setup extends PreRecheck, SymTransformer, SetupAPI:
 
         def paramsToCap(mt: Type)(using Context): Type = mt match
           case mt: MethodType =>
-            mt.derivedLambdaType(
-              paramInfos = mt.paramInfos.map(root.freshToCap),
-              resType = paramsToCap(mt.resType))
+            try
+              mt.derivedLambdaType(
+                paramInfos = mt.paramInfos.map(root.freshToCap),
+                resType = paramsToCap(mt.resType))
+            catch case ex: AssertionError =>
+              println(i"error while mapping params ${mt.paramInfos} of $sym")
+              throw ex
           case mt: PolyType =>
             mt.derivedLambdaType(resType = paramsToCap(mt.resType))
           case _ => mt

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

@@ -154,21 +154,6 @@ object root:
     override def eql(that: Annotation) = that match
       case Annot(kind) => this.kind eq kind
       case _ => false
-
-    /** Special treatment of `SubstBindingMaps` which can change the binder of
-     *  Result instances
-     */
-    override def mapWith(tm: TypeMap)(using Context) = kind match
-      case Kind.Result(binder) => tm match
-        case tm: Substituters.SubstBindingMap[MethodType] @unchecked if tm.from eq binder =>
-          derivedAnnotation(tm.to)
-        case tm: Substituters.SubstBindingsMap =>
-          var i = 0
-          while i < tm.from.length && (tm.from(i) ne binder) do i += 1
-          if i < tm.from.length then derivedAnnotation(tm.to(i).asInstanceOf[MethodType])
-          else this
-        case _ => this
-      case _ => this
   end Annot
 
   def cap(using Context): TermRef = defn.captureRoot.termRef
@@ -222,8 +207,7 @@ object root:
     override def apply(t: Type) =
       if variance <= 0 then t
       else t match
-        case t: CaptureRef if t.isCap =>
-          Fresh(origin)
+        case root(_) => assert(false)
         case t @ CapturingType(parent: TypeRef, _) if parent.symbol == defn.Caps_CapSet =>
           t
         case t @ CapturingType(_, _) =>
@@ -237,6 +221,11 @@ object root:
         case _ =>
           mapFollowingAliases(t)
 
+    override def mapCapability(c: CaptureRef, deep: Boolean): CaptureRef = c match
+      case c: CaptureRef if c.isCap => Fresh(origin)
+      case root(_) => c
+      case _ => super.mapCapability(c, deep)
+
     override def fuse(next: BiTypeMap)(using Context) = next match
       case next: Inverse => assert(false); Some(IdentityTypeMap)
       case _ => None
@@ -245,13 +234,14 @@ object root:
 
     class Inverse extends BiTypeMap, FollowAliasesMap:
       def apply(t: Type): Type = t match
-        case t @ Fresh(_) => cap
+        case root(_) => assert(false)
         case t @ CapturingType(_, refs) => mapOver(t)
         case _ => mapFollowingAliases(t)
 
-      override def fuse(next: BiTypeMap)(using Context) = next match
-        case next: CapToFresh => assert(false); Some(IdentityTypeMap)
-        case _ => None
+      override def mapCapability(c: CaptureRef, deep: Boolean): CaptureRef = c match
+        case c @ Fresh(_) => cap
+        case root(_) => c
+        case _ => super.mapCapability(c, deep)
 
       def inverse = thisMap
       override def toString = thisMap.toString + ".inverse"
@@ -283,9 +273,7 @@ object root:
       var localBinders: SimpleIdentitySet[MethodType] = SimpleIdentitySet.empty
 
       def apply(t: Type): Type = t match
-        case t @ Result(binder) =>
-          if localBinders.contains(binder) then t // keep bound references
-          else seen.getOrElseUpdate(t.annot, Fresh(origin)) // map free references to Fresh()
+        case root(_) => assert(false)
         case t: MethodType =>
           // skip parameters
           val saved = localBinders
@@ -298,6 +286,14 @@ object root:
         case _ =>
           mapOver(t)
 
+      override def mapCapability(c: CaptureRef, deep: Boolean) = c match
+        case t @ Result(binder) =>
+          if localBinders.contains(binder) then t // keep bound references
+          else seen.getOrElseUpdate(t.annot, Fresh(origin)) // map free references to Fresh()
+        case root(_) => c
+        case _ => super.mapCapability(c, deep)
+    end subst
+
     subst(tp)
   end resultToFresh
 
@@ -320,15 +316,7 @@ object root:
       private val seen = EqHashMap[CaptureRef, Result]()
 
       def apply(t: Type) = t match
-        case t: CaptureRef if t.isCapOrFresh =>
-          if variance > 0 then
-            seen.getOrElseUpdate(t, Result(mt))
-          else
-            if variance == 0 then
-              fail(em"""$tp captures the root capability `cap` in invariant position.
-                       |This capability cannot be converted to an existential in the result type of a function.""")
-            // we accept variance < 0, and leave the cap as it is
-            super.mapOver(t)
+        case root(_) => assert(false)
         case defn.FunctionNOf(args, res, contextual) if t.typeSymbol.name.isImpureFunction =>
           if variance > 0 then
             super.mapOver:
@@ -337,26 +325,48 @@ object root:
           else mapOver(t)
         case _ =>
           mapOver(t)
+
+      override def mapCapability(c: CaptureRef, deep: Boolean) = c match
+        case c: CaptureRef if c.isCapOrFresh =>
+          if variance > 0 then
+            seen.getOrElseUpdate(c, Result(mt))
+          else
+            if variance == 0 then
+              fail(em"""$tp captures the root capability `cap` in invariant position.
+                       |This capability cannot be converted to an existential in the result type of a function.""")
+            // we accept variance < 0, and leave the cap as it is
+            c
+        case root(_) => c
+        case _ =>
+          super.mapCapability(c, deep)
+
         //.showing(i"mapcap $t = $result")
       override def toString = "toVar"
 
       object inverse extends BiTypeMap:
         def apply(t: Type) = t match
-          case t @ Result(`mt`) =>
+          case root(_) => assert(false)
+          case _ => mapOver(t)
+        def inverse = toVar.this
+        override def toString = "toVar.inverse"
+
+        override def mapCapability(c: CaptureRef, deep: Boolean) = c match
+          case c @ Result(`mt`) =>
             // do a reverse getOrElseUpdate on `seen` to produce the
             // `Fresh` assosicated with `t`
             val it = seen.iterator
             var ref: CaptureRef | Null = null
             while it.hasNext && ref == null do
               val (k, v) = it.next
-              if v.annot eq t.annot then ref = k
+              if v eq c then ref = k
             if ref == null then
               ref = Fresh(Origin.Unknown)
-              seen(ref) = t
+              seen(ref) = c
             ref
-          case _ => mapOver(t)
-        def inverse = toVar.this
-        override def toString = "toVar.inverse"
+          case root(_) => c
+          case _ =>
+            super.mapCapability(c, deep)
+      end inverse
     end toVar
 
     toVar(tp)