Merge pull request #9948 from dotty-staging/refector-pickled-quotes-logic

Refactor pickled quotes logic

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 PickledQuote_make: Symbol = requiredMethod("scala.internal.quoted.PickledQuote.make")
+  @tu lazy val PickledQuote_unpickleExpr: Symbol = requiredMethod("scala.internal.quoted.PickledQuote.unpickleExpr")
+  @tu lazy val PickledQuote_unpickleType: Symbol = requiredMethod("scala.internal.quoted.PickledQuote.unpickleType")
 
   @tu lazy val EqlClass: ClassSymbol = requiredClass("scala.Eql")
     def Eql_eqlAny(using Context): TermSymbol = EqlClass.companionModule.requiredMethod(nme.eqlAny)

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

@@ -17,24 +17,22 @@ import dotty.tools.dotc.core.tasty.DottyUnpickler
 import dotty.tools.dotc.core.tasty.TreeUnpickler.UnpickleMode
 import dotty.tools.dotc.report
 
-import dotty.tools.tasty.TastyString
-
 import scala.reflect.ClassTag
 
-import scala.internal.quoted.Unpickler._
+import scala.internal.quoted.PickledQuote
 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`
       val pickled = pickle(tree)
-      TastyString.pickle(pickled)
+      scala.internal.quoted.TastyString.pickle(pickled)
     }
 
   /** Transform the expression into its fully spliced Tree */
@@ -52,27 +50,25 @@ 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(pickledQuote: PickledQuote)(using Context): Tree = {
     val unpickled = withMode(Mode.ReadPositions)(
-      unpickle(tastyBytes, splices, isType = false))
+      unpickle(pickledQuote, isType = false))
     val Inlined(call, Nil, expnasion) = unpickled
     val inlineCtx = inlineContext(call)
-    val expansion1 = spliceTypes(expnasion, splices)(using inlineCtx)
-    val expansion2 = spliceTerms(expansion1, splices)(using inlineCtx)
+    val expansion1 = spliceTypes(expnasion, pickledQuote)(using inlineCtx)
+    val expansion2 = spliceTerms(expansion1, pickledQuote)(using inlineCtx)
     cpy.Inlined(unpickled)(call, Nil, expansion2)
   }
 
   /** Unpickle the tree contained in the TastyType */
-  def unpickleType(tasty: PickledQuote, args: PickledArgs)(using Context): Tree = {
-    val tastyBytes = TastyString.unpickle(tasty)
+  def unpickleTypeTree(pickledQuote: PickledQuote)(using Context): Tree = {
     val unpickled = withMode(Mode.ReadPositions)(
-      unpickle(tastyBytes, args, isType = true))
-    spliceTypes(unpickled, args)
+      unpickle(pickledQuote, isType = true))
+    spliceTypes(unpickled, pickledQuote)
   }
 
   /** Replace all term holes with the spliced terms */
-  private def spliceTerms(tree: Tree, splices: PickledArgs)(using Context): Tree = {
+  private def spliceTerms(tree: Tree, pickledQuote: PickledQuote)(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 +77,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 = pickledQuote.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 +87,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 = pickledQuote.typeSplice(idx)(reifiedArgs)
             PickledQuotes.quotedTypeToTree(quotedType)
         case tree: Select =>
           // Retain selected members
@@ -127,15 +122,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, pickledQuote: PickledQuote)(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 = pickledQuote.typeSplice(idx)(args)
                 PickledQuotes.quotedTypeToTree(quotedType)
               case TypeBoundsTree(_, tpt, _) =>
                 tpt
@@ -181,7 +176,8 @@ 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(pickledQuote: PickledQuote, isType: Boolean)(using Context): Tree = {
+    val bytes = pickledQuote.bytes()
     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

@@ -15,7 +15,7 @@ import dotty.tools.dotc.core.Decorators._
 import scala.quoted.QuoteContext
 import scala.quoted.show.SyntaxHighlight
 
-import scala.internal.quoted.Unpickler
+import scala.internal.quoted.PickledQuote
 import scala.tasty.reflect._
 
 object QuoteContextImpl {
@@ -2613,12 +2613,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(pickledQuote: PickledQuote): Term =
+      PickledQuotes.unpickleTerm(pickledQuote)
 
-    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(pickledQuote: PickledQuote): TypeTree =
+      PickledQuotes.unpickleTypeTree(pickledQuote)
 
     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 unpickleMeth = if isType then defn.PickledQuote_unpickleType else defn.PickledQuote_unpickleExpr
         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 pickledQuote = ref(defn.PickledQuote_make).appliedTo(pickledQuoteStrings, splicesList)
+        ref(unpickleMeth).appliedToType(originalTp).appliedTo(pickledQuote)
       }
 
       def taggedType(sym: Symbol) = ref(defn.InternalQuotedTypeModule).select(sym.name.toTermName)

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

@@ -0,0 +1,45 @@
+package scala.internal.quoted
+
+import scala.quoted._
+import scala.internal.tasty.CompilerInterface.quoteContextWithCompilerInterface
+
+/** Pickled representation of a quoted expression or type */
+trait PickledQuote:
+
+  /** Bytes of the TASTy containing the quoted expression or type */
+  def bytes(): Array[Byte]
+
+  /** 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 PickledQuote:
+
+  def unpickleExpr[T](pickledQuote: PickledQuote): QuoteContext ?=> Expr[T] =
+    val qctx = quoteContextWithCompilerInterface(summon[QuoteContext])
+    val tree = qctx.reflect.unpickleTerm(pickledQuote)
+    new scala.internal.quoted.Expr(tree, qctx.hashCode).asInstanceOf[Expr[T]]
+
+  def unpickleType[T](pickledQuote: PickledQuote): QuoteContext ?=> Type[T] =
+    val qctx = quoteContextWithCompilerInterface(summon[QuoteContext])
+    val tree = qctx.reflect.unpickleTypeTree(pickledQuote)
+    new scala.internal.quoted.Type(tree, qctx.hashCode).asInstanceOf[Type[T]]
+
+  /** Create an instance of PickledExpr from encoded tasty and sequence of labmdas to fill holes
+   *
+   *  @param pickled: Bytes of tasty encoded using TastyString.pickle
+   *  @param seq: Sequence containing all the functions needed to fill the holes of the quote
+   */
+  def make(pickled: List[String], seq: Seq[Seq[Any /* Expr[Any] | Type[?] */] => Any]): PickledQuote =
+    // TODO: generate a more efficient representation
+    //       - avoid creation of lambdas
+    //       - use swich on id
+    new PickledQuote:
+      def bytes(): Array[Byte] =
+        TastyString.unpickle(pickled)
+      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/PickledQuote.scala

@@ -1,25 +1,23 @@
-package dotty.tools.tasty
+package scala.internal.quoted
 
 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
+  private inline 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/core/tasty/TastyString.scala renamed to library/src/scala/internal/quoted/TastyString.scala

@@ -2,7 +2,7 @@ package scala.internal.tasty
 
 import scala.quoted.QuoteContext
 import scala.tasty.reflect._
-import scala.internal.quoted.Unpickler
+import scala.internal.quoted.PickledQuote
 
 /** Part of the reflection interface that needs to be implemented by the compiler */
 trait CompilerInterface { self: scala.tasty.Reflection =>
@@ -12,14 +12,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(pickledQuote: PickledQuote): 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(pickledQuote: PickledQuote): TypeTree
 
   /** Pattern matches the scrutinee against the pattern and returns a tuple
    *  with the matched holes if successful.