Apply this adaptation also to assignments

Author: odersky

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

@@ -195,7 +195,7 @@ abstract class Recheck extends Phase, IdentityDenotTransformer:
 
     def recheckAssign(tree: Assign)(using Context): Type =
       val lhsType = recheck(tree.lhs)
-      recheck(tree.rhs, lhsType.widen)
+      recheckRHS(tree.rhs, lhsType.widen, tree.lhs.symbol)
       defn.UnitType
 
     def recheckBlock(stats: List[Tree], expr: Tree, pt: Type)(using Context): Type =

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

@@ -213,6 +213,12 @@ class CheckCaptures extends Recheck:
         interpolateVarsIn(tree.tpt)
         curEnv = saved
 
+    /** Capture check right hand side of the definition of `sym`, or of an assignment
+     *  to `sym`. If `sym` is a member of a final class with self type `cs T`, recheck `tree`
+     *  with an expected type that allows as a captureset all references in `cs`
+     *  except references that are covered by some parameter of `sym`. See #13657 for
+     *  a more detailed explanation of why we want to do this, and why it looks sound.
+     */
     override def recheckRHS(tree: Tree, pt: Type, sym: Symbol)(using Context): Type =
       val pt1 = pt match
         case CapturingType(core, refs, _)

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

@@ -0,0 +1,49 @@
+import language.experimental.saferExceptions
+import annotation.unchecked.uncheckedVariance
+
+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] =>
+
+  var forced = false
+  var cache: {this} LazyList[T @uncheckedVariance] = compiletime.uninitialized
+
+  private def force =
+    if !forced then
+      cache = xs()
+      forced = true
+    cache
+
+  def isEmpty = false
+  def head = x
+  def tail: {this} LazyList[T] = force
+
+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))
+
+class Ex1 extends Exception
+class Ex2 extends Exception
+
+def test(using cap1: CanThrow[Ex1], cap2: CanThrow[Ex2]) =
+  val xs = LazyCons(1, () => LazyNil)
+
+  def f(x: Int): Int throws Ex1 =
+    if x < 0 then throw Ex1()
+    x * x
+
+  val res = xs.map(f)
+  res: {cap1} LazyList[Int]
+