Place staged type captures in Quote AST

Author: nicolasstucki

compiler/src/dotty/tools/dotc/CompilationUnit.scala

@@ -157,7 +157,7 @@ object CompilationUnit {
       if tree.symbol.is(Flags.Inline) then
         containsInline = true
       tree match
-        case tpd.Quote(_) =>
+        case _: tpd.Quote =>
           containsQuote = true
         case tree: tpd.Apply if tree.symbol == defn.QuotedTypeModule_of =>
           containsQuote = true

compiler/src/dotty/tools/dotc/ast/Desugar.scala

@@ -1976,7 +1976,7 @@ object desugar {
         trees foreach collect
       case Block(Nil, expr) =>
         collect(expr)
-      case Quote(body) =>
+      case Quote(body, _) =>
         new UntypedTreeTraverser {
           def traverse(tree: untpd.Tree)(using Context): Unit = tree match {
             case Splice(expr) => collect(expr)

compiler/src/dotty/tools/dotc/ast/Trees.scala

@@ -690,9 +690,14 @@ object Trees {
    *  when type checking. TASTy files will not contain type quotes. Type quotes are used again
    *  in the `staging` phase to represent the reification of `Type.of[T]]`.
    *
+   *  Type tags `tags` are always empty before the `staging` phase. Tags for stage inconsistent
+   *  types are added in the `staging` phase to level 0 quotes. Tags for types that refer to
+   *  definitions in an outer quote are added in the `splicing` phase
+   *
    *  @param  body  The tree that was quoted
+   *  @param  tags  Term references to instances of `Type[T]` for `T`s that are used in the quote
    */
-  case class Quote[+T <: Untyped] private[ast] (body: Tree[T])(implicit @constructorOnly src: SourceFile)
+  case class Quote[+T <: Untyped] private[ast] (body: Tree[T], tags: List[Tree[T]])(implicit @constructorOnly src: SourceFile)
     extends TermTree[T] {
     type ThisTree[+T <: Untyped] = Quote[T]
 
@@ -1313,9 +1318,9 @@ object Trees {
         case tree: Inlined if (call eq tree.call) && (bindings eq tree.bindings) && (expansion eq tree.expansion) => tree
         case _ => finalize(tree, untpd.Inlined(call, bindings, expansion)(sourceFile(tree)))
       }
-      def Quote(tree: Tree)(body: Tree)(using Context): Quote = tree match {
-        case tree: Quote if (body eq tree.body) => tree
-        case _ => finalize(tree, untpd.Quote(body)(sourceFile(tree)))
+      def Quote(tree: Tree)(body: Tree, tags: List[Tree])(using Context): Quote = tree match {
+        case tree: Quote if (body eq tree.body) && (tags eq tree.tags) => tree
+        case _ => finalize(tree, untpd.Quote(body, tags)(sourceFile(tree)))
       }
       def Splice(tree: Tree)(expr: Tree)(using Context): Splice = tree match {
         case tree: Splice if (expr eq tree.expr) => tree
@@ -1558,8 +1563,8 @@ object Trees {
             case Thicket(trees) =>
               val trees1 = transform(trees)
               if (trees1 eq trees) tree else Thicket(trees1)
-            case tree @ Quote(body) =>
-              cpy.Quote(tree)(transform(body)(using quoteContext))
+            case Quote(body, tags) =>
+              cpy.Quote(tree)(transform(body)(using quoteContext), transform(tags))
             case tree @ Splice(expr) =>
               cpy.Splice(tree)(transform(expr)(using spliceContext))
             case tree @ Hole(isTerm, idx, args, content, tpt) =>
@@ -1703,8 +1708,8 @@ object Trees {
               this(this(x, arg), annot)
             case Thicket(ts) =>
               this(x, ts)
-            case Quote(body) =>
-              this(x, body)(using quoteContext)
+            case Quote(body, tags) =>
+              this(this(x, body)(using quoteContext), tags)
             case Splice(expr) =>
               this(x, expr)(using spliceContext)
             case Hole(_, _, args, content, tpt) =>

compiler/src/dotty/tools/dotc/ast/tpd.scala

@@ -170,8 +170,8 @@ object tpd extends Trees.Instance[Type] with TypedTreeInfo {
   def Inlined(call: Tree, bindings: List[MemberDef], expansion: Tree)(using Context): Inlined =
     ta.assignType(untpd.Inlined(call, bindings, expansion), bindings, expansion)
 
-  def Quote(body: Tree)(using Context): Quote =
-    untpd.Quote(body).withBodyType(body.tpe)
+  def Quote(body: Tree, tags: List[Tree])(using Context): Quote =
+    untpd.Quote(body, tags).withBodyType(body.tpe)
 
   def Splice(expr: Tree, tpe: Type)(using Context): Splice =
     untpd.Splice(expr).withType(tpe)

compiler/src/dotty/tools/dotc/ast/untpd.scala

@@ -397,7 +397,7 @@ object untpd extends Trees.Instance[Untyped] with UntypedTreeInfo {
   def SeqLiteral(elems: List[Tree], elemtpt: Tree)(implicit src: SourceFile): SeqLiteral = new SeqLiteral(elems, elemtpt)
   def JavaSeqLiteral(elems: List[Tree], elemtpt: Tree)(implicit src: SourceFile): JavaSeqLiteral = new JavaSeqLiteral(elems, elemtpt)
   def Inlined(call: tpd.Tree, bindings: List[MemberDef], expansion: Tree)(implicit src: SourceFile): Inlined = new Inlined(call, bindings, expansion)
-  def Quote(body: Tree)(implicit src: SourceFile): Quote = new Quote(body)
+  def Quote(body: Tree, tags: List[Tree])(implicit src: SourceFile): Quote = new Quote(body, tags)
   def Splice(expr: Tree)(implicit src: SourceFile): Splice = new Splice(expr)
   def TypeTree()(implicit src: SourceFile): TypeTree = new TypeTree()
   def InferredTypeTree()(implicit src: SourceFile): TypeTree = new InferredTypeTree()

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

@@ -211,6 +211,7 @@ object Phases {
     private var mySbtExtractDependenciesPhase: Phase = _
     private var myPicklerPhase: Phase = _
     private var myInliningPhase: Phase = _
+    private var myStagingPhase: Phase = _
     private var mySplicingPhase: Phase = _
     private var myFirstTransformPhase: Phase = _
     private var myCollectNullableFieldsPhase: Phase = _
@@ -235,6 +236,7 @@ object Phases {
     final def sbtExtractDependenciesPhase: Phase = mySbtExtractDependenciesPhase
     final def picklerPhase: Phase = myPicklerPhase
     final def inliningPhase: Phase = myInliningPhase
+    final def stagingPhase: Phase = myStagingPhase
     final def splicingPhase: Phase = mySplicingPhase
     final def firstTransformPhase: Phase = myFirstTransformPhase
     final def collectNullableFieldsPhase: Phase = myCollectNullableFieldsPhase
@@ -262,6 +264,7 @@ object Phases {
       mySbtExtractDependenciesPhase = phaseOfClass(classOf[sbt.ExtractDependencies])
       myPicklerPhase = phaseOfClass(classOf[Pickler])
       myInliningPhase = phaseOfClass(classOf[Inlining])
+      myStagingPhase = phaseOfClass(classOf[Staging])
       mySplicingPhase = phaseOfClass(classOf[Splicing])
       myFirstTransformPhase = phaseOfClass(classOf[FirstTransform])
       myCollectNullableFieldsPhase = phaseOfClass(classOf[CollectNullableFields])
@@ -449,6 +452,7 @@ object Phases {
   def sbtExtractDependenciesPhase(using Context): Phase = ctx.base.sbtExtractDependenciesPhase
   def picklerPhase(using Context): Phase                = ctx.base.picklerPhase
   def inliningPhase(using Context): Phase               = ctx.base.inliningPhase
+  def stagingPhase(using Context): Phase               = ctx.base.stagingPhase
   def splicingPhase(using Context): Phase               = ctx.base.splicingPhase
   def firstTransformPhase(using Context): Phase         = ctx.base.firstTransformPhase
   def refchecksPhase(using Context): Phase              = ctx.base.refchecksPhase

compiler/src/dotty/tools/dotc/core/tasty/TreePickler.scala

@@ -665,7 +665,7 @@ class TreePickler(pickler: TastyPickler) {
               pickleTree(hi)
               pickleTree(alias)
           }
-        case tree @ Quote(body) =>
+        case tree @ Quote(body, Nil) =>
           // TODO: Add QUOTE tag to TASTy
           assert(body.isTerm,
             """Quote with type should not be pickled.

compiler/src/dotty/tools/dotc/core/tasty/TreeUnpickler.scala

@@ -1269,7 +1269,7 @@ class TreeUnpickler(reader: TastyReader,
 
       def quotedExpr(fn: Tree, args: List[Tree]): Tree =
         val TypeApply(_, targs) = fn: @unchecked
-        untpd.Quote(args.head).withBodyType(targs.head.tpe)
+        untpd.Quote(args.head, Nil).withBodyType(targs.head.tpe)
 
       def splicedExpr(fn: Tree, args: List[Tree]): Tree =
         val TypeApply(_, targs) = fn: @unchecked

compiler/src/dotty/tools/dotc/inlines/Inliner.scala

@@ -827,7 +827,7 @@ class Inliner(val call: tpd.Tree)(using Context):
 
     override def typedQuote(tree: untpd.Quote, pt: Type)(using Context): Tree =
       super.typedQuote(tree, pt) match
-        case Quote(Splice(inner)) => inner
+        case Quote(Splice(inner), _) => inner
         case tree1 =>
           ctx.compilationUnit.needsStaging = true
           tree1

compiler/src/dotty/tools/dotc/parsing/Parsers.scala

@@ -1243,7 +1243,7 @@ object Parsers {
               }
             }
             in.nextToken()
-            Quote(t)
+            Quote(t, Nil)
           }
           else
             if !in.featureEnabled(Feature.symbolLiterals) then
@@ -2480,7 +2480,7 @@ object Parsers {
               val body =
                 if (in.token == LBRACKET) inBrackets(typ())
                 else stagedBlock()
-              Quote(body)
+              Quote(body, Nil)
             }
           }
         case NEW =>

compiler/src/dotty/tools/dotc/printing/RefinedPrinter.scala

@@ -726,11 +726,12 @@ class RefinedPrinter(_ctx: Context) extends PlainPrinter(_ctx) {
         "Thicket {" ~~ toTextGlobal(trees, "\n") ~~ "}"
       case MacroTree(call) =>
         keywordStr("macro ") ~ toTextGlobal(call)
-      case tree @ Quote(body) =>
+      case tree @ Quote(body, tags) =>
+        val tagsText = (keywordStr("<") ~ toTextGlobal(tags, ", ") ~ keywordStr(">")).provided(tree.tags.nonEmpty)
         val exprTypeText = (keywordStr("[") ~ toTextGlobal(tree.bodyType) ~ keywordStr("]")).provided(printDebug && tree.typeOpt.exists)
         val open = if (body.isTerm) keywordStr("{") else keywordStr("[")
         val close = if (body.isTerm) keywordStr("}") else keywordStr("]")
-        keywordStr("'") ~ exprTypeText ~ open ~ toTextGlobal(body) ~ close
+        keywordStr("'") ~ tagsText ~ exprTypeText ~ open ~ toTextGlobal(body) ~ close
       case Splice(expr) =>
         val spliceTypeText = (keywordStr("[") ~ toTextGlobal(tree.typeOpt) ~ keywordStr("]")).provided(printDebug && tree.typeOpt.exists)
         keywordStr("$") ~ spliceTypeText ~ keywordStr("{") ~ toTextGlobal(expr) ~ keywordStr("}")

compiler/src/dotty/tools/dotc/staging/CrossStageSafety.scala

@@ -16,28 +16,36 @@ import dotty.tools.dotc.util.Property
 import dotty.tools.dotc.util.Spans._
 import dotty.tools.dotc.util.SrcPos
 
-/** Checks that staging level consistency holds and heals staged types .
+/** Checks that staging level consistency holds and heals staged types.
  *
  *  Local term references are level consistent if and only if they are used at the same level as their definition.
  *
  *  Local type references can be used at the level of their definition or lower. If used used at a higher level,
  *  it will be healed if possible, otherwise it is inconsistent.
  *
- *  Type healing consists in transforming a level inconsistent type `T` into `summon[Type[T]].Underlying`.
+ *  Healing a type consists in replacing locally defined types defined at staging level 0 and used in higher levels.
+ *  For each type local `T` that is defined at level 0 and used in a quote, we summon a tag `t: Type[T]`. This `t`
+ *  tag must be defined at level 0. The tags will be listed in the `tags` of the level 0 quote (`'<t>{ ... }`) and
+ *  each reference to `T` will be replaced by `t.Underlying` in the body of the quote.
+ *
+ *  We delay the healing of types in quotes at level 1 or higher until those quotes reach level 0. At this point
+ *  more types will be statically known and fewer types will need to be healed. This also keeps the nested quotes
+ *  in their original form, we do not want macro users to see any artifacts of this phase in quoted expressions
+ *  they might inspect.
+ *
+ *  Type heal example:
  *
- *  As references to types do not necessarily have an associated tree it is not always possible to replace the types directly.
- *  Instead we always generate a type alias for it and place it at the start of the surrounding quote. This also avoids duplication.
- *  For example:
  *    '{
  *      val x: List[T] = List[T]()
+ *      '{ .. T .. }
  *      ()
  *    }
  *
  *  is transformed to
  *
- *    '{
- *      type t$1 = summon[Type[T]].Underlying
- *      val x: List[t$1] = List[t$1]();
+ *    '<t>{ // where `t` is a given term of type `Type[T]`
+ *      val x: List[t.Underlying] = List[t.Underlying]();
+ *      '{ .. t.Underlying .. }
  *      ()
  *     }
  *
@@ -56,11 +64,18 @@ class CrossStageSafety extends TreeMapWithStages {
       case tree: Quote =>
         if (ctx.property(InAnnotation).isDefined)
           report.error("Cannot have a quote in an annotation", tree.srcPos)
-        val body1 = transformQuoteBody(tree.body, tree.span)
-        val stripAnnotationsDeep: TypeMap = new TypeMap:
-          def apply(tp: Type): Type = mapOver(tp.stripAnnots)
-        val bodyType1 = healType(tree.srcPos)(stripAnnotationsDeep(tree.bodyType))
-        cpy.Quote(tree)(body1).withBodyType(bodyType1)
+
+        val tree1 =
+          val stripAnnotationsDeep: TypeMap = new TypeMap:
+            def apply(tp: Type): Type = mapOver(tp.stripAnnots)
+          val bodyType1 = healType(tree.srcPos)(stripAnnotationsDeep(tree.bodyType))
+          tree.withBodyType(bodyType1)
+
+        if level == 0 then
+          val (tags, body1) = inContextWithQuoteTypeTags { transform(tree1.body)(using quoteContext) }
+          cpy.Quote(tree1)(body1, tags)
+        else
+          super.transform(tree1)
 
       case CancelledSplice(tree) =>
         transform(tree) // Optimization: `${ 'x }` --> `x`
@@ -74,22 +89,18 @@ class CrossStageSafety extends TreeMapWithStages {
       case tree @ QuotedTypeOf(body) =>
         if (ctx.property(InAnnotation).isDefined)
           report.error("Cannot have a quote in an annotation", tree.srcPos)
-        body.tpe match
-          case DirectTypeOf(termRef) =>
-            // Optimization: `quoted.Type.of[x.Underlying](quotes)`  -->  `x`
-            ref(termRef).withSpan(tree.span)
-          case _ =>
-            transformQuoteBody(body, tree.span) match
-              case DirectTypeOf.Healed(termRef) =>
-                // Optimization: `quoted.Type.of[@SplicedType type T = x.Underlying; T](quotes)`  -->  `x`
-                ref(termRef).withSpan(tree.span)
-              case transformedBody =>
-                val quotes = transform(tree.args.head)
-                // `quoted.Type.of[<body>](quotes)`  --> `quoted.Type.of[<body2>](quotes)`
-                val TypeApply(fun, _) = tree.fun: @unchecked
-                if level != 0 then cpy.Apply(tree)(cpy.TypeApply(tree.fun)(fun, transformedBody :: Nil), quotes :: Nil)
-                else tpd.Quote(transformedBody).select(nme.apply).appliedTo(quotes).withSpan(tree.span)
 
+        if level == 0 then
+          val (tags, body1) = inContextWithQuoteTypeTags { transform(body)(using quoteContext) }
+          val quotes = transform(tree.args.head)
+          tags match
+            case tag :: Nil if body1.isType && body1.tpe =:= tag.tpe.select(tpnme.Underlying) =>
+              tag // Optimization: `quoted.Type.of[x.Underlying](quotes)`  -->  `x`
+            case _ =>
+              // `quoted.Type.of[<body>](<quotes>)` --> `'[<body1>].apply(<quotes>)`
+              tpd.Quote(body1, tags).select(nme.apply).appliedTo(quotes).withSpan(tree.span)
+        else
+          super.transform(tree)
       case _: DefDef if tree.symbol.isInlineMethod =>
         tree
 
@@ -137,17 +148,6 @@ class CrossStageSafety extends TreeMapWithStages {
         super.transform(tree)
   end transform
 
-  private def transformQuoteBody(body: Tree, span: Span)(using Context): Tree = {
-    val taggedTypes = new QuoteTypeTags(span)
-    val contextWithQuote =
-      if level == 0 then contextWithQuoteTypeTags(taggedTypes)(using quoteContext)
-      else quoteContext
-    val transformedBody = transform(body)(using contextWithQuote)
-    taggedTypes.getTypeTags match
-      case Nil  => transformedBody
-      case tags => tpd.Block(tags, transformedBody).withSpan(body.span)
-  }
-
   def transformTypeAnnotationSplices(tp: Type)(using Context) = new TypeMap {
     def apply(tp: Type): Type = tp match
       case tp: AnnotatedType =>
@@ -234,7 +234,7 @@ class CrossStageSafety extends TreeMapWithStages {
     def unapply(tree: Splice): Option[Tree] =
       def rec(tree: Tree): Option[Tree] = tree match
         case Block(Nil, expr) => rec(expr)
-        case Quote(inner) => Some(inner)
+        case Quote(inner, _) => Some(inner)
         case _ => None
       rec(tree.expr)
 }

compiler/src/dotty/tools/dotc/staging/DirectTypeOf.scala

@@ -25,7 +25,7 @@ class HealType(pos: SrcPos)(using Context) extends TypeMap {
    *
    *  If `T` is a reference to a type at the wrong level, try to heal it by replacing it with
    *  a type tag of type `quoted.Type[T]`.
-   *  The tag is generated by an instance of `QuoteTypeTags` directly if the splice is explicit
+   *  The tag is recorded by an instance of `QuoteTypeTags` directly if the splice is explicit
    *  or indirectly by `tryHeal`.
    */
   def apply(tp: Type): Type =
@@ -43,11 +43,9 @@ class HealType(pos: SrcPos)(using Context) extends TypeMap {
 
   private def healTypeRef(tp: TypeRef): Type =
     tp.prefix match
-      case NoPrefix if tp.typeSymbol.hasAnnotation(defn.QuotedRuntime_SplicedTypeAnnot) =>
-        tp
       case prefix: TermRef if tp.symbol.isTypeSplice =>
         checkNotWildcardSplice(tp)
-        if level == 0 then tp else getQuoteTypeTags.getTagRef(prefix)
+        if level == 0 then tp else getTagRef(prefix)
       case _: NamedType | _: ThisType | NoPrefix =>
         if levelInconsistentRootOfPath(tp).exists then
           tryHeal(tp)
@@ -58,7 +56,7 @@ class HealType(pos: SrcPos)(using Context) extends TypeMap {
 
   private object NonSpliceAlias:
     def unapply(tp: TypeRef)(using Context): Option[Type] = tp.underlying match
-      case TypeAlias(alias) if !tp.symbol.isTypeSplice && !tp.typeSymbol.hasAnnotation(defn.QuotedRuntime_SplicedTypeAnnot) => Some(alias)
+      case TypeAlias(alias) if !tp.symbol.isTypeSplice => Some(alias)
       case _ => None
 
   private def checkNotWildcardSplice(splice: TypeRef): Unit =
@@ -78,7 +76,7 @@ class HealType(pos: SrcPos)(using Context) extends TypeMap {
 
   /** Try to heal reference to type `T` used in a higher level than its definition.
    *  Returns a reference to a type tag generated by `QuoteTypeTags` that contains a
-   *  reference to a type alias containing the equivalent of `${summon[quoted.Type[T]]}`.
+   *  reference to a type alias containing the equivalent of `${summon[quoted.Type[T]]}.Underlying`.
    *  Emits an error if `T` cannot be healed and returns `T`.
    */
   protected def tryHeal(tp: TypeRef): Type = {
@@ -88,7 +86,7 @@ class HealType(pos: SrcPos)(using Context) extends TypeMap {
       case tp: TermRef =>
         ctx.typer.checkStable(tp, pos, "type witness")
         if levelOf(tp.symbol) > 0 then tp.select(tpnme.Underlying)
-        else getQuoteTypeTags.getTagRef(tp)
+        else getTagRef(tp)
       case _: SearchFailureType =>
         report.error(
           ctx.typer.missingArgMsg(tag, reqType, "")