Fix capture set healing

Capture set healing happens at the end of capture checking (in the `postCheck`
function), which checks the capture set in each type argument and widen the
ill-formed `TermParamRef`s by widening them. However, it is possible that a
capture set is healed first, and then get a `TermParamRef` propagated into it
later when we are healing another capture set. This lead to unsoundness.

This commit installs an handler on capture sets when healing them and will
inspect the newly added elements and heal these new elements as well.

Author: Linyxus

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

@@ -275,6 +275,11 @@ sealed abstract class CaptureSet extends Showable:
     if isUniversal then handler()
     this
 
+  /** Invoke handler on the elements to check wellformedness of the capture set */
+  def ensureWellformed(handler: List[CaptureRef] => Context ?=> Unit)(using Context): this.type =
+    handler(elems.toList)
+    this
+
   /** An upper approximation of this capture set, i.e. a constant set that is
    *  subcaptured by this set. If the current set is a variable
    *  it is the intersection of all upper approximations of known supersets
@@ -375,6 +380,9 @@ object CaptureSet:
     /** A handler to be invoked if the root reference `cap` is added to this set */
     var rootAddedHandler: () => Context ?=> Unit = () => ()
 
+    /** A handler to be invoked when new elems are added to this set */
+    var newElemAddedHandler: List[CaptureRef] => Context ?=> Unit = _ => ()
+
     var description: String = ""
 
     /** Record current elements in given VarState provided it does not yet
@@ -405,7 +413,8 @@ object CaptureSet:
       if !isConst && recordElemsState() then
         elems ++= newElems
         if isUniversal then rootAddedHandler()
-        // assert(id != 2 || elems.size != 2, this)
+        newElemAddedHandler(newElems.toList)
+        // assert(id != 5 || elems.size != 3, this)
         (CompareResult.OK /: deps) { (r, dep) =>
           r.andAlso(dep.tryInclude(newElems, this))
         }
@@ -425,6 +434,10 @@ object CaptureSet:
       rootAddedHandler = handler
       super.disallowRootCapability(handler)
 
+    override def ensureWellformed(handler: List[CaptureRef] => (Context) ?=> Unit)(using Context): this.type =
+      newElemAddedHandler = handler
+      super.ensureWellformed(handler)
+
     private var computingApprox = false
 
     /** Roughly: the intersection of all constant known supersets of this set.

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

@@ -976,8 +976,11 @@ class CheckCaptures extends Recheck, SymTransformer:
           recur(refs, Nil)
 
         private def healCaptureSet(cs: CaptureSet): Unit =
-          val toInclude = widenParamRefs(cs.elems.toList.filter(!isAllowed(_)).asInstanceOf)
-          toInclude.foreach(checkSubset(_, cs, tree.srcPos))
+          def avoidance(elems: List[CaptureRef])(using Context): Unit =
+            val toInclude = widenParamRefs(elems.filter(!isAllowed(_)).asInstanceOf)
+            //println(i"HEAL $cs by widening to $toInclude")
+            toInclude.foreach(checkSubset(_, cs, tree.srcPos))
+          cs.ensureWellformed(avoidance)
 
         private var allowed: SimpleIdentitySet[TermParamRef] = SimpleIdentitySet.empty
 

tests/neg-custom-args/captures/usingLogFile-alt.check

@@ -0,0 +1,7 @@
+-- Error: tests/neg-custom-args/captures/usingLogFile-alt.scala:18:2 ---------------------------------------------------
+18 |  usingFile(  // error
+   |  ^^^^^^^^^
+   |  Sealed type variable T cannot  be instantiated to box () => Unit since
+   |  that type captures the root capability `cap`.
+   |  This is often caused by a local capability in the body of method usingFile
+   |  leaking as part of its result.

tests/neg-custom-args/captures/usingLogFile-alt.scala

@@ -0,0 +1,23 @@
+// Reported in issue #17517
+
+import language.experimental.captureChecking
+import java.io.*
+
+object Test:
+  class Logger(f: OutputStream^):
+    def log(msg: String): Unit = ???
+
+  def usingFile[sealed T](name: String, op: OutputStream^ => T): T =
+    val f = new FileOutputStream(name)
+    val result = op(f)
+    f.close()
+    result
+
+  def usingLogger[sealed T](f: OutputStream^)(op: Logger^{f} => T): T = ???
+
+  usingFile(  // error
+    "foo",
+    file => {
+      usingLogger(file)(l => () => l.log("test"))
+    }
+  )

tests/neg-custom-args/captures/usingLogFile.check

@@ -45,10 +45,3 @@
    |                that type captures the root capability `cap`.
    |                This is often caused by a local capability in the body of method usingFile
    |                leaking as part of its result.
--- Error: tests/neg-custom-args/captures/usingLogFile.scala:72:6 -------------------------------------------------------
-72 |      usingLogger(_, l => () => l.log("test"))) // error
-   |      ^^^^^^^^^^^
-   |      Sealed type variable T cannot  be instantiated to box () => Unit since
-   |      that type captures the root capability `cap`.
-   |      This is often caused by a local capability in the body of method usingLogger
-   |      leaking as part of its result.

tests/neg-custom-args/captures/usingLogFile.scala

@@ -69,5 +69,5 @@ object Test4:
 
   def test =
     val later = usingFile("logfile",            // error
-      usingLogger(_, l => () => l.log("test"))) // error
+      usingLogger(_, l => () => l.log("test"))) // ok, since we can widen `l` to `file` instead of to `cap`
     later()