From e5987d66f01aa65d82423f1ba4109ea4e3c9eacb Mon Sep 17 00:00:00 2001
From: odersky <odersky@gmail.com>
Date: Fri, 23 Aug 2024 19:41:07 +0200
Subject: [PATCH] Check all top-level covariant capture sets in
 checkNotUniversal

Fixes #21401
---
 .../dotty/tools/dotc/cc/CheckCaptures.scala   | 29 ++++++++++++-------
 tests/neg-custom-args/captures/i21401.check   | 10 +++++++
 tests/neg-custom-args/captures/i21401.scala   | 19 ++++++++++++
 3 files changed, 47 insertions(+), 11 deletions(-)
 create mode 100644 tests/neg-custom-args/captures/i21401.check
 create mode 100644 tests/neg-custom-args/captures/i21401.scala

diff --git a/compiler/src/dotty/tools/dotc/cc/CheckCaptures.scala b/compiler/src/dotty/tools/dotc/cc/CheckCaptures.scala
index 27a3d6024b65..384c6e1f29ef 100644
--- a/compiler/src/dotty/tools/dotc/cc/CheckCaptures.scala
+++ b/compiler/src/dotty/tools/dotc/cc/CheckCaptures.scala
@@ -978,21 +978,28 @@ class CheckCaptures extends Recheck, SymTransformer:
         case _: RefTree | _: Apply | _: TypeApply => tree.symbol.unboxesResult
         case _: Try => true
         case _ => false
-      def checkNotUniversal(tp: Type): Unit = tp.widenDealias match
-        case wtp @ CapturingType(parent, refs) =>
-          refs.disallowRootCapability { () =>
-            report.error(
-              em"""The expression's type $wtp is not allowed to capture the root capability `cap`.
-                  |This usually means that a capability persists longer than its allowed lifetime.""",
-              tree.srcPos)
-          }
-          checkNotUniversal(parent)
-        case _ =>
+
+      object checkNotUniversal extends TypeTraverser:
+        def traverse(tp: Type) =
+         tp.dealias match
+          case wtp @ CapturingType(parent, refs) =>
+            if variance > 0 then
+              refs.disallowRootCapability: () =>
+                def part = if wtp eq tpe.widen then "" else i" in its part $wtp"
+                report.error(
+                  em"""The expression's type ${tpe.widen} is not allowed to capture the root capability `cap`$part.
+                    |This usually means that a capability persists longer than its allowed lifetime.""",
+                  tree.srcPos)
+            if !wtp.isBoxed then traverse(parent)
+          case tp =>
+            traverseChildren(tp)
+
       if !ccConfig.useSealed
           && !tpe.hasAnnotation(defn.UncheckedCapturesAnnot)
           && needsUniversalCheck
+          && tpe.widen.isValueType
       then
-        checkNotUniversal(tpe)
+        checkNotUniversal.traverse(tpe.widen)
       super.recheckFinish(tpe, tree, pt)
     end recheckFinish
 
diff --git a/tests/neg-custom-args/captures/i21401.check b/tests/neg-custom-args/captures/i21401.check
new file mode 100644
index 000000000000..e9a5fbd4678c
--- /dev/null
+++ b/tests/neg-custom-args/captures/i21401.check
@@ -0,0 +1,10 @@
+-- Error: tests/neg-custom-args/captures/i21401.scala:15:22 ------------------------------------------------------------
+15 |  val a = usingIO[IO^](x => x) // error: The expression's type IO^ is not allowed to capture the root capability `cap`
+   |          ^^^^^^^^^^^^^^^^^^^^
+   |          The expression's type box IO^ is not allowed to capture the root capability `cap`.
+   |          This usually means that a capability persists longer than its allowed lifetime.
+-- Error: tests/neg-custom-args/captures/i21401.scala:16:70 ------------------------------------------------------------
+16 |  val leaked: [R, X <: Boxed[IO^] -> R] -> (op: X) -> R = usingIO[Res](mkRes) // error: The expression's type Res is not allowed to capture the root capability `cap` in its part box IO^
+   |                                                          ^^^^^^^^^^^^^^^^^^^
+   |              The expression's type Res is not allowed to capture the root capability `cap` in its part box IO^.
+   |              This usually means that a capability persists longer than its allowed lifetime.
diff --git a/tests/neg-custom-args/captures/i21401.scala b/tests/neg-custom-args/captures/i21401.scala
new file mode 100644
index 000000000000..07d407a79809
--- /dev/null
+++ b/tests/neg-custom-args/captures/i21401.scala
@@ -0,0 +1,19 @@
+import language.experimental.captureChecking
+
+trait IO:
+  def println(s: String): Unit
+def usingIO[R](op: IO^ => R): R = ???
+
+case class Boxed[+T](unbox: T)
+
+type Res = [R, X <: Boxed[IO^] -> R] -> (op: X) -> R
+def mkRes(x: IO^): Res =
+  [R, X <: Boxed[IO^] -> R] => (op: X) =>
+    val op1: Boxed[IO^] -> R = op
+    op1(Boxed[IO^](x))
+def test2() =
+  val a = usingIO[IO^](x => x) // error: The expression's type IO^ is not allowed to capture the root capability `cap`
+  val leaked: [R, X <: Boxed[IO^] -> R] -> (op: X) -> R = usingIO[Res](mkRes) // error: The expression's type Res is not allowed to capture the root capability `cap` in its part box IO^
+  val x: Boxed[IO^] = leaked[Boxed[IO^], Boxed[IO^] -> Boxed[IO^]](x => x)
+  val y: IO^{x*} = x.unbox
+  y.println("boom")