Fix inferred capture-polymorphic function literals

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

@@ -8,6 +8,7 @@ import Names.{Name, TermName}
 import ast.{tpd, untpd}
 import Decorators.*, NameOps.*
 import config.Printers.capt
+import util.{NoSourcePosition, SrcPos}
 import util.Property.Key
 import tpd.*
 import Annotations.Annotation
@@ -898,17 +899,111 @@ extension (tp: AnnotatedType) {
 /** A prototype that indicates selection */
 class PathSelectionProto(val selector: Symbol, val pt: Type, val tree: Tree) extends typer.ProtoTypes.WildcardSelectionProto
 
-/** 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).
+/** A TypeMap that strips `@retains` annotations from inferred types,
+ *  preserving only those whose elements describe a stable, capture-polymorphic
+ *  interface (refs to capset binders in scope). Inferred types reach the
+ *  pickler with retains that the user never wrote — this cleanup makes them
+ *  consistent with what the user can actually rely on.
+ *
+ *  When applied to the inferred type of a capture-polymorphic poly-function
+ *  literal, an unsupported retain is reported as an implementation
+ *  restriction. The supported shape is documented at the report site.
+ *
+ *  @param reportPos  The source position to use in case of illegal signatures.
+ *                    Defaults to `NoSourcePosition` for callers that don't need
+ *                    to report (UnApply / Bind cleanups in PostTyper).
  */
-class CleanupRetains(using Context) extends TypeMap:
+class CleanupRetains(reportPos: SrcPos = NoSourcePosition, ownerWalk: Boolean = true)(using Context) extends TypeMap:
+  // Capset-binding TypeLambdas in scope. Used to recognize valid retain
+  // elements during descent and to detect a capture-polymorphic context.
+  //
+  // A poly-fn literal `{ [C^] => x => body }` desugars to a `DefDef $anonfun`
+  // carrying the [C^] PolyType on its symbol. PostTyper visits two inferred
+  // TypeTrees per literal: the outer val/def tpt sees the whole type so [C^]
+  // is reached by descent; the inner $anonfun return tpt sees only a fragment
+  // and never crosses [C^]. Owner walking covers that fragment case by
+  // seeding the binder from `ctx.owner.info` — gated on anon-fn owners so we
+  // don't pull enclosing PolyTypes into the scope of an unrelated regular
+  // val (e.g. one inside `extension [T, C^]`).
+  //
+  // `ownerWalk = false` is for callers that just want plain erasure
+  // (e.g. inlining) and shouldn't interact with the literal-tpt machinery.
+  private var capSetBinders: List[TypeLambda] =
+    if ownerWalk && Feature.ccEnabled && ctx.owner.isAnonymousFunction then
+      def enclosing(sym: Symbol): List[TypeLambda] =
+        if !sym.exists || sym.is(Package) then Nil
+        else sym.info match
+          case lt: TypeLambda if lt.paramRefs.exists(_.derivesFromCapSet) =>
+            lt :: enclosing(sym.owner)
+          case _ => enclosing(sym.owner)
+      enclosing(ctx.owner)
+    else Nil
+
+  // At most one implementation-restriction error per type tree.
+  private var reported = false
+
+  // Whether to keep a retain element. `parent` is the annotated type.
+  private def isPreserved(elem: Type, parent: Type): Boolean = elem match
+    case ref: TypeParamRef =>
+      // Capset param of an in-scope TypeLambda.
+      ref.derivesFromCapSet && capSetBinders.exists(_ eq ref.binder)
+    case ref: TypeRef =>
+      // Capset typedef aliasing an in-scope binder — required for curried
+      // literals where inner-fragment visits see the outer binder via a
+      // TypeRef (e.g. `TypeRef(NoPrefix, type C)`) rather than as a
+      // TypeParamRef. `frozen_=:=` matches it against the binder's paramRef
+      // without leaking constraints into type inference.
+      ref.derivesFromCapSet && capSetBinders.exists(_.paramRefs.exists(_ frozen_=:= ref))
+    case ref: TermRef =>
+      // caps.any only in capset bounds (synthesized `<: CapSet^{any}` for `[C^]`).
+      ref.isCapsAnyRef && parent.derivesFromCapSet
+    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 AnnotatedType(this(parent), RetainingAnnotation(annot.symbol.asClass, defn.NothingType))
+        else
+          val elems = annot.retainedType.retainedElementsRaw
+          val keep = elems.nonEmpty && elems.forall(isPreserved(_, parent))
+          // Suppress only the narrow case where the entire retains is a single
+          // synthesized TypeBounds — `nonDependentResultApprox` (Types.scala)
+          // produces these when collapsing a dependent retain into a `range(...)`
+          // and an invariant position turns it back into bounds. Anything else
+          // gets reported, even if some bad element is a TypeBounds.
+          val isSyntheticOnly = elems match
+            case (_: TypeBounds) :: Nil => true
+            case _ => false
+          if !keep && capSetBinders.nonEmpty && !reported && reportPos.span.exists && !isSyntheticOnly then
+            elems.find(e => !isPreserved(e, parent)) match
+              case Some(bad) =>
+                reported = true
+                report.restrictionError(
+                  em"""inferred capture-polymorphic function literals can only retain capture-set parameters from enclosing type lambdas; found $bad.
+                      |Define a class with a polymorphic apply method to retain arbitrary capture sets.""",
+                  reportPos)
+              case None =>
+          AnnotatedType(this(parent),
+            if keep then annot
+            else RetainingAnnotation(annot.symbol.asClass, defn.NothingType))
       else this(parent)
+    case tp: TypeLambda if tp.paramRefs.exists(_.derivesFromCapSet) =>
+      // A second capset TypeLambda nested inside another would crash cc later
+      // (Setup's lambda recursion + SubstBindingMap on a Var holding outer
+      // binder refs hits an invariant in CaptureSet.map). Reject it here.
+      // Descend with empty `capSetBinders` so retains in the inner lambda
+      // erase to Nothing — this avoids the Const{C, D}-with-rebound-C shape
+      // that triggers the downstream crash.
+      val saved = capSetBinders
+      val nested = capSetBinders.nonEmpty
+      if nested && !reported && reportPos.span.exists then
+        reported = true
+        report.restrictionError(
+          em"""nested capture-set type-lambda binders aren't supported in inferred capture-polymorphic function literals.
+              |Combine binders into a single section like `[C^, D^, ...]`, or define a class with a polymorphic apply method.""",
+          reportPos)
+      capSetBinders = if nested then Nil else tp :: capSetBinders
+      try mapOver(tp) finally capSetBinders = saved
     case _ => mapOver(tp)
 
 /** A base class for extractors that match annotated types with a specific
@@ -1016,4 +1111,3 @@ abstract class DeepTypeAccumulator[T](using Context) extends TypeAccumulator[T]:
       case _ =>
         foldOver(acc, t)
 end DeepTypeAccumulator
-

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

@@ -338,14 +338,38 @@ class Setup extends PreRecheck, SymTransformer, SetupAPI:
 
       def innerApply(tp: Type) =
         val tp1 = tp match
+          case AnnotatedType(parent, annot: RetainingAnnotation) =>
+            val elems = annot.retainedType.retainedElementsRaw
+            def isRetainedParamRef(elem: Type): Boolean = elem match
+              case _: ParamRef => true
+              // CleanupRetains may also preserve capset TypeRefs (e.g. to a
+              // type member of a synthetic poly-function class that aliases an
+              // outer apply method's type parameter).
+              case ref: TypeRef => ref.derivesFromCapSet
+              case _ => false
+            def promote(caps: List[Capability]) =
+              // Refs preserved by CleanupRetains are part of the inferred
+              // polymorphic function interface. Keep them as a Const
+              // CapturingType and strip the empty Var that `apply(parent)`
+              // would have layered underneath.
+              val parent1 = apply(parent) match
+                case CapturingType(p, refs) if refs.elems.isEmpty && !refs.isConst => p
+                case other => other
+              CapturingType(parent1, CaptureSet(caps*))
+            try
+              if elems.nonEmpty && elems.forall(isRetainedParamRef) then
+                promote(elems.map(_.toCapability))
+              else apply(parent)
+            catch case _: IllegalCaptureRef =>
+              apply(parent)
           case AnnotatedType(parent, annot)
           if annot.symbol.isRetains || annot.symbol == defn.InferredAnnot =>
-            // Drop explicit retains and @inferred annotations
+            // Drop non-RetainingAnnotation retains (e.g. pickle-read) and @inferred.
             apply(parent)
           case tp: TypeLambda =>
-            // Don't recurse into parameter bounds, just cleanup any stray retains annotations
+            // Leave parameter bounds alone; CleanupRetains already filtered them.
             tp.derivedLambdaType(
-              paramInfos = tp.paramInfos.mapConserve(_.dropAllRetains.bounds),
+              paramInfos = tp.paramInfos,
               resType = this(tp.resType))
           case tp @ RefinedType(parent, rname, rinfo) =>
             val saved = refiningNames
@@ -879,6 +903,9 @@ class Setup extends PreRecheck, SymTransformer, SetupAPI:
    *   - If type is a capturing type that has already a capture set variable or has
    *     the universal capture set, it does not need a variable.
    */
+  private def isNonEmptyParamRefSet(refs: CaptureSet)(using Context): Boolean =
+    !refs.elems.isEmpty && refs.elems.forall(_.core.isInstanceOf[ParamRef])
+
   def needsVariable(tp: Type)(using Context): Boolean =
     tp.typeParams.isEmpty && tp.match
       case tp: (TypeRef | AppliedType) =>
@@ -899,6 +926,9 @@ 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
+        // Non-empty Const sets that contain only parameter refs must stay Const:
+        // a Var's elements don't rewrite under SubstParamsMap. See i25830.
+        && !isNonEmptyParamRefSet(refs)
       case AnnotatedType(parent, _) =>
         needsVariable(parent)
       case _ =>

compiler/src/dotty/tools/dotc/inlines/Inlines.scala

@@ -102,8 +102,12 @@ object Inlines:
    */
   def inlineCall(tree: Tree)(using Context): Tree = ctx.profiler.onInlineCall(tree.symbol):
 
-    /** Strip @retains annotations from inferred types in the call tree */
-    val stripRetains = CleanupRetains()
+    /** Strip @retains annotations from inferred types in the call tree.
+     *  Setup expects retains in inferred types to have been cleaned; types
+     *  arriving here from inlining bypass PostTyper's CleanupRetains, so we
+     *  apply a plain (non-poly-literal-aware) cleanup pass here.
+     */
+    val stripRetains = CleanupRetains(ownerWalk = false)
     val stripper = new TreeTypeMap(
       treeMap = {
         case tree: InferredTypeTree =>
@@ -127,7 +131,7 @@ object Inlines:
       // not to provoke cyclic references. See i16116 for a test case.
       CrossVersionChecks.checkRef(tree0.symbol, tree0.srcPos)
 
-    if tree0.symbol.isConstructor then return tree // error already reported for the inline constructor definition
+    if tree0.symbol.isConstructor then return tree0 // error already reported for the inline constructor definition
 
     /** Set the position of all trees logically contained in the expansion of
      *  inlined call `call` to the position of `call`. This transform is necessary
@@ -175,17 +179,17 @@ object Inlines:
         tree
     }
 
-    // assertAllPositioned(tree0)   // debug
-    val tree1 = liftBindings(tree0, identity)
+    // assertAllPositioned(tree)   // debug
+    val tree1 = liftBindings(tree, identity)
     val tree2  =
       if bindings.nonEmpty then
-        cpy.Block(tree0)(bindings.toList, inlineCall(tree1))
+        cpy.Block(tree)(bindings.toList, inlineCall(tree1))
       else if enclosingInlineds.length < ctx.settings.XmaxInlines.value && !reachedInlinedTreesLimit then
         val body =
-          try bodyToInline(tree0.symbol) // can typecheck the tree and thereby produce errors
+          try bodyToInline(tree.symbol) // can typecheck the tree and thereby produce errors
           catch case _: MissingInlineInfo =>
             throw CyclicReference(ctx.owner)
-        new InlineCall(tree0).expand(body)
+        new InlineCall(tree).expand(body)
       else
         ctx.base.stopInlining = true
         val (reason, setting) =

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

@@ -16,7 +16,7 @@ import Symbols.*, NameOps.*
 import ContextFunctionResults.annotateContextResults
 import config.Printers.typr
 import config.Feature
-import util.{SrcPos, Stats}
+import util.{SrcPos, Stats, NoSourcePosition}
 import reporting.*
 import NameKinds.{WildcardParamName, TempResultName}
 import typer.Applications.{spread, HasSpreads}
@@ -30,6 +30,19 @@ import ast.TreeInfo
 object PostTyper {
   val name: String = "posttyper"
   val description: String = "additional checks and cleanups after type checking"
+
+  /** Detect whether `rhs` (or, transitively, any closure body it nests into)
+   *  is a capture-polymorphic poly-fn literal — i.e. a closure whose
+   *  synthesized `$anonfun` carries a `PolyType` with at least one capset
+   *  binder. The recursion catches cases like `(i: Int) => { [C^] => ... }`,
+   *  where the user-written literal is the body of an outer Function1.
+   */
+  def isCapsetPolyFunLiteralRhs(rhs: tpd.Tree)(using Context): Boolean = rhs match
+    case tpd.closureDef(dd) =>
+      dd.symbol.info match
+        case pt: PolyType => pt.paramRefs.exists(_.derivesFromCapSet)
+        case _ => isCapsetPolyFunLiteralRhs(dd.rhs)
+    case _ => false
 }
 
 /** A macro transform that runs immediately after typer and that performs the following functions:
@@ -107,6 +120,14 @@ class PostTyper extends MacroTransform with InfoTransformer { thisPhase =>
     initContextCalled = true
     compilingScala2StdLib = Feature.shouldBehaveAsScala2(using ctx)
 
+  /** The set of as-yet untransformed TypeTrees that are the inferred tpt of a val/def
+   *  whose RHS is a capture-polymorphic poly-function literal. Used to direct
+   *  `cc.CleanupRetains` whether to enable the implementation-restriction error:
+   *  it should fire only on user-written literals, not on inferred poly-fn types
+   *  that come from method calls or aliases.
+   */
+  private val hasCapsetPolyFunLiteralRhs = mutable.Set[Tree]()
+
   val superAcc: SuperAccessors = new SuperAccessors(thisPhase)
   val synthMbr: SyntheticMembers = new SyntheticMembers(thisPhase)
   val beanProps: BeanProperties = new BeanProperties(thisPhase)
@@ -585,12 +606,27 @@ class PostTyper extends MacroTransform with InfoTransformer { thisPhase =>
           val tree1 = cpy.ValDef(tree)(tpt = makeOverrideTypeDeclared(tree.symbol, tree.tpt))
           if tree1.removeAttachment(desugar.UntupledParam).isDefined then
             checkStableSelection(tree.rhs)
+          tree1.tpt match
+            case tpt: TypeTree if tpt.isInferred && PostTyper.isCapsetPolyFunLiteralRhs(tree.rhs) =>
+              hasCapsetPolyFunLiteralRhs += tpt
+            case _ =>
           processValOrDefDef(super.transform(tree1))
         case tree: DefDef =>
           registerIfHasMacroAnnotations(tree)
           Checking.checkPolyFunctionType(tree.tpt)
           annotateContextResults(tree)
           val tree1 = cpy.DefDef(tree)(tpt = makeOverrideTypeDeclared(tree.symbol, tree.tpt))
+          // Only attach for user-named defs. Synthetic anon-funs (closure
+          // bodies) shouldn't carry the marker — that would fire restriction
+          // errors for literals nested inside Function1s even when the
+          // enclosing val has an explicit type ascription.
+          tree1.tpt match
+            case tpt: TypeTree
+            if !tree.symbol.isAnonymousFunction
+               && tpt.isInferred
+               && PostTyper.isCapsetPolyFunLiteralRhs(tree.rhs) =>
+              hasCapsetPolyFunLiteralRhs += tpt
+            case _ =>
           processValOrDefDef(superAcc.wrapDefDef(tree1)(super.transform(tree1).asInstanceOf[DefDef]))
         case tree: TypeDef =>
           registerIfHasMacroAnnotations(tree)
@@ -677,7 +713,11 @@ class PostTyper extends MacroTransform with InfoTransformer { thisPhase =>
               report.error(em"type ${alias.tpe} outside bounds $bounds", tree.srcPos)
           super.transform(tree)
         case tree: TypeTree =>
-          val tpe = if tree.isInferred then CleanupRetains()(tree.tpe) else tree.tpe
+          def reportPos: SrcPos =
+            if hasCapsetPolyFunLiteralRhs.remove(tree) && Feature.ccEnabled
+            then tree.srcPos
+            else NoSourcePosition
+          val tpe = if tree.isInferred then CleanupRetains(reportPos)(tree.tpe) else tree.tpe
           tree.withType(transformAnnotsIn(tpe))
         case Typed(Ident(nme.WILDCARD), _) =>
           withMode(Mode.Pattern)(super.transform(tree))

tests/neg-custom-args/captures/i25830-unsupported.scala

@@ -0,0 +1,56 @@
+import language.experimental.captureChecking
+import caps.*
+
+class File extends SharedCapability
+
+// TODO these cases are currently restricted by the implementation, but we should consider allowing some of them in the future. See i25830 for details.
+// The general case for capture-polymorphic lambdas is non-trivial to implement and we want to get the implementation right before allowing more cases.
+// On the other hand, they are currently rarely needed in practice and the workaround is not too bad, so we can live with the restriction for now.
+
+def mixedExternal() =
+  val external = File()
+  val f = // error
+    { [C^] => (xs: List[File^{C, external}]) => xs }
+
+def externalOnly() =
+  val external = File()
+  val f = // error
+    { [C^] => (xs: List[File^{external}]) => xs }
+
+def mixedExternalWithLowerBoundedParam() =
+  val external = File()
+  val f = // error
+    { [C^, D^ >: {C}] => (xs: List[File^{D, external}]) => xs }
+
+def mixedExternalInLaterParamList() =
+  val external = File()
+  val f = // error
+    { [C^] => (xs: List[File^{C}]) => (ys: List[File^{C, external}]) => ys }
+
+def enclosingParam(external: File^) =
+  val f = // error
+    { [C^] => (xs: List[File^{C}]) => (ys: List[File^{external}]) => ys }
+
+def unsupportedDef() =
+  val external = File()
+  def f = // error
+    { [C^] => (xs: List[File^{C, external}]) => xs }
+
+def unsupportedInsideAnonymousFunction() =
+  List(File()).map: external =>
+    val f = // error
+      { [C^] => (xs: List[File^{C}]) => (ys: List[File^{external}]) => ys }
+
+def externalInBound() =
+  val external = File()
+  val f = // error
+    { [C^, D^ <: {C, external}] => (xs: List[File^{D}]) => xs }
+
+def nestedCapsetBinders() =
+  val f = // error
+    { [C^] => (xs: List[File^{C}]) => [D^] => (ys: List[File^{C, D}]) => ys }
+
+def literalNestedInFunction1() =
+  val external = File()
+  val f = // error
+    (i: Int) => { [C^] => (xs: List[File^{C, external}]) => xs }