Skip to content

Refactor pickled quotes logic #9948

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions compiler/src/dotty/tools/dotc/core/Definitions.scala
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
38 changes: 17 additions & 21 deletions compiler/src/dotty/tools/dotc/quoted/PickledQuotes.scala
Original file line number Diff line number Diff line change
Expand Up @@ -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 */
Expand All @@ -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) =>
Expand All @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
11 changes: 5 additions & 6 deletions compiler/src/dotty/tools/dotc/quoted/QuoteContextImpl.scala
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -2607,12 +2607,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)
Expand Down
6 changes: 4 additions & 2 deletions compiler/src/dotty/tools/dotc/transform/ReifyQuotes.scala
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
28 changes: 0 additions & 28 deletions library/src-bootstrapped/scala/internal/quoted/Unpickler.scala

This file was deleted.

45 changes: 45 additions & 0 deletions library/src/scala/internal/quoted/PickledQuote.scala
Original file line number Diff line number Diff line change
@@ -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]
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In the original design, the interface defines abstract methods such as unpickle(). Will it be more flexible to keep the original design without committing to a concrete representation in the interface?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We have the same flexibility with both designs. In this one we just need to add can add a PickledQuote.unpickleV2 if we need to change the unpickling logic without breaking the old one. Though make is the one that will most likely evolve as it contains the byte representation logic and the hole logic.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The old PickledExpr interface was guided by an attempt to have a user-facing interface for pickled quotes. But we could provide the same abstraction and internally use PickledQuote to achieve the same.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As we put the protocol methods in a single trait. What will happen if we put the methods in CompilerInterface similar to the original design before this PR? Will that have the same benefit?

trait CompilerInterface {
  type PickledQuote
  type PickledArgs
  def unpickleExpr(repr: PickledQuote, args: PickledArgs): Term
  def unpickleType(repr: PickledQuote, args: PickledArgs): TypeTree
}

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How would one create a PickledQuote in that design? We need to create it with some concrete class that exists in the library, but that design seems to indicate that it is not known in the library.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just repeat the message in the email, it seems I miss something about the following point:

And we still need an implementation for each one of those on the library side as is the current case.

I think qctx.reflect.pickTasty(...) will remove the need to put concrete implementation in stdlib.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Indeed, it could all be implemented on the compiler. But that means that each improvement will be painfully complex in the compiler (i.e. the reason why we have not been able to improve it so far) and would have to add a new method on the interface on each change. I don't see how that would scale. The version I propose has the option of adding new static helper method if it is simpler or lets the code be generated without the CompilerInterface which is more flexible.

If we do it your way we also need to add the static helper methods as we do not have access to pickTasty method directly because we don't have access to the qctx where the transformation happens.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If we do it your way we also need to add the static helper methods as we do not have access to pickTasty method directly because we don't have access to the qctx where the transformation happens.

This seems to be an implementation detail, as a QuoteContext must be available in scope. Currently, in the implementation, it seems each quoted tree is associated with a QuoteContext, thus it's possible to get that in later compiler phases. It seems to me that polluting the API design with accidental implementation details is not a good idea.

But there might be other aspects that I'm not aware of in the design. Therefore, feel free to choose whatever you think makes more sense technically.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This design also generated reduces the code size of each quote by putting common code in those static methods. We do not want to generate the following code (or equivalent) at each place a quote is load.

val qctx = quoteContextWithCompilerInterface(summon[QuoteContext])
val tree = qctx.reflect.unpickleTerm(pickledQuote)
new scala.internal.quoted.Expr(tree, qctx.hashCode).asInstanceOf[Expr[T]]

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would use this design as I can see some clear benefits for the optimization I have in mind. We can always evolve the API the way you said later if we need to do a major change.


/** 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[?]]
Original file line number Diff line number Diff line change
@@ -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))
Expand Down
10 changes: 5 additions & 5 deletions library/src/scala/internal/tasty/CompilerInterface.scala
Original file line number Diff line number Diff line change
Expand Up @@ -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 =>
Expand All @@ -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.
Expand Down