From 7beff3a1fd2139bcc8c03c4dc9b30d5216969cd9 Mon Sep 17 00:00:00 2001
From: odersky <odersky@gmail.com>
Date: Sat, 4 Nov 2023 11:29:22 +0100
Subject: [PATCH] Simplify CapturedVars phase

No need for a separate entry in the context's store; we can keep everything in the phase
itself, which is more efficient and modular.
---
 .../tools/dotc/transform/CapturedVars.scala   | 55 +++++++------------
 1 file changed, 20 insertions(+), 35 deletions(-)

diff --git a/compiler/src/dotty/tools/dotc/transform/CapturedVars.scala b/compiler/src/dotty/tools/dotc/transform/CapturedVars.scala
index a018bbd1a3ac..cab61a6cce4c 100644
--- a/compiler/src/dotty/tools/dotc/transform/CapturedVars.scala
+++ b/compiler/src/dotty/tools/dotc/transform/CapturedVars.scala
@@ -13,25 +13,20 @@ import core.NameKinds.TempResultName
 import core.Constants._
 import util.Store
 import dotty.tools.uncheckedNN
-
-import scala.compiletime.uninitialized
+import ast.tpd.*
+import compiletime.uninitialized
 
 /** This phase translates variables that are captured in closures to
  *  heap-allocated refs.
  */
 class CapturedVars extends MiniPhase with IdentityDenotTransformer:
   thisPhase =>
-  import ast.tpd._
 
   override def phaseName: String = CapturedVars.name
 
   override def description: String = CapturedVars.description
 
-  private[this] var Captured: Store.Location[util.ReadOnlySet[Symbol]] = uninitialized
-  private def captured(using Context) = ctx.store(Captured)
-
-  override def initContext(ctx: FreshContext): Unit =
-    Captured = ctx.addLocation(util.ReadOnlySet.empty)
+  private var captured = util.HashSet[Symbol]()
 
   private class RefInfo(using Context) {
     /** The classes for which a Ref type exists. */
@@ -57,33 +52,10 @@ class CapturedVars extends MiniPhase with IdentityDenotTransformer:
     myRefInfo.uncheckedNN
   }
 
-  private class CollectCaptured extends TreeTraverser {
-    private val captured = util.HashSet[Symbol]()
-    def traverse(tree: Tree)(using Context) = tree match {
-      case id: Ident =>
-        val sym = id.symbol
-        if (sym.is(Mutable, butNot = Method) && sym.owner.isTerm) {
-          val enclMeth = ctx.owner.enclosingMethod
-          if (sym.enclosingMethod != enclMeth) {
-            report.log(i"capturing $sym in ${sym.enclosingMethod}, referenced from $enclMeth")
-            captured += sym
-          }
-        }
-      case _ =>
-        traverseChildren(tree)
-    }
-    def runOver(tree: Tree)(using Context): util.ReadOnlySet[Symbol] = {
-      traverse(tree)
-      captured
-    }
-  }
-
-  override def prepareForUnit(tree: Tree)(using Context): Context = {
-    val captured = atPhase(thisPhase) {
-      CollectCaptured().runOver(ctx.compilationUnit.tpdTree)
-    }
-    ctx.fresh.updateStore(Captured, captured)
-  }
+  override def prepareForUnit(tree: Tree)(using Context): Context =
+    captured.clear()
+    atPhase(thisPhase)(CapturedVars.collect(captured)).traverse(tree)
+    ctx
 
   /** The {Volatile|}{Int|Double|...|Object}Ref class corresponding to the class `cls`,
     *  depending on whether the reference should be @volatile
@@ -143,3 +115,16 @@ class CapturedVars extends MiniPhase with IdentityDenotTransformer:
 object CapturedVars:
   val name: String = "capturedVars"
   val description: String = "represent vars captured by closures as heap objects"
+
+  def collect(captured: util.HashSet[Symbol]): TreeTraverser = new:
+    def traverse(tree: Tree)(using Context) = tree match
+      case id: Ident =>
+        val sym = id.symbol
+        if sym.is(Mutable, butNot = Method) && sym.owner.isTerm then
+          val enclMeth = ctx.owner.enclosingMethod
+          if sym.enclosingMethod != enclMeth then
+            report.log(i"capturing $sym in ${sym.enclosingMethod}, referenced from $enclMeth")
+            captured += sym
+      case _ =>
+        traverseChildren(tree)
+end CapturedVars