Special rule for {this} in capture sets of class members

compiler/src/dotty/tools/dotc/cc/CaptureSet.scala

@@ -49,7 +49,14 @@ sealed abstract class CaptureSet extends Showable:
   /** Is this capture set definitely non-empty? */
   final def isNotEmpty: Boolean = !elems.isEmpty
 
-  /** Cast to variable. @pre: @isConst */
+  /** Cast to Const. @pre: isConst */
+  def asConst: Const = this match
+    case c: Const => c
+    case v: Var =>
+      assert(v.isConst)
+      Const(v.elems)
+
+  /** Cast to variable. @pre: !isConst */
   def asVar: Var =
     assert(!isConst)
     asInstanceOf[Var]

compiler/src/dotty/tools/dotc/transform/Recheck.scala

@@ -196,12 +196,15 @@ abstract class Recheck extends Phase, IdentityDenotTransformer:
         bindType
 
     def recheckValDef(tree: ValDef, sym: Symbol)(using Context): Unit =
-      if !tree.rhs.isEmpty then recheck(tree.rhs, sym.info)
+      if !tree.rhs.isEmpty then recheckRHS(tree.rhs, sym.info, sym)
 
     def recheckDefDef(tree: DefDef, sym: Symbol)(using Context): Unit =
       val rhsCtx = linkConstructorParams(sym).withOwner(sym)
       if !tree.rhs.isEmpty && !sym.isInlineMethod && !sym.isEffectivelyErased then
-        inContext(rhsCtx) { recheck(tree.rhs, recheck(tree.tpt)) }
+        inContext(rhsCtx) { recheckRHS(tree.rhs, recheck(tree.tpt), sym) }
+
+    def recheckRHS(tree: Tree, pt: Type, sym: Symbol)(using Context): Type =
+      recheck(tree, pt)
 
     def recheckTypeDef(tree: TypeDef, sym: Symbol)(using Context): Type =
       recheck(tree.rhs)

compiler/src/dotty/tools/dotc/typer/CheckCaptures.scala

@@ -322,6 +322,22 @@ class CheckCaptures extends Recheck:
         interpolateVarsIn(tree.tpt)
         curEnv = saved
 
+    override def recheckRHS(tree: Tree, pt: Type, sym: Symbol)(using Context): Type =
+      val pt1 = pt match
+        case CapturingType(core, refs, _)
+        if sym.owner.isClass && !sym.owner.isExtensibleClass
+            && refs.elems.contains(sym.owner.thisType) =>
+          val paramCaptures =
+            sym.paramSymss.flatten.foldLeft(CaptureSet.empty) { (cs, p) =>
+              val pcs = p.info.captureSet
+              (cs ++ (if pcs.isConst then pcs else CaptureSet.universal)).asConst
+            }
+          val declaredCaptures = sym.owner.asClass.givenSelfType.captureSet
+          pt.derivedCapturingType(core, refs ++ (declaredCaptures -- paramCaptures))
+        case _ =>
+          pt
+      recheck(tree, pt1)
+
     override def recheckClassDef(tree: TypeDef, impl: Template, cls: ClassSymbol)(using Context): Type =
       for param <- cls.paramGetters do
         if param.is(Private) && !param.info.captureSet.isAlwaysEmpty then

tests/neg-custom-args/captures/lazylists1.check

@@ -0,0 +1,7 @@
+-- [E007] Type Mismatch Error: tests/neg-custom-args/captures/lazylists1.scala:25:63 -----------------------------------
+25 |      def concat(other: {f} LazyList[A]): {this} LazyList[A] = ??? : ({xs, f} LazyList[A]) // error
+   |                                                               ^^^^^^^^^^^^^^^^^^^^^^^^^^^
+   |                                                               Found:    {xs, f} LazyList[A]
+   |                                                               Required: {Mapped.this, xs} LazyList[A]
+
+longer explanation available when compiling with `-explain`

tests/neg-custom-args/captures/lazylists1.scala

@@ -0,0 +1,27 @@
+class CC
+type Cap = {*} CC
+
+trait LazyList[+A]:
+  this: ({*} LazyList[A]) =>
+
+  def isEmpty: Boolean
+  def head: A
+  def tail: {this} LazyList[A]
+
+object LazyNil extends LazyList[Nothing]:
+  def isEmpty: Boolean = true
+  def head = ???
+  def tail = ???
+
+extension [A](xs: {*} LazyList[A])
+  def map[B](f: {*} A => B): {xs, f} LazyList[B] =
+    final class Mapped extends LazyList[B]:
+      this: ({xs, f} Mapped) =>
+
+      def isEmpty = false
+      def head: B = f(xs.head)
+      def tail: {this} LazyList[B] = xs.tail.map(f)  // OK
+      def drop(n: Int): {this} LazyList[B] = ??? : ({xs, f} LazyList[B]) // OK
+      def concat(other: {f} LazyList[A]): {this} LazyList[A] = ??? : ({xs, f} LazyList[A]) // error
+    new Mapped
+

