Merge pull request #6963 from dotty-staging/fix-#5376

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): PickledQuote = {
     if (ctx.reporter.hasErrors) Nil
     else {
       assert(!tree.isInstanceOf[Hole]) // Should not be pickled as it represents `'{$x}` which should be optimized to `x`
@@ -33,33 +35,12 @@ 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)
-  }
+  def quotedExprToTree[T](expr: quoted.Expr[T])(implicit ctx: Context): Tree =
+    healOwner(expr.asInstanceOf[TastyTreeExpr[Tree]].tree)
 
   /** 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 = {
@@ -72,15 +53,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: PickledQuote, 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: PickledQuote, 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,25 +1,25 @@
 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
 
+import scala.runtime.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 an Seq of String */
-  def pickle(bytes: Array[Byte]): Pickled = {
+  /** Encode TASTY bytes into a List of String */
+  def pickle(bytes: Array[Byte]): PickledQuote = {
     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: PickledQuote): 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,11 @@
+package dotty.tools.dotc.quoted
+
+import dotty.tools.dotc.core.Contexts.Context
+import dotty.tools.dotc.tastyreflect.ReflectionImpl
+
+object QuoteContext {
+
+  def apply() given Context: scala.quoted.QuoteContext =
+    new scala.quoted.QuoteContext(ReflectionImpl(the[Context]))
+
+}

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

@@ -16,6 +16,7 @@ import dotty.tools.dotc.parsing.Parsers.Parser
 import dotty.tools.dotc.typer.Implicits.{AmbiguousImplicits, DivergingImplicit, NoMatchingImplicits, SearchFailure, SearchFailureType}
 import dotty.tools.dotc.util.{SourceFile, SourcePosition, Spans}
 
+import scala.runtime.quoted.Unpickler
 import scala.tasty.reflect.CompilerInterface
 
 class ReflectionCompilerInterface(val rootContext: core.Contexts.Context) extends CompilerInterface {
@@ -28,6 +29,16 @@ class ReflectionCompilerInterface(val rootContext: core.Contexts.Context) extend
   def rootPosition: util.SourcePosition =
     tastyreflect.MacroExpansion.position.getOrElse(SourcePosition(rootContext.source, Spans.NoSpan))
 
+  //
+  // QUOTE UNPICKLING
+  //
+
+  def unpickleExpr(repr: Unpickler.PickledQuote, args: Unpickler.PickledExprArgs): scala.quoted.Expr[_] =
+    new scala.internal.quoted.TastyTreeExpr(PickledQuotes.unpickleExpr(repr, args))
+
+  def unpickleType(repr: Unpickler.PickledQuote, args: Unpickler.PickledTypeArgs): scala.quoted.Type[_] =
+    new scala.internal.quoted.TreeType(PickledQuotes.unpickleType(repr, args))
+
   //
   // 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/dotty/internal/StringContextMacro.scala

@@ -66,7 +66,7 @@ object StringContextMacro {
     }
     def splitParts(seq: Expr[Seq[String]]) = (seq, seq) match {
       case (ExprSeq(p1), ConstSeq(p2)) => Some((p1.toList, p2.toList))
-      case _ => notStatic
+      case (_, _) => notStatic
     }
     strCtxExpr match {
       case '{ StringContext($parts: _*) } => splitParts(parts)

library/src/scala/quoted/Expr.scala

@@ -79,12 +79,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/Type.scala

@@ -68,16 +68,9 @@ package quoted {
 
 package internal {
   package quoted {
-    import scala.quoted.Type
-    import scala.runtime.quoted.Unpickler.Pickled
-
-    /** A Type backed by a pickled TASTY tree */
-    final class TastyType[T](val tasty: Pickled, val args: Seq[Any]) extends Type[T] {
-      override def toString(): String = s"Type(<pickled tasty>)"
-    }
 
     /** An Type backed by a tree */
-    final class TreeType[Tree](val typeTree: Tree) extends Type[Any] {
+    final class TreeType[Tree](val typeTree: Tree) extends scala.quoted.Type[Any] {
       override def toString: String = s"Type(<tasty tree>)"
     }
 

library/src/scala/runtime/quoted/Unpickler.scala

@@ -1,24 +1,25 @@
-package scala.runtime.quoted
+package scala.runtime.quoted // TODO move to scala.internal.quoted
 
-import scala.internal.quoted.{TastyExpr, TastyType}
 import scala.quoted.{Expr, QuoteContext, Type}
 
 /** Provides methods to unpickle `Expr` and `Type` trees. */
 object Unpickler {
 
-  /** Representation of pickled trees. For now a List[String],
-   *  but it should be changed to some kind of TASTY bundle.
-   */
-  type Pickled = List[String]
+  type PickledQuote = List[String]
+  type PickledExprArgs = Seq[Seq[Any] => ((given QuoteContext => Expr[Any]) | Type[_])]
+  type PickledTypeArgs = Seq[Seq[Any] => Type[_]]
 
   /** Unpickle `repr` which represents a pickled `Expr` tree,
    *  replacing splice nodes with `args`
    */
-  def unpickleExpr[T](repr: Pickled, args: Seq[Seq[Any] => ((given QuoteContext => Expr[Any]) | Type[_])]): given QuoteContext => Expr[T] = new TastyExpr[T](repr, args)
+  def unpickleExpr[T](repr: PickledQuote, args: PickledExprArgs): given QuoteContext => Expr[T] =
+    the[QuoteContext].tasty.internal.unpickleExpr(repr, args).asInstanceOf[Expr[T]]
 
   /** Unpickle `repr` which represents a pickled `Type` tree,
    *  replacing splice nodes with `args`
    */
-  def unpickleType[T](repr: Pickled, args: Seq[Seq[Any] => Type[_]]): given QuoteContext => Type[T] = new TastyType[T](repr, args)
+  def unpickleType[T](repr: PickledQuote, args: PickledTypeArgs): given QuoteContext => Type[T] =
+    the[QuoteContext].tasty.internal.unpickleType(repr, args).asInstanceOf[Type[T]]
+
 
 }

library/src/scala/tasty/Reflection.scala

@@ -1,5 +1,6 @@
 package scala.tasty
 
+import scala.quoted.QuoteContext
 import scala.tasty.reflect._
 
 class Reflection(private[scala] val internal: CompilerInterface)

library/src/scala/tasty/reflect/Internal.scala renamed to library/src/scala/tasty/reflect/CompilerInterface.scala

@@ -1,6 +1,8 @@
-package scala.tasty.reflect
+package scala.tasty.reflect // TODO move to scala.internal.tasty.reflect
 
 import scala.quoted.QuoteContext
+import scala.tasty.Reflection
+import scala.runtime.quoted.Unpickler
 
 /** Tasty reflect abstract types
  *
@@ -129,6 +131,20 @@ trait CompilerInterface {
 
   def settings: Settings
 
+  //
+  // QUOTE UNPICKLING
+  //
+
+  /** Unpickle `repr` which represents a pickled `Expr` tree,
+   *  replacing splice nodes with `args`
+   */
+  def unpickleExpr(repr: Unpickler.PickledQuote, args: Unpickler.PickledExprArgs): scala.quoted.Expr[_]
+
+  /** Unpickle `repr` which represents a pickled `Type` tree,
+   *  replacing splice nodes with `args`
+   */
+  def unpickleType(repr: Unpickler.PickledQuote, args: Unpickler.PickledTypeArgs): scala.quoted.Type[_]
+
   //
   // CONTEXT
   //

tests/patmat/i6255b.check

@@ -1 +0,0 @@
-2: Pattern Match Exhaustivity: _: Expr[Int]

tests/run-macros/quote-matching-optimize-1.check

@@ -18,15 +18,15 @@ Result: ()
 
 Original: scala.List.apply[scala.Int]((1, 2, 3: scala.<repeated>[scala.Int])).map[scala.Int, scala.collection.immutable.List[scala.Int]](((a: scala.Int) => a.*(2)))(scala.collection.immutable.List.canBuildFrom[scala.Int]).map[java.lang.String, scala.collection.immutable.List[java.lang.String]](((b: scala.Int) => b.toString()))(scala.collection.immutable.List.canBuildFrom[java.lang.String])
 Optimized: scala.List.apply[scala.Int]((1, 2, 3: scala.<repeated>[scala.Int])).map[java.lang.String, scala.collection.immutable.List[java.lang.String]](((x: scala.Int) => {
-  val x$5: scala.Int = x.*(2)
-  x$5.toString()
+  val x$1: scala.Int = x.*(2)
+  x$1.toString()
 }))(scala.collection.immutable.List.canBuildFrom[java.lang.String])
 Result: List(2, 4, 6)
 
 Original: scala.List.apply[scala.Int]((55, 67, 87: scala.<repeated>[scala.Int])).map[scala.Char, scala.collection.immutable.List[scala.Char]](((a: scala.Int) => a.toChar))(scala.collection.immutable.List.canBuildFrom[scala.Char]).map[java.lang.String, scala.collection.immutable.List[java.lang.String]](((b: scala.Char) => b.toString()))(scala.collection.immutable.List.canBuildFrom[java.lang.String])
 Optimized: scala.List.apply[scala.Int]((55, 67, 87: scala.<repeated>[scala.Int])).map[java.lang.String, scala.collection.immutable.List[java.lang.String]](((x: scala.Int) => {
-  val x$10: scala.Char = x.toChar
-  x$10.toString()
+  val x$2: scala.Char = x.toChar
+  x$2.toString()
 }))(scala.collection.immutable.List.canBuildFrom[java.lang.String])
 Result: List(7, C, W)
 

tests/run-macros/quote-matching-optimize-2.check

@@ -18,15 +18,15 @@ Result: ()
 
 Original: ls.map[scala.Int, scala.collection.immutable.List[scala.Int]](((a: scala.Int) => a.*(2)))(scala.collection.immutable.List.canBuildFrom[scala.Int]).map[java.lang.String, scala.collection.immutable.List[java.lang.String]](((b: scala.Int) => b.toString()))(scala.collection.immutable.List.canBuildFrom[java.lang.String])
 Optimized: ls.map[java.lang.String, scala.collection.immutable.List[java.lang.String]](((x: scala.Int) => {
-  val x$5: scala.Int = x.*(2)
-  x$5.toString()
+  val x$1: scala.Int = x.*(2)
+  x$1.toString()
 }))(scala.collection.immutable.List.canBuildFrom[java.lang.String])
 Result: List(2, 4, 6)
 
 Original: ls.map[scala.Char, scala.collection.immutable.List[scala.Char]](((a: scala.Int) => a.toChar))(scala.collection.immutable.List.canBuildFrom[scala.Char]).map[java.lang.String, scala.collection.immutable.List[java.lang.String]](((b: scala.Char) => b.toString()))(scala.collection.immutable.List.canBuildFrom[java.lang.String])
 Optimized: ls.map[java.lang.String, scala.collection.immutable.List[java.lang.String]](((x: scala.Int) => {
-  val x$10: scala.Char = x.toChar
-  x$10.toString()
+  val x$2: scala.Char = x.toChar
+  x$2.toString()
 }))(scala.collection.immutable.List.canBuildFrom[java.lang.String])
 Result: List(
, , )
 

tests/run-with-compiler/i5376.check

@@ -0,0 +1 @@
+1.+(1).+(2).+(3)

tests/run-with-compiler/i5376.scala

@@ -0,0 +1,13 @@
+import scala.quoted._
+
+object Test {
+  implicit val toolbox: scala.quoted.Toolbox = scala.quoted.Toolbox.make(getClass.getClassLoader)
+
+  def main(args: Array[String]): Unit = withQuoteContext {
+    var e = '{1}
+    e = '{$e + 1}
+    e = '{$e + 2}
+    e = '{$e + 3}
+    println(e.show)
+  }
+}

tests/run-with-compiler/i5434.check

@@ -0,0 +1,5 @@
+start
+start e
+splice
+end e
+end