From 593872a3fcafa2ac2dc76d1a8a226247988b8f1a Mon Sep 17 00:00:00 2001
From: noti0na1 <noti0na1@users.noreply.github.com>
Date: Fri, 1 Nov 2024 13:22:47 +0100
Subject: [PATCH] Handle CapsOf in more places

---
 .../src/dotty/tools/dotc/cc/CaptureOps.scala  |  2 +-
 .../src/dotty/tools/dotc/cc/CaptureRef.scala  |  2 ++
 .../dotty/tools/dotc/cc/CheckCaptures.scala   |  2 +-
 .../src/dotty/tools/dotc/core/Types.scala     |  4 ++-
 .../captures/capture-poly.scala               | 25 +++++++++++++++++++
 5 files changed, 32 insertions(+), 3 deletions(-)
 create mode 100644 tests/neg-custom-args/captures/capture-poly.scala

diff --git a/compiler/src/dotty/tools/dotc/cc/CaptureOps.scala b/compiler/src/dotty/tools/dotc/cc/CaptureOps.scala
index aad6ca8ddeac..a9272b73e605 100644
--- a/compiler/src/dotty/tools/dotc/cc/CaptureOps.scala
+++ b/compiler/src/dotty/tools/dotc/cc/CaptureOps.scala
@@ -198,7 +198,7 @@ extension (tp: Type)
       || tp.isRootCapability
       ) && !tp.symbol.isOneOf(UnstableValueFlags)
     case tp: TypeRef =>
-      tp.symbol.isAbstractOrParamType && tp.derivesFrom(defn.Caps_CapSet)
+      tp.symbol.isType && tp.derivesFrom(defn.Caps_CapSet)
     case tp: TypeParamRef =>
       tp.derivesFrom(defn.Caps_CapSet)
     case AnnotatedType(parent, annot) =>
diff --git a/compiler/src/dotty/tools/dotc/cc/CaptureRef.scala b/compiler/src/dotty/tools/dotc/cc/CaptureRef.scala
index 590beda42903..199114880c2b 100644
--- a/compiler/src/dotty/tools/dotc/cc/CaptureRef.scala
+++ b/compiler/src/dotty/tools/dotc/cc/CaptureRef.scala
@@ -140,6 +140,8 @@ trait CaptureRef extends TypeProxy, ValueType:
         case ReachCapability(x1) => x1.subsumes(y.stripReach)
         case x: TermRef => viaInfo(x.info)(subsumingRefs(_, y))
         case x: TermParamRef => subsumesExistentially(x, y)
+        case x: TypeRef if x.symbol.info.derivesFrom(defn.Caps_CapSet) =>
+          x.captureSetOfInfo.elems.exists(_.subsumes(y))
         case x: TypeRef => assumedContainsOf(x).contains(y)
         case _ => false
   end subsumes
diff --git a/compiler/src/dotty/tools/dotc/cc/CheckCaptures.scala b/compiler/src/dotty/tools/dotc/cc/CheckCaptures.scala
index 4d905a5df4ab..eefce44e1628 100644
--- a/compiler/src/dotty/tools/dotc/cc/CheckCaptures.scala
+++ b/compiler/src/dotty/tools/dotc/cc/CheckCaptures.scala
@@ -132,7 +132,7 @@ object CheckCaptures:
       elem match
         case CapsOfApply(arg) =>
           def isLegalCapsOfArg =
-            arg.symbol.isAbstractOrParamType && arg.symbol.info.derivesFrom(defn.Caps_CapSet)
+            arg.symbol.isType && arg.symbol.info.derivesFrom(defn.Caps_CapSet)
           if !isLegalCapsOfArg then
             report.error(
               em"""$arg is not a legal prefix for `^` here,
diff --git a/compiler/src/dotty/tools/dotc/core/Types.scala b/compiler/src/dotty/tools/dotc/core/Types.scala
index d8a4453325f2..14dc4f0dd0e5 100644
--- a/compiler/src/dotty/tools/dotc/core/Types.scala
+++ b/compiler/src/dotty/tools/dotc/core/Types.scala
@@ -4007,7 +4007,7 @@ object Types extends TypeUtils {
                 (compute(status, parent, theAcc) /: refs.elems) {
                   (s, ref) => ref.stripReach match
                     case tp: TermParamRef if tp.binder eq thisLambdaType => combine(s, CaptureDeps)
-                    case _ => s
+                    case tp => combine(s, compute(status, tp, theAcc))
                 }
               case _ =>
                 if tp.annot.refersToParamOf(thisLambdaType) then TrueDeps
@@ -6077,6 +6077,8 @@ object Types extends TypeUtils {
         case tp: CaptureRef =>
           if tp.isTrackableRef then tp
           else ensureTrackable(tp.underlying)
+        case tp: TypeAlias =>
+          ensureTrackable(tp.alias)
         case _ =>
           assert(false, i"not a trackable captureRef ref: $result, ${result.underlyingIterator.toList}")
       ensureTrackable(result)
diff --git a/tests/neg-custom-args/captures/capture-poly.scala b/tests/neg-custom-args/captures/capture-poly.scala
new file mode 100644
index 000000000000..a3a7a4c2a3d7
--- /dev/null
+++ b/tests/neg-custom-args/captures/capture-poly.scala
@@ -0,0 +1,25 @@
+import caps.*
+
+trait Foo extends Capability
+
+trait CaptureSet:
+  type C <: CapSet^
+
+def capturePoly[C^](a: Foo^{C^}): Foo^{C^} = a
+def capturePoly2(c: CaptureSet)(a: Foo^{c.C^}): Foo^{c.C^} = a
+
+def test =
+  val x: Foo^ = ???
+  val y: Foo^ = ???
+
+  object X extends CaptureSet:
+    type C = CapSet^{x}
+
+  val z1: Foo^{X.C^} = x
+  val z2: Foo^{X.C^} = y // error
+
+  val z3: Foo^{x} = capturePoly(x)
+  val z4: Foo^{x} = capturePoly(y) // error
+
+  val z5: Foo^{x} = capturePoly2(X)(x)
+  val z6: Foo^{x} = capturePoly2(X)(y) // error
\ No newline at end of file