From c2265771133aaf84708f132f81280894a9984643 Mon Sep 17 00:00:00 2001
From: Martin Odersky <odersky@gmail.com>
Date: Sat, 19 Jun 2021 12:56:46 +0200
Subject: [PATCH 01/14] Refactor tests and fix isGlobal

---
 .../tools/dotc/typer/CheckCaptures.scala      |  1 +
 .../dotty/tools/dotc/CompilationTests.scala   |  2 +-
 tests/neg-custom-args/captures/io.scala       | 21 +++++++++++++++++++
 .../captures/boxmap.scala                     |  0
 .../captures/capt-depfun.scala                |  0
 .../captures/capt1.scala                      |  0
 .../captures/list-encoding.scala              |  0
 .../captures/try.scala                        |  0
 .../captures/try3.scala                       |  0
 9 files changed, 23 insertions(+), 1 deletion(-)
 create mode 100644 tests/neg-custom-args/captures/io.scala
 rename tests/{pos-special => pos-custom-args}/captures/boxmap.scala (100%)
 rename tests/{pos-special => pos-custom-args}/captures/capt-depfun.scala (100%)
 rename tests/{pos-special => pos-custom-args}/captures/capt1.scala (100%)
 rename tests/{pos-special => pos-custom-args}/captures/list-encoding.scala (100%)
 rename tests/{pos-special => pos-custom-args}/captures/try.scala (100%)
 rename tests/{pos-special => pos-custom-args}/captures/try3.scala (100%)

diff --git a/compiler/src/dotty/tools/dotc/typer/CheckCaptures.scala b/compiler/src/dotty/tools/dotc/typer/CheckCaptures.scala
index 4fceb72046ef..dce20d5a1d54 100644
--- a/compiler/src/dotty/tools/dotc/typer/CheckCaptures.scala
+++ b/compiler/src/dotty/tools/dotc/typer/CheckCaptures.scala
@@ -124,6 +124,7 @@ class CheckCaptures extends RefineTypes:
               val isGlobal = ref match
                 case ref: TypeRef => ref.isRootCapability
                 case ref: TermRef => ref.prefix != NoPrefix && ref.symbol.hasAnnotation(defn.AbilityAnnot)
+                case _ => false
               val what = if ref.isRootCapability then "universal" else "global"
               if isGlobal then
                 val notAllowed = i" is not allowed to capture the $what capability $ref"
