Refactor pickled quotes logic

Introduce a PickledExpr, PickledType and PickledSplice abstractions.
This will allow the underlying representation of the pickled quote or quote arguments to change in a future version if needed.

Author: nicolasstucki

compiler/src/dotty/tools/dotc/core/Definitions.scala

@@ -837,8 +837,9 @@ class Definitions {
 
   @tu lazy val TastyReflectionClass: ClassSymbol = requiredClass("scala.tasty.Reflection")
 
-  @tu lazy val Unpickler_unpickleExpr: Symbol = requiredMethod("scala.internal.quoted.Unpickler.unpickleExpr")
-  @tu lazy val Unpickler_unpickleType: Symbol = requiredMethod("scala.internal.quoted.Unpickler.unpickleType")
+  @tu lazy val PickledExpr_make: Symbol = requiredMethod("scala.internal.quoted.PickledExpr.make")
+  @tu lazy val PickledType_make: Symbol = requiredMethod("scala.internal.quoted.PickledType.make")
+  @tu lazy val PickledSplices_make: Symbol = requiredMethod("scala.internal.quoted.PickledSplices.make")
 
   @tu lazy val EqlClass: ClassSymbol = requiredClass("scala.Eql")
     def Eql_eqlAny(using Context): TermSymbol = EqlClass.companionModule.requiredMethod(nme.eqlAny)

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

@@ -4,22 +4,20 @@ import java.io._
 import java.util.Base64
 import java.nio.charset.StandardCharsets.UTF_8
 
-import scala.internal.quoted.Unpickler.PickledQuote
-
 /** Utils for String representation of TASTY */
 object TastyString {
 
   // Max size of a string literal in the bytecode
   private final val maxStringSize = 65535
 
   /** Encode TASTY bytes into a List of String */
-  def pickle(bytes: Array[Byte]): PickledQuote = {
+  def pickle(bytes: Array[Byte]): List[String] = {
     val str = new String(Base64.getEncoder().encode(bytes), UTF_8)
     str.toSeq.sliding(maxStringSize, maxStringSize).map(_.unwrap).toList
   }
 
   /** Decode the List of Strings into TASTY bytes */
-  def unpickle(strings: PickledQuote): Array[Byte] = {
+  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/quoted/PickledQuotes.scala

@@ -21,15 +21,17 @@ import dotty.tools.tasty.TastyString
 
 import scala.reflect.ClassTag
 
-import scala.internal.quoted.Unpickler._
+import scala.internal.quoted.PickledExpr
+import scala.internal.quoted.PickledSplices
+import scala.internal.quoted.PickledType
 import scala.quoted.QuoteContext
 import scala.collection.mutable
 
 object PickledQuotes {
   import tpd._
 
   /** Pickle the tree of the quote into strings */
-  def pickleQuote(tree: Tree)(using Context): PickledQuote =
+  def pickleQuote(tree: Tree)(using 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`
@@ -52,8 +54,7 @@ object PickledQuotes {
   }
 
   /** Unpickle the tree contained in the TastyExpr */
-  def unpickleExpr(tasty: PickledQuote, splices: PickledArgs)(using Context): Tree = {
-    val tastyBytes = TastyString.unpickle(tasty)
+  def unpickleTerm(tastyBytes: Array[Byte], splices: PickledSplices)(using Context): Tree = {
     val unpickled = withMode(Mode.ReadPositions)(
       unpickle(tastyBytes, splices, isType = false))
     val Inlined(call, Nil, expnasion) = unpickled
@@ -64,15 +65,14 @@ object PickledQuotes {
   }
 
   /** Unpickle the tree contained in the TastyType */
-  def unpickleType(tasty: PickledQuote, args: PickledArgs)(using Context): Tree = {
-    val tastyBytes = TastyString.unpickle(tasty)
+  def unpickleTypeTree(tastyBytes: Array[Byte], splices: PickledSplices)(using Context): Tree = {
     val unpickled = withMode(Mode.ReadPositions)(
-      unpickle(tastyBytes, args, isType = true))
-    spliceTypes(unpickled, args)
+      unpickle(tastyBytes, splices, isType = true))
+    spliceTypes(unpickled, splices)
   }
 
   /** Replace all term holes with the spliced terms */
-  private def spliceTerms(tree: Tree, splices: PickledArgs)(using Context): Tree = {
+  private def spliceTerms(tree: Tree, splices: PickledSplices)(using Context): Tree = {
     val evaluateHoles = new TreeMap {
       override def transform(tree: tpd.Tree)(using Context): tpd.Tree = tree match {
         case Hole(isTerm, idx, args) =>
@@ -81,8 +81,7 @@ object PickledQuotes {
             else new scala.internal.quoted.Type(arg, QuoteContextImpl.scopeId)
           }
           if isTerm then
-            val splice1 = splices(idx).asInstanceOf[Seq[Any] => QuoteContext ?=> quoted.Expr[?]]
-            val quotedExpr = splice1(reifiedArgs)(using dotty.tools.dotc.quoted.QuoteContextImpl())
+            val quotedExpr = splices.exprSplice(idx)(reifiedArgs)(dotty.tools.dotc.quoted.QuoteContextImpl())
             val filled = PickledQuotes.quotedExprToTree(quotedExpr)
 
             // We need to make sure a hole is created with the source file of the surrounding context, even if
@@ -92,7 +91,7 @@ object PickledQuotes {
           else
             // Replaces type holes generated by ReifyQuotes (non-spliced types).
             // These are types defined in a quote and used at the same level in a nested quote.
-            val quotedType = splices(idx).asInstanceOf[Seq[Any] => quoted.Type[?]](reifiedArgs)
+            val quotedType = splices.typeSplice(idx)(reifiedArgs)
             PickledQuotes.quotedTypeToTree(quotedType)
         case tree: Select =>
           // Retain selected members
@@ -127,15 +126,15 @@ object PickledQuotes {
   }
 
   /** Replace all type holes generated with the spliced types */
-  private def spliceTypes(tree: Tree, splices: PickledArgs)(using Context): Tree = {
+  private def spliceTypes(tree: Tree, splices: PickledSplices)(using Context): Tree = {
     tree match
       case Block(stat :: rest, expr1) if stat.symbol.hasAnnotation(defn.InternalQuoted_QuoteTypeTagAnnot) =>
         val typeSpliceMap = (stat :: rest).iterator.map {
           case tdef: TypeDef =>
             assert(tdef.symbol.hasAnnotation(defn.InternalQuoted_QuoteTypeTagAnnot))
             val tree = tdef.rhs match
               case TypeBoundsTree(_, Hole(_, idx, args), _) =>
-                val quotedType = splices(idx).asInstanceOf[Seq[Any] => quoted.Type[?]](args)
+                val quotedType = splices.typeSplice(idx)(args)
                 PickledQuotes.quotedTypeToTree(quotedType)
               case TypeBoundsTree(_, tpt, _) =>
                 tpt
@@ -181,7 +180,7 @@ object PickledQuotes {
   }
 
   /** Unpickle TASTY bytes into it's tree */
-  private def unpickle(bytes: Array[Byte], splices: Seq[Any], isType: Boolean)(using Context): Tree = {
+  private def unpickle(bytes: Array[Byte], splices: PickledSplices, isType: Boolean)(using Context): Tree = {
     quotePickling.println(s"**** unpickling quote from TASTY\n${new TastyPrinter(bytes).printContents()}")
 
     val mode = if (isType) UnpickleMode.TypeTree else UnpickleMode.Term

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

@@ -16,7 +16,9 @@ import dotty.tools.dotc.typer.Implicits
 import scala.quoted.QuoteContext
 import scala.quoted.show.SyntaxHighlight
 
-import scala.internal.quoted.Unpickler
+import scala.internal.quoted.PickledExpr
+import scala.internal.quoted.PickledSplices
+import scala.internal.quoted.PickledType
 import scala.tasty.reflect._
 
 object QuoteContextImpl {
@@ -2604,12 +2606,11 @@ class QuoteContextImpl private (ctx: Context) extends QuoteContext:
     private def withDefaultPos[T <: Tree](fn: Context ?=> T): T =
       fn(using ctx.withSource(rootPosition.source)).withSpan(rootPosition.span)
 
+    def unpickleTerm(bytes: Array[Byte], splices: PickledSplices): Term =
+      PickledQuotes.unpickleTerm(bytes, splices)
 
-    def unpickleExpr(repr: Unpickler.PickledQuote, args: Unpickler.PickledArgs): Term =
-      PickledQuotes.unpickleExpr(repr, args)
-
-    def unpickleType(repr: Unpickler.PickledQuote, args: Unpickler.PickledArgs): TypeTree =
-      PickledQuotes.unpickleType(repr, args)
+    def unpickleTypeTree(bytes: Array[Byte], splices: PickledSplices): TypeTree =
+      PickledQuotes.unpickleTypeTree(bytes, splices)
 
     def termMatch(scrutinee: Term, pattern: Term): Option[Tuple] =
       treeMatch(scrutinee, pattern)

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

@@ -163,10 +163,12 @@ class ReifyQuotes extends MacroTransform {
       }
 
       def pickleAsTasty() = {
-        val meth = if isType then defn.Unpickler_unpickleType else defn.Unpickler_unpickleExpr
+        val meth = if isType then defn.PickledType_make else defn.PickledExpr_make
         val pickledQuoteStrings = liftList(PickledQuotes.pickleQuote(body).map(x => Literal(Constant(x))), defn.StringType)
+        // TODO: generate an instance of PickledSplices directly instead of passing through a List
         val splicesList = liftList(splices, defn.FunctionType(1).appliedTo(defn.SeqType.appliedTo(defn.AnyType), defn.AnyType))
-        ref(meth).appliedToType(originalTp).appliedTo(pickledQuoteStrings, splicesList)
+        val pickledSplices = ref(defn.PickledSplices_make).appliedTo(splicesList)
+        ref(meth).appliedToType(originalTp).appliedTo(pickledQuoteStrings, pickledSplices).select("unpickle".toTermName).appliedToArgs(Nil)
       }
 
       def taggedType(sym: Symbol) = ref(defn.InternalQuotedTypeModule).select(sym.name.toTermName)

library/src-bootstrapped/scala/internal/quoted/Unpickler.scala

@@ -0,0 +1,18 @@
+package scala.internal.quoted
+
+import scala.quoted._
+import scala.internal.tasty.CompilerInterface.quoteContextWithCompilerInterface
+
+/** An Expression that is pickled in a  */
+trait PickledExpr[+T]:
+  def unpickle(): QuoteContext ?=> Expr[T]
+
+object PickledExpr:
+
+  def make[T](pickled: List[String], splices: PickledSplices): PickledExpr[T] =
+    new PickledExpr[T]:
+      def unpickle(): QuoteContext ?=> Expr[T] =
+        val bytes = TastyString.unpickle(pickled)
+        val qctx = quoteContextWithCompilerInterface(summon[QuoteContext])
+        val tree = qctx.tasty.unpickleTerm(bytes, splices)
+        new scala.internal.quoted.Expr(tree, qctx.hashCode).asInstanceOf[Expr[T]]

library/src/scala/internal/quoted/PickledExpr.scala

@@ -0,0 +1,23 @@
+package scala.internal.quoted
+
+import scala.quoted._
+
+/** Provider of expressions and types that will fill the holes a pickled quote */
+trait PickledSplices:
+
+  /** Expression that will fill the hole `Hole(<idx> | <args>*)` */
+  def exprSplice(idx: Int)(args: Seq[Any /* Expr[Any] | Type[?] */])(qctx: QuoteContext): Expr[Any]
+
+  /** Type that will fill the hole `Hole(<idx> | <args>*)` */
+  def typeSplice(idx: Int)(args: Seq[Any /* Expr[Any] | Type[?] */]): Type[?]
+
+object PickledSplices:
+  // TODO: generate a more efficient representation
+  //       - avoid creation of lambdas
+  //       - use swich on id
+  def make(seq: Seq[Seq[Any /* Expr[Any] | Type[?] */] => Any]): PickledSplices =
+    new PickledSplices:
+      def exprSplice(idx: Int)(args: Seq[Any])(qctx: QuoteContext): Expr[Any] =
+        seq(idx)(args).asInstanceOf[QuoteContext => Expr[Any]](qctx)
+      def typeSplice(idx: Int)(args: Seq[Any]): Type[?] =
+        seq(idx)(args).asInstanceOf[Type[?]]

library/src/scala/internal/quoted/PickledSplices.scala

@@ -0,0 +1,17 @@
+package scala.internal.quoted
+
+import scala.quoted._
+import scala.internal.tasty.CompilerInterface.quoteContextWithCompilerInterface
+
+trait PickledType[T <: AnyKind]:
+  def unpickle(): QuoteContext ?=> Type[T]
+
+object PickledType:
+
+  def make[T <: AnyKind](pickled: List[String], splices: PickledSplices): PickledType[T] =
+    new PickledType[T]:
+      def unpickle(): QuoteContext ?=> Type[T] =
+        val bytes = TastyString.unpickle(pickled)
+        val qctx = quoteContextWithCompilerInterface(summon[QuoteContext])
+        val tree = qctx.tasty.unpickleTypeTree(bytes, splices)
+        new scala.internal.quoted.Type(tree, qctx.hashCode).asInstanceOf[Type[T]]

library/src/scala/internal/quoted/PickledType.scala

@@ -0,0 +1,25 @@
+package scala.internal.quoted
+
+import java.io._
+import java.util.Base64
+import java.nio.charset.StandardCharsets.UTF_8
+
+/** Utils for String representation of TASTY */
+object TastyString {
+
+  // Max size of a string literal in the bytecode
+  private final val maxStringSize = 65535
+
+  /** 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.toSeq.sliding(maxStringSize, maxStringSize).map(_.unwrap).toList
+  }
+
+  /** 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))
+  }
+}

library/src/scala/internal/quoted/TastyString.scala

@@ -2,7 +2,9 @@ package scala.internal.tasty
 
 import scala.quoted.QuoteContext
 import scala.tasty.reflect._
-import scala.internal.quoted.Unpickler
+import scala.internal.quoted.PickledExpr
+import scala.internal.quoted.PickledSplices
+import scala.internal.quoted.PickledType
 
 /** Part of the reflection interface that needs to be implemented by the compiler */
 trait CompilerInterface { self: scala.tasty.Reflection =>
@@ -12,14 +14,14 @@ trait CompilerInterface { self: scala.tasty.Reflection =>
   //////////////////////
 
   /** Unpickle `repr` which represents a pickled `Expr` tree,
-   *  replacing splice nodes with `args`
+   *  replacing splice nodes with `holes`
    */
-  def unpickleExpr(repr: Unpickler.PickledQuote, args: Unpickler.PickledArgs): Term
+  def unpickleTerm(bytes: Array[Byte], splices: PickledSplices): Term
 
   /** Unpickle `repr` which represents a pickled `Type` tree,
-   *  replacing splice nodes with `args`
+   *  replacing splice nodes with `holes`
    */
-  def unpickleType(repr: Unpickler.PickledQuote, args: Unpickler.PickledArgs): TypeTree
+  def unpickleTypeTree(bytes: Array[Byte], splices: PickledSplices): TypeTree
 
   /** Pattern matches the scrutinee against the pattern and returns a tuple
    *  with the matched holes if successful.