Fix type inference for capture-polymorphic lambdas

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

@@ -900,15 +900,49 @@ class PathSelectionProto(val selector: Symbol, val pt: Type, val tree: Tree) ext
 
 /** Drop retains annotations in the inferred type if CC is not enabled
  *  or transform them into retains annotations with Nothing (i.e. empty set) as
- *   argument if CC is enabled (we need to do that to keep by-name status).
+ *  argument if CC is enabled (we need to do that to keep by-name status).
+ *  Retains are preserved as-is when they reference a named capture-set type parameter
+ *  symbol (e.g. `T^{C}` where `C: CapSet^`): those references are load-bearing for
+ *  capture-polymorphic lambdas and cannot be recovered after being rewritten to
+ *  `Nothing`. Retains whose refs are anonymous TypeParamRefs — e.g. the bound of an
+ *  un-named polymorphic lambda type parameter — are still cleaned up, since keeping
+ *  them would leave orphan parameter references in the annotation tree when pickled.
+ *  See i25830.
  */
 class CleanupRetains(using Context) extends TypeMap:
+
+  // LambdaTypes entered during the traversal of the *current* tpe only —
+  // outer binders from the surrounding source are not here.
+  private var binders: List[LambdaType] = Nil
+
+  private def isLocalCapSetParam(tp: Type): Boolean = tp match
+    // Proper scope check: the ref's binder must sit inside the current tpe.
+    case ref: TypeParamRef =>
+      ref.derivesFromCapSet && binders.contains(ref.binder)
+    // A TypeRef has no .binder, so two heuristics stand in:
+    //  - owner is an anonymous function: the ref points at an enclosing
+    //    $anonfun's C (curried poly-lambda value, inner closure tpt).
+    //  - binders.nonEmpty: the tpe itself has a LambdaType structure, so
+    //    preserving an outer cap-set ref inside it is load-bearing.
+    // Refs to a named method's type param sitting in a non-poly inferred
+    // tpe stay erased — see nicolas1.scala, cap-paramlists5.scala.
+    case ref: TypeRef if ref.derivesFromCapSet =>
+      val sym = ref.symbol
+      sym.isType && (binders.nonEmpty || sym.owner.isAnonymousFunction)
+    case _ => false
+
   def apply(tp: Type): Type = tp match
     case tp @ AnnotatedType(parent, annot: RetainingAnnotation) =>
       if Feature.ccEnabled then
         if annot.symbol == defn.RetainsCapAnnot then tp
+        else if annot.retainedType.retainedElementsRaw.exists(isLocalCapSetParam) then
+          tp.derivedAnnotatedType(this(parent), annot)
         else AnnotatedType(this(parent), RetainingAnnotation(annot.symbol.asClass, defn.NothingType))
       else this(parent)
