From 5e7de2e3462824289a9372137b5a0eb98aaea915 Mon Sep 17 00:00:00 2001
From: odersky <odersky@gmail.com>
Date: Mon, 12 Aug 2024 21:02:00 +0200
Subject: [PATCH] Expand Capability types T to T^ only if not explicit capture
 set is given

---
 .../dotty/tools/dotc/cc/CheckCaptures.scala   |  6 ++---
 compiler/src/dotty/tools/dotc/cc/Setup.scala  | 22 +++++++++++++++----
 .../dotty/tools/dotc/core/Definitions.scala   |  3 +++
 .../pos-custom-args/captures/cap-refine.scala | 12 ++++++++++
 4 files changed, 36 insertions(+), 7 deletions(-)
 create mode 100644 tests/pos-custom-args/captures/cap-refine.scala

diff --git a/compiler/src/dotty/tools/dotc/cc/CheckCaptures.scala b/compiler/src/dotty/tools/dotc/cc/CheckCaptures.scala
index 51cf362ca667..27a3d6024b65 100644
--- a/compiler/src/dotty/tools/dotc/cc/CheckCaptures.scala
+++ b/compiler/src/dotty/tools/dotc/cc/CheckCaptures.scala
@@ -632,7 +632,7 @@ class CheckCaptures extends Recheck, SymTransformer:
         def addParamArgRefinements(core: Type, initCs: CaptureSet): (Type, CaptureSet) =
           var refined: Type = core
           var allCaptures: CaptureSet =
-            if core.derivesFromCapability then CaptureSet.universal else initCs
+            if core.derivesFromCapability then defn.universalCSImpliedByCapability else initCs
           for (getterName, argType) <- mt.paramNames.lazyZip(argTypes) do
             val getter = cls.info.member(getterName).suchThat(_.isRefiningParamAccessor).symbol
             if !getter.is(Private) && getter.hasTrackedParts then
@@ -809,7 +809,7 @@ class CheckCaptures extends Recheck, SymTransformer:
         val localSet = capturedVars(sym)
         if !localSet.isAlwaysEmpty then
           curEnv = Env(sym, EnvKind.Regular, localSet, curEnv)
-          
+
         // ctx with AssumedContains entries for each Contains parameter
         val bodyCtx =
           var ac = CaptureSet.assumedContains
@@ -828,7 +828,7 @@ class CheckCaptures extends Recheck, SymTransformer:
               interpolateVarsIn(tree.tpt)
             curEnv = saved
     end recheckDefDef
-    
+
     /** If val or def definition with inferred (result) type is visible
      *  in other compilation units, check that the actual inferred type
      *  conforms to the expected type where all inferred capture sets are dropped.
diff --git a/compiler/src/dotty/tools/dotc/cc/Setup.scala b/compiler/src/dotty/tools/dotc/cc/Setup.scala
index f578d10702e9..1c3326c5cc12 100644
--- a/compiler/src/dotty/tools/dotc/cc/Setup.scala
+++ b/compiler/src/dotty/tools/dotc/cc/Setup.scala
@@ -299,16 +299,29 @@ class Setup extends PreRecheck, SymTransformer, SetupAPI:
           CapturingType(fntpe, cs, boxed = false)
         else fntpe
 
+      def stripImpliedCaptureSet(tp: Type): Type = tp match
+        case tp @ CapturingType(parent, refs)
+        if (refs eq defn.universalCSImpliedByCapability) && !tp.isBoxedCapturing =>
+          parent
+        case tp @ CapturingType(parent, refs) =>
+         // println(i"other $tp")
+          tp
+        case _ =>
+          tp
+
       def apply(t: Type) =
         t match
           case t @ CapturingType(parent, refs) =>
-            t.derivedCapturingType(this(parent), refs)
+            t.derivedCapturingType(stripImpliedCaptureSet(this(parent)), refs)
+              //.showing(i"CAPT $result")
           case t @ AnnotatedType(parent, ann) =>
             val parent1 = this(parent)
             if ann.symbol.isRetains then
+              val parent2 = stripImpliedCaptureSet(parent1)
               for tpt <- tptToCheck do
-                checkWellformedLater(parent1, ann.tree, tpt)
-              CapturingType(parent1, ann.tree.toCaptureSet)
+                checkWellformedLater(parent2, ann.tree, tpt)
+              CapturingType(parent2, ann.tree.toCaptureSet)
+                //.showing(i"CAPT2 $result")
             else
               t.derivedAnnotatedType(parent1, ann)
           case throwsAlias(res, exc) =>
@@ -316,7 +329,8 @@ class Setup extends PreRecheck, SymTransformer, SetupAPI:
           case t =>
             // Map references to capability classes C to C^
             if t.derivesFromCapability && !t.isSingleton && t.typeSymbol != defn.Caps_Exists
-            then CapturingType(t, CaptureSet.universal, boxed = false)
+            then CapturingType(t, defn.universalCSImpliedByCapability, boxed = false)
+              //.showing(i"gen cap $result")
             else normalizeCaptures(mapOver(t))
     end toCapturing
 
diff --git a/compiler/src/dotty/tools/dotc/core/Definitions.scala b/compiler/src/dotty/tools/dotc/core/Definitions.scala
index 8981aa4aa6ac..f95bb3cea351 100644
--- a/compiler/src/dotty/tools/dotc/core/Definitions.scala
+++ b/compiler/src/dotty/tools/dotc/core/Definitions.scala
@@ -1005,6 +1005,9 @@ class Definitions {
     @tu lazy val Caps_ContainsTrait: TypeSymbol = CapsModule.requiredType("Contains")
     @tu lazy val Caps_containsImpl: TermSymbol = CapsModule.requiredMethod("containsImpl")
 
+  /** The same as CaptureSet.universal but generated implicitly for references of Capability subtypes */
+  @tu lazy val universalCSImpliedByCapability = CaptureSet(captureRoot.termRef)
+
   @tu lazy val PureClass: Symbol = requiredClass("scala.Pure")
 
   // Annotation base classes
diff --git a/tests/pos-custom-args/captures/cap-refine.scala b/tests/pos-custom-args/captures/cap-refine.scala
new file mode 100644
index 000000000000..ed0b4d018b88
--- /dev/null
+++ b/tests/pos-custom-args/captures/cap-refine.scala
@@ -0,0 +1,12 @@
+//> using options -Werror
+import caps.Capability
+
+trait Buffer[T] extends Capability:
+  def append(x: T): this.type
+
+def f(buf: Buffer[Int]) =
+  val buf1 = buf.append(1).append(2)
+  val buf2: Buffer[Int]^{buf1} = buf1
+
+
+