Fix #5376, fix #5434: Unpickle quotes eagerly

Author: nicolasstucki

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

@@ -19,11 +19,13 @@ import dotty.tools.dotc.tastyreflect.ReflectionImpl
 import scala.internal.quoted._
 import scala.reflect.ClassTag
 
+import scala.runtime.quoted.Unpickler._
+
 object PickledQuotes {
   import tpd._
 
   /** Pickle the tree of the quote into strings */
-  def pickleQuote(tree: Tree)(implicit ctx: Context): scala.runtime.quoted.Unpickler.Pickled = {
+  def pickleQuote(tree: Tree)(implicit ctx: Context): List[String] = {
     if (ctx.reporter.hasErrors) Nil
     else {
       assert(!tree.isInstanceOf[Hole]) // Should not be pickled as it represents `'{$x}` which should be optimized to `x`
@@ -34,34 +36,15 @@ 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[_] =>
-      val unpickled = unpickleExpr(expr)
-      /** Force unpickling of the tree, removes the spliced type `@quotedTypeTag type` definitions and dealiases references to `@quotedTypeTag type` */
-      val forceAndCleanArtefacts = new TreeMap {
-        override def transform(tree: tpd.Tree)(implicit ctx: Context): tpd.Tree = tree match {
-          case Block(stat :: rest, expr1) if stat.symbol.hasAnnotation(defn.InternalQuoted_QuoteTypeTagAnnot) =>
-            assert(rest.forall { case tdef: TypeDef => tdef.symbol.hasAnnotation(defn.InternalQuoted_QuoteTypeTagAnnot) })
-            transform(expr1)
-          case tree => super.transform(tree).withType(dealiasTypeTags(tree.tpe))
-        }
-      }
-      forceAndCleanArtefacts.transform(unpickled)
     case expr: TastyTreeExpr[Tree] @unchecked => healOwner(expr.tree)
     case expr: FunctionAppliedTo[_] =>
-      functionAppliedTo(quotedExprToTree(expr.f), expr.args.map(arg => quotedExprToTree(arg(new scala.quoted.QuoteContext(ReflectionImpl(ctx))))).toList)
+      def qctx = dotty.tools.dotc.quoted.QuoteContext()
+      functionAppliedTo(quotedExprToTree(expr.f), expr.args.map(arg => quotedExprToTree(arg(qctx))).toList)
   }
 
   /** Transform the expression into its fully spliced TypeTree */
-  def quotedTypeToTree(expr: quoted.Type[_])(implicit ctx: Context): Tree = expr match {
-    case expr: TastyType[_] =>
-      unpickleType(expr) match {
-        case Block(aliases, tpt) =>
-          // `@quoteTypeTag type` aliasses are not required after unpickling
-          tpt
-        case tpt => tpt
-      }
-    case expr: TreeType[Tree] @unchecked => healOwner(expr.typeTree)
-  }
+  def quotedTypeToTree(expr: quoted.Type[_])(implicit ctx: Context): Tree =
+    healOwner(expr.asInstanceOf[TreeType[Tree]].typeTree)
 
   private def dealiasTypeTags(tp: Type)(implicit ctx: Context): Type = new TypeMap() {
     override def apply(tp: Type): Type = {
@@ -74,15 +57,31 @@ object PickledQuotes {
   }.apply(tp)
 
   /** Unpickle the tree contained in the TastyExpr */
-  private def unpickleExpr(expr: TastyExpr[_])(implicit ctx: Context): Tree = {
-    val tastyBytes = TastyString.unpickle(expr.tasty)
-    unpickle(tastyBytes, expr.args, isType = false)(ctx.addMode(Mode.ReadPositions))
+  def unpickleExpr(tasty: PickledExpr, args: PickledExprArgs)(implicit ctx: Context): Tree = {
+    val tastyBytes = TastyString.unpickle(tasty)
+    val unpickled = unpickle(tastyBytes, args, isType = false)(ctx.addMode(Mode.ReadPositions))
+    /** Force unpickling of the tree, removes the spliced type `@quotedTypeTag type` definitions and dealiases references to `@quotedTypeTag type` */
+    val forceAndCleanArtefacts = new TreeMap {
+      override def transform(tree: tpd.Tree)(implicit ctx: Context): tpd.Tree = tree match {
+        case Block(stat :: rest, expr1) if stat.symbol.hasAnnotation(defn.InternalQuoted_QuoteTypeTagAnnot) =>
+          assert(rest.forall { case tdef: TypeDef => tdef.symbol.hasAnnotation(defn.InternalQuoted_QuoteTypeTagAnnot) })
+          transform(expr1)
+        case tree => super.transform(tree).withType(dealiasTypeTags(tree.tpe))
+      }
+    }
+    forceAndCleanArtefacts.transform(unpickled)
   }
 
   /** Unpickle the tree contained in the TastyType */
-  private def unpickleType(ttpe: TastyType[_])(implicit ctx: Context): Tree = {
-    val tastyBytes = TastyString.unpickle(ttpe.tasty)
-    unpickle(tastyBytes, ttpe.args, isType = true)(ctx.addMode(Mode.ReadPositions))
+  def unpickleType(tasty: PickledType, args: PickledTypeArgs)(implicit ctx: Context): Tree = {
+    val tastyBytes = TastyString.unpickle(tasty)
+    val unpickled = unpickle(tastyBytes, args, isType = true)(ctx.addMode(Mode.ReadPositions))
+    unpickled match {
+      case Block(aliases, tpt) =>
+        // `@quoteTypeTag type` aliasses are not required after unpickling
+        tpt
+      case tpt => tpt
+    }
   }
 
   // TASTY picklingtests/pos/quoteTest.scala

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

@@ -1,7 +1,5 @@
 package dotty.tools.dotc.core.tasty
 
-import scala.runtime.quoted.Unpickler.Pickled
-
 import java.io._
 import java.util.Base64
 import java.nio.charset.StandardCharsets.UTF_8
@@ -12,14 +10,14 @@ object TastyString {
   // Max size of a string literal in the bytecode
   private final val maxStringSize = 65535
 
-  /** Encode TASTY bytes into an Seq of String */
-  def pickle(bytes: Array[Byte]): Pickled = {
+  /** Encode TASTY bytes into a List of String */
+  def pickle(bytes: Array[Byte]): List[String] = {
     val str = new String(Base64.getEncoder().encode(bytes), UTF_8)
     str.sliding(maxStringSize, maxStringSize).toList
   }
 
-  /** Decode the TASTY String into TASTY bytes */
-  def unpickle(strings: Pickled): Array[Byte] = {
+  /** Decode the List of Strings into TASTY bytes */
+  def unpickle(strings: List[String]): Array[Byte] = {
     val string = new StringBuilder
     strings.foreach(string.append)
     Base64.getDecoder().decode(string.result().getBytes(UTF_8))

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

@@ -1281,7 +1281,7 @@ class TreeUnpickler(reader: TastyReader,
         PickledQuotes.quotedTypeToTree(quotedType)
       } else {
         val splice1 = splice.asInstanceOf[Seq[Any] => given scala.quoted.QuoteContext => quoted.Expr[_]]
-        val quotedExpr = splice1(reifiedArgs) given new scala.quoted.QuoteContext(tastyreflect.ReflectionImpl(ctx))
+        val quotedExpr = splice1(reifiedArgs) given dotty.tools.dotc.quoted.QuoteContext()
         PickledQuotes.quotedExprToTree(quotedExpr)
       }
       // We need to make sure a hole is created with the source file of the surrounding context, even if

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

@@ -64,7 +64,8 @@ class QuoteCompiler 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(exprUnit.exprBuilder.apply(new QuoteContext(ReflectionImpl(ctx))))(ctx.withOwner(meth))
+          val qctx = dotty.tools.dotc.quoted.QuoteContext()
+          val quoted = PickledQuotes.quotedExprToTree(exprUnit.exprBuilder.apply(qctx))(ctx.withOwner(meth))
 
           getLiteral(quoted) match {
             case Some(value) =>

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

@@ -0,0 +1,35 @@
+package dotty.tools.dotc.quoted
+
+import dotty.tools.dotc.core.Contexts.Context
+import dotty.tools.dotc.core.Contexts.Context
+import dotty.tools.dotc.core.quoted.PickledQuotes
+import dotty.tools.dotc.tastyreflect.ReflectionImpl
+
+import scala.quoted.{Expr, Type}
+import scala.runtime.quoted.Unpickler._
+
+class QuoteContext private (val tasty: scala.tasty.Reflection) extends scala.quoted.QuoteContext {
+
+  def unpickleExpr[T](pickledExpr: PickledExpr, args: PickledExprArgs): given scala.quoted.QuoteContext => Expr[T] = given qctx => {
+    new scala.internal.quoted.TastyTreeExpr(
+      PickledQuotes.unpickleExpr(pickledExpr, args) given qctx.tasty.rootContext.asInstanceOf[Context]
+    ).asInstanceOf[Expr[T]]
+  }
+
+  def unpickleType[T](pickledType: PickledType, args: PickledTypeArgs): given scala.quoted.QuoteContext => Type[T] = given qctx => {
+    new scala.internal.quoted.TreeType(
+      PickledQuotes.unpickleType(pickledType, args) given qctx.tasty.rootContext.asInstanceOf[Context]
+    ).asInstanceOf[Type[T]]
+  }
+
+}
+
+
+object QuoteContext {
+
+  def apply() given Context: QuoteContext = apply(ReflectionImpl(the[Context]))
+
+  def apply(tastyInstance: scala.tasty.Reflection): QuoteContext { val tasty: tastyInstance.type } =
+    new QuoteContext(tastyInstance).asInstanceOf[QuoteContext { val tasty: tastyInstance.type }]
+
+}

compiler/src/dotty/tools/dotc/tastyreflect/KernelImpl.scala

@@ -27,6 +27,10 @@ class KernelImpl(val rootContext: core.Contexts.Context) extends Kernel {
   def rootPosition: util.SourcePosition =
     tastyreflect.MacroExpansion.position.getOrElse(SourcePosition(rootContext.source, Spans.NoSpan))
 
+  def QuoteContext_from_Reflection(tastyInstance: scala.tasty.Reflection): scala.quoted.QuoteContext { val tasty: tastyInstance.type } = {
+    dotty.tools.dotc.quoted.QuoteContext(tastyInstance)
+  }
+
   //
   // CONTEXT
   //

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

@@ -25,6 +25,8 @@ import dotty.tools.repl.AbstractFileClassLoader
 
 import scala.reflect.ClassTag
 
+import dotty.tools.dotc.quoted.QuoteContext
+
 /** Utility class to splice quoted expressions */
 object Splicer {
   import tpd._
@@ -42,7 +44,7 @@ object Splicer {
       try {
         // Some parts of the macro are evaluated during the unpickling performed in quotedExprToTree
         val interpretedExpr = interpreter.interpret[scala.quoted.QuoteContext => scala.quoted.Expr[Any]](tree)
-        interpretedExpr.fold(tree)(macroClosure => PickledQuotes.quotedExprToTree(macroClosure(new scala.quoted.QuoteContext(ReflectionImpl(ctx)))))
+        interpretedExpr.fold(tree)(macroClosure => PickledQuotes.quotedExprToTree(macroClosure(QuoteContext())))
       }
       catch {
         case ex: StopInterpretation =>
@@ -261,7 +263,7 @@ object Splicer {
       args.toSeq
 
     private def interpretQuoteContext()(implicit env: Env): Object =
-      new scala.quoted.QuoteContext(ReflectionImpl(ctx))
+      QuoteContext()
 
     private def interpretedStaticMethodCall(moduleClass: Symbol, fn: Symbol)(implicit env: Env): List[Object] => Object = {
       val (inst, clazz) =

library/src-bootstrapped/scala/tasty/reflect/TreeUtils.scala

@@ -286,7 +286,7 @@ trait TreeUtils
   /** Bind the `rhs` to a `val` and use it in `body` */
   def let(rhs: Term)(body: Ident => Term): Term = {
     import scala.quoted.QuoteContext
-    given as QuoteContext = new QuoteContext(this)
+    given as QuoteContext = QuoteContext.from(self)
     type T // TODO probably it is better to use the Sealed contruct rather than let the user create their own existential type
     implicit val rhsTpe: quoted.Type[T] = rhs.tpe.seal.asInstanceOf[quoted.Type[T]]
     val rhsExpr = rhs.seal.cast[T]

library/src/scala/quoted/Expr.scala

@@ -75,12 +75,7 @@ package quoted {
 package internal {
   package quoted {
 
-    import scala.quoted._
-
-    /** An Expr backed by a pickled TASTY tree */
-    final class TastyExpr[+T](val tasty: scala.runtime.quoted.Unpickler.Pickled, val args: Seq[Any]) extends Expr[T] {
-      override def toString: String = s"Expr(<pickled tasty>)"
-    }
+    import scala.quoted.{Expr, QuoteContext}
 
     /** An Expr backed by a tree. Only the current compiler trees are allowed.
      *

library/src/scala/quoted/QuoteContext.scala

@@ -2,6 +2,8 @@ package scala.quoted
 
 import scala.quoted.show.SyntaxHighlight
 
+import scala.runtime.quoted.Unpickler._
+
 /** Quotation context provided by a macro expansion or in the scope of `scala.quoted.run`.
  *  Used to perform all operations on quoted `Expr` or `Type`.
  *
@@ -10,7 +12,10 @@ import scala.quoted.show.SyntaxHighlight
  *
  *  @param tasty Typed AST API. Usage: `def f(qctx: QuoteContext) = { import qctx.tasty._; ... }`.
  */
-class QuoteContext(val tasty: scala.tasty.Reflection) {
+trait QuoteContext {
+
+  /** Typed AST API. Usage: `def f(qctx: QuoteContext) = { import qctx.tasty._; ... }`. */
+  val tasty: scala.tasty.Reflection
 
   def show(expr: Expr[_], syntaxHighlight: SyntaxHighlight): String = {
     import tasty._
@@ -46,8 +51,18 @@ class QuoteContext(val tasty: scala.tasty.Reflection) {
     tasty.warning(msg, expr.unseal.pos) given rootContext
   }
 
+  /** Unpickle `pickledExpr` which represents a pickled `Expr` tree
+   */
+  private[scala] def unpickleExpr[T](pickledExpr: PickledExpr, args: PickledExprArgs): given QuoteContext => Expr[T]
+
+  /** Unpickle `pickledType` which represents a pickled `Type` tree
+   */
+  private[scala] def unpickleType[T](pickledType: PickledType, args: PickledTypeArgs): given QuoteContext => Type[T]
+
 }
 
 object QuoteContext {
   def macroContext: QuoteContext = throw new Exception("Not in inline macro.")
+  def from(tastyInstance: scala.tasty.Reflection): QuoteContext { val tasty: tastyInstance.type } =
+    tastyInstance.kernel.QuoteContext_from_Reflection(tastyInstance)
 }