Fix SubstBindings BiTypeMap logic

 1. Also apply `SubstBindings` map to root annotations
    (previously it was only single `SubstBinding` maps).
 2. Don't preserve BiTypeMaps for sets that are set up for
    the first time.

Author: odersky

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

@@ -120,6 +120,11 @@ class CCState:
 
   private var capIsRoot: Boolean = false
 
+  /** If true, apply a BiTypeMap also to elements added to the set in the future
+   *  (and use its inverse when back-progating).
+   */
+  private var mapFutureElems = true
+
   var iterCount = 1
 
 object CCState:
@@ -178,9 +183,25 @@ object CCState:
       try op finally ccs.capIsRoot = saved
     else op
 
+  /** Don't map future elements in this `op` */
+  inline def withoutMappedFutureElems[T](op: => T)(using Context): T =
+    val ccs = ccState
+    val saved = ccs.mapFutureElems
+    ccs.mapFutureElems = false
+    try op finally ccs.mapFutureElems = saved
+
   /** Is `caps.cap` a root capability that is allowed to subsume other capabilities? */
   def capIsRoot(using Context): Boolean = ccState.capIsRoot
 
+  /** When mapping a capture set with a BiTypeMap, should we create a BiMapped set
+   *  so that future elements can also be mapped, and elements added to the BiMapped
+   *  are back-propagated? Turned off when creating capture set variables for the
+   *  first time, since we then do not want to change the binder to the original type
+   *  without capture sets when back propagating. Error case where this shows:
+   *  pos-customargs/captures/lists.scala, method m2c.
+   */
+  def mapFutureElems(using Context) = ccState.mapFutureElems
+
   /** The currently opened existential scopes */
   def openExistentialScopes(using Context): List[MethodType] = ccState.openExistentialScopes
 

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

@@ -322,13 +322,14 @@ sealed abstract class CaptureSet extends Showable:
         if isConst then
           if mappedElems == elems then this
           else Const(mappedElems)
-        else
+        else if CCState.mapFutureElems then
           def unfused = BiMapped(asVar, tm, mappedElems)
           this match
             case self: BiMapped => self.bimap.fuse(tm) match
               case Some(fused: BiTypeMap) => BiMapped(self.source, fused, mappedElems)
               case _ => unfused
             case _ => unfused
+        else this
       case tm: IdentityCaptRefMap =>
         this
       case tm: AvoidMap if this.isInstanceOf[HiddenSet] =>
@@ -732,11 +733,13 @@ object CaptureSet:
      *  is not derived from some other variable.
      */
     protected def ids(using Context): String =
+      def descr = getClass.getSimpleName.nn.take(1)
       val trail = this.match
-        case dv: DerivedVar => dv.source.ids
-        case _ => ""
-      val descr = getClass.getSimpleName.nn.take(1)
-      s"$id$descr$trail"
+        case dv: DerivedVar =>
+          def summary = if ctx.settings.YccVerbose.value then dv.summarize else descr
+          s"$summary${dv.source.ids}"
+        case _ => descr
+      s"$id$trail"
     override def toString = s"Var$id$elems"
   end Var
 

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

@@ -302,9 +302,10 @@ class Setup extends PreRecheck, SymTransformer, SetupAPI:
             apply(parent)
           case tp: TypeLambda =>
             // Don't recurse into parameter bounds, just cleanup any stray retains annotations
-            tp.derivedLambdaType(
-              paramInfos = tp.paramInfos.mapConserve(_.dropAllRetains.bounds),
-              resType = this(tp.resType))
+            withoutMappedFutureElems:
+              tp.derivedLambdaType(
+                paramInfos = tp.paramInfos.mapConserve(_.dropAllRetains.bounds),
+                resType = this(tp.resType))
           case _ =>
             mapFollowingAliases(tp)
         addVar(

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

@@ -112,6 +112,11 @@ object root:
       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
@@ -317,7 +322,12 @@ object root:
       case t: (LazyRef | TypeVar) =>
         mapConserveSuper(t)
       case _ =>
-        if keepAliases then mapOver(t) else mapFollowingAliases(t)
+        try
+          if keepAliases then mapOver(t)
+          else mapFollowingAliases(t)
+        catch case ex: AssertionError =>
+          println(i"error while mapping $t")
+          throw ex
   end toResultInResults
 
   /** If `refs` contains an occurrence of `cap` or `cap.rd`, the current context

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

@@ -4219,6 +4219,35 @@ object Types extends TypeUtils {
           }
       mt
     }
+
+    /** Not safe to use in general: Check that all references to an enclosing
+     *  TermParamRef name point to that TermParamRef
+     */
+    def checkValid2(mt: MethodType)(using Context): mt.type = {
+      var t = new TypeTraverser:
+        val ps = mt.paramNames.zip(mt.paramRefs).toMap
+        def traverse(t: Type) =
+          t match
+            case CapturingType(p, refs) =>
+              def checkRefs(refs: CaptureSet) =
+                for elem <- refs.elems do
+                  elem match
+                    case elem: TermParamRef =>
+                      val elemName = elem.binder.paramNames(elem.paramNum)
+                      //assert(elemName.toString != "f")
+                      ps.get(elemName) match
+                        case Some(elemRef) => assert(elemRef eq elem, i"bad $mt")
+                        case _ =>
+                    case root.Result(binder) if binder ne mt =>
+                      assert(binder.paramNames.toList != mt.paramNames.toList, i"bad $mt")
+                    case _ =>
+              checkRefs(refs)
+              traverse(p)
+            case _ =>
+              traverseChildren(t)
+      t.traverse(mt.resType)
+      mt
+    }
   }
 
   object MethodType extends MethodTypeCompanion("MethodType") {

tests/pos-custom-args/captures/lists.scala

@@ -45,9 +45,7 @@ def test(c: Cap, d: Cap, e: Cap) =
   def m2 = [A, B] =>
       (f: A => B) => (xs: LIST[A]) => xs.map(f)
 
-//  def m2c: [A, B] -> (f: A => B) -> LIST[A] ->{f} LIST[B] = m2
-//  !!! m2 has a bad type due to spurious level error. Need to follow up and
-//  fix this.
+  def m2c: [A, B] -> (f: A => B) -> LIST[A] ->{f} LIST[B] = m2
 
   def eff[A](x: A) = if x == e then x else x