tests/neg-custom-args/captures/lazylists2.check

@@ -0,0 +1,45 @@
+-- [E163] Declaration Error: tests/neg-custom-args/captures/lazylists2.scala:50:10 -------------------------------------
+50 |      def tail: {xs, f} LazyList[B] = xs.tail.map(f) // error
+   |          ^
+   |          error overriding method tail in trait LazyList of type => {Mapped.this} LazyList[B];
+   |            method tail of type => {xs, f} LazyList[B] has incompatible type
+
+longer explanation available when compiling with `-explain`
+-- [E007] Type Mismatch Error: tests/neg-custom-args/captures/lazylists2.scala:18:4 ------------------------------------
+18 |    final class Mapped extends LazyList[B]:  // error
+   |    ^
+   |    Found:    {f, xs} LazyList[B]
+   |    Required: {f} LazyList[B]
+19 |      this: ({xs, f} Mapped) =>
+20 |      def isEmpty = false
+21 |      def head: B = f(xs.head)
+22 |      def tail: {this} LazyList[B] = xs.tail.map(f)
+23 |    new Mapped
+
+longer explanation available when compiling with `-explain`
+-- [E007] Type Mismatch Error: tests/neg-custom-args/captures/lazylists2.scala:27:4 ------------------------------------
+27 |    final class Mapped extends LazyList[B]:  // error
+   |    ^
+   |    Found:    {f, xs} LazyList[B]
+   |    Required: {xs} LazyList[B]
+28 |      this: ({xs, f} Mapped) =>
+29 |      def isEmpty = false
+30 |      def head: B = f(xs.head)
+31 |      def tail: {this} LazyList[B] = xs.tail.map(f)
+32 |    new Mapped
+
+longer explanation available when compiling with `-explain`
+-- [E007] Type Mismatch Error: tests/neg-custom-args/captures/lazylists2.scala:41:48 -----------------------------------
+41 |      def tail: {this} LazyList[B] = xs.tail.map(f) // error
+   |                                     ^^^^^^^^^^^^^^
+   |                                     Found:    {f} LazyList[B]
+   |                                     Required: {xs} LazyList[B]
+
+longer explanation available when compiling with `-explain`
+-- [E007] Type Mismatch Error: tests/neg-custom-args/captures/lazylists2.scala:59:48 -----------------------------------
+59 |      def tail: {this} LazyList[B] = xs.tail.map(f) // error
+   |                                     ^^^^^^^^^^^^^^
+   |                                     Found:    {f} LazyList[B]
+   |                                     Required: {Mapped.this} LazyList[B]
+
+longer explanation available when compiling with `-explain`

tests/neg-custom-args/captures/lazylists2.scala

@@ -0,0 +1,64 @@
+class CC
+type Cap = {*} CC
+
+trait LazyList[+A]:
+  this: ({*} LazyList[A]) =>
+
+  def isEmpty: Boolean
+  def head: A
+  def tail: {this} LazyList[A]
+
+object LazyNil extends LazyList[Nothing]:
+  def isEmpty: Boolean = true
+  def head = ???
+  def tail = ???
+
+extension [A](xs: {*} LazyList[A])
+  def map[B](f: {*} A => B): {f} LazyList[B] =
+    final class Mapped extends LazyList[B]:  // error
+      this: ({xs, f} Mapped) =>
+
+      def isEmpty = false
+      def head: B = f(xs.head)
+      def tail: {this} LazyList[B] = xs.tail.map(f)
+    new Mapped
+
+  def map2[B](f: {*} A => B): {xs} LazyList[B] =
+    final class Mapped extends LazyList[B]:  // error
+      this: ({xs, f} Mapped) =>
+
+      def isEmpty = false
+      def head: B = f(xs.head)
+      def tail: {this} LazyList[B] = xs.tail.map(f)
+    new Mapped
+
+  def map3[B](f: {*} A => B): {xs} LazyList[B] =
+    final class Mapped extends LazyList[B]:
+      this: ({xs} Mapped) =>
+
+      def isEmpty = false
+      def head: B = f(xs.head)
+      def tail: {this} LazyList[B] = xs.tail.map(f) // error
+    new Mapped
+
+  def map4[B](f: {*} A => B): {xs} LazyList[B] =
+    final class Mapped extends LazyList[B]:
+      this: ({xs, f} Mapped) =>
+
+      def isEmpty = false
+      def head: B = f(xs.head)
+      def tail: {xs, f} LazyList[B] = xs.tail.map(f) // error
+    new Mapped
+
+  def map5[B](f: {*} A => B): LazyList[B] =
+    class Mapped extends LazyList[B]:
+      this: ({xs, f} Mapped) =>
+
+      def isEmpty = false
+      def head: B = f(xs.head)
+      def tail: {this} LazyList[B] = xs.tail.map(f) // error
+    class Mapped2 extends Mapped:
+      this: Mapped =>
+    new Mapped2
+
+

