Use a separate phase to add try owners

I tried several other strategies, but nothing worked.

 - update after Setup: This means we cannot to levelOwner during setup.
 - running levelOwner after Setup causes CyclicErrors when transforming symbols
   in Setup
 - having a seperate notion of ccOwner maintained in a CCState table does
   not work since there are too many hidden uses of owner everywhere that
   would ignore that table.

Author: odersky

compiler/src/dotty/tools/dotc/Compiler.scala

@@ -82,6 +82,7 @@ class Compiler {
          new PatternMatcher) ::      // Compile pattern matches
     List(new TestRecheck.Pre) ::     // Test only: run rechecker, enabled under -Yrecheck-test
     List(new TestRecheck) ::         // Test only: run rechecker, enabled under -Yrecheck-test
+    List(new cc.AddTryOwners) ::     // Add symbols as owners of try blocks, enabled under captureChecking
     List(new cc.Setup) ::            // Preparations for check captures phase, enabled under captureChecking
     List(new cc.CheckCaptures) ::    // Check captures, enabled under captureChecking
     List(new ElimOpaque,             // Turn opaque into normal aliases

compiler/src/dotty/tools/dotc/ast/tpd.scala

@@ -1284,6 +1284,21 @@ object tpd extends Trees.Instance[Type] with TypedTreeInfo {
       !(sym.is(Method) && sym.info.isInstanceOf[MethodOrPoly]) // if is a method it is parameterless
   }
 
+  /** A tree traverser that generates the the same import contexts as original typer for statements.
+   *  TODO: Should we align TreeMapWithPreciseStatContexts and also keep track of exprOwners?
+   */
+  abstract class TreeTraverserWithPreciseImportContexts extends TreeTraverser:
+    override def apply(x: Unit, trees: List[Tree])(using Context): Unit =
+      def recur(trees: List[Tree]): Unit = trees match
+        case (imp: Import) :: rest =>
+          traverse(rest)(using ctx.importContext(imp, imp.symbol))
+        case tree :: rest =>
+          traverse(tree)
+          traverse(rest)
+        case Nil =>
+      recur(trees)
+  end TreeTraverserWithPreciseImportContexts
+
   extension (xs: List[tpd.Tree])
     def tpes: List[Type] = xs match {
       case x :: xs1 => x.tpe :: xs1.tpes

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

@@ -0,0 +1,45 @@
+package dotty.tools.dotc
+package cc
+
+import core.*
+import DenotTransformers.IdentityDenotTransformer
+import Phases.Phase
+import Contexts.*, Symbols.*, Flags.*, Types.*
+import config.Feature
+import ast.tpd
+import StdNames.nme
+import Decorators.i
+
+object AddTryOwners:
+  val name: String = "addTryOwners"
+  val description: String = "add symbols for try blocks in preparation of capture checking"
+
+class AddTryOwners extends Phase, IdentityDenotTransformer:
+  thisPhase =>
+
+  import tpd.*
+
+  override def phaseName: String = AddTryOwners.name
+  override def description: String = AddTryOwners.description
+
+  override def isRunnable(using Context) =
+    super.isRunnable && Feature.ccEnabledSomewhere
+
+  override def isCheckable = false
+
+  def run(using Context): Unit =
+    val addTryOwners = new TreeTraverserWithPreciseImportContexts:
+      def traverse(tree: Tree)(using Context): Unit = tree match
+        case tree @ Try(expr, cases, finalizer) if Feature.enabled(Feature.saferExceptions) =>
+          val tryOwner = newSymbol(ctx.owner, nme.TRY_BLOCK, SyntheticMethod, MethodType(Nil, defn.UnitType))
+          ccState.tryBlockOwner(tree) = tryOwner
+          expr.changeOwnerAfter(ctx.owner, tryOwner, thisPhase)
+          inContext(ctx.withOwner(tryOwner)):
+            traverse(expr)
+          traverse(cases)
+          traverse(finalizer)
+        case _ =>
+          traverseChildren(tree)
+    addTryOwners.traverse(ctx.compilationUnit.tpdTree)
+end AddTryOwners
+

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

@@ -380,18 +380,18 @@ extension (sym: Symbol)
    *   - _root_
    */
   def levelOwner(using Context): Symbol =
-    def recur(sym: Symbol)(using Context): Symbol =
+    def recur(sym: Symbol): Symbol =
       if !sym.exists || sym.isRoot || sym.isStaticOwner then defn.RootClass
       else if sym.isLevelOwner then sym
       else recur(sym.owner)
-    recur(sym)(using ctx.withPhase(Phases.checkCapturesPhase))
+    recur(sym)
 
   /** The level owner enclosing `sym` which has the given name, or NoSymbol if none exists.
    *  If name refers to a val that has a closure as rhs, we return the closure as level
    *  owner.
    */
   def levelOwnerNamed(name: String)(using Context): Symbol =
-    def recur(sym: Symbol, prev: Symbol)(using Context): Symbol =
+    def recur(sym: Symbol, prev: Symbol): Symbol =
       if sym.name.toString == name then
         if sym.isLevelOwner then sym
         else if sym.isTerm && !sym.isOneOf(Method | Module) && prev.exists then prev
@@ -401,7 +401,7 @@ extension (sym: Symbol)
       else
         val prev1 = if sym.isAnonymousFunction && sym.isLevelOwner then sym else NoSymbol
         recur(sym.owner, prev1)
-    recur(sym, NoSymbol)(using ctx.withPhase(Phases.checkCapturesPhase))
+    recur(sym, NoSymbol)
       .showing(i"find outer $sym [ $name ] = $result", capt)
 
   /** The nesting level of `sym` for the purposes of `cc`,

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

@@ -182,6 +182,7 @@ class CheckCaptures extends Recheck, SymTransformer:
   def phaseName: String = "cc"
 
   override def isRunnable(using Context) = super.isRunnable && Feature.ccEnabledSomewhere
+  override def firstPrepPhase = preRecheckPhase.prev.asInstanceOf[AddTryOwners]
 
   def newRechecker()(using Context) = CaptureChecker(ctx)
 

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

@@ -402,31 +402,17 @@ class Setup extends PreRecheck, SymTransformer, SetupAPI:
       def inverse = thisMap
   end SubstParams
 
-  /** If the outer context directly enclosing the definition of `sym`
-   *  has a <try block> owner, that owner, otherwise null.
-   */
-  def newOwnerFor(sym: Symbol)(using Context): Symbol =
-    var octx = ctx
-    while octx.owner == sym do octx = octx.outer
-    if octx.owner.name == nme.TRY_BLOCK then octx.owner else sym.maybeOwner
-
   /** Update info of `sym` for CheckCaptures phase only */
   private def updateInfo(sym: Symbol, info: Type)(using Context) =
-    sym.updateInfo(thisPhase, info, newFlagsFor(sym), newOwnerFor(sym))
+    sym.updateInfo(thisPhase, info, newFlagsFor(sym))
     sym.namedType match
       case ref: CaptureRef => ref.invalidateCaches() // TODO: needed?
       case _ =>
 
-  /** Update the owner of `sym` to `newOwnerFor(sym)`.
-   *  This can happen because of inserted try block owners.
-   */
-  private def updateOwner(sym: Symbol)(using Context) =
-    updateInfo(sym, sym.info)
-
   extension (sym: Symbol) def nextInfo(using Context): Type =
     atPhase(thisPhase.next)(sym.info)
 
-  def setupTraverser(recheckDef: DefRecheck) = new TreeTraverser:
+  def setupTraverser(recheckDef: DefRecheck) = new TreeTraverserWithPreciseImportContexts:
 
     def traverse(tree: Tree)(using Context): Unit =
       tree match
@@ -435,7 +421,6 @@ class Setup extends PreRecheck, SymTransformer, SetupAPI:
           if isExcluded(meth) then
             return
 
-          updateOwner(meth)
           inContext(ctx.withOwner(meth)):
             paramss.foreach(traverse)
             transformTT(tpt, boxed = false,
@@ -456,12 +441,8 @@ class Setup extends PreRecheck, SymTransformer, SetupAPI:
             case _ => false
 
           val sym = tree.symbol
-          val newOwner =
-            if sym.isOneOf(TermParamOrAccessor) then ctx.owner
-            else
-              updateOwner(sym)
-              sym
-          inContext(ctx.withOwner(newOwner)):
+          val defCtx = if sym.isOneOf(TermParamOrAccessor) then ctx else ctx.withOwner(sym)
+          inContext(defCtx):
             val rootsNeedMap =
               // need to run before cc so that rhsClosure is defined without transforming
               // symbols to cc, since rhsClosure is used in that transformation.
@@ -493,32 +474,14 @@ class Setup extends PreRecheck, SymTransformer, SetupAPI:
             transformTT(arg, boxed = true, exact = false, rootTarget = ctx.owner) // type arguments in type applications are boxed
 
         case tree: TypeDef if tree.symbol.isClass =>
-          val cls = tree.symbol
-          updateOwner(cls)
-          if cls.is(ModuleClass) then updateOwner(cls.sourceModule)
-          inContext(ctx.withOwner(cls)):
-            traverseChildren(tree)
-
-        case tree: Try if Feature.enabled(Feature.saferExceptions) =>
-          val tryOwner = newSymbol(ctx.owner, nme.TRY_BLOCK, SyntheticMethod, MethodType(Nil, defn.UnitType))
-          ccState.isLevelOwner(tryOwner) = true
-          ccState.tryBlockOwner(tree) = tryOwner
-          inContext(ctx.withOwner(tryOwner)):
+          inContext(ctx.withOwner(tree.symbol)):
             traverseChildren(tree)
 
         case _ =>
           traverseChildren(tree)
       postProcess(tree)
     end traverse
 
-    override def apply(x: Unit, trees: List[Tree])(using Context): Unit = trees match
-      case (imp: Import) :: rest =>
-        traverse(rest)(using ctx.importContext(imp, imp.symbol))
-      case tree :: rest =>
-        traverse(tree)
-        traverse(rest)
-      case Nil =>
-
     extension (sym: Symbol)(using Context) def unlessStatic =
       val owner = sym.levelOwner
       if sym.isStaticOwner then NoSymbol else owner
@@ -644,9 +607,6 @@ class Setup extends PreRecheck, SymTransformer, SetupAPI:
                 val modul = cls.sourceModule
                 updateInfo(modul, CapturingType(modul.info, newSelfType.captureSet))
                 modul.termRef.invalidateCaches()
-            else
-              updateOwner(cls)
-              if cls.is(ModuleClass) then updateOwner(cls.sourceModule)
           case _ =>
             val info = tree.symbol.info
             val newInfo = transformExplicitType(info, rootTarget = ctx.owner.unlessStatic)

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

@@ -138,18 +138,27 @@ abstract class Recheck extends Phase, SymTransformer:
   import ast.tpd.*
   import Recheck.*
 
+  /** The phase before rechecking, used to setup symbol infos. */
   def preRecheckPhase = this.prev.asInstanceOf[PreRecheck]
 
+  /** The first phase that pepares for rechecking. This is usually preRecheckPhase
+   *  but could also be before. Updated symbols will snap back to their
+   *  denotations at firestPrepPhase after rechecking.
+   */
+  def firstPrepPhase: Phase = preRecheckPhase
+
   override def changesBaseTypes: Boolean = true
 
   override def isCheckable = false
     // TODO: investigate what goes wrong we Ycheck directly after rechecking.
     // One failing test is pos/i583a.scala
 
-  /** Change any `ResetPrivate` flags back to `Private` */
+  /** Change denotation back to what it was before (pre-)rechecking` */
   def transformSym(symd: SymDenotation)(using Context): SymDenotation =
     val sym = symd.symbol
-    if sym.isUpdatedAfter(preRecheckPhase) then atPhase(preRecheckPhase)(sym.denot)
+    def updatedAfter(p: Phase): Boolean =
+      sym.isUpdatedAfter(p) || p != preRecheckPhase && updatedAfter(p.next)
+    if updatedAfter(firstPrepPhase) then atPhase(firstPrepPhase)(sym.denot)
     else symd
 
   def run(using Context): Unit =

tests/neg-custom-args/captures/real-try.check

@@ -24,13 +24,11 @@
    |    cannot be included in outer capture set ?, defined at level 1 in method test
    |
    | longer explanation available when compiling with `-explain`
--- [E007] Type Mismatch Error: tests/neg-custom-args/captures/real-try.scala:31:4 --------------------------------------
+-- Error: tests/neg-custom-args/captures/real-try.scala:31:21 ----------------------------------------------------------
 31 |    Cell(() => foo(1))// // error
-   |    ^^^^^^^^^^^^^^^^^^
-   |    Found:    Cell[box () ->{cap[<try block>]} Unit]^?
-   |    Required: Cell[box () ->? Unit]^?
+   |                     ^
+   |(canThrow$4 : CanThrow[Ex1 | Ex2]^{cap[<try block>]}) cannot be referenced here; it is not included in the allowed capture set ?
+   |of an enclosing function literal with expected type box () ->? Unit
    |
-   |    Note that reference (cap[<try block>] : caps.Cap), defined at level 2
-   |    cannot be included in outer capture set ?, defined at level 1 in method test
-   |
-   | longer explanation available when compiling with `-explain`
+   |Note that reference (canThrow$4 : CanThrow[Ex1 | Ex2]^{cap[<try block>]}), defined at level 2
+   |cannot be included in outer capture set ?, defined at level 1 in method test