+    case tp: LambdaType =>
+      val saved = binders
+      binders = tp :: binders
+      try mapOver(tp) finally binders = saved
     case _ => mapOver(tp)
 
 /** A base class for extractors that match annotated types with a specific

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

@@ -338,14 +338,39 @@ class Setup extends PreRecheck, SymTransformer, SetupAPI:
 
       def innerApply(tp: Type) =
         val tp1 = tp match
+          case AnnotatedType(parent, annot: RetainingAnnotation)
+          if annot.retainedType.retainedElementsRaw.exists(_.derivesFromCapSet) =>
+            // Preserve retains annotations that reference capture-set type parameters
+            // (e.g. `File^{C}` where `C: CapSet^`). Process the parent but strip the
+            // fresh empty variable that `addVar` would attach, then re-attach the
+            // original retains set so the reference to `C` survives later type-arg
+            // substitution. Without this, `{C}` would be silently erased to `{}`.
+            // See i25830.
+            val parent1 = apply(parent) match
+              case CapturingType(p, refs) if refs.elems.isEmpty && !refs.isConst => p
+              case other => other
+            // `toCaptureSet` can throw when an element isn't a well-formed capability
+            // — mirrors `decomposeCapturingType` in CapturingType.scala. Fall back to
+            // the plain parent in that case.
+            try CapturingType(parent1, annot.toCaptureSet)
+            catch case _: IllegalCaptureRef => parent1
           case AnnotatedType(parent, annot)
           if annot.symbol.isRetains || annot.symbol == defn.InferredAnnot =>
             // Drop explicit retains and @inferred annotations
             apply(parent)
           case tp: TypeLambda =>
-            // Don't recurse into parameter bounds, just cleanup any stray retains annotations
+            // Don't recurse into parameter bounds, just cleanup any stray retains
+            // annotations. Retains on CapSet-derived parts are meaningful (they encode
+            // the `CapSet^{any}` upper bound of a capture-set type parameter `[C^]`)
+            // and must be preserved so that `{C}` is a valid reference in the body.
+            // See i25830.
+            def cleanBound(bnd: Type): Type = bnd match
+              case bnd: TypeBounds =>
+                bnd.derivedTypeBounds(cleanBound(bnd.lo), cleanBound(bnd.hi))
+              case _ =>
+                if bnd.derivesFromCapSet then bnd else bnd.dropAllRetains
             tp.derivedLambdaType(
-              paramInfos = tp.paramInfos.mapConserve(_.dropAllRetains.bounds),
+              paramInfos = tp.paramInfos.mapConserve(cleanBound(_).bounds),
               resType = this(tp.resType))
           case tp @ RefinedType(parent, rname, rinfo) =>
             val saved = refiningNames
@@ -899,6 +924,10 @@ class Setup extends PreRecheck, SymTransformer, SetupAPI:
         needsVariable(parent)
         && refs.isConst       // if refs is a variable, no need to add another
         && !refs.isUniversal  // if refs is {caps.any}, an added variable would not change anything
+        && !refs.elems.exists(_.coreType.derivesFromCapSet)
+          // Don't wrap a const set that references a capture-set type parameter in a
+          // fresh variable: the resulting Var's elements would no longer track the
+          // original ParamRef across type-argument substitution. See i25830.
       case AnnotatedType(parent, _) =>
         needsVariable(parent)
       case _ =>

tests/pos-custom-args/captures/i25830-mixed-scopes.scala

@@ -0,0 +1,86 @@
+import language.experimental.captureChecking
+import caps.*
+
+class File extends SharedCapability
+
+// Every local poly lambda below mentions at least one *outer* capture-set
+// parameter in its own signature, and varies the lambda's own binder shape —
+// multiple cap-sets, interleaved with plain types, nested poly, etc.
+// Shapes where the local lambda does not reference anything from the outer
+// scope are covered by tests/pos-custom-args/captures/i25830.scala.
+
+// (1) val-form of the `def g[C^] = (xs) => xs.head` TODO line from
+//     tests/neg-custom-args/captures/use-capset.scala. The def form still
+//     errors (separate @use/classifier issue); the val form works because
+//     the lambda binds its own C.
+def useCapsetVal() =
+  val g = { [C^] => (xs: List[File^{C}]) => xs.head }
+  val io = File()
+  val _ : File^{io} = g[{io}](List[File^{io}](io))
+
+// (2) Baseline: simple local `[C^]` inside a def with `[OuterC^]`, the
+//     lambda's signature mentions both.
+def baseline[OuterC^](a: File^{OuterC}) =
+  val mk = { [C^] => (x: File^{C}, y: File^{OuterC}) => x }
+  val b = File()
+  val _ : File^{b} = mk[{b}](b, a)
+
+// (3) Fully interleaved local binders `[T, C^, U, D^]` around an outer cap-set.
+def interleaved_local[OuterC^](a: File^{OuterC}) =
+  val mk = { [T, C^, U, D^] =>
+    (t: T, x: File^{C}, u: U, y: File^{D}, z: File^{OuterC}) => (t, u, x)
+  }
+  val b = File(); val c = File()
+  val _ = mk[Int, {b}, String, {c}](1, b, "s", c, a)
+
+// (4) Outer itself mixes plain + cap-set; local lambda does too.
+def mixed_both_sides[T, OuterC^](t: T, a: File^{OuterC}) =
+  val mk = { [U, C^] =>
+    (u: U, x: File^{C}, t2: T, y: File^{OuterC}) => (u, x, t2)
+  }
+  val b = File()
+  val _ = mk[Boolean, {b}](true, b, t, a)
+
+// (5) Class field: lambda's interleaved signature mentions the class's
+//     cap-set parameter.
+class Holder[OuterC^](val outer: File^{OuterC}):
+  val mk = { [T, C^, U] =>
+    (t: T, x: File^{C}, u: U, y: File^{OuterC}) => x
+  }
+
+def useHolder() =
+  val a = File(); val b = File()
+  val h = Holder[{a}](a)
+  val _ : File^{b} = h.mk[Int, {b}, String](1, b, "s", a)
+
+// (6) Trait abstract member + subclass override, both mentioning the
+//     enclosing cap-set parameter. Exercises the explicit-tpt path.
+trait Ops[OuterC^]:
+  val mk: [T, C^] -> (t: T, x: File^{C}, y: File^{OuterC}) -> File^{C}
+
+class OpsImpl[X^](x0: File^{X}) extends Ops[X]:
+  val mk = { [T, C^] => (t: T, x: File^{C}, y: File^{X}) => x }
+
+def useOps() =
+  val a = File(); val b = File()
+  val ops = OpsImpl[{a}](a)
+  val _ : File^{b} = ops.mk[Int, {b}](7, b, a)
+
+// (7) Nested outer scopes (class + inner def) each contribute a cap-set;
+//     the local lambda mentions both.
+class Outer[X^](val xr: File^{X}):
+  def inner[T, Y^](yr: File^{Y})(t0: T) =
+    val mk = { [U, C^] =>
+      (u: U, x: File^{C}, a: File^{X}, b: File^{Y}, t: T) => x
+    }
+    val c = File()
+    val _ : File^{c} = mk[String, {c}]("s", c, xr, yr, t0)
+
+// (8) Nested poly lambda: the body of the outer lambda is itself a poly
+//     lambda, with interleaved binders at each level, and the inner level
+//     mentions the enclosing method's OuterC.
+def nested_poly_interleaved[OuterC^](a: File^{OuterC}) =
+  val mk = { [T, C^] => (t: T, x: File^{C}) =>
+    { [U, D^] => (u: U, y: File^{D}, z: File^{OuterC}) => (t, x) } }
+  val b = File(); val c = File()
+  val _ = mk[Int, {b}](1, b)[String, {c}]("s", c, a)

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

@@ -0,0 +1,10 @@
+import language.experimental.captureChecking
+import caps.*
+
+class File extends SharedCapability
+
+@main def test =
+  val convert = { [C^] => (xs: List[File^{C}]) => xs.map(_ => ()) }
+  val x = File()
+  val files: List[File^{x}] = List(x)
+  val result = convert[{x}](files)