From 9c9653819e9473c0cc44a2a7cbc1cd7441e194ce Mon Sep 17 00:00:00 2001
From: odersky <odersky@gmail.com>
Date: Mon, 3 Jul 2023 14:36:42 +0200
Subject: [PATCH] Follow span of function types when computing captureSetOfInfo

---
 .../src/dotty/tools/dotc/cc/CaptureSet.scala  | 63 +++++++++++--------
 .../dotty/tools/dotc/cc/CheckCaptures.scala   |  3 +-
 .../src/dotty/tools/dotc/core/Types.scala     |  2 +-
 .../captures/dependent-pure.scala             |  5 ++
 4 files changed, 45 insertions(+), 28 deletions(-)
 create mode 100644 tests/pos-custom-args/captures/dependent-pure.scala

diff --git a/compiler/src/dotty/tools/dotc/cc/CaptureSet.scala b/compiler/src/dotty/tools/dotc/cc/CaptureSet.scala
index fd07133b55e1..9de68220d2cf 100644
--- a/compiler/src/dotty/tools/dotc/cc/CaptureSet.scala
+++ b/compiler/src/dotty/tools/dotc/cc/CaptureSet.scala
@@ -846,34 +846,45 @@ object CaptureSet:
   /** The capture set of the type underlying a CaptureRef */
   def ofInfo(ref: CaptureRef)(using Context): CaptureSet = ref match
     case ref: TermRef if ref.isRootCapability => ref.singletonCaptureSet
-    case _ => ofType(ref.underlying)
+    case _ => ofType(ref.underlying, followResult = true)
 
   /** Capture set of a type */
-  def ofType(tp: Type)(using Context): CaptureSet =
-    def recur(tp: Type): CaptureSet = tp.dealias match
-      case tp: TermRef =>
-        tp.captureSet
-      case tp: TermParamRef =>
-        tp.captureSet
-      case _: TypeRef =>
-        if tp.classSymbol.hasAnnotation(defn.CapabilityAnnot) then universal else empty
-      case _: TypeParamRef =>
-        empty
-      case CapturingType(parent, refs) =>
-        recur(parent) ++ refs
-      case AppliedType(tycon, args) =>
-        val cs = recur(tycon)
-        tycon.typeParams match
-          case tparams @ (LambdaParam(tl, _) :: _) => cs.substParams(tl, args)
-          case _ => cs
-      case tp: TypeProxy =>
-        recur(tp.underlying)
-      case AndType(tp1, tp2) =>
-        recur(tp1) ** recur(tp2)
-      case OrType(tp1, tp2) =>
-        recur(tp1) ++ recur(tp2)
-      case _ =>
-        empty
+  def ofType(tp: Type, followResult: Boolean)(using Context): CaptureSet =
+    def recur(tp: Type): CaptureSet = trace(i"ofType $tp, ${tp.getClass} $followResult", show = true):
+      tp.dealias match
+        case tp: TermRef =>
+          tp.captureSet
+        case tp: TermParamRef =>
+          tp.captureSet
+        case _: TypeRef =>
+          if tp.classSymbol.hasAnnotation(defn.CapabilityAnnot) then universal else empty
+        case _: TypeParamRef =>
+          empty
+        case CapturingType(parent, refs) =>
+          recur(parent) ++ refs
+        case tpd @ RefinedType(parent, _, rinfo: MethodType)
+        if followResult && defn.isFunctionType(tpd) =>
+          ofType(parent, followResult = false)                 // pick up capture set from parent type
+          ++ (recur(rinfo.resType)                             // add capture set of result
+          -- CaptureSet(rinfo.paramRefs.filter(_.isTracked)*)) // but disregard bound parameters
+        case tpd @ AppliedType(tycon, args) =>
+          if followResult && defn.isNonRefinedFunction(tpd) then
+            recur(args.last)
+              // must be (pure) FunctionN type since ImpureFunctions have already
+              // been eliminated in selector's dealias. Use capture set of result.
+          else
+            val cs = recur(tycon)
+            tycon.typeParams match
+              case tparams @ (LambdaParam(tl, _) :: _) => cs.substParams(tl, args)
+              case _ => cs
+        case tp: TypeProxy =>
+          recur(tp.underlying)
+        case AndType(tp1, tp2) =>
+          recur(tp1) ** recur(tp2)
+        case OrType(tp1, tp2) =>
+          recur(tp1) ++ recur(tp2)
+        case _ =>
+          empty
     recur(tp)
       .showing(i"capture set of $tp = $result", capt)
 
diff --git a/compiler/src/dotty/tools/dotc/cc/CheckCaptures.scala b/compiler/src/dotty/tools/dotc/cc/CheckCaptures.scala
index d186e5c605bb..0f4486c86d68 100644
--- a/compiler/src/dotty/tools/dotc/cc/CheckCaptures.scala
+++ b/compiler/src/dotty/tools/dotc/cc/CheckCaptures.scala
@@ -854,8 +854,9 @@ class CheckCaptures extends Recheck, SymTransformer:
       actual match
         case ref: CaptureRef if ref.isTracked =>
           actualw match
-            case CapturingType(p, refs) =>
+            case CapturingType(p, refs) if ref.singletonCaptureSet.mightSubcapture(refs) =>
               actualw = actualw.derivedCapturingType(p, ref.singletonCaptureSet)
+                .showing(i"improve $actualw to $result", capt)
                 // given `a: C T`, improve `C T` to `{a} T`
             case _ =>
         case _ =>
diff --git a/compiler/src/dotty/tools/dotc/core/Types.scala b/compiler/src/dotty/tools/dotc/core/Types.scala
index e3aaa235858b..a68986ea9d18 100644
--- a/compiler/src/dotty/tools/dotc/core/Types.scala
+++ b/compiler/src/dotty/tools/dotc/core/Types.scala
@@ -1561,7 +1561,7 @@ object Types {
     }
 
     /** The capture set of this type. Overridden and cached in CaptureRef */
-    def captureSet(using Context): CaptureSet = CaptureSet.ofType(this)
+    def captureSet(using Context): CaptureSet = CaptureSet.ofType(this, followResult = false)
 
     // ----- Normalizing typerefs over refined types ----------------------------
 
diff --git a/tests/pos-custom-args/captures/dependent-pure.scala b/tests/pos-custom-args/captures/dependent-pure.scala
new file mode 100644
index 000000000000..ad10d9590f25
--- /dev/null
+++ b/tests/pos-custom-args/captures/dependent-pure.scala
@@ -0,0 +1,5 @@
+import language.experimental.captureChecking
+class ContextCls
+type Context = ContextCls^
+
+class Filtered(p: (c: Context) ?-> () ->{c} Boolean) extends Pure