Extract StagingLevel from TreeMapWithStages

Author: nicolasstucki

compiler/src/dotty/tools/dotc/inlines/PrepareInlineable.scala

@@ -21,7 +21,7 @@ import transform.{AccessProxies, PCPCheckAndHeal, Splicer}
 import transform.SymUtils.*
 import config.Printers.inlining
 import util.Property
-import dotty.tools.dotc.transform.TreeMapWithStages._
+import dotty.tools.dotc.staging.StagingLevel.freshStagingLevelContext
 
 object PrepareInlineable {
   import tpd._
@@ -293,7 +293,7 @@ object PrepareInlineable {
           if (code.symbol.flags.is(Inline))
             report.error("Macro cannot be implemented with an `inline` method", code.srcPos)
           Splicer.checkValidMacroBody(code)
-          new PCPCheckAndHeal(freshStagingContext).transform(body) // Ignore output, only check PCP
+          (new PCPCheckAndHeal).transform(body)(using freshStagingLevelContext) // Ignore output, only check PCP
         case Block(List(stat), Literal(Constants.Constant(()))) => checkMacro(stat)
         case Block(Nil, expr) => checkMacro(expr)
         case Typed(expr, _) => checkMacro(expr)

compiler/src/dotty/tools/dotc/quoted/Interpreter.scala

@@ -25,7 +25,6 @@ import dotty.tools.dotc.core.Symbols._
 import dotty.tools.dotc.core.TypeErasure
 import dotty.tools.dotc.core.Types._
 import dotty.tools.dotc.quoted._
-import dotty.tools.dotc.transform.TreeMapWithStages._
 import dotty.tools.dotc.typer.ImportInfo.withRootImports
 import dotty.tools.dotc.util.SrcPos
 import dotty.tools.dotc.reporting.Message

compiler/src/dotty/tools/dotc/staging/StagingLevel.scala

@@ -0,0 +1,38 @@
+package dotty.tools.dotc.staging
+
+import dotty.tools.dotc.ast.{TreeMapWithImplicits, tpd}
+import dotty.tools.dotc.config.Printers.staging
+import dotty.tools.dotc.core.Decorators._
+import dotty.tools.dotc.core.Contexts._
+import dotty.tools.dotc.core.StagingContext._
+import dotty.tools.dotc.core.Symbols._
+import dotty.tools.dotc.util.Property
+
+import scala.collection.mutable
+
+object StagingLevel {
+
+  /** A key to be used in a context property that caches the `levelOf` mapping */
+  private val LevelOfKey = new Property.Key[mutable.HashMap[Symbol, Int]]
+
+  /** Initial context for a StagingTransformer transformation. */
+  def freshStagingLevelContext(using Context): Context =
+    ctx.fresh.setProperty(LevelOfKey, new mutable.HashMap[Symbol, Int])
+
+  /** The quotation level of the definition of the locally defined symbol */
+  def levelOf(sym: Symbol)(using Context): Int =
+    ctx.property(LevelOfKey).get.getOrElse(sym, 0)
+
+  def removeLevelOf(sym: Symbol)(using Context): Unit =
+    val levelOfMap = ctx.property(LevelOfKey).get
+    levelOfMap -= sym
+
+  /** Enter staging level of symbol defined by `tree` */
+  def markSymbol(sym: Symbol)(using Context): Boolean =
+    val levelOfMap = ctx.property(LevelOfKey).get
+    if level != 0 && !levelOfMap.contains(sym) then
+      levelOfMap(sym) = level
+      true
+    else
+      false
+}

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

@@ -17,11 +17,10 @@ import dotty.tools.dotc.transform.SymUtils._
 import dotty.tools.dotc.typer.Checking
 import dotty.tools.dotc.typer.Implicits.SearchFailureType
 import dotty.tools.dotc.core.Annotations._
+import dotty.tools.dotc.staging.StagingLevel.*
 
 import dotty.tools.dotc.util.Property
 
-import scala.annotation.constructorOnly
-
 /** Checks that the Phase Consistency Principle (PCP) holds and heals types.
  *
  *  Local term references are phase consistent if and only if they are used at the same level as their definition.
@@ -48,7 +47,7 @@ import scala.annotation.constructorOnly
  *     }
  *
  */
-class PCPCheckAndHeal(@constructorOnly ictx: Context) extends TreeMapWithStages(ictx) with Checking {
+class PCPCheckAndHeal extends TreeMapWithStages with Checking {
   import tpd._
 
   private val InAnnotation = Property.Key[Unit]()

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

@@ -19,7 +19,7 @@ import scala.collection.mutable
 import dotty.tools.dotc.core.Annotations._
 import dotty.tools.dotc.core.StdNames._
 import dotty.tools.dotc.quoted._
-import dotty.tools.dotc.transform.TreeMapWithStages._
+import dotty.tools.dotc.staging.StagingLevel.freshStagingLevelContext
 import dotty.tools.dotc.inlines.Inlines
 
 import scala.annotation.constructorOnly
@@ -93,7 +93,7 @@ class PickleQuotes extends MacroTransform {
       case _ =>
 
   override def run(using Context): Unit =
-    if (ctx.compilationUnit.needsStaging) super.run(using freshStagingContext)
+    if (ctx.compilationUnit.needsStaging) super.run(using freshStagingLevelContext)
 
   protected def newTransformer(using Context): Transformer = new Transformer {
     override def transform(tree: tpd.Tree)(using Context): tpd.Tree =

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

@@ -17,7 +17,6 @@ import dotty.tools.dotc.core.Annotations._
 import dotty.tools.dotc.core.Names._
 import dotty.tools.dotc.core.StdNames._
 import dotty.tools.dotc.quoted._
-import dotty.tools.dotc.transform.TreeMapWithStages._
 
 import scala.annotation.constructorOnly
 

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

@@ -21,7 +21,7 @@ import dotty.tools.dotc.core.Annotations._
 import dotty.tools.dotc.core.Names._
 import dotty.tools.dotc.core.StdNames._
 import dotty.tools.dotc.quoted._
-import dotty.tools.dotc.transform.TreeMapWithStages._
+import dotty.tools.dotc.staging.StagingLevel.freshStagingLevelContext
 import dotty.tools.dotc.config.ScalaRelease.*
 
 import scala.annotation.constructorOnly
@@ -77,7 +77,7 @@ class Splicing extends MacroTransform:
 
   override def run(using Context): Unit =
     if ctx.compilationUnit.needsStaging then
-      super.run(using freshStagingContext)
+      super.run(using freshStagingLevelContext)
 
   protected def newTransformer(using Context): Transformer = Level0QuoteTransformer
 
@@ -246,7 +246,7 @@ class Splicing extends MacroTransform:
         if tree.symbol == defn.QuotedTypeModule_of && containsCapturedType(tpt.tpe) =>
           val newContent = capturedPartTypes(tpt)
           newContent match
-            case block: Block => 
+            case block: Block =>
               inContext(ctx.withSource(tree.source)) {
                 Apply(TypeApply(typeof, List(newContent)), List(quotes)).withSpan(tree.span)
               }
@@ -342,7 +342,7 @@ class Splicing extends MacroTransform:
       val bindingSym = refBindingMap.getOrElseUpdate(tree.symbol, (tree, newBinding))._2
       ref(bindingSym)
 
-    private def newQuotedTypeClassBinding(tpe: Type)(using Context) = 
+    private def newQuotedTypeClassBinding(tpe: Type)(using Context) =
       newSymbol(
         spliceOwner,
         UniqueName.fresh(nme.Type).toTermName,
@@ -376,7 +376,7 @@ class Splicing extends MacroTransform:
       tpt match
         case block: Block =>
           cpy.Block(block)(newHealedTypes ::: block.stats, TypeTree(captured))
-        case _ => 
+        case _ =>
           if newHealedTypes.nonEmpty then
             cpy.Block(tpt)(newHealedTypes, TypeTree(captured))
           else

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

@@ -11,7 +11,7 @@ import dotty.tools.dotc.core.Symbols._
 import dotty.tools.dotc.core.Types._
 import dotty.tools.dotc.util.SrcPos
 import dotty.tools.dotc.transform.SymUtils._
-import dotty.tools.dotc.transform.TreeMapWithStages._
+import dotty.tools.dotc.staging.StagingLevel.*
 
 
 
@@ -35,7 +35,7 @@ class Staging extends MacroTransform {
       // Recheck that PCP holds but do not heal any inconsistent types as they should already have been heald
       tree match {
         case PackageDef(pid, _) if tree.symbol.owner == defn.RootClass =>
-          val checker = new PCPCheckAndHeal(freshStagingContext) {
+          val checker = new PCPCheckAndHeal {
             override protected def tryHeal(sym: Symbol, tp: TypeRef, pos: SrcPos)(using Context): TypeRef = {
               def symStr =
                 if (sym.is(ModuleClass)) sym.sourceModule.show
@@ -51,7 +51,7 @@ class Staging extends MacroTransform {
               tp
             }
           }
-          checker.transform(tree)
+          checker.transform(tree)(using freshStagingLevelContext)
         case _ =>
       }
 
@@ -66,11 +66,11 @@ class Staging extends MacroTransform {
     }
 
   override def run(using Context): Unit =
-    if (ctx.compilationUnit.needsStaging) super.run(using freshStagingContext)
+    if (ctx.compilationUnit.needsStaging) super.run(using freshStagingLevelContext)
 
   protected def newTransformer(using Context): Transformer = new Transformer {
     override def transform(tree: tpd.Tree)(using Context): tpd.Tree =
-      new PCPCheckAndHeal(ctx).transform(tree)
+      (new PCPCheckAndHeal).transform(tree)
   }
 }
 

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

@@ -8,33 +8,25 @@ import dotty.tools.dotc.core.Contexts._
 import dotty.tools.dotc.core.StagingContext._
 import dotty.tools.dotc.core.Symbols._
 import dotty.tools.dotc.util.Property
+import dotty.tools.dotc.staging.StagingLevel
 
 import scala.collection.mutable
-import scala.annotation.constructorOnly
 
 /** The main transformer class
  *  @param  level      the current level, where quotes add one and splices subtract one level.
  *                     The initial level is 0, a level `l` where `l > 0` implies code has been quoted `l` times
  *                     and `l == -1` is code inside a top level splice (in an inline method).
  *  @param  levels     a stacked map from symbols to the levels in which they were defined
  */
-abstract class TreeMapWithStages(@constructorOnly ictx: Context) extends TreeMapWithImplicits {
-
+abstract class TreeMapWithStages extends TreeMapWithImplicits {
   import tpd._
-  import TreeMapWithStages._
-
-  /** A map from locally defined symbols to their definition quotation level */
-  private[this] val levelOfMap: mutable.HashMap[Symbol, Int] = ictx.property(LevelOfKey).get
 
   /** A stack of entered symbols, to be unwound after scope exit */
   private[this] var enteredSyms: List[Symbol] = Nil
 
   /** If we are inside a quote or a splice */
   private[this] var inQuoteOrSplice = false
 
-  /** The quotation level of the definition of the locally defined symbol */
-  protected def levelOf(sym: Symbol): Int = levelOfMap.getOrElse(sym, 0)
-
   /** Locally defined symbols seen so far by `StagingTransformer.transform` */
   protected def localSymbols: List[Symbol] = enteredSyms
 
@@ -43,8 +35,7 @@ abstract class TreeMapWithStages(@constructorOnly ictx: Context) extends TreeMap
 
   /** Enter staging level of symbol defined by `tree` */
   private def markSymbol(sym: Symbol)(using Context): Unit =
-    if level != 0 && !levelOfMap.contains(sym) then
-      levelOfMap(sym) = level
+    if StagingLevel.markSymbol(sym) then
       enteredSyms = sym :: enteredSyms
 
   /** Enter staging level of symbol defined by `tree`, if applicable. */
@@ -79,7 +70,7 @@ abstract class TreeMapWithStages(@constructorOnly ictx: Context) extends TreeMap
         try super.transform(tree)
         finally
           while (enteredSyms ne lastEntered) {
-            levelOfMap -= enteredSyms.head
+            StagingLevel.removeLevelOf(enteredSyms.head)
             enteredSyms = enteredSyms.tail
           }
 
@@ -150,15 +141,3 @@ abstract class TreeMapWithStages(@constructorOnly ictx: Context) extends TreeMap
       }
     }
 }
-
-
-object TreeMapWithStages {
-
-  /** A key to be used in a context property that caches the `levelOf` mapping */
-  private val LevelOfKey = new Property.Key[mutable.HashMap[Symbol, Int]]
-
-  /** Initial context for a StagingTransformer transformation. */
-  def freshStagingContext(using Context): Context =
-    ctx.fresh.setProperty(LevelOfKey, new mutable.HashMap[Symbol, Int])
-
-}