Simplify Quote/Splice transformations in CrossStageSafety

Author: nicolasstucki

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

@@ -714,6 +714,17 @@ object Trees {
         defn.FunctionType(1, isContextual = true)
           .appliedTo(defn.QuotesClass.typeRef, exprType)
       withType(quoteType)
+
+    /** Cancel this Quote if it contains a Splice */
+    def cancelled(using Context): Option[Tree[T]] =
+      def rec(tree: Tree[T]): Option[Tree[T]] = tree match
+        case Block(Nil, expr) => rec(expr)
+        case Splice(inner) =>
+          // Optimization: `'{ $x }` --> `x`
+          // and adapt the refinement of `Quotes { type reflect: ... } ?=> Expr[T]`
+          Some(inner)
+        case _ => None
+      rec(body)
   }
 
   /** A tree representing a splice `${ expr }`
@@ -730,6 +741,16 @@ object Trees {
   case class Splice[+T <: Untyped] private[ast] (expr: Tree[T])(implicit @constructorOnly src: SourceFile)
     extends TermTree[T] {
     type ThisTree[+T <: Untyped] = Splice[T]
+
+    /** Cancel this Splice if it contains a Quote */
+    def cancelled(using Context): Option[Tree[T]] =
+      def rec(tree: Tree[T]): Option[Tree[T]] = tree match
+        case Block(Nil, expr1) => rec(expr1)
+        case Quote(body) =>
+          // Optimization: `${ 'x }` --> `x`
+          Some(body)
+        case _ => None
+      rec(expr)
   }
 
   /** A type tree that represents an existing or inferred type */

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

@@ -53,6 +53,14 @@ class CrossStageSafety extends TreeMapWithStages {
     else if !inQuoteOrSpliceScope then
       checkAnnotations(tree)
       tree match
+        case tree @ Quote(quotedTree) =>
+          tree.cancelled match
+            case Some(tree1) => transform(tree1.asInstance(tree.tpe)) // TODO is this asInstance still needed?
+            case None => transformQuote(quotedTree, tree)
+        case tree @ Splice(splicedTree) =>
+          tree.cancelled match
+            case Some(tree1) => transform(tree1)
+            case None => transformSplice(splicedTree, tree)
         case tree @ QuotedTypeOf(quotedTree) =>
           transformQuotedType(quotedTree, tree)
         case _ => super.transform(tree)
@@ -95,6 +103,14 @@ class CrossStageSafety extends TreeMapWithStages {
       case tree: TypeDef if tree.symbol.is(Case) && level > 0 =>
         report.error(reporting.CaseClassInInlinedCode(tree), tree)
         super.transform(tree)
+      case tree @ Quote(quotedTree) =>
+        tree.cancelled match
+          case Some(tree1) => transform(tree1.asInstance(tree.tpe)) // TODO is this asInstance still needed?
+          case None => transformQuote(quotedTree, tree)
+      case tree @ Splice(splicedTree) =>
+        tree.cancelled match
+          case Some(tree1) => transform(tree1)
+          case None => transformSplice(splicedTree, tree)
       case tree @ QuotedTypeOf(quotedTree) =>
         transformQuotedType(quotedTree, tree)
       case tree @ SplicedType(splicedTree) =>
@@ -104,7 +120,7 @@ class CrossStageSafety extends TreeMapWithStages {
     }
 
   /** Transform quoted trees while maintaining level correctness */
-  override protected def transformQuote(body: Tree, quote: Quote)(using Context): Tree = {
+  private def transformQuote(body: Tree, quote: Quote)(using Context): Tree = {
     if (ctx.property(InAnnotation).isDefined)
       report.error("Cannot have a quote in an annotation", quote.srcPos)
     val transformedBody = transformQuoteBody(body, quote.span)
@@ -150,7 +166,7 @@ class CrossStageSafety extends TreeMapWithStages {
    *  - If inside inlined code, expand the macro code.
    *  - If inside of a macro definition, check the validity of the macro.
    */
-  protected def transformSplice(body: Tree, splice: Splice)(using Context): Tree = {
+  private def transformSplice(body: Tree, splice: Splice)(using Context): Tree = {
     val body1 = transform(body)(using spliceContext)
     val tpe1 =
       if level == 0 then splice.tpe

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

@@ -6,7 +6,6 @@ import dotty.tools.dotc.config.Printers.staging
 import dotty.tools.dotc.core.Decorators._
 import dotty.tools.dotc.core.Contexts._
 import dotty.tools.dotc.core.Symbols._
-import dotty.tools.dotc.util.Property
 import dotty.tools.dotc.staging.StagingLevel.*
 
 import scala.collection.mutable
@@ -15,45 +14,11 @@ import scala.collection.mutable
 abstract class TreeMapWithStages extends TreeMapWithImplicits {
   import tpd._
 
-  /** Transform the quote `quote` which contains the quoted `body`.
-   *
-   *  - `quoted.runtime.Expr.quote[T](<body0>)`  --> `quoted.runtime.Expr.quote[T](<body>)`
-   */
-  protected def transformQuote(body: Tree, quote: Quote)(using Context): Tree =
-    cpy.Quote(quote)(body)
-
-  /** Transform the expression splice `splice` which contains the spliced `body`. */
-  protected def transformSplice(body: Tree, splice: Splice)(using Context): Tree
-
   override def transform(tree: Tree)(using Context): Tree =
     if (tree.source != ctx.source && tree.source.exists)
       transform(tree)(using ctx.withSource(tree.source))
-    else reporting.trace(i"StagingTransformer.transform $tree at $level", staging, show = true) {
-      def dropEmptyBlocks(tree: Tree): Tree = tree match {
-        case Block(Nil, expr) => dropEmptyBlocks(expr)
-        case _ => tree
-      }
-
+    else reporting.trace(i"TreeMapWithStages.transform $tree at $level", staging, show = true) {
       tree match {
-        case tree @ Quote(quotedTree) =>
-          dropEmptyBlocks(quotedTree) match {
-            case Splice(t) =>
-              // Optimization: `'{ $x }` --> `x`
-              // and adapt the refinement of `Quotes { type reflect: ... } ?=> Expr[T]`
-              transform(t).asInstance(tree.tpe)
-            case _ =>
-              transformQuote(quotedTree, tree)
-          }
-
-        case tree @ Splice(splicedTree) =>
-          dropEmptyBlocks(splicedTree) match {
-            case Quote(t) =>
-              // Optimization: `${ 'x }` --> `x`
-              transform(t)
-            case _ =>
-              transformSplice(splicedTree, tree)
-          }
-
         case Block(stats, _) =>
           val defSyms = stats.collect { case defTree: DefTree => defTree.symbol }
           super.transform(tree)(using symbolsInCurrentLevel(defSyms))