Track variable insytantiations in capture sets

If a variable in a capture set gets instantiated, the capture set
has to update itself as well.

Author: odersky

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

@@ -11,22 +11,36 @@ import reporting.trace
 import printing.{Showable, Printer}
 import printing.Texts.*
 
-case class CaptureSet private (elems: CaptureSet.Refs) extends Showable:
+case class CaptureSet private (elems0: CaptureSet.Refs) extends Showable:
   import CaptureSet.*
 
-  def isEmpty: Boolean = elems.isEmpty
-  def nonEmpty: Boolean = !isEmpty
+  def isEmpty(using Context): Boolean = elems.isEmpty
+  def nonEmpty(using Context): Boolean = !isEmpty
 
-  def ++ (that: CaptureSet): CaptureSet =
+  private var isProvisional = true
+  private var myElems: CaptureSet.Refs = elems0
+
+  def elems(using Context): CaptureSet.Refs =
+    if isProvisional then
+      isProvisional = false
+      myElems.foreach {
+        case tv: TypeVar =>
+          if tv.isInstantiated then myElems = myElems - tv ++ tv.inst.captureSet.elems
+          else isProvisional = true
+        case _ =>
+      }
+    myElems
+
+  def ++ (that: CaptureSet)(using Context): CaptureSet =
     if this.isEmpty then that
     else if that.isEmpty then this
-    else CaptureSet(elems ++ that.elems)
+    else CaptureSet(myElems ++ that.elems)
 
-  def + (ref: CaptureRef) =
+  def + (ref: CaptureRef)(using Context) =
     if elems.contains(ref) then this
     else CaptureSet(elems + ref)
 
-  def intersect (that: CaptureSet): CaptureSet =
+  def intersect (that: CaptureSet)(using Context): CaptureSet =
     CaptureSet(this.elems.intersect(that.elems))
 
   /** {x} <:< this   where <:< is subcapturing */
@@ -46,10 +60,10 @@ case class CaptureSet private (elems: CaptureSet.Refs) extends Showable:
       case ref => ref.singletonCaptureSet
     }
 
-  override def toString = elems.toString
+  override def toString = myElems.toString
 
   override def toText(printer: Printer): Text =
-    Str("{") ~ Text(elems.toList.map(printer.toTextCaptureRef), ", ") ~ Str("}")
+    Str("{") ~ Text(myElems.toList.map(printer.toTextCaptureRef), ", ") ~ Str("}")
 
 object CaptureSet:
   type Refs = SimpleIdentitySet[CaptureRef]

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

@@ -4706,7 +4706,11 @@ object Types {
 
     override def captureSetOfInfo(using Context): CaptureSet = instanceOpt match
       case ref: CaptureRef => ref.captureSetOfInfo
-      case tp => tp.captureSet
+      case _ => underlying.captureSet
+
+    override def captureSet(using Context): CaptureSet =
+      if isInstantiated then inst.captureSet
+      else super.captureSet
 
     // Object members
 

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

@@ -0,0 +1,33 @@
+import language.experimental.erasedDefinitions
+import annotation.ability
+import java.io.IOException
+
+class CanThrow[E] extends Retains[*]
+type Top  = Any retains *
+infix type ==> [A, B] = (A => B) retains *
+
+def handle[E <: Exception, T <: Top](op: (CanThrow[E] ?=> T))(handler: (E => T) retains T): T =
+  val x: CanThrow[E] = ???
+  try op(using x)
+  catch case ex: E => handler(ex)
+
+def raise[E <: Exception](ex: E)(using CanThrow[E]): Nothing =
+  throw ex
+
+def test2: Int ==> Int =
+  def f(a: Boolean): Boolean => CanThrow[IOException] ?=> Int ==> Int =
+    handle {
+      if !a then raise(IOException())
+      (b: Boolean) => (_: CanThrow[IOException]) ?=>
+        if !b then raise(IOException())
+        (x: Int) => 1
+    } {
+      ex => (b: Boolean) => (_: CanThrow[IOException]) ?=> (x: Int) => -1
+    }
+  handle {  // error: inferred type argument ((Int => Int) retains *) is not allowed to capture the universal capability *
+    val g = f(true)
+    g(false) // would raise an uncaught exception
+    f(true)(false) // would raise an uncaught exception
+  } {
+    ex => (x: Int) => -1
+  }