diff --git a/compiler/test/dotty/tools/dotc/CompilationTests.scala b/compiler/test/dotty/tools/dotc/CompilationTests.scala
index 6b0c77d77b32..757ebf8310b9 100644
--- a/compiler/test/dotty/tools/dotc/CompilationTests.scala
+++ b/compiler/test/dotty/tools/dotc/CompilationTests.scala
@@ -33,13 +33,13 @@ class CompilationTests {
       compileFilesInDir("tests/pos-special/sourcepath/outer", defaultOptions.and("-sourcepath", "tests/pos-special/sourcepath")),
       compileFile("tests/pos-special/sourcepath/outer/nested/Test4.scala", defaultOptions.and("-sourcepath", "tests/pos-special/sourcepath")),
       compileFilesInDir("tests/pos-special/fatal-warnings", defaultOptions.and("-Xfatal-warnings", "-deprecation", "-feature")),
-      compileFilesInDir("tests/pos-special/captures", defaultOptions.and("-Ycc")),
       compileFile("tests/pos-special/avoid-warn-deprecation.scala", defaultOptions.and("-Xfatal-warnings", "-feature")),
       compileFilesInDir("tests/pos-special/spec-t5545", defaultOptions),
       compileFilesInDir("tests/pos-special/strawman-collections", allowDeepSubtypes),
       compileFilesInDir("tests/pos-special/isInstanceOf", allowDeepSubtypes.and("-Xfatal-warnings")),
       compileFilesInDir("tests/new", defaultOptions),
       compileFilesInDir("tests/pos-scala2", scala2CompatMode),
+      compileFilesInDir("tests/pos-custom-args/captures", defaultOptions.and("-Ycc")),
       compileFilesInDir("tests/pos-custom-args/erased", defaultOptions.and("-language:experimental.erasedDefinitions")),
       compileFilesInDir("tests/pos", defaultOptions.and("-Ysafe-init")),
       compileFilesInDir("tests/pos-deep-subtype", allowDeepSubtypes),
diff --git a/tests/neg-custom-args/captures/io.scala b/tests/neg-custom-args/captures/io.scala
new file mode 100644
index 000000000000..a9636b045694
--- /dev/null
+++ b/tests/neg-custom-args/captures/io.scala
@@ -0,0 +1,21 @@
+sealed trait IO:
+  def puts(msg: Any): Unit = println(msg)
+
+def test1 =
+  val IO : IO retains * = new IO {}
+  def foo = IO.puts("hello")
+  val x : () => Unit = () => foo  // error: Found: (() => Unit) retains IO; Required: () => Unit
+
+def test2 =
+  val IO : IO retains * = new IO {}
+  def puts(msg: Any, io: IO retains *) = println(msg)
+  def foo() = puts("hello", IO)
+  val x : () => Unit = () => foo()  // error: Found: (() => Unit) retains IO; Required: () => Unit
+
+type Capability[T] = T retains *
+
+def test3 =
+  val IO : Capability[IO] = new IO {}
+  def puts(msg: Any, io: Capability[IO]) = println(msg)
+  def foo() = puts("hello", IO)
+  val x : () => Unit = () => foo()  // error: Found: (() => Unit) retains IO; Required: () => Unit
diff --git a/tests/pos-special/captures/boxmap.scala b/tests/pos-custom-args/captures/boxmap.scala
similarity index 100%
rename from tests/pos-special/captures/boxmap.scala
rename to tests/pos-custom-args/captures/boxmap.scala
diff --git a/tests/pos-special/captures/capt-depfun.scala b/tests/pos-custom-args/captures/capt-depfun.scala
similarity index 100%
rename from tests/pos-special/captures/capt-depfun.scala
rename to tests/pos-custom-args/captures/capt-depfun.scala
diff --git a/tests/pos-special/captures/capt1.scala b/tests/pos-custom-args/captures/capt1.scala
similarity index 100%
rename from tests/pos-special/captures/capt1.scala
rename to tests/pos-custom-args/captures/capt1.scala
diff --git a/tests/pos-special/captures/list-encoding.scala b/tests/pos-custom-args/captures/list-encoding.scala
similarity index 100%
rename from tests/pos-special/captures/list-encoding.scala
rename to tests/pos-custom-args/captures/list-encoding.scala
diff --git a/tests/pos-special/captures/try.scala b/tests/pos-custom-args/captures/try.scala
similarity index 100%
rename from tests/pos-special/captures/try.scala
rename to tests/pos-custom-args/captures/try.scala
diff --git a/tests/pos-special/captures/try3.scala b/tests/pos-custom-args/captures/try3.scala
similarity index 100%
rename from tests/pos-special/captures/try3.scala
rename to tests/pos-custom-args/captures/try3.scala

From acc93bb1639deacd32b8b0f46ced82c30d51dc23 Mon Sep 17 00:00:00 2001
From: Martin Odersky <odersky@gmail.com>
Date: Sat, 19 Jun 2021 13:00:58 +0200
Subject: [PATCH 02/14] Fix printing of poly function types

---
 .../tools/dotc/printing/RefinedPrinter.scala  | 37 ++++++++++++++-----
 tests/neg/polymorphic-functions1.check        |  7 ++++
 tests/neg/polymorphic-functions1.scala        |  1 +
 3 files changed, 35 insertions(+), 10 deletions(-)
 create mode 100644 tests/neg/polymorphic-functions1.check
 create mode 100644 tests/neg/polymorphic-functions1.scala

diff --git a/compiler/src/dotty/tools/dotc/printing/RefinedPrinter.scala b/compiler/src/dotty/tools/dotc/printing/RefinedPrinter.scala
index 54de6a7e13fc..e17798e0ad5d 100644
--- a/compiler/src/dotty/tools/dotc/printing/RefinedPrinter.scala
+++ b/compiler/src/dotty/tools/dotc/printing/RefinedPrinter.scala
@@ -158,14 +158,25 @@ class RefinedPrinter(_ctx: Context) extends PlainPrinter(_ctx) {
         argStr ~ " " ~ arrow(isGiven) ~ " " ~ argText(args.last)
       }
 
-    def toTextDependentFunction(appType: MethodType): Text =
-      "("
-      ~ keywordText("erased ").provided(appType.isErasedMethod)
-      ~ paramsText(appType)
-      ~ ") "
-      ~ arrow(appType.isImplicitMethod)
-      ~ " "
-      ~ toText(appType.resultType)
+    def toTextMethodAsFunction(info: Type): Text = info match
+      case info: MethodType =>
+        "("
+        ~ keywordText("erased ").provided(info.isErasedMethod)
+        ~ ( if info.isParamDependent || info.isResultDependent
+            then paramsText(info)
+            else argsText(info.paramInfos)
+          )
+        ~ ") "
+        ~ arrow(info.isImplicitMethod)
+        ~ " "
+        ~ toTextMethodAsFunction(info.resultType)
+      case info: PolyType =>
+        "["
+        ~ paramsText(info)
+        ~ "] => "
+        ~ toTextMethodAsFunction(info.resultType)
+      case _ =>
+        toText(info)
 
     def isInfixType(tp: Type): Boolean = tp match
       case AppliedType(tycon, args) =>
@@ -229,8 +240,10 @@ class RefinedPrinter(_ctx: Context) extends PlainPrinter(_ctx) {
       if !printDebug && appliedText(tp.asInstanceOf[HKLambda].resType).isEmpty =>
         // don't eta contract if the application would be printed specially
         toText(tycon)
-      case tp: RefinedType if defn.isFunctionType(tp) && !printDebug =>
-        toTextDependentFunction(tp.refinedInfo.asInstanceOf[MethodType])
+      case tp: RefinedType
+      if (defn.isFunctionType(tp) || (tp.parent.typeSymbol eq defn.PolyFunctionClass))
+          && !printDebug =>
+        toTextMethodAsFunction(tp.refinedInfo)
       case tp: TypeRef =>
         if (tp.symbol.isAnonymousClass && !showUniqueIds)
           toText(tp.info)
@@ -244,6 +257,10 @@ class RefinedPrinter(_ctx: Context) extends PlainPrinter(_ctx) {
       case ErasedValueType(tycon, underlying) =>
         "ErasedValueType(" ~ toText(tycon) ~ ", " ~ toText(underlying) ~ ")"
       case tp: ClassInfo =>
+        if tp.cls.derivesFrom(defn.PolyFunctionClass) then
+          tp.member(nme.apply).info match
+            case info: PolyType => return toTextMethodAsFunction(info)
+            case _ =>
         toTextParents(tp.parents) ~~ "{...}"
       case JavaArrayType(elemtp) =>
         toText(elemtp) ~ "[]"
diff --git a/tests/neg/polymorphic-functions1.check b/tests/neg/polymorphic-functions1.check
new file mode 100644
index 000000000000..b9459340fac7
--- /dev/null
+++ b/tests/neg/polymorphic-functions1.check
@@ -0,0 +1,7 @@
+-- [E007] Type Mismatch Error: tests/neg/polymorphic-functions1.scala:1:53 ---------------------------------------------
+1 |val f: [T] => (x: T) => x.type = [T] => (x: Int) => x // error
+  |                                                     ^
+  |                                                     Found:    [T] => (Int) => Int
+  |                                                     Required: [T] => (x: T) => x.type
+
+longer explanation available when compiling with `-explain`
diff --git a/tests/neg/polymorphic-functions1.scala b/tests/neg/polymorphic-functions1.scala
new file mode 100644
index 000000000000..de887f3b8c50
--- /dev/null
+++ b/tests/neg/polymorphic-functions1.scala
@@ -0,0 +1 @@
+val f: [T] => (x: T) => x.type = [T] => (x: Int) => x // error

From 87c92d145a15d06ff6d7a1a437dc28ac8920d50a Mon Sep 17 00:00:00 2001
From: Martin Odersky <odersky@gmail.com>
Date: Sat, 19 Jun 2021 14:56:48 +0200
Subject: [PATCH 03/14] Allow capturing types as CFT protos

---
 compiler/src/dotty/tools/dotc/core/Definitions.scala | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/compiler/src/dotty/tools/dotc/core/Definitions.scala b/compiler/src/dotty/tools/dotc/core/Definitions.scala
index 25dc0602ae76..397392dbec84 100644
--- a/compiler/src/dotty/tools/dotc/core/Definitions.scala
+++ b/compiler/src/dotty/tools/dotc/core/Definitions.scala
@@ -1565,7 +1565,7 @@ class Definitions {
    *   - the upper bound of a TypeParamRef in the current constraint
    */
   def asContextFunctionType(tp: Type)(using Context): Type =
-    tp.stripTypeVar.dealias match
+    tp.stripped.dealias match
       case tp1: TypeParamRef if ctx.typerState.constraint.contains(tp1) =>
         asContextFunctionType(TypeComparer.bounds(tp1).hiBound)
       case tp1 =>

From 651b79183fef333538c94506ecfdba0ca9cffc75 Mon Sep 17 00:00:00 2001
From: Martin Odersky <odersky@gmail.com>
Date: Sat, 19 Jun 2021 14:58:26 +0200
Subject: [PATCH 04/14] Don't print trees twice after refiner phases

---
 compiler/src/dotty/tools/dotc/Run.scala                 | 4 ++--
 compiler/src/dotty/tools/dotc/typer/CheckCaptures.scala | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/compiler/src/dotty/tools/dotc/Run.scala b/compiler/src/dotty/tools/dotc/Run.scala
index 0afea7988958..f4d0c4973e18 100644
--- a/compiler/src/dotty/tools/dotc/Run.scala
+++ b/compiler/src/dotty/tools/dotc/Run.scala
@@ -9,7 +9,7 @@ import Types._
 import Scopes._
 import Names.Name
 import Denotations.Denotation
-import typer.Typer
+import typer.{Typer, RefineTypes}
 import typer.ImportInfo._
 import Decorators._
 import io.{AbstractFile, PlainFile, VirtualFile}
@@ -204,7 +204,7 @@ class Run(comp: Compiler, ictx: Context) extends ImplicitRunInfo with Constraint
             val profileBefore = profiler.beforePhase(phase)
             units = phase.runOn(units)
             profiler.afterPhase(phase, profileBefore)
-            if (ctx.settings.Xprint.value.containsPhase(phase))
+            if ctx.settings.Xprint.value.containsPhase(phase) && !phase.isInstanceOf[RefineTypes] then
               for (unit <- units)
                 lastPrintedTree =
                   printTree(lastPrintedTree)(using ctx.fresh.setPhase(phase.next).setCompilationUnit(unit))
diff --git a/compiler/src/dotty/tools/dotc/typer/CheckCaptures.scala b/compiler/src/dotty/tools/dotc/typer/CheckCaptures.scala
index dce20d5a1d54..e38d590c1e0c 100644
--- a/compiler/src/dotty/tools/dotc/typer/CheckCaptures.scala
+++ b/compiler/src/dotty/tools/dotc/typer/CheckCaptures.scala
@@ -11,7 +11,7 @@ import StdNames._
 import Decorators._
 import ProtoTypes._
 import Inferencing.isFullyDefined
-import config.Printers.refinr
+import config.Printers.capt
 import ast.{tpd, untpd, Trees}
 import NameKinds.{DocArtifactName, OuterSelectName, DefaultGetterName}
 import Trees._
@@ -78,7 +78,7 @@ class CheckCaptures extends RefineTypes:
     override def typedClosure(tree: untpd.Closure, pt: Type)(using Context): Tree =
       super.typedClosure(tree, pt) match
         case tree1: Closure =>
-          refinr.println(i"typing closure ${tree1.meth.symbol} with fvs ${capturedVars(tree1.meth.symbol)}")
+          capt.println(i"typing closure ${tree1.meth.symbol} with fvs ${capturedVars(tree1.meth.symbol)}")
           tree1.withType(tree1.tpe.capturing(capturedVars(tree1.meth.symbol)))
         case tree1 => tree1
 

From f5d25dae25807c46bbfad3460c1f80bba60ce6a1 Mon Sep 17 00:00:00 2001
From: Martin Odersky <odersky@gmail.com>
Date: Sat, 19 Jun 2021 15:00:30 +0200
Subject: [PATCH 05/14] Shallow capture sets

---
 .../dotty/tools/dotc/core/CaptureSet.scala    | 70 ++++++++-----------
 1 file changed, 28 insertions(+), 42 deletions(-)

diff --git a/compiler/src/dotty/tools/dotc/core/CaptureSet.scala b/compiler/src/dotty/tools/dotc/core/CaptureSet.scala
index cf1d1aae21fa..8210863ce9a0 100644
--- a/compiler/src/dotty/tools/dotc/core/CaptureSet.scala
+++ b/compiler/src/dotty/tools/dotc/core/CaptureSet.scala
@@ -35,7 +35,16 @@ case class CaptureSet private (elems: CaptureSet.Refs) extends Showable:
     myClosure
 
   def ++ (that: CaptureSet): CaptureSet =
-    CaptureSet(elems ++ that.elems)
+    if this.isEmpty then that
+    else if that.isEmpty then this
+    else CaptureSet(elems ++ that.elems)
+
+  def + (ref: CaptureRef) =
+    if elems.contains(ref) then this
+    else CaptureSet(elems + ref)
+
+  def intersect (that: CaptureSet): CaptureSet =
+    CaptureSet(this.elems.intersect(that.elems))
 
   /** {x} <:< this   where <:< is subcapturing */
   def accountsFor(x: CaptureRef)(using Context) =
@@ -82,46 +91,23 @@ object CaptureSet:
     css.foldLeft(empty)(_ ++ _)
 
   def ofType(tp: Type)(using Context): CaptureSet =
-    val collect = new TypeAccumulator[Refs]:
-      var localBinders: SimpleIdentitySet[BindingType] = SimpleIdentitySet.empty
-      var seenLazyRefs: SimpleIdentitySet[LazyRef] = SimpleIdentitySet.empty
-      def apply(elems: Refs, tp: Type): Refs = trace(i"capt $elems, $tp", capt, show = true) {
-        tp match
-          case tp: NamedType =>
-            if variance < 0 then elems
-            else elems ++ tp.captureSet.elems
-          case tp: ParamRef =>
-            if variance < 0 || localBinders.contains(tp.binder) then elems
-            else elems ++ tp.captureSet.elems
-          case tp: LambdaType =>
-            localBinders += tp
-            try apply(elems, tp.resultType)
-            finally localBinders -= tp
-          case AndType(tp1, tp2) =>
-            val elems1 = apply(SimpleIdentitySet.empty, tp1)
-            val elems2 = apply(SimpleIdentitySet.empty, tp2)
-            elems ++ elems1.intersect(elems2)
-          case CapturingType(parent, ref) =>
-            val elems1 = apply(elems, parent)
-            if variance >= 0 then elems1 + ref else elems1
-          case TypeBounds(_, hi) =>
-            apply(elems, hi)
-          case tp: ClassInfo =>
-            elems ++ ofClass(tp, Nil).elems
-          case tp: LazyRef =>
-            if seenLazyRefs.contains(tp)
-              || tp.evaluating  // shapeless gets an assertion error without this test
-            then elems
-            else
-              seenLazyRefs += tp
-              foldOver(elems, tp)
-//          case tp: MatchType =>
-//            val normed = tp.tryNormalize
-//            if normed.exists then apply(elems, normed) else foldOver(elems, tp)
-          case _ =>
-            foldOver(elems, tp)
-      }
-
-    CaptureSet(collect(empty.elems, tp))
+    def recur(tp: Type): CaptureSet = tp match
+      case tp: NamedType =>
+        tp.captureSet
+      case tp: ParamRef =>
+        tp.captureSet
+      case CapturingType(parent, ref) =>
+        recur(parent) + ref
+      case tp: TypeProxy =>
+        recur(tp.underlying)
+      case AndType(tp1, tp2) =>
+        recur(tp1).intersect(recur(tp2))
+      case OrType(tp1, tp2) =>
+        recur(tp1) ++ recur(tp2)
+      case tp: ClassInfo =>
+        ofClass(tp, Nil)
+      case _ =>
+        empty
+    recur(tp)
       .showing(i"capture set of $tp = $result", capt)
 

From 9b7fe8b46e43a933c19b735746d95623fd4692df Mon Sep 17 00:00:00 2001
From: Martin Odersky <odersky@gmail.com>
Date: Sat, 19 Jun 2021 16:26:09 +0200
Subject: [PATCH 06/14] Fix condition in TreeChecker

The previous condition did not work for qualifiers with the
```
PolyFunction { ... } retains ...
```
---
 compiler/src/dotty/tools/dotc/transform/TreeChecker.scala | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/compiler/src/dotty/tools/dotc/transform/TreeChecker.scala b/compiler/src/dotty/tools/dotc/transform/TreeChecker.scala
index 8972a1e12ddd..b932ea865ac5 100644
--- a/compiler/src/dotty/tools/dotc/transform/TreeChecker.scala
+++ b/compiler/src/dotty/tools/dotc/transform/TreeChecker.scala
@@ -375,14 +375,14 @@ class TreeChecker extends Phase with SymTransformer {
       val tpe = tree.typeOpt
 
       // Polymorphic apply methods stay structural until Erasure
-      val isPolyFunctionApply = (tree.name eq nme.apply) && (tree.qualifier.typeOpt <:< defn.PolyFunctionType)
+      val isPolyFunctionApply = (tree.name eq nme.apply) && tree.qualifier.typeOpt.derivesFrom(defn.PolyFunctionClass)
       // Outer selects are pickled specially so don't require a symbol
       val isOuterSelect = tree.name.is(OuterSelectName)
       val isPrimitiveArrayOp = ctx.erasedTypes && nme.isPrimitiveName(tree.name)
       if !(tree.isType || isPolyFunctionApply || isOuterSelect || isPrimitiveArrayOp) then
         val denot = tree.denot
         assert(denot.exists, i"Selection $tree with type $tpe does not have a denotation")
-        assert(denot.symbol.exists, i"Denotation $denot of selection $tree with type $tpe does not have a symbol")
+        assert(denot.symbol.exists, i"Denotation $denot of selection $tree with type $tpe does not have a symbol, ${tree.qualifier.typeOpt}")
 
       val sym = tree.symbol
       val symIsFixed = tpe match {

From 3cf68e3bdfad1d81467816e3890f0b40542fac0a Mon Sep 17 00:00:00 2001
From: Martin Odersky <odersky@gmail.com>
Date: Sat, 19 Jun 2021 16:35:38 +0200
Subject: [PATCH 07/14] Fix precedence of dependent and poly function printing

---
 .../tools/dotc/printing/RefinedPrinter.scala  | 32 +++++++++++--------
 1 file changed, 18 insertions(+), 14 deletions(-)

diff --git a/compiler/src/dotty/tools/dotc/printing/RefinedPrinter.scala b/compiler/src/dotty/tools/dotc/printing/RefinedPrinter.scala
index e17798e0ad5d..bcc473fb450b 100644
--- a/compiler/src/dotty/tools/dotc/printing/RefinedPrinter.scala
+++ b/compiler/src/dotty/tools/dotc/printing/RefinedPrinter.scala
@@ -160,21 +160,25 @@ class RefinedPrinter(_ctx: Context) extends PlainPrinter(_ctx) {
 
     def toTextMethodAsFunction(info: Type): Text = info match
       case info: MethodType =>
-        "("
-        ~ keywordText("erased ").provided(info.isErasedMethod)
-        ~ ( if info.isParamDependent || info.isResultDependent
-            then paramsText(info)
-            else argsText(info.paramInfos)
-          )
-        ~ ") "
-        ~ arrow(info.isImplicitMethod)
-        ~ " "
-        ~ toTextMethodAsFunction(info.resultType)
+        changePrec(GlobalPrec) {
+          "("
+          ~ keywordText("erased ").provided(info.isErasedMethod)
+          ~ ( if info.isParamDependent || info.isResultDependent
+              then paramsText(info)
+              else argsText(info.paramInfos)
+            )
+          ~ ") "
+          ~ arrow(info.isImplicitMethod)
+          ~ " "
+          ~ toTextMethodAsFunction(info.resultType)
+        }
       case info: PolyType =>
-        "["
-        ~ paramsText(info)
-        ~ "] => "
-        ~ toTextMethodAsFunction(info.resultType)
+        changePrec(GlobalPrec) {
+          "["
+          ~ paramsText(info)
+          ~ "] => "
+          ~ toTextMethodAsFunction(info.resultType)
+        }
       case _ =>
         toText(info)
 

From 25b23c1d3b1f8f1bc212e22ea7344dde573ba74a Mon Sep 17 00:00:00 2001
From: Martin Odersky <odersky@gmail.com>
Date: Sun, 20 Jun 2021 10:04:47 +0200
Subject: [PATCH 08/14] Fix capture sets of applied types

A substitution was missing for type parameters that appear in the capture set
of the constructor.
---
 .../src/dotty/tools/dotc/core/CaptureSet.scala     | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/compiler/src/dotty/tools/dotc/core/CaptureSet.scala b/compiler/src/dotty/tools/dotc/core/CaptureSet.scala
index 8210863ce9a0..d28d271fc3b9 100644
--- a/compiler/src/dotty/tools/dotc/core/CaptureSet.scala
+++ b/compiler/src/dotty/tools/dotc/core/CaptureSet.scala
@@ -54,6 +54,15 @@ case class CaptureSet private (elems: CaptureSet.Refs) extends Showable:
   def <:< (that: CaptureSet)(using Context): Boolean =
     elems.isEmpty || elems.forall(that.accountsFor)
 
+  def flatMap(f: CaptureRef => CaptureSet)(using Context): CaptureSet =
+    (empty /: elems)((cs, ref) => cs ++ f(ref))
+
+  def substParams(tl: BindingType, to: List[Type])(using Context) =
+    flatMap {
+      case ref: ParamRef if ref.binder eq tl => to(ref.paramNum).captureSet
+      case ref => ref.singletonCaptureSet
+    }
+
   override def toString = elems.toString
 
   override def toText(printer: Printer): Text =
@@ -98,6 +107,11 @@ object CaptureSet:
         tp.captureSet
       case CapturingType(parent, ref) =>
         recur(parent) + ref
+      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) =>

From 27c9fc53fbff333798f1914ab8e3390b5d6aeff6 Mon Sep 17 00:00:00 2001
From: Martin Odersky <odersky@gmail.com>
Date: Sun, 20 Jun 2021 10:06:35 +0200
Subject: [PATCH 09/14] Fix decomposeProto

Strip off capturing types before decomposing a function prototype
---
 compiler/src/dotty/tools/dotc/typer/Typer.scala | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/compiler/src/dotty/tools/dotc/typer/Typer.scala b/compiler/src/dotty/tools/dotc/typer/Typer.scala
index 111f5f42a525..0fc519ace9f0 100644
--- a/compiler/src/dotty/tools/dotc/typer/Typer.scala
+++ b/compiler/src/dotty/tools/dotc/typer/Typer.scala
@@ -1127,7 +1127,7 @@ class Typer extends Namer
         case _ => mapOver(t)
     }
 
-    val pt1 = pt.stripTypeVar.dealias
+    val pt1 = pt.stripped.dealias
     if (pt1 ne pt1.dropDependentRefinement)
        && defn.isContextFunctionType(pt1.nonPrivateMember(nme.apply).info.finalResultType)
     then

From 988c4c054b120bfbb3c3a50dc2686699dda9c702 Mon Sep 17 00:00:00 2001
From: Martin Odersky <odersky@gmail.com>
Date: Sun, 20 Jun 2021 10:07:27 +0200
Subject: [PATCH 10/14] Adapt tests to shallow capture sets

---
 tests/neg-custom-args/captures/boxmap.scala   |  2 +-
 tests/neg-custom-args/captures/capt1.check    | 42 +++++++++++++++++++
 tests/neg-custom-args/captures/capt1.scala    |  3 +-
 tests/neg-custom-args/captures/try2.check     |  4 +-
 tests/neg-custom-args/captures/try2.scala     |  2 +-
 tests/pos-custom-args/captures/boxmap.scala   |  4 +-
 tests/pos-custom-args/captures/capt1.scala    |  7 ++--
 .../captures/list-encoding.scala              |  2 +-
 8 files changed, 54 insertions(+), 12 deletions(-)
 create mode 100644 tests/neg-custom-args/captures/capt1.check

diff --git a/tests/neg-custom-args/captures/boxmap.scala b/tests/neg-custom-args/captures/boxmap.scala
index d61f2ca0e357..cb77c6794cb1 100644
--- a/tests/neg-custom-args/captures/boxmap.scala
+++ b/tests/neg-custom-args/captures/boxmap.scala
@@ -3,7 +3,7 @@ class Cap extends Retains[*]
 
 infix type ==> [A, B] = (A => B) retains *
 
-type Box[+T <: Top] = [K <: Top] => (T ==> K) => K
+type Box[+T <: Top] = ([K <: Top] => (T ==> K) => K) retains T
 
 def box[T <: Top](x: T): Box[T] =
   [K <: Top] => (k: T ==> K) => k(x)
diff --git a/tests/neg-custom-args/captures/capt1.check b/tests/neg-custom-args/captures/capt1.check
new file mode 100644
index 000000000000..15ba35a52e91
--- /dev/null
+++ b/tests/neg-custom-args/captures/capt1.check
@@ -0,0 +1,42 @@
+-- [E007] Type Mismatch Error: tests/neg-custom-args/captures/capt1.scala:3:2 ------------------------------------------
+3 |  () => if x == null then y else y  // error
+  |  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  |  Found:    (() => C) retains x
+  |  Required: () => C
+
+longer explanation available when compiling with `-explain`
+-- [E007] Type Mismatch Error: tests/neg-custom-args/captures/capt1.scala:6:2 ------------------------------------------
+6 |  () => if x == null then y else y  // error
+  |  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  |  Found:    (() => C) retains x
+  |  Required: Any
+
+longer explanation available when compiling with `-explain`
+-- [E007] Type Mismatch Error: tests/neg-custom-args/captures/capt1.scala:14:2 -----------------------------------------
+14 |  f  // error
+   |  ^
+   |  Found:    (Int => Int) retains x
+   |  Required: Any
+
+longer explanation available when compiling with `-explain`
+-- [E007] Type Mismatch Error: tests/neg-custom-args/captures/capt1.scala:23:3 -----------------------------------------
+23 |  F(22)  // error
+   |  ^^^^^
+   |  Found:    F retains x
+   |  Required: A
+
+longer explanation available when compiling with `-explain`
+-- [E007] Type Mismatch Error: tests/neg-custom-args/captures/capt1.scala:27:40 ----------------------------------------
+27 |    def m() = if x == null then y else y  // error
+   |                                        ^
+   |                                        Found:    A {...} retains x
+   |                                        Required: A
+
+longer explanation available when compiling with `-explain`
+-- [E007] Type Mismatch Error: tests/neg-custom-args/captures/capt1.scala:32:24 ----------------------------------------
+32 |  val z2 = h[() => Cap](() => x)(() => C())  // error
+   |                        ^^^^^^^
+   |                        Found:    (() => Cap) retains x
+   |                        Required: () => Cap
+
+longer explanation available when compiling with `-explain`
diff --git a/tests/neg-custom-args/captures/capt1.scala b/tests/neg-custom-args/captures/capt1.scala
index a7f2b0925dca..6f9c02d1540f 100644
--- a/tests/neg-custom-args/captures/capt1.scala
+++ b/tests/neg-custom-args/captures/capt1.scala
@@ -30,5 +30,6 @@ def foo() =
   val x: C retains * = ???
   def h[X <:Top](a: X)(b: X) = a
   val z2 = h[() => Cap](() => x)(() => C())  // error
-  val z3 = h(() => x)(() => C())  // error
+  val z3 = h[(() => Cap) retains x.type](() => x)(() => C())  // ok
+  val z4 = h[(() => Cap) retains x.type](() => x)(() => C())  // what was inferred for z3
 
diff --git a/tests/neg-custom-args/captures/try2.check b/tests/neg-custom-args/captures/try2.check
index 55af935635e7..a73ee901406d 100644
--- a/tests/neg-custom-args/captures/try2.check
+++ b/tests/neg-custom-args/captures/try2.check
@@ -29,10 +29,10 @@ longer explanation available when compiling with `-explain`
 -- Error: tests/neg-custom-args/captures/try2.scala:24:28 --------------------------------------------------------------
 24 |  val a = handle[Exception, CanThrow[Exception]] { // error
    |                            ^^^^^^^^^^^^^^^^^^^
-   |                            type argument is not allowed to capture the global capability (canThrow : List[*])
+   |                            type argument is not allowed to capture the global capability (canThrow : *)
 -- Error: tests/neg-custom-args/captures/try2.scala:36:11 --------------------------------------------------------------
 36 |  val xx = handle { // error
    |           ^^^^^^
-   |inferred type argument ((() => Int) retains canThrow) is not allowed to capture the global capability (canThrow : List[*])
+   |inferred type argument ((() => Int) retains canThrow) is not allowed to capture the global capability (canThrow : *)
    |
    |The inferred arguments are: [Exception, ((() => Int) retains canThrow)]
diff --git a/tests/neg-custom-args/captures/try2.scala b/tests/neg-custom-args/captures/try2.scala
index 3b7bee8224a9..469d9cf8d2f2 100644
--- a/tests/neg-custom-args/captures/try2.scala
+++ b/tests/neg-custom-args/captures/try2.scala
@@ -1,7 +1,7 @@
 import language.experimental.erasedDefinitions
 import annotation.ability
 
-@ability erased val canThrow: List[*] = ???
+@ability erased val canThrow: * = ???
 
 class CanThrow[E <: Exception] extends Retains[canThrow.type]
 type Top  = Any retains *
diff --git a/tests/pos-custom-args/captures/boxmap.scala b/tests/pos-custom-args/captures/boxmap.scala
index adda37f6c03a..50a84e5c6ae5 100644
--- a/tests/pos-custom-args/captures/boxmap.scala
+++ b/tests/pos-custom-args/captures/boxmap.scala
@@ -3,7 +3,7 @@ class Cap extends Retains[*]
 
 infix type ==> [A, B] = (A => B) retains *
 
-type Box[+T <: Top] = [K <: Top] => (T ==> K) => K
+type Box[+T <: Top] = ([K <: Top] => (T ==> K) => K) retains T
 
 def box[T <: Top](x: T): Box[T] =
   [K <: Top] => (k: T ==> K) => k(x)
@@ -17,5 +17,5 @@ def lazymap[A <: Top, B <: Top](b: Box[A])(f: A ==> B): (() => Box[B]) retains b
 def test[A <: Top, B <: Top] =
   def lazymap[A <: Top, B <: Top](b: Box[A])(f: A ==> B) =
     () => b[Box[B]]((x: A) => box(f(x)))
-  val x: (b: Box[A]) => (f: A ==> B) => (() => Box[B]) retains b.type | f.type = lazymap[A, B]
+  val x: (b: Box[A]) => ((f: A ==> B) => (() => Box[B]) retains b.type | f.type) retains b.type = lazymap[A, B]
   ()
diff --git a/tests/pos-custom-args/captures/capt1.scala b/tests/pos-custom-args/captures/capt1.scala
index cada7aa41de5..b311e1bf9d8f 100644
--- a/tests/pos-custom-args/captures/capt1.scala
+++ b/tests/pos-custom-args/captures/capt1.scala
@@ -1,7 +1,7 @@
 class C
 type Cap = C retains *
 type Top = Any retains *
-def f1(c: Cap): () => C retains c.type = () => c // ok
+def f1(c: Cap): (() => c.type) retains c.type = () => c // ok
 
 def f2: Int =
   val g: (Boolean => Int) retains * = ???
@@ -18,11 +18,10 @@ def foo() =
   val x: C retains * = ???
   val y: C retains x.type = x
   val x2: (() => C) retains x.type = ???
-  val y2: () => C retains x.type = x2
+  val y2: (() => C retains x.type) retains x.type = x2
 
-  val z1: () => Cap = f1(x)
+  val z1: (() => Cap) retains * = f1(x)
   def h[X <:Top](a: X)(b: X) = a
 
   val z2 =
     if x == null then () => x else () => C()
-
diff --git a/tests/pos-custom-args/captures/list-encoding.scala b/tests/pos-custom-args/captures/list-encoding.scala
index eda95898de65..91cf6d92c08f 100644
--- a/tests/pos-custom-args/captures/list-encoding.scala
+++ b/tests/pos-custom-args/captures/list-encoding.scala
@@ -5,7 +5,7 @@ type Op[T <: Top, C <: Top] =
   ((v: T) => ((s: C) => C) retains *) retains *
 
 type List[T <: Top] =
-  [C <: Top] => (op: Op[T, C]) => ((s: C) => C) retains op.type
+  ([C <: Top] => (op: Op[T, C]) => ((s: C) => C) retains op.type) retains T
 
 def nil[T <: Top]: List[T] =
   [C <: Top] => (op: Op[T, C]) => (s: C) => s

From 6977ad147a577af7bf3e3b873408d010f25df16f Mon Sep 17 00:00:00 2001
From: Martin Odersky <odersky@gmail.com>
Date: Mon, 21 Jun 2021 18:07:12 +0200
Subject: [PATCH 11/14] Handle TypeVars in CaptureSets

---
 .../dotty/tools/dotc/core/CaptureSet.scala    | 21 +--------------
 .../src/dotty/tools/dotc/core/Types.scala     | 27 ++++++++++++++++---
 .../tools/dotc/typer/CheckCaptures.scala      | 10 ++++---
 tests/neg-custom-args/captures/try3.scala     |  2 +-
 4 files changed, 32 insertions(+), 28 deletions(-)

diff --git a/compiler/src/dotty/tools/dotc/core/CaptureSet.scala b/compiler/src/dotty/tools/dotc/core/CaptureSet.scala
index d28d271fc3b9..6f7d03b92733 100644
--- a/compiler/src/dotty/tools/dotc/core/CaptureSet.scala
+++ b/compiler/src/dotty/tools/dotc/core/CaptureSet.scala
@@ -17,23 +17,6 @@ case class CaptureSet private (elems: CaptureSet.Refs) extends Showable:
   def isEmpty: Boolean = elems.isEmpty
   def nonEmpty: Boolean = !isEmpty
 
-  private var myClosure: Refs | Null = null
-
-  def closure(using Context): Refs =
-    if myClosure == null then
-      var cl = elems
-      var seen: Refs = SimpleIdentitySet.empty
-      while
-        val prev = cl
-        for ref <- cl do
-          if !seen.contains(ref) then
-            seen += ref
-            cl = cl ++ ref.captureSetOfInfo.elems
-        prev ne cl
-      do ()
-      myClosure = cl
-    myClosure
-
   def ++ (that: CaptureSet): CaptureSet =
     if this.isEmpty then that
     else if that.isEmpty then this
@@ -101,9 +84,7 @@ object CaptureSet:
 
   def ofType(tp: Type)(using Context): CaptureSet =
     def recur(tp: Type): CaptureSet = tp match
-      case tp: NamedType =>
-        tp.captureSet
-      case tp: ParamRef =>
+      case tp: CaptureRef =>
         tp.captureSet
       case CapturingType(parent, ref) =>
         recur(parent) + ref
diff --git a/compiler/src/dotty/tools/dotc/core/Types.scala b/compiler/src/dotty/tools/dotc/core/Types.scala
index b8dd2434a071..5620e76ed0b0 100644
--- a/compiler/src/dotty/tools/dotc/core/Types.scala
+++ b/compiler/src/dotty/tools/dotc/core/Types.scala
@@ -3620,7 +3620,7 @@ object Types {
           case tp: TermParamRef if tp.binder eq thisLambdaType => TrueDeps
           case tp: CapturingType =>
             val status1 = compute(status, tp.parent, theAcc)
-            tp.ref match
+            tp.ref.stripTypeVar match
               case tp: TermParamRef if tp.binder eq thisLambdaType => combine(status1, CaptureDeps)
               case _ => status1
           case _: ThisType | _: BoundType | NoPrefix => status
@@ -4505,9 +4505,10 @@ object Types {
    *  @param  origin        The parameter that's tracked by the type variable.
    *  @param  creatorState  The typer state in which the variable was created.
    */
-  final class TypeVar private(initOrigin: TypeParamRef, creatorState: TyperState, nestingLevel: Int) extends CachedProxyType with ValueType {
+  final class TypeVar private(initOrigin: TypeParamRef, creatorState: TyperState, nestingLevel: Int)
+  extends CachedProxyType, CaptureRef {
 
-    private var currentOrigin = initOrigin
+   private var currentOrigin = initOrigin
 
     def origin: TypeParamRef = currentOrigin
 
@@ -4689,6 +4690,26 @@ object Types {
       if (inst.exists) inst else origin
     }
 
+    // Capture ref methods
+
+    def canBeTracked(using Context): Boolean = underlying match
+      case ref: CaptureRef => ref.canBeTracked
+      case _ => false
+
+    override def normalizedRef(using Context): CaptureRef = instanceOpt match
+      case ref: CaptureRef => ref
+      case _ => this
+
+    override def singletonCaptureSet(using Context) = instanceOpt match
+      case ref: CaptureRef => ref.singletonCaptureSet
+      case _ => super.singletonCaptureSet
+
+    override def captureSetOfInfo(using Context): CaptureSet = instanceOpt match
+      case ref: CaptureRef => ref.captureSetOfInfo
+      case tp => tp.captureSet
+
+    // Object members
+
     override def computeHash(bs: Binders): Int = identityHash(bs)
     override def equals(that: Any): Boolean = this.eq(that.asInstanceOf[AnyRef])
 
diff --git a/compiler/src/dotty/tools/dotc/typer/CheckCaptures.scala b/compiler/src/dotty/tools/dotc/typer/CheckCaptures.scala
index e38d590c1e0c..98d9d1d4d979 100644
--- a/compiler/src/dotty/tools/dotc/typer/CheckCaptures.scala
+++ b/compiler/src/dotty/tools/dotc/typer/CheckCaptures.scala
@@ -99,10 +99,12 @@ class CheckCaptures extends RefineTypes:
   def checkWellFormed(whole: Type, pos: SrcPos)(using Context): Unit =
     def checkRelativeVariance(mt: MethodType) = new TypeTraverser:
       def traverse(tp: Type): Unit = tp match
-        case CapturingType(parent, ref @ TermParamRef(`mt`, _)) =>
-          if variance <= 0 then
-            val direction = if variance < 0 then "contra" else "in"
-            report.error(em"captured reference $ref appears ${direction}variantly in type $whole", pos)
+        case CapturingType(parent, ref) =>
+          ref.stripTypeVar match
+            case ref @ TermParamRef(`mt`, _) if variance <= 0 =>
+              val direction = if variance < 0 then "contra" else "in"
+              report.error(em"captured reference $ref appears ${direction}variantly in type $whole", pos)
+            case _ =>
           traverse(parent)
         case _ =>
           traverseChildren(tp)
diff --git a/tests/neg-custom-args/captures/try3.scala b/tests/neg-custom-args/captures/try3.scala
index f90ad6aff0aa..ece7870ffecc 100644
--- a/tests/neg-custom-args/captures/try3.scala
+++ b/tests/neg-custom-args/captures/try3.scala
@@ -3,7 +3,7 @@ import java.io.IOException
 class CanThrow[E] extends Retains[*]
 type Top  = Any retains *
 
-def handle[E <: Exception, T <: Top](op: CanThrow[E] ?=> T)(handler: E => T): T =
+def handle[E <: Exception, T <: Top](op: (CanThrow[E] ?=> T) retains T)(handler: (E => T) retains T): T =
   val x: CanThrow[E] = ???
   try op(using x)
   catch case ex: E => handler(ex)

From cceff77ca13f17f6177c3d361d7ea9fca66f0cfb Mon Sep 17 00:00:00 2001
From: Martin Odersky <odersky@gmail.com>
Date: Mon, 21 Jun 2021 18:22:58 +0200
Subject: [PATCH 12/14] Assume capture sets of result type in enclosing
 function

---
 .../tools/dotc/config/ScalaSettings.scala     |   1 +
 .../tools/dotc/typer/CheckCaptures.scala      | 101 ++++++++++++++++++
 .../dotty/tools/dotc/typer/TypeAssigner.scala |  10 +-
 .../src/dotty/tools/dotc/typer/Typer.scala    |  16 +--
 tests/neg-custom-args/captures/boxmap.check   |   2 +-
 tests/neg-custom-args/captures/capt1.check    |   9 +-
 tests/neg-custom-args/captures/capt1.scala    |   2 +-
 .../captures/try3-abbrev.scala                |  26 +++++
 .../captures/capt1-abbrev.scala               |  27 +++++
 9 files changed, 174 insertions(+), 20 deletions(-)
 create mode 100644 tests/neg-custom-args/captures/try3-abbrev.scala
 create mode 100644 tests/pos-custom-args/captures/capt1-abbrev.scala

diff --git a/compiler/src/dotty/tools/dotc/config/ScalaSettings.scala b/compiler/src/dotty/tools/dotc/config/ScalaSettings.scala
index 1392baa80a75..094e685ddc75 100644
--- a/compiler/src/dotty/tools/dotc/config/ScalaSettings.scala
+++ b/compiler/src/dotty/tools/dotc/config/ScalaSettings.scala
@@ -203,6 +203,7 @@ trait AllScalaSettings extends CommonScalaSettings { self: Settings.SettingGroup
   val YrequireTargetName: Setting[Boolean] = BooleanSetting("-Yrequire-targetName", "Warn if an operator is defined without a @targetName annotation")
   val YrefineTypes: Setting[Boolean] = BooleanSetting("-Yrefine-types", "Run experimental type refiner (test only)")
   val Ycc: Setting[Boolean] = BooleanSetting("-Ycc", "Check captured references")
+  val YccNoAbbrev: Setting[Boolean] = BooleanSetting("-Ycc-no-abbrev", "Used in conjunction with -Ycc, suppress type abbreviations")
 
   /** Area-specific debug output */
   val YexplainLowlevel: Setting[Boolean] = BooleanSetting("-Yexplain-lowlevel", "When explaining type errors, show types at a lower level.")
diff --git a/compiler/src/dotty/tools/dotc/typer/CheckCaptures.scala b/compiler/src/dotty/tools/dotc/typer/CheckCaptures.scala
index 98d9d1d4d979..5a33ab726a9f 100644
--- a/compiler/src/dotty/tools/dotc/typer/CheckCaptures.scala
+++ b/compiler/src/dotty/tools/dotc/typer/CheckCaptures.scala
@@ -152,4 +152,105 @@ class CheckCaptures extends RefineTypes:
   def postRefinerCheck(tree: tpd.Tree)(using Context): Unit =
     PostRefinerCheck.traverse(tree)
 
+
+object CheckCaptures:
+  import ast.tpd.*
+
+  def expandFunctionTypes(using Context) =
+    ctx.settings.Ycc.value && !ctx.settings.YccNoAbbrev.value && !ctx.isAfterTyper
+
+  object FunctionTypeTree:
+    def unapply(tree: Tree)(using Context): Option[(List[Type], Type)] =
+      if defn.isFunctionType(tree.tpe) then
+        tree match
+          case AppliedTypeTree(tycon: TypeTree, args) =>
+            Some((args.init.tpes, args.last.tpe))
+          case RefinedTypeTree(_, (appDef: DefDef) :: Nil) if appDef.span == tree.span =>
+            appDef.symbol.info match
+              case mt: MethodType => Some((mt.paramInfos, mt.resultType))
+              case _ => None
+          case _ =>
+            None
+      else None
+
+  object CapturingTypeTree:
+    def unapply(tree: Tree)(using Context): Option[(Tree, Tree, CaptureRef)] = tree match
+      case AppliedTypeTree(tycon, parent :: _ :: Nil)
+      if tycon.symbol == defn.Predef_retainsType =>
+        tree.tpe match
+          case CapturingType(_, ref) => Some((tycon, parent, ref))
+          case _ => None
+      case _ => None
+
+  def addRetains(tree: Tree, ref: CaptureRef)(using Context): Tree =
+    untpd.AppliedTypeTree(
+        TypeTree(defn.Predef_retainsType.typeRef), List(tree, TypeTree(ref)))
+      .withType(CapturingType(tree.tpe, ref))
+      .showing(i"add inferred capturing $result", capt)
+
+  /**  Under -Ycc but not -Ycc-no-abbrev, if `tree` represents a function type
+   *   `(ARGS) => T` where T is tracked and all ARGS are pure, expand it to
+   *   `(ARGS) => T retains CS` where CS is the capture set of `T`. These synthesized
+   *   additions will be removed again if the function type is wrapped in an
+   *   explicit `retains` type.
+   */
+  def addResultCaptures(tree: Tree)(using Context): Tree =
+    if expandFunctionTypes then
+      tree match
+        case FunctionTypeTree(argTypes, resType) =>
+          val cs = resType.captureSet
+          if cs.nonEmpty && argTypes.forall(_.captureSet.isEmpty)
+          then (tree /: cs.elems)(addRetains)
+          else tree
+        case _ =>
+          tree
+    else tree
+
+  private def addCaptures(tp: Type, refs: Type)(using Context): Type = refs match
+    case ref: CaptureRef => CapturingType(tp, ref)
+    case OrType(refs1, refs2) => addCaptures(addCaptures(tp, refs1), refs2)
+    case _ => tp
+
+  /** @pre: `tree is a tree of the form `T retains REFS`.
+   *  Return the same tree with `parent1` instead of `T` with its type
+   *  recomputed accordingly.
+   */
+  private def derivedCapturingTree(tree: AppliedTypeTree, parent1: Tree)(using Context): AppliedTypeTree =
+    tree match
+      case AppliedTypeTree(tycon, parent :: (rest @ (refs :: Nil))) if parent ne parent1 =>
+        cpy.AppliedTypeTree(tree)(tycon, parent1 :: rest)
+          .withType(addCaptures(parent1.tpe, refs.tpe))
+      case _ =>
+        tree
+
+  private def stripCaptures(tree: Tree, ref: CaptureRef)(using Context): Tree = tree match
+    case tree @ AppliedTypeTree(tycon, parent :: refs :: Nil) if tycon.symbol == defn.Predef_retainsType =>
+      val parent1 = stripCaptures(parent, ref)
+      val isSynthetic = tycon.isInstanceOf[TypeTree]
+      if isSynthetic then
+        parent1.showing(i"drop inferred capturing $tree => $result", capt)
+      else
+        if parent1.tpe.captureSet.accountsFor(ref) then
+          report.warning(
+            em"redundant capture: $parent1 already contains $ref with capture set ${ref.captureSet} in its capture set ${parent1.tpe.captureSet}",
+            tree.srcPos)
+        derivedCapturingTree(tree, parent1)
+    case _ => tree
+
+  private def stripCaptures(tree: Tree, refs: Type)(using Context): Tree = refs match
+    case ref: CaptureRef => stripCaptures(tree, ref)
+    case OrType(refs1, refs2) => stripCaptures(stripCaptures(tree, refs1), refs2)
+    case _ => tree
+
+  /** If this is a tree of the form `T retains REFS`,
+   *   - strip any synthesized captures directly in T;
+   *   - warn if a reference in REFS is accounted for by the capture set of the remaining type
+   */
+  def refineNestedCaptures(tree: AppliedTypeTree)(using Context): AppliedTypeTree = tree match
+    case AppliedTypeTree(tycon, parent :: (rest @ (refs :: Nil))) if tycon.symbol == defn.Predef_retainsType =>
+      derivedCapturingTree(tree, stripCaptures(parent, refs.tpe))
+    case _ =>
+      tree
+
 end CheckCaptures
+
diff --git a/compiler/src/dotty/tools/dotc/typer/TypeAssigner.scala b/compiler/src/dotty/tools/dotc/typer/TypeAssigner.scala
index b083361fe6b2..32a1200c5c4d 100644
--- a/compiler/src/dotty/tools/dotc/typer/TypeAssigner.scala
+++ b/compiler/src/dotty/tools/dotc/typer/TypeAssigner.scala
@@ -12,6 +12,7 @@ import config.Printers.typr
 import ast.Trees._
 import NameOps._
 import ProtoTypes._
+import CheckCaptures.refineNestedCaptures
 import collection.mutable
 import reporting._
 import Checking.{checkNoPrivateLeaks, checkNoWildcard}
@@ -189,8 +190,6 @@ trait TypeAssigner {
     def captType(tp: Type, refs: Type): Type = refs match
       case ref: NamedType =>
         if ref.isTracked then
-          if tp.captureSet.accountsFor(ref) then
-            report.warning(em"redundant capture: $tp already contains $ref with capture set ${ref.captureSet} in its capture set ${tp.captureSet}", tree.srcPos)
           CapturingType(tp, ref)
         else
           val reason =
@@ -479,7 +478,7 @@ trait TypeAssigner {
     tree.withType(RecType.closeOver(rt => refined.substThis(refineCls, rt.recThis)))
   }
 
-  def assignType(tree: untpd.AppliedTypeTree, tycon: Tree, args: List[Tree])(using Context): AppliedTypeTree = {
+  def assignType(tree: untpd.AppliedTypeTree, tycon: Tree, args: List[Tree])(using Context): AppliedTypeTree =
     assert(!hasNamedArg(args) || ctx.reporter.errorsReported, tree)
     val tparams = tycon.tpe.typeParams
     val ownType =
@@ -487,8 +486,9 @@ trait TypeAssigner {
         wrongNumberOfTypeArgs(tycon.tpe, tparams, args, tree.srcPos)
       else
         processAppliedType(tree, tycon.tpe.appliedTo(args.tpes))
-    tree.withType(ownType)
-  }
+    val tree1 = tree.withType(ownType)
+    if ctx.settings.Ycc.value then refineNestedCaptures(tree1)
+    else tree1
 
   def assignType(tree: untpd.LambdaTypeTree, tparamDefs: List[TypeDef], body: Tree)(using Context): LambdaTypeTree =
     tree.withType(HKTypeLambda.fromParams(tparamDefs.map(_.symbol.asType), body.tpe))
diff --git a/compiler/src/dotty/tools/dotc/typer/Typer.scala b/compiler/src/dotty/tools/dotc/typer/Typer.scala
index 0fc519ace9f0..f020ea8f802c 100644
--- a/compiler/src/dotty/tools/dotc/typer/Typer.scala
+++ b/compiler/src/dotty/tools/dotc/typer/Typer.scala
@@ -28,6 +28,7 @@ import Checking._
 import Inferencing._
 import Dynamic.isDynamicExpansion
 import EtaExpansion.etaExpand
+import CheckCaptures.addResultCaptures
 import TypeComparer.CompareResult
 import util.Spans._
 import util.common._
@@ -1249,13 +1250,14 @@ class Typer extends Namer
       RefinedTypeTree(core, List(appDef), ctx.owner.asClass)
     end typedDependent
 
-    args match {
-      case ValDef(_, _, _) :: _ =>
-        typedDependent(args.asInstanceOf[List[untpd.ValDef]])(
-          using ctx.fresh.setOwner(newRefinedClassSymbol(tree.span)).setNewScope)
-      case _ =>
-        propagateErased(
-          typed(cpy.AppliedTypeTree(tree)(untpd.TypeTree(funCls.typeRef), args :+ body), pt))
+    addResultCaptures {
+      args match
+        case ValDef(_, _, _) :: _ =>
+          typedDependent(args.asInstanceOf[List[untpd.ValDef]])(
+            using ctx.fresh.setOwner(newRefinedClassSymbol(tree.span)).setNewScope)
+        case _ =>
+          propagateErased(
+            typed(cpy.AppliedTypeTree(tree)(untpd.TypeTree(funCls.typeRef), args :+ body), pt))
     }
   }
 
diff --git a/tests/neg-custom-args/captures/boxmap.check b/tests/neg-custom-args/captures/boxmap.check
index bbf1879464de..fa9189457180 100644
--- a/tests/neg-custom-args/captures/boxmap.check
+++ b/tests/neg-custom-args/captures/boxmap.check
@@ -2,7 +2,7 @@
 15 |  () => b[Box[B]]((x: A) => box(f(x)))  // error
    |  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    |  Found:    (() => Box[B]) retains b retains f
-   |  Required: () => Box[B]
+   |  Required: (() => Box[B]) retains B
    |
    |  where:    B is a type in method lazymap with bounds <: Top
 
diff --git a/tests/neg-custom-args/captures/capt1.check b/tests/neg-custom-args/captures/capt1.check
index 15ba35a52e91..6a52a1cd7481 100644
--- a/tests/neg-custom-args/captures/capt1.check
+++ b/tests/neg-custom-args/captures/capt1.check
@@ -33,10 +33,7 @@ longer explanation available when compiling with `-explain`
    |                                        Required: A
 
 longer explanation available when compiling with `-explain`
--- [E007] Type Mismatch Error: tests/neg-custom-args/captures/capt1.scala:32:24 ----------------------------------------
+-- Error: tests/neg-custom-args/captures/capt1.scala:32:13 -------------------------------------------------------------
 32 |  val z2 = h[() => Cap](() => x)(() => C())  // error
-   |                        ^^^^^^^
-   |                        Found:    (() => Cap) retains x
-   |                        Required: () => Cap
-
-longer explanation available when compiling with `-explain`
+   |             ^^^^^^^^^
+   |             type argument is not allowed to capture the universal capability *
diff --git a/tests/neg-custom-args/captures/capt1.scala b/tests/neg-custom-args/captures/capt1.scala
index 6f9c02d1540f..9f68a09cbab8 100644
--- a/tests/neg-custom-args/captures/capt1.scala
+++ b/tests/neg-custom-args/captures/capt1.scala
@@ -30,6 +30,6 @@ def foo() =
   val x: C retains * = ???
   def h[X <:Top](a: X)(b: X) = a
   val z2 = h[() => Cap](() => x)(() => C())  // error
-  val z3 = h[(() => Cap) retains x.type](() => x)(() => C())  // ok
+  val z3 = h(() => x)(() => C())  // ok
   val z4 = h[(() => Cap) retains x.type](() => x)(() => C())  // what was inferred for z3
 
diff --git a/tests/neg-custom-args/captures/try3-abbrev.scala b/tests/neg-custom-args/captures/try3-abbrev.scala
new file mode 100644
index 000000000000..f90ad6aff0aa
--- /dev/null
+++ b/tests/neg-custom-args/captures/try3-abbrev.scala
@@ -0,0 +1,26 @@
+import java.io.IOException
+
+class CanThrow[E] extends Retains[*]
+type Top  = Any retains *
+
+def handle[E <: Exception, T <: Top](op: CanThrow[E] ?=> T)(handler: E => 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
+
+@main def Test: Int =
+  def f(a: Boolean) =
+    handle { // error
+      if !a then raise(IOException())
+      (b: Boolean) =>
+        if !b then raise(IOException())
+        0
+    } {
+      ex => (b: Boolean) => -1
+    }
+  val g = f(true)
+  g(false) // would raise an uncaught exception
+  f(true)(false) // would raise an uncaught exception
diff --git a/tests/pos-custom-args/captures/capt1-abbrev.scala b/tests/pos-custom-args/captures/capt1-abbrev.scala
new file mode 100644
index 000000000000..c2b72bb21b70
--- /dev/null
+++ b/tests/pos-custom-args/captures/capt1-abbrev.scala
@@ -0,0 +1,27 @@
+class C
+type Cap = C retains *
+type Top = Any retains *
+def f1(c: Cap): () => c.type = () => c // ok
+
+def f2: Int =
+  val g: (Boolean => Int) retains * = ???
+  val x = g(true)
+  x
+
+def f3: Int =
+  def g: (Boolean => Int) retains * = ???
+  def h = g
+  val x = g.apply(true)
+  x
+
+def foo() =
+  val x: C retains * = ???
+  val y: C retains x.type = x
+  val x2: (() => C) retains x.type = ???
+  val y2: () => C retains x.type = x2
+
+  val z1: (() => Cap) retains * = f1(x)
+  def h[X <:Top](a: X)(b: X) = a
+
+  val z2 =
+    if x == null then () => x else () => C()

From 7a282ad41657614289ba5b682212f678fa327edf Mon Sep 17 00:00:00 2001
From: Martin Odersky <odersky@gmail.com>
Date: Mon, 21 Jun 2021 18:24:46 +0200
Subject: [PATCH 13/14] Drop unused val in appliedTo

---
 compiler/src/dotty/tools/dotc/core/TypeApplications.scala | 1 -
 1 file changed, 1 deletion(-)

diff --git a/compiler/src/dotty/tools/dotc/core/TypeApplications.scala b/compiler/src/dotty/tools/dotc/core/TypeApplications.scala
index aa380b574b98..968fffe82809 100644
--- a/compiler/src/dotty/tools/dotc/core/TypeApplications.scala
+++ b/compiler/src/dotty/tools/dotc/core/TypeApplications.scala
@@ -310,7 +310,6 @@ class TypeApplications(val self: Type) extends AnyVal {
    */
   final def appliedTo(args: List[Type])(using Context): Type = {
     record("appliedTo")
-    val typParams = self.typeParams
     val stripped = self.stripTypeVar
     val dealiased = stripped.safeDealias
     if (args.isEmpty || ctx.erasedTypes) self

From 59151bdb6c2237fce7d0e6b07b6abfcd799be5bc Mon Sep 17 00:00:00 2001
From: Martin Odersky <odersky@gmail.com>
Date: Mon, 21 Jun 2021 18:25:18 +0200
Subject: [PATCH 14/14] Print `retains` type trees infix

---
 .../src/dotty/tools/dotc/printing/RefinedPrinter.scala   | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/compiler/src/dotty/tools/dotc/printing/RefinedPrinter.scala b/compiler/src/dotty/tools/dotc/printing/RefinedPrinter.scala
index bcc473fb450b..37fd56931a87 100644
--- a/compiler/src/dotty/tools/dotc/printing/RefinedPrinter.scala
+++ b/compiler/src/dotty/tools/dotc/printing/RefinedPrinter.scala
@@ -527,13 +527,16 @@ class RefinedPrinter(_ctx: Context) extends PlainPrinter(_ctx) {
       case RefinedTypeTree(tpt, refines) =>
         toTextLocal(tpt) ~ " " ~ blockText(refines)
       case AppliedTypeTree(tpt, args) =>
-        if (tpt.symbol == defn.orType && args.length == 2)
+        if tpt.symbol == defn.orType && args.length == 2 then
           changePrec(OrTypePrec) { toText(args(0)) ~ " | " ~ atPrec(OrTypePrec + 1) { toText(args(1)) } }
-        else if (tpt.symbol == defn.andType && args.length == 2)
+        else if tpt.symbol == defn.andType && args.length == 2 then
           changePrec(AndTypePrec) { toText(args(0)) ~ " & " ~ atPrec(AndTypePrec + 1) { toText(args(1)) } }
+        else if tpt.symbol == defn.Predef_retainsType && args.length == 2 then
+          changePrec(InfixPrec) { toText(args(0)) ~ " retains " ~ toText(args(1)) }
         else if defn.isFunctionClass(tpt.symbol)
             && tpt.isInstanceOf[TypeTree] && tree.hasType && !printDebug
-        then changePrec(GlobalPrec) { toText(tree.typeOpt) }
+        then
+          changePrec(GlobalPrec) { toText(tree.typeOpt) }
         else args match
           case arg :: _ if arg.isTerm =>
             toTextLocal(tpt) ~ "(" ~ Text(args.map(argText), ", ") ~ ")"