Move quote tagging to Staging

Fix #8302, fix #8052 and fix #7521.

Author: nicolasstucki

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

@@ -1,7 +1,12 @@
 package dotty.tools.dotc.core
 
+import dotty.tools.dotc.ast.tpd
+import dotty.tools.dotc.core.Contexts._
+import dotty.tools.dotc.core.Types._
+import dotty.tools.dotc.core.Symbols._
 import dotty.tools.dotc.core.Contexts._
 import dotty.tools.dotc.util.Property
+import dotty.tools.dotc.transform.PCPCheckAndHeal
 
 import scala.collection.mutable
 
@@ -10,6 +15,8 @@ object StagingContext {
   /** A key to be used in a context property that tracks the quoteation level */
   private val QuotationLevel = new Property.Key[Int]
 
+  private val TaggedTypes = new Property.Key[PCPCheckAndHeal.QuoteTypeTags]
+
   /** All enclosing calls that are currently inlined, from innermost to outermost. */
   def level(implicit ctx: Context): Int =
     ctx.property(QuotationLevel).getOrElse(0)
@@ -21,5 +28,12 @@ object StagingContext {
   /** Context with a decremented quotation level. */
   def spliceContext(implicit ctx: Context): Context =
     ctx.fresh.setProperty(QuotationLevel, level - 1)
+
+  def contextWithQuoteTypeTags(taggedTypes: PCPCheckAndHeal.QuoteTypeTags)(implicit ctx: Context) =
+    ctx.fresh.setProperty(TaggedTypes, taggedTypes)
+
+  def getQuoteTypeTags(implicit ctx: Context): PCPCheckAndHeal.QuoteTypeTags =
+    ctx.property(TaggedTypes).get
+
 }
 

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

@@ -52,7 +52,8 @@ object PickledQuotes {
   private def dealiasTypeTags(tp: Type)(implicit ctx: Context): Type = new TypeMap() {
     override def apply(tp: Type): Type = {
       val tp1 = tp match {
-        case tp: TypeRef if tp.typeSymbol.hasAnnotation(defn.InternalQuoted_QuoteTypeTagAnnot) => tp.dealias
+        case tp: TypeRef if tp.typeSymbol.hasAnnotation(defn.InternalQuoted_QuoteTypeTagAnnot) =>
+          tp.symbol.info.hiBound
         case _ => tp
       }
       mapOver(tp1)

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

@@ -21,6 +21,7 @@ import dotty.tools.dotc.transform.TreeMapWithStages._
 import dotty.tools.dotc.typer.Checking
 import dotty.tools.dotc.typer.Implicits.SearchFailureType
 import dotty.tools.dotc.typer.Inliner
+import dotty.tools.dotc.core.Annotations._
 
 import scala.collection.mutable
 import dotty.tools.dotc.util.SourcePosition
@@ -54,10 +55,26 @@ class PCPCheckAndHeal(@constructorOnly ictx: Context) extends TreeMapWithStages(
 
   /** Transform quoted trees while maintaining phase correctness */
   override protected def transformQuotation(body: Tree, quote: Tree)(implicit ctx: Context): Tree = {
+    val taggedTypes = new PCPCheckAndHeal.QuoteTypeTags(quote.span)(using ctx)
+    given Context =
+      if level(ctx) == 0 then contextWithQuoteTypeTags(taggedTypes)(ctx)
+      else ctx
+
     if (ctx.property(InAnnotation).isDefined)
       ctx.error("Cannot have a quote in an annotation", quote.sourcePos)
-    val body1 = transform(body)(quoteContext)
-    super.transformQuotation(body1, quote)
+
+    val body1 =
+      transform(body)(quoteContext)
+
+    val body2 =
+      if level != 0 then body1
+      else
+        val tags = taggedTypes.getTypeTags
+        if tags.isEmpty then body1
+        else tpd.Block(taggedTypes.getTypeTags, body1).withSpan(body.span)
+
+    super.transformQuotation(body2, quote)
+
   }
 
   /** Transform splice
@@ -73,7 +90,9 @@ class PCPCheckAndHeal(@constructorOnly ictx: Context) extends TreeMapWithStages(
         // internal.Quoted.expr[F[T]](... T ...)  -->  internal.Quoted.expr[F[$t]](... T ...)
         val tp = checkType(splice.sourcePos).apply(splice.tpe.widenTermRefExpr)
         cpy.Apply(splice)(cpy.TypeApply(fun)(fun.fun, tpd.TypeTree(tp) :: Nil), body1 :: Nil)
-      case splice: Select => cpy.Select(splice)(body1, splice.name)
+      case splice: Select =>
+        val tagRef = getQuoteTypeTags.getTagRef(splice.qualifier.tpe.asInstanceOf[TermRef])
+        ref(tagRef).withSpan(splice.span)
     }
   }
 
@@ -120,11 +139,14 @@ class PCPCheckAndHeal(@constructorOnly ictx: Context) extends TreeMapWithStages(
         case tp: TypeRef if tp.symbol.isSplice =>
           if (tp.isTerm)
             ctx.error(i"splice outside quotes", pos)
-          tp
+          if level > 0 then getQuoteTypeTags.getTagRef(tp.prefix.asInstanceOf[TermRef])
+          else tp
         case tp: TypeRef if tp.symbol == defn.QuotedTypeClass.typeParams.head =>
-          // Adapt direct references to the type of the type parameter T of a quoted.Type[T].
-          // Replace it with a properly encoded type splice. This is the normal for expected for type splices.
-          tp.prefix.select(tpnme.splice)
+          if level > 0 then
+            // Adapt direct references to the type of the type parameter T of a quoted.Type[T].
+            // Replace it with a properly encoded type splice. This is the normal for expected for type splices.
+            getQuoteTypeTags.getTagRef(tp.prefix.asInstanceOf[TermRef])
+          else tp
         case tp: NamedType =>
           checkSymLevel(tp.symbol, tp, pos) match {
             case Some(tpRef) => tpRef.tpe
@@ -201,6 +223,7 @@ class PCPCheckAndHeal(@constructorOnly ictx: Context) extends TreeMapWithStages(
             sym.isClass // reference to this in inline methods
           )
     case None =>
+      (sym.hasAnnotation(defn.InternalQuoted_QuoteTypeTagAnnot) && level > 0) ||
       sym.is(Package) || sym.owner.isStaticOwner || levelOK(sym.owner)
   }
 
@@ -212,10 +235,11 @@ class PCPCheckAndHeal(@constructorOnly ictx: Context) extends TreeMapWithStages(
   protected def tryHeal(sym: Symbol, tp: TypeRef, pos: SourcePosition)(implicit ctx: Context): Option[Tree] = {
     val reqType = defn.QuotedTypeClass.typeRef.appliedTo(tp)
     val tag = ctx.typer.inferImplicitArg(reqType, pos.span)
+
     tag.tpe match
       case tp: TermRef =>
         checkStable(tp, pos)
-        Some(tag.select(tpnme.splice))
+        Some(ref(getQuoteTypeTags.getTagRef(tp)))
       case _: SearchFailureType =>
         levelError(sym, tp, pos,
                     i"""
@@ -242,3 +266,34 @@ class PCPCheckAndHeal(@constructorOnly ictx: Context) extends TreeMapWithStages(
   }
 }
 
+object PCPCheckAndHeal {
+  import tpd._
+
+  class QuoteTypeTags(span: Span)(using ctx: Context) {
+
+    private val tags = collection.mutable.LinkedHashMap.empty[Symbol, TypeDef]
+
+    def getTagRef(spliced: TermRef): TypeRef = {
+      val typeDef = tags.getOrElseUpdate(spliced.symbol, mkTagSymbolAndAssignType(spliced))
+      typeDef.symbol.typeRef
+    }
+
+    def getTypeTags: List[TypeDef] = tags.valuesIterator.toList
+
+    private def mkTagSymbolAndAssignType(spliced: TermRef): TypeDef = {
+      val splicedTree = tpd.ref(spliced).withSpan(span)
+      val rhs = splicedTree.select(tpnme.splice).withSpan(span)
+      val alias = ctx.typeAssigner.assignType(untpd.TypeBoundsTree(rhs, rhs), rhs, rhs, EmptyTree)
+      val local = ctx.newSymbol(
+        owner = ctx.owner,
+        name = UniqueName.fresh((splicedTree.symbol.name.toString + "$_").toTermName).toTypeName,
+        flags = Synthetic,
+        info = TypeAlias(splicedTree.tpe.select(tpnme.splice)),
+        coord = span).asType
+      local.addAnnotation(Annotation(defn.InternalQuoted_QuoteTypeTagAnnot))
+      ctx.typeAssigner.assignType(untpd.TypeDef(local.name, alias), local)
+    }
+
+  }
+
+}

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

@@ -119,57 +119,6 @@ class ReifyQuotes extends MacroTransform {
       new QuoteReifier(this, capturers, nestedEmbedded, ctx.owner)(ctx)
     }
 
-    /** Assuming <expr> contains types `${<tag1>}, ..., ${<tagN>}`, the expression
-     *
-     *      { @quoteTypeTag type <Type1> = ${<tag1>}
-     *        ...
-     *        @quoteTypeTag type <TypeN> = ${<tagN>}
-     *        <expr>
-     *      }
-     *
-     *  references to `TypeI` in `expr` are rewired to point to the locally
-     *  defined versions. As a side effect, prepend the expressions `tag1, ..., `tagN`
-     *  as splices.
-     */
-    private def addTags(expr: Tree)(implicit ctx: Context): Tree = {
-
-      def mkTagSymbolAndAssignType(spliced: TermRef): TypeDef = {
-        val splicedTree = tpd.ref(spliced).withSpan(expr.span)
-        val rhs = transform(splicedTree.select(tpnme.splice))
-        val alias = ctx.typeAssigner.assignType(untpd.TypeBoundsTree(rhs, rhs), rhs, rhs, EmptyTree)
-        val local = ctx.newSymbol(
-          owner = ctx.owner,
-          name = UniqueName.fresh((splicedTree.symbol.name.toString + "$_").toTermName).toTypeName,
-          flags = Synthetic,
-          info = TypeAlias(splicedTree.tpe.select(tpnme.splice)),
-          coord = spliced.termSymbol.coord).asType
-        local.addAnnotation(Annotation(defn.InternalQuoted_QuoteTypeTagAnnot))
-        ctx.typeAssigner.assignType(untpd.TypeDef(local.name, alias), local)
-      }
-
-      val tagDefCache = new mutable.LinkedHashMap[Symbol, TypeDef]()
-
-      def typeTagMap = new TypeMap() {
-        def apply(tp: Type): Type = tp match {
-          case tp: TypeRef if tp.symbol.isSplice =>
-            tp.prefix match {
-              case prefix: TermRef =>
-                val tagDef = tagDefCache.getOrElseUpdate(prefix.symbol, mkTagSymbolAndAssignType(prefix))
-                tagDef.symbol.typeRef
-            }
-          case AnnotatedType(parent, _) =>
-            apply(parent) // Only keep the Annotated tree
-          case _ =>
-            mapOver(tp)
-        }
-      }
-
-      val tagedTree = new TreeTypeMap(typeMap = typeTagMap).apply(expr)
-
-      if (tagDefCache.isEmpty) expr
-      else Block(tagDefCache.valuesIterator.toList, tagedTree)
-    }
-
     /** Split `body` into a core and a list of embedded splices.
      *  Then if inside a splice, make a hole from these parts.
      *  If outside a splice, generate a call tp `scala.quoted.Unpickler.unpickleType` or
@@ -229,7 +178,7 @@ class ReifyQuotes extends MacroTransform {
       def pickleAsTasty() = {
         val meth =
           if (isType) ref(defn.Unpickler_unpickleType).appliedToType(originalTp)
-          else ref(defn.Unpickler_unpickleExpr).appliedToType(originalTp.widen)
+          else ref(defn.Unpickler_unpickleExpr).appliedToType(originalTp.widen.dealias)
         val spliceResType =
           if (isType) defn.QuotedTypeClass.typeRef.appliedTo(WildcardType)
           else defn.FunctionType(1, isContextual = true).appliedTo(defn.QuoteContextClass.typeRef, defn.QuotedExprClass.typeRef.appliedTo(defn.AnyType)) | defn.QuotedTypeClass.typeRef.appliedTo(WildcardType)
@@ -365,10 +314,10 @@ class ReifyQuotes extends MacroTransform {
       level == 1 && levelOf(sym).contains(1) && capturers.contains(sym)
 
     /** Transform `tree` and return the resulting tree and all `embedded` quotes
-     *  or splices as a pair, after performing the `addTags` transform.
+     *  or splices as a pair.
      */
     private def splitQuote(tree: Tree)(implicit ctx: Context): (Tree, List[Tree]) = {
-      val tree1 = addTags(transform(tree))
+      val tree1 = stipTypeAnnotations(transform(tree))
       (tree1, embedded.getTrees)
     }
 
@@ -377,6 +326,16 @@ class ReifyQuotes extends MacroTransform {
       (tree1, embedded.getTrees)
     }
 
+    private def stipTypeAnnotations(tree: Tree)(using Context): Tree = {
+      def typeTagMap = new TypeMap() {
+        def apply(tp: Type): Type =
+          tp match
+            case AnnotatedType(parent, _) => apply(parent) // Only keep the Annotated tree
+            case _ => mapOver(tp)
+      }
+      new TreeTypeMap(typeMap = typeTagMap).apply(tree)
+    }
+
     /** Register `body` as an `embedded` quote or splice
      *  and return a hole with `splices` as arguments and the given type `tpe`.
      */

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

@@ -48,7 +48,9 @@ class Staging extends MacroTransform {
                 if (sym.is(ModuleClass)) sym.sourceModule.show
                 else i"${sym.name}.this"
               val errMsg = s"\nin ${ctx.owner.fullName}"
-              assert(false,
+              assert(
+                ctx.owner.hasAnnotation(defn.InternalQuoted_QuoteTypeTagAnnot) ||
+                (sym.isType && levelOf(sym).getOrElse(0) > 0),
                 em"""access to $symStr from wrong staging level:
                     | - the definition is at level ${levelOf(sym).getOrElse(0)},
                     | - but the access is at level $level.$errMsg""")

compiler/src/dotty/tools/dotc/typer/QuotesAndSplices.scala

@@ -81,7 +81,7 @@ trait QuotesAndSplices {
           if (c.owner eq c.outer.owner) markAsMacro(c.outer)
           else if (c.owner.isInlineMethod) c.owner.setFlag(Macro)
           else if (!c.outer.owner.is(Package)) markAsMacro(c.outer)
-          else assert(false) // Did not find inline def to mark as macro
+          else assert(ctx.reporter.hasErrors) // Did not find inline def to mark as macro
         markAsMacro(ctx)
       }
 

compiler/test/dotc/pos-test-pickling.blacklist

@@ -11,6 +11,7 @@ t3612.scala
 reference
 scala-days-2019-slides
 i7048e.scala
+i8052.scala
 
 # Stale symbol: package object scala
 seqtype-cycle

library/src/scala/quoted/util/ExprMap.scala

@@ -64,9 +64,11 @@ trait ExprMap {
           val tp = tpt.tpe match
             // TODO improve code
             case AppliedType(TypeRef(ThisType(TypeRef(NoPrefix(), "scala")), "<repeated>"), List(tp0: Type)) =>
+              // TODO rewrite without using quotes
               type T
-              val a = tp0.seal.asInstanceOf[quoted.Type[T]]
-              '[Seq[$a]].unseal.tpe
+              val qtp: quoted.Type[T] = tp0.seal.asInstanceOf[quoted.Type[T]]
+              given qtp.type = qtp
+              '[Seq[T]].unseal.tpe
             case tp => tp
           Typed.copy(tree)(transformTerm(expr, tp), transformTypeTree(tpt))
         case tree: NamedArg =>