tests/pos-custom-args/captures/lazylists-mono.scala

@@ -0,0 +1,27 @@
+class CC
+type Cap = {*} CC
+
+//-------------------------------------------------
+
+def test(E: Cap) =
+
+  trait LazyList[+A]:
+    protected def contents: {E} () => (A, {E} LazyList[A])
+    def isEmpty: Boolean
+    def head: A = contents()._1
+    def tail: {E} LazyList[A] = contents()._2
+
+  class LazyCons[+A](override val contents: {E} () => (A, {E} LazyList[A]))
+  extends LazyList[A]:
+    def isEmpty: Boolean = false
+
+  object LazyNil extends LazyList[Nothing]:
+    def contents: {E} () => (Nothing, LazyList[Nothing]) = ???
+    def isEmpty: Boolean = true
+
+  extension [A](xs: {E} LazyList[A])
+    def map[B](f: {E} A => B): {E} LazyList[B] =
+      if xs.isEmpty then LazyNil
+      else
+        val cons = () => (f(xs.head), xs.tail.map(f))
+        LazyCons(cons)

tests/pos-custom-args/captures/lazylists.scala

@@ -0,0 +1,42 @@
+class CC
+type Cap = {*} CC
+
+trait LazyList[+A]:
+  this: ({*} LazyList[A]) =>
+
+  def isEmpty: Boolean
+  def head: A
+  def tail: {this} LazyList[A]
+
+object LazyNil extends LazyList[Nothing]:
+  def isEmpty: Boolean = true
+  def head = ???
+  def tail = ???
+
+extension [A](xs: {*} LazyList[A])
+  def map[B](f: {*} A => B): {xs, f} LazyList[B] =
+    final class Mapped extends LazyList[B]:
+      this: ({xs, f} Mapped) =>
+
+      def isEmpty = false
+      def head: B = f(xs.head)
+      def tail: {this} LazyList[B] = xs.tail.map(f)  // OK
+      def concat(other: {f} LazyList[A]): {this, f} LazyList[A] = ??? : ({xs, f} LazyList[A]) // OK
+    if xs.isEmpty then LazyNil
+    else new Mapped
+
+def test(cap1: Cap, cap2: Cap) =
+  def f(x: String): String = if cap1 == cap1 then "" else "a"
+  def g(x: String): String = if cap2 == cap2 then "" else "a"
+
+  val xs =
+    class Initial extends LazyList[String]:
+      this: ({cap1} Initial) =>
+
+      def isEmpty = false
+      def head = f("")
+      def tail = LazyNil
+    new Initial
+  val xsc: {cap1} LazyList[String] = xs
+  val ys = xs.map(g)
+  val ysc: {cap1, cap2} LazyList[String] = ys

tests/pos-custom-args/captures/lazylists1.scala

@@ -0,0 +1,35 @@
+class CC
+type Cap = {*} CC
+
+trait LazyList[+A]:
+  this: ({*} LazyList[A]) =>
+
+  def isEmpty: Boolean
+  def head: A
+  def tail: {this} LazyList[A]
+
+object LazyNil extends LazyList[Nothing]:
+  def isEmpty: Boolean = true
+  def head = ???
+  def tail = ???
+
+final class LazyCons[+T](val x: T, val xs: {*} () => {*} LazyList[T]) extends LazyList[T]:
+  this: ({*} LazyList[T]) =>
+
+  def isEmpty = false
+  def head = x
+  def tail: {this} LazyList[T] = xs()
+
+extension [A](xs: {*} LazyList[A])
+  def map[B](f: {*} A => B): {xs, f} LazyList[B] =
+    if xs.isEmpty then LazyNil
+    else LazyCons(f(xs.head), () => xs.tail.map(f))
+
+def test(cap1: Cap, cap2: Cap) =
+  def f(x: String): String = if cap1 == cap1 then "" else "a"
+  def g(x: String): String = if cap2 == cap2 then "" else "a"
+
+  val xs = LazyCons("", () => if f("") == f("") then LazyNil else LazyNil)
+  val xsc: {cap1} LazyList[String] = xs
+  val ys = xs.map(g)
+  val ysc: {cap1, cap2} LazyList[String] = ys