Change owners only once for nested unpickled quotes

Author: nicolasstucki

compiler/src/dotty/tools/dotc/core/quoted/PickledQuotes.scala

@@ -19,6 +19,8 @@ import scala.quoted.Exprs._
 
 import scala.reflect.ClassTag
 
+import scala.collection.mutable
+
 object PickledQuotes {
   import tpd._
 
@@ -34,43 +36,58 @@ object PickledQuotes {
   }
 
   /** Transform the expression into its fully spliced Tree */
-  def quotedExprToTree[T](expr: quoted.Expr[T])(implicit ctx: Context): Tree = expr match {
-    case expr: TastyExpr[_] => unpickleExpr(expr)
+  def quotedExprToTree[T](expr: quoted.Expr[T], nested: Boolean)(implicit ctx: Context): Tree = expr match {
+    case expr: TastyExpr[_] => unpickleExpr(expr, nested)
     case expr: LiftedExpr[T] =>
       expr.value match {
         case value: Class[_] => ref(defn.Predef_classOf).appliedToType(classToType(value))
-        case value=> Literal(Constant(value))
+        case value => Literal(Constant(value))
       }
     case expr: TreeExpr[Tree] @unchecked => expr.tree
     case expr: FunctionAppliedTo[_, _] =>
-      functionAppliedTo(quotedExprToTree(expr.f), quotedExprToTree(expr.x))
+      functionAppliedTo(quotedExprToTree(expr.f, nested), quotedExprToTree(expr.x, nested))
   }
 
   /** Transform the expression into its fully spliced TypeTree */
-  def quotedTypeToTree(expr: quoted.Type[_])(implicit ctx: Context): Tree = expr match {
-    case expr: TastyType[_] => unpickleType(expr)
+  def quotedTypeToTree(expr: quoted.Type[_], nested: Boolean)(implicit ctx: Context): Tree = expr match {
+    case expr: TastyType[_] => unpickleType(expr, nested)
     case expr: TaggedType[_] => classTagToTypeTree(expr.ct)
     case expr: TreeType[Tree] @unchecked => expr.tree
   }
 
   /** Unpickle the tree contained in the TastyExpr */
-  private def unpickleExpr(expr: TastyExpr[_])(implicit ctx: Context): Tree = {
+  private def unpickleExpr(expr: TastyExpr[_], nested: Boolean)(implicit ctx: Context): Tree = {
     val tastyBytes = TastyString.unpickle(expr.tasty)
     val unpickled = unpickle(tastyBytes, expr.args)
     unpickled match {
       case PackageDef(_, (vdef: ValDef) :: Nil) =>
-        vdef.rhs.changeOwner(vdef.symbol, ctx.owner)
+        changeQuoteOwners(vdef.rhs, nested)
     }
   }
 
   /** Unpickle the tree contained in the TastyType */
-  private def unpickleType(ttpe: TastyType[_])(implicit ctx: Context): Tree = {
+  private def unpickleType(ttpe: TastyType[_], nested: Boolean)(implicit ctx: Context): Tree = {
     val tastyBytes = TastyString.unpickle(ttpe.tasty)
     val unpickled = unpickle(tastyBytes, ttpe.args)
     unpickled match {
       case PackageDef(_, (vdef: ValDef) :: Nil) =>
-        vdef.rhs.asInstanceOf[TypeApply].args.head
-          .changeOwner(vdef.symbol, ctx.owner)
+        changeQuoteOwners(vdef.rhs.asInstanceOf[TypeApply].args.head, nested)
+    }
+  }
+
+  private def changeQuoteOwners(tree: Tree, nested: Boolean)(implicit ctx: Context): Tree = {
+    if (nested) tree
+    else {
+      val set = mutable.HashSet.empty[Symbol]
+      new TreeTraverser {
+        override def traverse(tree: tpd.Tree)(implicit ctx: Context): Unit = {
+          if (tree.symbol.exists && tree.symbol.owner.name == "$quote".toTermName)
+            set += tree.symbol.owner
+          traverseChildren(tree)
+        }
+      }.traverse(tree)
+      val owners = set.toList
+      new TreeTypeMap(oldOwners = owners, newOwners = owners.map(_ => ctx.owner)).apply(tree)
     }
   }
 

compiler/src/dotty/tools/dotc/core/tasty/TreeUnpickler.scala

@@ -1145,12 +1145,12 @@ class TreeUnpickler(reader: TastyReader,
         val quotedType =
           if (reifiedArgs.isEmpty) splice.asInstanceOf[quoted.Type[_]]
           else splice.asInstanceOf[Seq[Any] => quoted.Type[_]](reifiedArgs)
-        PickledQuotes.quotedTypeToTree(quotedType)
+        PickledQuotes.quotedTypeToTree(quotedType, nested = true)
       } else {
         val quotedExpr =
           if (reifiedArgs.isEmpty) splice.asInstanceOf[quoted.Expr[_]]
           else splice.asInstanceOf[Seq[Any] => quoted.Expr[_]](reifiedArgs)
-        PickledQuotes.quotedExprToTree(quotedExpr)
+        PickledQuotes.quotedExprToTree(quotedExpr, nested = true)
       }
     }
 

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

@@ -61,7 +61,7 @@ class ExprCompiler(directory: AbstractFile) extends Compiler {
         case exprUnit: ExprCompilationUnit =>
           val tree =
             if (putInClass) inClass(exprUnit.expr)
-            else PickledQuotes.quotedExprToTree(exprUnit.expr)
+            else PickledQuotes.quotedExprToTree(exprUnit.expr, nested = false)
           val source = new SourceFile("", Seq())
           CompilationUnit.mkCompilationUnit(source, tree, forceTrees = true)
       }
@@ -80,7 +80,7 @@ class ExprCompiler(directory: AbstractFile) extends Compiler {
       cls.enter(ctx.newDefaultConstructor(cls), EmptyScope)
       val meth = ctx.newSymbol(cls, nme.apply, Method, ExprType(defn.AnyType), coord = pos).entered
 
-      val quoted = PickledQuotes.quotedExprToTree(expr)(ctx.withOwner(meth))
+      val quoted = PickledQuotes.quotedExprToTree(expr, nested = false)(ctx.withOwner(meth))
 
       val run = DefDef(meth, quoted)
       val classTree = ClassDef(cls, DefDef(cls.primaryConstructor.asTerm), run :: Nil)

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

@@ -39,7 +39,7 @@ object Splicer {
     val liftedArgs = getLiftedArgs(call, bindings)
     val interpreter = new Interpreter(pos)
     val interpreted = interpreter.interpretCallToSymbol[Seq[Any] => Object](call.symbol)
-    interpreted.flatMap(lambda => evaluateLambda(lambda, liftedArgs, pos)).fold(tree)(PickledQuotes.quotedExprToTree)
+    interpreted.flatMap(lambda => evaluateLambda(lambda, liftedArgs, pos)).fold(tree)(PickledQuotes.quotedExprToTree(_, nested = false))
   }
 
   /** Given the inline code and bindings, compute the lifted arguments that will be used to execute the macro