Avoid transforming the body in optimization case

nicolasstucki · nicolasstucki · commit a7a59c74de32 · 2023-03-14T16:07:06.000+01:00
diff --git a/compiler/src/dotty/tools/dotc/staging/PCPCheckAndHeal.scala b/compiler/src/dotty/tools/dotc/staging/PCPCheckAndHeal.scala
@@ -107,49 +107,52 @@ class PCPCheckAndHeal 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 tp @ TypeRef(x: TermRef, _) if tp.symbol == defn.QuotedType_splice =>
+          // Optimization: `quoted.Type.of[x.Underlying](quotes)`  -->  `x`
+          ref(x).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 HealedDirectQuotedTypeRef(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)
+
   }
 
+  private object HealedDirectQuotedTypeRef:
+    /** Matches `quoted.Type.of[@SplicedType type T = x.Underlying; T](quotes)` 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 =>
+              tdef.rhs.tpe.hiBound match
+                case tp @ TypeRef(x: TermRef, _) if tp.symbol == defn.QuotedType_splice => Some(x)
+                case _ => None
+            case _ => None
+        case _ => None
+
   /** Transform splice
    *  - If inside a quote, transform the contents of the splice.
    *  - If inside inlined code, expand the macro code.
