Merge pull request #10238 from dotty-staging/encode-unpickle-call-directly-on-context

Encode unpickle call directly on context

Author: nicolasstucki

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

@@ -795,6 +795,7 @@ class Definitions {
   @tu lazy val QuotedExprModule: Symbol = QuotedExprClass.companionModule
 
   @tu lazy val QuoteContextClass: ClassSymbol = requiredClass("scala.quoted.QuoteContext")
+  @tu lazy val QuoteContextInternalClass: ClassSymbol = requiredClass("scala.internal.quoted.QuoteContextInternal")
 
   @tu lazy val LiftableModule: Symbol = requiredModule("scala.quoted.Liftable")
     @tu lazy val LiftableModule_BooleanLiftable: Symbol = LiftableModule.requiredMethod("BooleanLiftable")
@@ -837,8 +838,6 @@ class Definitions {
   @tu lazy val TastyReflectionClass: ClassSymbol = requiredClass("scala.tasty.Reflection")
 
   @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/QuoteContextImpl.scala

@@ -2631,13 +2631,13 @@ class QuoteContextImpl private (ctx: Context) extends QuoteContext, scala.intern
 
   end reflect
 
-  def unpickleExpr(pickledQuote: PickledQuote): scala.quoted.Expr[Any] =
+  def unpickleExpr[T](pickledQuote: PickledQuote): scala.quoted.Expr[T] =
     val tree = PickledQuotes.unpickleTerm(pickledQuote)(using reflect.rootContext)
-    new scala.internal.quoted.Expr(tree, hash)
+    new scala.internal.quoted.Expr(tree, hash).asInstanceOf[scala.quoted.Expr[T]]
 
-  def unpickleType(pickledQuote: PickledQuote): scala.quoted.Type[?] =
+  def unpickleType[T <: AnyKind](pickledQuote: PickledQuote): scala.quoted.Type[T] =
     val tree = PickledQuotes.unpickleTypeTree(pickledQuote)(using reflect.rootContext)
-    new scala.internal.quoted.Type(tree, hash)
+    new scala.internal.quoted.Type(tree, hash).asInstanceOf[scala.quoted.Type[T]]
 
   def exprMatch(scrutinee: scala.quoted.Expr[Any], pattern: scala.quoted.Expr[Any]): Option[Tuple] =
     treeMatch(scrutinee.unseal(using this), pattern.unseal(using this))

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

@@ -155,10 +155,9 @@ class ReifyQuotes extends MacroTransform {
        */
       def pickleAsLiteral(lit: Literal) = {
         val exprType = defn.QuotedExprClass.typeRef.appliedTo(body.tpe)
-        val tpe = MethodType(defn.QuoteContextClass.typeRef :: Nil, exprType)
-        val meth = newSymbol(ctx.owner, UniqueName.fresh(nme.ANON_FUN), Synthetic | Method, tpe)
-        def mkConst(tss: List[List[Tree]]) = {
-          val reflect = tss.head.head.select("reflect".toTermName)
+        val lambdaTpe = MethodType(defn.QuoteContextClass.typeRef :: Nil, exprType)
+        def mkConst(ts: List[Tree]) = {
+          val reflect = ts.head.select("reflect".toTermName)
           val typeName = body.tpe.typeSymbol.name
           val literalValue =
             if lit.const.tag == Constants.NullTag || lit.const.tag == Constants.UnitTag then Nil
@@ -167,7 +166,7 @@ class ReifyQuotes extends MacroTransform {
           val literal = reflect.select("Literal".toTermName).select(nme.apply).appliedTo(constant)
           reflect.select("TreeMethods".toTermName).select("asExpr".toTermName).appliedTo(literal).asInstance(exprType)
         }
-        Closure(meth, mkConst).withSpan(body.span)
+        Lambda(lambdaTpe, mkConst).withSpan(body.span)
       }
 
       def pickleAsValue(lit: Literal) = {
@@ -190,13 +189,30 @@ class ReifyQuotes extends MacroTransform {
         }
       }
 
+      /** Encode quote using QuoteContextInternal.{unpickleExpr, unpickleType}
+       *
+       *  Generate the code
+       *  ```scala
+       *    qctx => qctx.asInstanceOf[QuoteContextInternal].<unpickleExpr|unpickleType>[<type>](
+       *      <pickledQuote>
+       *    )
+       *  ```
+       *  this closure is always applied directly to the actual context and the BetaReduce phase removes it.
+       */
       def pickleAsTasty() = {
-        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))
         val pickledQuote = ref(defn.PickledQuote_make).appliedTo(pickledQuoteStrings, splicesList)
-        ref(unpickleMeth).appliedToType(originalTp).appliedTo(pickledQuote)
+        val quoteClass = if isType then defn.QuotedTypeClass else defn.QuotedExprClass
+        val quotedType = quoteClass.typeRef.appliedTo(originalTp)
+        val lambdaTpe = MethodType(defn.QuoteContextClass.typeRef :: Nil, quotedType)
+        def callUnpickle(ts: List[Tree]) = {
+          val qctx = ts.head.asInstance(defn.QuoteContextInternalClass.typeRef)
+          val unpickleMethName = if isType then "unpickleType" else "unpickleExpr"
+          qctx.select(unpickleMethName.toTermName).appliedToType(originalTp).appliedTo(pickledQuote)
+        }
+        Lambda(lambdaTpe, callUnpickle).withSpan(body.span)
       }
 
       /** Encode quote using Reflection.TypeRepr.typeConstructorOf
@@ -212,14 +228,13 @@ class ReifyQuotes extends MacroTransform {
       def taggedType() =
         val typeType = defn.QuotedTypeClass.typeRef.appliedTo(body.tpe)
         val classTree = TypeApply(ref(defn.Predef_classOf.termRef), body :: Nil)
-        val tpe = MethodType(defn.QuoteContextClass.typeRef :: Nil, typeType)
-        val meth = newSymbol(ctx.owner, UniqueName.fresh(nme.ANON_FUN), Synthetic | Method, tpe)
-        def mkConst(tss: List[List[Tree]]) = {
-          val reflect = tss.head.head.select("reflect".toTermName)
+        val lambdaTpe = MethodType(defn.QuoteContextClass.typeRef :: Nil, typeType)
+        def callTypeConstructorOf(ts: List[Tree]) = {
+          val reflect = ts.head.select("reflect".toTermName)
           val typeRepr = reflect.select("TypeRepr".toTermName).select("typeConstructorOf".toTermName).appliedTo(classTree)
           reflect.select("TypeReprMethods".toTermName).select("asType".toTermName).appliedTo(typeRepr).asInstance(typeType)
         }
-        Closure(meth, mkConst).withSpan(body.span)
+        Lambda(lambdaTpe, callTypeConstructorOf).withSpan(body.span)
 
       if (isType) {
         if (splices.isEmpty && body.symbol.isPrimitiveValueClass) taggedType()

library/src/scala/internal/quoted/PickledQuote.scala

@@ -16,12 +16,6 @@ trait PickledQuote:
 
 object PickledQuote:
 
-  def unpickleExpr[T](pickledQuote: PickledQuote): QuoteContext ?=> Expr[T] =
-    qctx.asInstanceOf[QuoteContextInternal].unpickleExpr(pickledQuote).asInstanceOf[Expr[T]]
-
-  def unpickleType[T](pickledQuote: PickledQuote): QuoteContext ?=> Type[T] =
-    qctx.asInstanceOf[QuoteContextInternal].unpickleType(pickledQuote).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

library/src/scala/internal/quoted/QuoteContextInternal.scala

@@ -10,12 +10,12 @@ trait QuoteContextInternal { self: scala.quoted.QuoteContext =>
   /** Unpickle `repr` which represents a pickled `Expr` tree,
    *  replacing splice nodes with `holes`
    */
-  def unpickleExpr(pickledQuote: PickledQuote): scala.quoted.Expr[Any]
+  def unpickleExpr[T](pickledQuote: PickledQuote): scala.quoted.Expr[T]
 
   /** Unpickle `repr` which represents a pickled `Type` tree,
    *  replacing splice nodes with `holes`
    */
-  def unpickleType(pickledQuote: PickledQuote): scala.quoted.Type[?]
+  def unpickleType[T <: AnyKind](pickledQuote: PickledQuote): scala.quoted.Type[T]
 
   /** Pattern matches the scrutinee against the pattern and returns a tuple
    *  with the matched holes if successful.