Move quote tagging to Staging

Fix #8302, fix #8052 and fix #7521

Author: nicolasstucki

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

@@ -1,8 +1,13 @@
 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.ast.tpd
 import dotty.tools.dotc.util.Property
+import dotty.tools.dotc.transform.PCPCheckAndHeal
 
 import scala.collection.mutable
 
@@ -16,6 +21,8 @@ object StagingContext {
    */
   private val QuoteContextStack = new Property.Key[List[tpd.Tree]]
 
+  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)
@@ -34,6 +41,12 @@ object StagingContext {
   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
+
   /** Context with a decremented quotation level and pops the Some of top of the quote context stack or None if the stack is empty.
    *  The quotation stack could be empty if we are in a top level splice or an eroneous splice directly witin a top level splice.
    */

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,21 @@ 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)
+
     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 contextWithQuote =
+      if level == 0 then contextWithQuoteTypeTags(taggedTypes)(quoteContext)
+      else quoteContext
+    val body1 = transform(body)(contextWithQuote)
+    val body2 =
+      taggedTypes.getTypeTags match
+        case Nil  => body1
+        case tags => tpd.Block(tags, body1).withSpan(body.span)
+
+    super.transformQuotation(body2, quote)
   }
 
   /** Transform splice
@@ -73,7 +85,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) :: qctx :: 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 +134,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 form 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,7 +218,9 @@ class PCPCheckAndHeal(@constructorOnly ictx: Context) extends TreeMapWithStages(
             sym.isClass // reference to this in inline methods
           )
     case None =>
-      sym.is(Package) || sym.owner.isStaticOwner || levelOK(sym.owner)
+      sym.is(Package) || sym.owner.isStaticOwner ||
+      (sym.hasAnnotation(defn.InternalQuoted_QuoteTypeTagAnnot) && level > 0) ||
+      levelOK(sym.owner)
   }
 
   /** Try to heal reference to type `T` used in a higher level than its definition.
@@ -212,10 +231,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 +262,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 pickledQuoteStrings = liftList(PickledQuotes.pickleQuote(body).map(x => Literal(Constant(x))), defn.StringType)
         val splicesList = liftList(splices, defn.FunctionType(1).appliedTo(defn.SeqType.appliedTo(defn.AnyType), defn.AnyType))
         meth.appliedTo(pickledQuoteStrings, splicesList)
@@ -362,10 +311,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)
     }
 
@@ -374,6 +323,9 @@ class ReifyQuotes extends MacroTransform {
       (tree1, embedded.getTrees)
     }
 
+    private def stipTypeAnnotations(tree: Tree)(using Context): Tree =
+      new TreeTypeMap(typeMap = _.stripAnnots).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

@@ -87,7 +87,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 =>

tests/run-staging/quote-macro-in-splice/quoted_1.scala renamed to tests/disabled/run-staging/quote-macro-in-splice/quoted_1.scala

@@ -9,7 +9,7 @@ abstract class Test {
 
   def foo(using QuoteContext): Expr[Any] = {
 
-    val r = '{Option.empty[T]} // error
+    val r = '{Option.empty[T]} // error: is not stable
 
     {
       val t: Test = this
@@ -22,14 +22,14 @@ abstract class Test {
     {
       val r1 = '{Option.empty[${T}]} // works
       val r2 = '{Option.empty[List[${T}]]} // works
-      // val r3 = '{summon[Type[${T}]]} // access to Test.this from wrong staging level
-      val r4 = '{summon[${T} <:< Any]} // error
+      val r3 = '{summon[Type[${T}]]} // error: is not stable
+      val r4 = '{summon[${T} <:< Any]} // error: is not stable
     }
 
     {
-      val s = '{Option.empty[${T}]}
+      val s = '{Option.empty[${T}]} // works
       val r = '{identity($s)} // works
-      val r2 = '{identity(${s: Expr[Option[T]]})} // error // error
+      val r2 = '{identity(${s: Expr[Option[T]]})} // error // error : is not stable
     }
 
     r

tests/run-staging/quote-macro-in-splice/quoted_2.scala renamed to tests/disabled/run-staging/quote-macro-in-splice/quoted_2.scala

@@ -0,0 +1,19 @@
+import scala.deriving._
+import scala.quoted._
+import scala.quoted.matching._
+
+object Macro2 {
+  trait TC[T] {
+    def test(): Unit
+  }
+
+  object TC {
+    def derived[T: Type](ev: Expr[Mirror.Of[T]])(using qctx: QuoteContext): Expr[TC[T]] = '{
+      new TC[T] {
+        def encode(): Unit = $ev match {
+          case '{ $m: Mirror.ProductOf[T] } => ??? // error
+        }
+      }
+    }
+  }
+}

tests/neg/i7048e.scala

@@ -0,0 +1,10 @@
+import scala.quoted._
+import scala.annotation.StaticAnnotation
+
+object Test {
+  inline def quote[T]: Unit = ${ quoteImpl[T] }
+  def quoteImpl[T: Type](using qctx: QuoteContext): Expr[Unit] = '{
+     class Annot extends StaticAnnotation
+     var test: T @Annot = ???
+  }
+}

tests/neg/i8052.scala

@@ -0,0 +1,21 @@
+import scala.deriving._
+import scala.quoted._
+import scala.quoted.matching._
+
+object Macro2 {
+  trait TC[T] {
+    def test(): Unit
+  }
+
+  object TC {
+    def derived[T: Type](ev: Expr[Mirror.Of[T]])(using qctx: QuoteContext): Expr[TC[T]] = '{
+      new TC[T] {
+        def encode(): Unit = ${
+          ev match {
+            case '{ $m: Mirror.ProductOf[T] } => ???
+          }
+        }
+      }
+    }
+  }
+}

tests/pos/i7521.scala

@@ -0,0 +1,8 @@
+import scala.quoted._
+def foo[T](using qctx: QuoteContext, tpe: Type[T]): Expr[Any] =
+  '{
+    type TT = T
+    val t = '[TT]
+    ???
+  }
+

tests/pos/i8052.scala

@@ -0,0 +1,6 @@
+import scala.quoted._
+
+class Test {
+  def fold[W: Type](s: Expr[W]): QuoteContext ?=> Expr[W] =
+    '{ ???; $s }
+}

tests/pos/i8302.scala

@@ -1,2 +1,2 @@
 null.asInstanceOf[java.lang.Object]
-null.asInstanceOf[scala.List[java.lang.Object]]
+null.asInstanceOf[scala.collection.immutable.List[java.lang.Object]]