Refactor checks

Author: nicolasstucki

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

@@ -38,12 +38,10 @@ class PCPCheckAndHeal(@constructorOnly ictx: Context) extends TreeMapWithStages(
 
   private val InAnnotation = Property.Key[Unit]()
 
-  private[this] var inQuoteOrSplice = false
-
   override def transform(tree: Tree)(implicit ctx: Context): Tree =
     if (tree.source != ctx.source && tree.source.exists)
       transform(tree)(ctx.withSource(tree.source))
-    else if !inQuoteOrSplice then
+    else if !isInQuoteOrSplice then
       checkAnnotations(tree)
       super.transform(tree)
     else tree match {
@@ -54,11 +52,9 @@ class PCPCheckAndHeal(@constructorOnly ictx: Context) extends TreeMapWithStages(
         else TypeTree(healedType).withSpan(tree.span)
 
       case _: This =>
-        if level != -1 && level != levelOf(tree.symbol) then
-          levelError(tree.symbol, tree.tpe, tree.sourcePos, "")
-        tree
+        tree.withType(healTermType(tree.sourcePos)(tree.tpe))
       case _: Ident =>
-        checkTermLevel(tree).withType(healTermType(tree.sourcePos)(tree.tpe))
+        tree.withType(healTermType(tree.sourcePos)(tree.tpe))
 
       // Remove inline defs in quoted code. Already fully inlined.
       case tree: DefDef if tree.symbol.is(Inline) && level > 0 => EmptyTree
@@ -79,16 +75,14 @@ class PCPCheckAndHeal(@constructorOnly ictx: Context) extends TreeMapWithStages(
   /** Transform quoted trees while maintaining phase correctness */
   override protected def transformQuotation(body: Tree, quote: Tree)(implicit ctx: Context): Tree = {
     val taggedTypes = new PCPCheckAndHeal.QuoteTypeTags(quote.span)(using ctx)
+
     if (ctx.property(InAnnotation).isDefined)
       ctx.error("Cannot have a quote in an annotation", quote.sourcePos)
 
     val contextWithQuote =
       if level == 0 then contextWithQuoteTypeTags(taggedTypes)(quoteContext)
       else quoteContext
-    val inQuoteOrSpliceOld = inQuoteOrSplice
-    inQuoteOrSplice = true
     val body1 = transform(body)(contextWithQuote)
-    inQuoteOrSplice = inQuoteOrSpliceOld
     val body2 =
       taggedTypes.getTypeTags match
         case Nil  => body1
@@ -103,10 +97,7 @@ class PCPCheckAndHeal(@constructorOnly ictx: Context) extends TreeMapWithStages(
    *  - If inside of a macro definition, check the validity of the macro.
    */
   protected def transformSplice(body: Tree, splice: Tree)(implicit ctx: Context): Tree = {
-    val inQuoteOrSpliceOld = inQuoteOrSplice
-    inQuoteOrSplice = true
     val body1 = transform(body)(spliceContext)
-    inQuoteOrSplice = inQuoteOrSpliceOld
     splice match {
       case Apply(fun @ TypeApply(_, _ :: Nil), _) if splice.isTerm =>
         // Type of the splice itsel must also be healed
@@ -154,9 +145,8 @@ class PCPCheckAndHeal(@constructorOnly ictx: Context) extends TreeMapWithStages(
               tryHeal(tp.symbol, tp, pos).getOrElse(tp)
             case _ =>
               mapOver(tp)
-        case tp: ThisType if level != levelOf(tp.cls) =>
-          if level != -1 then // Ok after inlined
-            assert(levelError(tp.cls, tp, pos, "").isEmpty)
+        case tp: ThisType if level != -1 && level != levelOf(tp.cls) =>
+          levelError(tp.cls, tp, pos, "")
           tp
         case tp: AnnotatedType =>
           val newAnnotTree = transform(tp.annot.tree)
@@ -170,22 +160,17 @@ class PCPCheckAndHeal(@constructorOnly ictx: Context) extends TreeMapWithStages(
       tp match
         case tp @ TypeRef(NoPrefix, _) if level > levelOf(tp.symbol) =>
           tryHeal(tp.symbol, tp, pos).getOrElse(tp)
+        case tp @ TermRef(NoPrefix, _) if !tp.symbol.isStatic && level != levelOf(tp.symbol) =>
+          levelError(tp.symbol, tp, pos, "")
+          tp
+        case tp: ThisType if level != -1 && level != levelOf(tp.cls) =>
+          levelError(tp.cls, tp, pos, "")
+          tp
         case _ =>
           if tp.typeSymbol.is(Package) then tp
           else mapOver(tp)
   }
 
-  private def checkTermLevel(tree: Tree)(using Context): tree.type = {
-    new TypeTraverser {
-      def traverse(tp: Type): Unit =
-        tp match
-          case tp @ TermRef(NoPrefix, _) if !tp.symbol.isStatic && level != levelOf(tp.symbol) =>
-            levelError(tree.symbol, tree.tpe, tree.sourcePos, "")
-          case _ => traverseChildren(tp)
-    }.traverse(tree.tpe)
-    tree
-  }
-
   // TODO return new type or else old?
   protected def tryHeal(sym: Symbol, tp: TypeRef, pos: SourcePosition)(implicit ctx: Context): Option[TypeRef] = {
     val reqType = defn.QuotedTypeClass.typeRef.appliedTo(tp)

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

@@ -35,17 +35,23 @@ abstract class TreeMapWithStages(@constructorOnly ictx: Context) extends TreeMap
   import TreeMapWithStages._
 
   /** A map from locally defined symbols to their definition quotation level */
-  private val levelOfMap: mutable.HashMap[Symbol, Int] = ictx.property(LevelOfKey).get
+  private[this] val levelOfMap: mutable.HashMap[Symbol, Int] = ictx.property(LevelOfKey).get
 
   /** A stack of entered symbols, to be unwound after scope exit */
-  private var enteredSyms: List[Symbol] = Nil
+  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)
 
   /** Localy defined symbols seen so far by `StagingTransformer.transform` */
   protected def localSymbols: List[Symbol] = enteredSyms
 
+  /** If we are inside a quote or a splice */
+  protected def isInQuoteOrSplice: Boolean = inQuoteOrSplice
+
   /** Enter staging level of symbol defined by `tree` */
   private def markSymbol(sym: Symbol)(implicit ctx: Context): Unit =
     if (!levelOfMap.contains(sym))
@@ -88,19 +94,25 @@ abstract class TreeMapWithStages(@constructorOnly ictx: Context) extends TreeMap
       tree match {
 
         case Quoted(quotedTree) =>
-          dropEmptyBlocks(quotedTree) match {
+          val old = inQuoteOrSplice
+          inQuoteOrSplice = true
+          try dropEmptyBlocks(quotedTree) match {
             case Spliced(t) =>
               // '{ $x } --> x
               // and adapt the refinment of `QuoteContext { type tasty: ... } ?=> Expr[T]`
               transform(t).asInstance(tree.tpe)
             case _ => transformQuotation(quotedTree, tree)
           }
+          finally inQuoteOrSplice = old
 
         case tree @ Spliced(splicedTree) =>
-          dropEmptyBlocks(splicedTree) match {
+          val old = inQuoteOrSplice
+          inQuoteOrSplice = true
+          try dropEmptyBlocks(splicedTree) match {
             case Quoted(t) => transform(t) // ${ 'x } --> x
             case _ => transformSplice(splicedTree, tree)
           }
+          finally inQuoteOrSplice = old
 
         case Block(stats, _) =>
           val last = enteredSyms