Avoid transforming the body in optimization case

nicolasstucki · nicolasstucki · commit d479de600520 · 2023-03-22T08:01:32.000+01:00
diff --git a/compiler/src/dotty/tools/dotc/staging/CrossStageSafety.scala b/compiler/src/dotty/tools/dotc/staging/CrossStageSafety.scala
@@ -107,47 +107,37 @@ class CrossStageSafety extends TreeMapWithStages {
     val stripAnnotsDeep: TypeMap = new TypeMap:
       def apply(tp: Type): Type = mapOver(tp.stripAnnots)
 
-    val contextWithQuote =
-      if level == 0 then contextWithQuoteTypeTags(taggedTypes)(using quoteContext)
-      else quoteContext
-    val body1 = transform(body)(using contextWithQuote)
-    val body2 =
+    def transformBody() =
+      val contextWithQuote =
+        if level == 0 then contextWithQuoteTypeTags(taggedTypes)(using quoteContext)
+        else quoteContext
+      val transformedBody = transform(body)(using contextWithQuote)
       taggedTypes.getTypeTags match
-        case Nil  => body1
-        case tags => tpd.Block(tags, body1).withSpan(body.span)
+        case Nil  => transformedBody
+        case tags => tpd.Block(tags, transformedBody).withSpan(body.span)
 
     if body.isTerm then
+      val transformedBody = transformBody()
       // `quoted.runtime.Expr.quote[T](<body>)`  --> `quoted.runtime.Expr.quote[T2](<body2>)`
       val TypeApply(fun, targs) = quote.fun: @unchecked
       val targs2 = targs.map(targ => TypeTree(healType(quote.fun.srcPos)(stripAnnotsDeep(targ.tpe))))
-      cpy.Apply(quote)(cpy.TypeApply(quote.fun)(fun, targs2), body2 :: Nil)
+      cpy.Apply(quote)(cpy.TypeApply(quote.fun)(fun, targs2), transformedBody :: Nil)
     else
-      object DirectTypeOfRef:
-        def unapply(body: Tree): Option[Tree] =
-          body.tpe match
-            case tp @ TypeRef(x: TermRef, _) if tp.symbol == defn.QuotedType_splice =>
-              // Optimization: `quoted.Type.of[x.Underlying](quotes)`  -->  `x`
-              Some(ref(x).withSpan(quote.span))
-            case _ =>
-              body2 match
-                case Block(List(tdef: TypeDef), tpt: TypeTree) =>
-                  tpt.tpe match
-                    case tpe: TypeRef if tpe.typeSymbol == tdef.symbol =>
-                      tdef.rhs.tpe.hiBound match
-                        case tp @ TypeRef(x: TermRef, _) if tp.symbol == defn.QuotedType_splice =>
-                          // Optimization: `quoted.Type.of[@SplicedType type T = x.Underlying; T](quotes)`  -->  `x`
-                          Some(ref(x).withSpan(quote.span))
-                        case _ => None
-                    case _ => None
-                case _ => None
-
-      body match
-        case DirectTypeOfRef(ref) => ref
+      body.tpe match
+        case DirectTypeOf(termRef) =>
+          // Optimization: `quoted.Type.of[x.Underlying](quotes)`  -->  `x`
+          ref(termRef).withSpan(quote.span)
         case _ =>
-          val quotes = quote.args.mapConserve(transform)
-          // `quoted.Type.of[<body>](quotes)`  --> `quoted.Type.of[<body2>](quotes)`
-          val TypeApply(fun, _) = quote.fun: @unchecked
-          cpy.Apply(quote)(cpy.TypeApply(quote.fun)(fun, body2 :: Nil), quotes)
+          transformBody() match
+            case DirectTypeOf.Healed(termRef) =>
+              // Optimization: `quoted.Type.of[@SplicedType type T = x.Underlying; T](quotes)`  -->  `x`
+              ref(termRef).withSpan(quote.span)
+            case transformedBody =>
+              val quotes = quote.args.mapConserve(transform)
+              // `quoted.Type.of[<body>](quotes)`  --> `quoted.Type.of[<body2>](quotes)`
+              val TypeApply(fun, _) = quote.fun: @unchecked
+              cpy.Apply(quote)(cpy.TypeApply(quote.fun)(fun, transformedBody :: Nil), quotes)
+
   }
 
   /** Transform splice
diff --git a/compiler/src/dotty/tools/dotc/staging/DirectTypeOf.scala b/compiler/src/dotty/tools/dotc/staging/DirectTypeOf.scala
@@ -0,0 +1,25 @@
+package dotty.tools.dotc.staging
+
+import dotty.tools.dotc.ast.{tpd, untpd}
+import dotty.tools.dotc.core.Contexts._
+import dotty.tools.dotc.core.Symbols._
+import dotty.tools.dotc.core.Types._
+
+object DirectTypeOf:
+  import tpd.*
+
+  /** Matches `x.Underlying` and extracts the TermRef to `x` */
+  def unapply(tpe: Type)(using Context): Option[TermRef] = tpe match
+    case tp @ TypeRef(x: TermRef, _) if tp.symbol == defn.QuotedType_splice => Some(x)
+    case _ => None
+
+  object Healed:
+    /** Matches `{ @SplicedType type T = x.Underlying; T }` and extracts the TermRef to `x` */
+    def unapply(body: Tree)(using Context): Option[TermRef] =
+      body match
+        case Block(List(tdef: TypeDef), tpt: TypeTree) =>
+          tpt.tpe match
+            case tpe: TypeRef if tpe.typeSymbol == tdef.symbol =>
+              DirectTypeOf.unapply(tdef.rhs.tpe.hiBound)
+            case _ => None
+        case _ => None
