Use polyfunction for quote hole contents

This avoids leaking references to types defined within the quotes when the
contents of the holes are extracted from the quote.

Author: nicolasstucki

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

@@ -89,7 +89,7 @@ class Compiler {
          new ExplicitOuter,          // Add accessors to outer classes from nested ones.
          new ExplicitSelf,           // Make references to non-trivial self types explicit as casts
          new StringInterpolatorOpt,  // Optimizes raw and s and f string interpolators by rewriting them to string concatenations or formats
-         new DropBreaks) ::          // Optimize local Break throws by rewriting them 
+         new DropBreaks) ::          // Optimize local Break throws by rewriting them
     List(new PruneErasedDefs,        // Drop erased definitions from scopes and simplify erased expressions
          new UninitializedDefs,      // Replaces `compiletime.uninitialized` by `_`
          new InlinePatterns,         // Remove placeholders of inlined patterns

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

@@ -1374,54 +1374,57 @@ object Types {
             case Atoms.Unknown => Atoms.Unknown
         case _ => Atoms.Unknown
 
-    private def dealias1(keep: AnnotatedType => Context ?=> Boolean, keepOpaques: Boolean)(using Context): Type = this match {
+    private def dealias1(keep: AnnotatedType => Context ?=> Boolean, keepOpaques: Boolean, keepSplicedTypes: Boolean)(using Context): Type = this match {
       case tp: TypeRef =>
         if (tp.symbol.isClass) tp
         else tp.info match {
-          case TypeAlias(alias) if !(keepOpaques && tp.symbol.is(Opaque)) =>
-            alias.dealias1(keep, keepOpaques)
+          case TypeAlias(alias) if !(keepOpaques && tp.symbol.is(Opaque)) && !(keepSplicedTypes && tp.symbol.hasAnnotation(defn.QuotedRuntime_SplicedTypeAnnot)) =>
+            alias.dealias1(keep, keepOpaques, keepSplicedTypes)
           case _ => tp
         }
       case app @ AppliedType(tycon, _) =>
-        val tycon1 = tycon.dealias1(keep, keepOpaques)
-        if (tycon1 ne tycon) app.superType.dealias1(keep, keepOpaques)
+        val tycon1 = tycon.dealias1(keep, keepOpaques, keepSplicedTypes)
+        if (tycon1 ne tycon) app.superType.dealias1(keep, keepOpaques, keepSplicedTypes)
         else this
       case tp: TypeVar =>
         val tp1 = tp.instanceOpt
-        if (tp1.exists) tp1.dealias1(keep, keepOpaques) else tp
+        if (tp1.exists) tp1.dealias1(keep, keepOpaques, keepSplicedTypes) else tp
       case tp: AnnotatedType =>
-        val parent1 = tp.parent.dealias1(keep, keepOpaques)
+        val parent1 = tp.parent.dealias1(keep, keepOpaques, keepSplicedTypes)
         tp match
           case tp @ CapturingType(parent, refs) =>
             tp.derivedCapturingType(parent1, refs)
           case _ =>
             if keep(tp) then tp.derivedAnnotatedType(parent1, tp.annot)
             else parent1
       case tp: LazyRef =>
-        tp.ref.dealias1(keep, keepOpaques)
+        tp.ref.dealias1(keep, keepOpaques, keepSplicedTypes)
       case _ => this
     }
 
     /** Follow aliases and dereferences LazyRefs, annotated types and instantiated
      *  TypeVars until type is no longer alias type, annotated type, LazyRef,
      *  or instantiated type variable.
      */
-    final def dealias(using Context): Type = dealias1(keepNever, keepOpaques = false)
+    final def dealias(using Context): Type = dealias1(keepNever, keepOpaques = false, keepSplicedTypes = false)
 
     /** Follow aliases and dereferences LazyRefs and instantiated TypeVars until type
      *  is no longer alias type, LazyRef, or instantiated type variable.
      *  Goes through annotated types and rewraps annotations on the result.
      */
-    final def dealiasKeepAnnots(using Context): Type = dealias1(keepAlways, keepOpaques = false)
+    final def dealiasKeepAnnots(using Context): Type = dealias1(keepAlways, keepOpaques = false, keepSplicedTypes = false)
 
     /** Like `dealiasKeepAnnots`, but keeps only refining annotations */
-    final def dealiasKeepRefiningAnnots(using Context): Type = dealias1(keepIfRefining, keepOpaques = false)
+    final def dealiasKeepRefiningAnnots(using Context): Type = dealias1(keepIfRefining, keepOpaques = false, keepSplicedTypes = false)
 
     /** Follow non-opaque aliases and dereferences LazyRefs, annotated types and instantiated
      *  TypeVars until type is no longer alias type, annotated type, LazyRef,
      *  or instantiated type variable.
      */
-    final def dealiasKeepOpaques(using Context): Type = dealias1(keepNever, keepOpaques = true)
+    final def dealiasKeepOpaques(using Context): Type = dealias1(keepNever, keepOpaques = true, keepSplicedTypes = false)
+
+    /** Like `dealias`, but keeps aliases marked with `@SplicedType` */
+    final def dealiasKeepSpliceTypes(using Context): Type = dealias1(keepNever, keepOpaques = true, keepSplicedTypes = true)
 
     /** Approximate this type with a type that does not contain skolem types. */
     final def deskolemized(using Context): Type =
@@ -1449,7 +1452,7 @@ object Types {
     def tryNormalize(using Context): Type = NoType
 
     private def widenDealias1(keep: AnnotatedType => Context ?=> Boolean)(using Context): Type = {
-      val res = this.widen.dealias1(keep, keepOpaques = false)
+      val res = this.widen.dealias1(keep, keepOpaques = false, keepSplicedTypes = false)
       if (res eq this) res else res.widenDealias1(keep)
     }
 

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

@@ -44,7 +44,7 @@ import dotty.tools.dotc.util.SrcPos
  *
  */
 class PCPCheckAndHeal extends TreeMapWithStages {
-  import tpd._
+  import tpd.*
 
   private val InAnnotation = Property.Key[Unit]()
 
@@ -107,30 +107,37 @@ class PCPCheckAndHeal extends TreeMapWithStages {
     val stripAnnotsDeep: TypeMap = new TypeMap:
       def apply(tp: Type): Type = mapOver(tp.stripAnnots)
 
-    val contextWithQuote =
-      if level == 0 then contextWithQuoteTypeTags(taggedTypes)(using quoteContext)
-      else quoteContext
-    val body1 = transform(body)(using contextWithQuote)
-    val body2 =
+    def transformBody() =
+      val contextWithQuote =
+        if level == 0 then contextWithQuoteTypeTags(taggedTypes)(using quoteContext)
+        else quoteContext
+      val transformedBody = transform(body)(using contextWithQuote)
       taggedTypes.getTypeTags match
-        case Nil  => body1
-        case tags => tpd.Block(tags, body1).withSpan(body.span)
+        case Nil  => transformedBody
+        case tags => tpd.Block(tags, transformedBody).withSpan(body.span)
 
     if body.isTerm then
+      val transformedBody = transformBody()
       // `quoted.runtime.Expr.quote[T](<body>)`  --> `quoted.runtime.Expr.quote[T2](<body2>)`
       val TypeApply(fun, targs) = quote.fun: @unchecked
       val targs2 = targs.map(targ => TypeTree(healType(quote.fun.srcPos)(stripAnnotsDeep(targ.tpe))))
-      cpy.Apply(quote)(cpy.TypeApply(quote.fun)(fun, targs2), body2 :: Nil)
+      cpy.Apply(quote)(cpy.TypeApply(quote.fun)(fun, targs2), transformedBody :: Nil)
     else
-      val quotes = quote.args.mapConserve(transform)
       body.tpe match
         case tp @ TypeRef(x: TermRef, _) if tp.symbol == defn.QuotedType_splice =>
           // Optimization: `quoted.Type.of[x.Underlying](quotes)`  -->  `x`
-          ref(x)
+          ref(x).withSpan(quote.span)
         case _ =>
-          // `quoted.Type.of[<body>](quotes)`  --> `quoted.Type.of[<body2>](quotes)`
-          val TypeApply(fun, _) = quote.fun: @unchecked
-          cpy.Apply(quote)(cpy.TypeApply(quote.fun)(fun, body2 :: Nil), quotes)
+          transformBody() match
+            case HealedDirectQuotedTypeRef(termRef) =>
+              // Optimization: `quoted.Type.of[@SplicedType type T = x.Underlying; T](quotes)`  -->  `x`
+              tpd.ref(termRef).withSpan(quote.span)
+            case transformedBody =>
+              val quotes = quote.args.mapConserve(transform)
+              // `quoted.Type.of[<body>](quotes)`  --> `quoted.Type.of[<body2>](quotes)`
+              val TypeApply(fun, _) = quote.fun: @unchecked
+              cpy.Apply(quote)(cpy.TypeApply(quote.fun)(fun, transformedBody :: Nil), quotes)
+
   }
 
   /** Transform splice
@@ -236,4 +243,5 @@ class PCPCheckAndHeal extends TreeMapWithStages {
           | - but the access is at level $level.$hint""", pos)
     tp
   }
+
 }

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

@@ -117,7 +117,7 @@ object BetaReduce:
           case ref @ TypeRef(NoPrefix, _) =>
             ref.symbol
           case _ =>
-            val binding = TypeDef(newSymbol(ctx.owner, tparam.name, EmptyFlags, targ.tpe, coord = targ.span)).withSpan(targ.span)
+            val binding = TypeDef(newSymbol(ctx.owner, tparam.name, EmptyFlags, TypeAlias(targ.tpe), coord = targ.span)).withSpan(targ.span)
             bindings += binding
             binding.symbol
 

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

@@ -48,7 +48,7 @@ class ExpandSAMs extends MiniPhase:
       tpt.tpe match {
         case NoType =>
           tree // it's a plain function
-        case tpe if defn.isContextFunctionType(tpe) =>
+        case tpe if defn.isFunctionOrPolyType(tpe) =>
           tree
         case tpe @ SAMType(_) if tpe.isRef(defn.PartialFunctionClass) =>
           val tpe1 = checkRefinements(tpe, fn)

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

@@ -314,17 +314,28 @@ object PickleQuotes {
               defn.QuotedExprClass.typeRef.appliedTo(defn.AnyType)),
             args =>
               val cases = termSplices.map { case (splice, idx) =>
-                val defn.FunctionOf(argTypes, defn.FunctionOf(quotesType :: _, _, _), _) = splice.tpe: @unchecked
+                val (typeParamCount, argTypes, quotesType) = splice.tpe match
+                  case defn.FunctionOf(argTypes, defn.FunctionOf(quotesType :: _, _, _), _) => (0, argTypes, quotesType)
+                  case RefinedType(polyFun, nme.apply, pt @ PolyType(tparams, _)) if polyFun.typeSymbol.derivesFrom(defn.PolyFunctionClass) =>
+                    pt.instantiate(pt.paramInfos.map(_.hi)) match
+                      case MethodTpe(_, argTypes, defn.FunctionOf(quotesType :: _, _, _)) =>
+                        (tparams.size, argTypes, quotesType)
+
                 val rhs = {
                   val spliceArgs = argTypes.zipWithIndex.map { (argType, i) =>
                     args(1).select(nme.apply).appliedTo(Literal(Constant(i))).asInstance(argType)
                   }
                   val Block(List(ddef: DefDef), _) = splice: @unchecked
-                  // TODO: beta reduce inner closure? Or wait until BetaReduce phase?
-                  BetaReduce(
-                    splice
-                      .select(nme.apply).appliedToArgs(spliceArgs))
-                      .select(nme.apply).appliedTo(args(2).asInstance(quotesType))
+
+                  val typeArgs = ddef.symbol.info match
+                    case pt: PolyType => pt.paramInfos
+                    case _ => Nil
+
+                  val sel1 = splice.changeOwner(ddef.symbol.owner, ctx.owner).select(nme.apply)
+                  val appTpe = if typeParamCount == 0 then sel1 else sel1.appliedToTypes(List.fill(typeParamCount)(defn.AnyType))
+                  val app1 = appTpe.appliedToArgs(spliceArgs)
+                  val sel2 = app1.select(nme.apply)
+                  sel2.appliedTo(args(2).asInstance(quotesType))
                 }
                 CaseDef(Literal(Constant(idx)), EmptyTree, rhs)
               }
@@ -333,18 +344,31 @@ object PickleQuotes {
                 case _ => Match(args(0).annotated(New(ref(defn.UncheckedAnnot.typeRef))), cases)
           )
 
+      def dealiasSplicedTypes(tp: Type) = new TypeMap {
+        def apply(tp: Type): Type = tp match
+          case tp: TypeRef if tp.typeSymbol.hasAnnotation(defn.QuotedRuntime_SplicedTypeAnnot) =>
+            val TypeAlias(alias) = tp.info: @unchecked
+            alias
+          case tp1 => mapOver(tp)
+      }.apply(tp)
+
+      val adaptedType =
+        if isType then dealiasSplicedTypes(originalTp)
+        else originalTp
+
       val quoteClass = if isType then defn.QuotedTypeClass else defn.QuotedExprClass
-      val quotedType = quoteClass.typeRef.appliedTo(originalTp)
+      val quotedType = quoteClass.typeRef.appliedTo(adaptedType)
       val lambdaTpe = MethodType(defn.QuotesClass.typeRef :: Nil, quotedType)
       val unpickleMeth =
         if isType then defn.QuoteUnpickler_unpickleTypeV2
         else defn.QuoteUnpickler_unpickleExprV2
       val unpickleArgs =
         if isType then List(pickledQuoteStrings, types)
         else List(pickledQuoteStrings, types, termHoles)
+
       quotes
         .asInstance(defn.QuoteUnpicklerClass.typeRef)
-        .select(unpickleMeth).appliedToType(originalTp)
+        .select(unpickleMeth).appliedToType(adaptedType)
         .appliedToArgs(unpickleArgs).withSpan(body.span)
     }
 

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

@@ -25,6 +25,7 @@ import dotty.tools.dotc.staging.PCPCheckAndHeal
 import dotty.tools.dotc.staging.QuoteContext.*
 import dotty.tools.dotc.staging.StagingLevel.*
 import dotty.tools.dotc.staging.QuoteTypeTags
+import dotty.tools.dotc.staging.HealedDirectQuotedTypeRef
 
 import scala.annotation.constructorOnly
 
@@ -133,7 +134,7 @@ class Splicing extends MacroTransform:
             case None =>
               val holeIdx = numHoles
               numHoles += 1
-              val hole = tpd.Hole(false, holeIdx, Nil, ref(qual), TypeTree(tp))
+              val hole = tpd.Hole(false, holeIdx, Nil, ref(qual), TypeTree(tp.dealias))
               typeHoles.put(qual.symbol, hole)
               hole
           cpy.TypeDef(tree)(rhs = hole)
@@ -155,7 +156,7 @@ class Splicing extends MacroTransform:
 
     private def transformAnnotations(tree: DefTree)(using Context): Unit =
       tree.symbol.annotations = tree.symbol.annotations.mapconserve { annot =>
-        val newAnnotTree = transform(annot.tree)(using ctx.withOwner(tree.symbol))
+        val newAnnotTree = transform(annot.tree)
         if (annot.tree == newAnnotTree) annot
         else ConcreteAnnotation(newAnnotTree)
       }
@@ -199,10 +200,56 @@ class Splicing extends MacroTransform:
       val newTree = transform(tree)
       val (refs, bindings) = refBindingMap.values.toList.unzip
       val bindingsTypes = bindings.map(_.termRef.widenTermRefExpr)
-      val methType = MethodType(bindingsTypes, newTree.tpe)
+      val types = bindingsTypes.collect {
+        case AppliedType(tycon, List(arg: TypeRef)) if tycon.derivesFrom(defn.QuotedTypeClass) => arg
+      }
+      val newTypeParams = types.map { tpe =>
+        newSymbol(
+          spliceOwner,
+          (tpe.symbol.name.toString + "$tpe").toTypeName,
+          Param,
+          TypeBounds.empty
+        )
+      }
+      val methType =
+        if types.nonEmpty then
+          PolyType(types.map(tp => (tp.symbol.name.toString + "$").toTypeName))(
+            pt => types.map(_ => TypeBounds.empty),
+            pt => {
+               val tpParamMap = new TypeMap {
+                private val mapping = types.map(_.typeSymbol).zip(pt.paramRefs).toMap
+                def apply(tp: Type): Type = tp match
+                  case tp: TypeRef => mapping.getOrElse(tp.typeSymbol, tp)
+                  case tp => mapOver(tp)
+               }
+               MethodType(bindingsTypes.map(tpParamMap), tpParamMap(newTree.tpe))
+            }
+          )
+        else MethodType(bindingsTypes, newTree.tpe)
       val meth = newSymbol(spliceOwner, nme.ANON_FUN, Synthetic | Method, methType)
-      val ddef = DefDef(meth, List(bindings), newTree.tpe, newTree.changeOwner(ctx.owner, meth))
-      val fnType = defn.FunctionType(bindings.size, isContextual = false).appliedTo(bindingsTypes :+ newTree.tpe)
+
+      def substituteTypes(tree: Tree): Tree =
+        if types.nonEmpty then
+          TreeTypeMap(
+            typeMap = new TypeMap {
+              def apply(tp: Type): Type = tp match
+                case tp @ TypeRef(x: TermRef, _) if tp.symbol == defn.QuotedType_splice => tp
+                case tp: TypeRef if tp.typeSymbol.hasAnnotation(defn.QuotedRuntime_SplicedTypeAnnot) => tp
+                case _: TypeRef =>
+                  val idx = types.indexWhere(_ =:= tp) // TODO performance
+                  if idx == -1 then mapOver(tp)
+                  else newTypeParams(idx).typeRef
+                case _ => mapOver(tp)
+            }
+          ).transform(tree)
+        else tree
+      val paramss =
+        if types.nonEmpty then List(newTypeParams, bindings)
+        else List(bindings)
+      val ddef = substituteTypes(DefDef(meth, paramss, newTree.tpe, newTree.changeOwner(ctx.owner, meth)))
+      val fnType =
+        if types.isEmpty then defn.FunctionType(bindings.size, isContextual = false).appliedTo(bindingsTypes :+ newTree.tpe) // FIXME add type parameter?
+        else RefinedType(defn.PolyFunctionType, nme.apply, methType)
       val closure = Block(ddef :: Nil, Closure(Nil, ref(meth), TypeTree(fnType)))
       tpd.Hole(true, holeIdx, refs, closure, TypeTree(tpe))
 
@@ -248,6 +295,9 @@ class Splicing extends MacroTransform:
         if tree.symbol == defn.QuotedTypeModule_of && containsCapturedType(tpt.tpe) =>
           val newContent = capturedPartTypes(tpt)
           newContent match
+            case HealedDirectQuotedTypeRef(termRef) =>
+              // Optimization: `quoted.Type.of[@SplicedType type T = x.Underlying; T](quotes)`  -->  `x`
+              tpd.ref(termRef).withSpan(tpt.span)
             case block: Block =>
               inContext(ctx.withSource(tree.source)) {
                 Apply(TypeApply(typeof, List(newContent)), List(quotes)).withSpan(tree.span)
@@ -347,7 +397,7 @@ class Splicing extends MacroTransform:
     private def newQuotedTypeClassBinding(tpe: Type)(using Context) =
       newSymbol(
         spliceOwner,
-        UniqueName.fresh(nme.Type).toTermName,
+        UniqueName.fresh(tpe.typeSymbol.name.toTermName),
         Param,
         defn.QuotedTypeClass.typeRef.appliedTo(tpe),
       )

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

@@ -533,6 +533,11 @@ object TreeChecker {
         i"owner chain = ${tree.symbol.ownersIterator.toList}%, %, ctxOwners = ${ctx.outersIterator.map(_.owner).toList}%, %")
     }
 
+    override def typedTypeDef(tdef: untpd.TypeDef, sym: Symbol)(using Context): Tree = {
+      assert(sym.info.isInstanceOf[ClassInfo] || sym.info.isInstanceOf[TypeBounds], i"wrong type, expect a template or type bounds for ${sym.fullName}, but found: ${sym.info}")
+      super.typedTypeDef(tdef, sym)
+    }
+
     override def typedClassDef(cdef: untpd.TypeDef, cls: ClassSymbol)(using Context): Tree = {
       val TypeDef(_, impl @ Template(constr, _, _, _)) = cdef: @unchecked
       assert(cdef.symbol == cls)
@@ -665,6 +670,9 @@ object TreeChecker {
       else assert(tpt.typeOpt =:= pt)
 
       // Check that the types of the args conform to the types of the contents of the hole
+      val typeArgsTypes = args.collect { case arg if arg.isType =>
+        arg.typeOpt
+      }
       val argQuotedTypes = args.map { arg =>
         if arg.isTerm then
           val tpe = arg.typeOpt.widenTermRefExpr match
@@ -682,7 +690,28 @@ object TreeChecker {
       val contextualResult =
         defn.FunctionOf(List(defn.QuotesClass.typeRef), expectedResultType, isContextual = true)
       val expectedContentType =
-        defn.FunctionOf(argQuotedTypes, contextualResult)
+        if typeArgsTypes.isEmpty then defn.FunctionOf(argQuotedTypes, contextualResult)
+        else RefinedType(defn.PolyFunctionType, nme.apply, PolyType(typeArgsTypes.map(_ => TypeBounds.empty))(pt =>
+          val tpParamMap = new TypeMap {
+            private val mapping = typeArgsTypes.map(_.typeSymbol).zip(pt.paramRefs).toMap
+            def apply(tp: Type): Type = tp match
+              case tp: TypeRef => mapping.getOrElse(tp.typeSymbol, tp)
+              case tp => mapOver(tp)
+          }
+          MethodType(
+          args.zipWithIndex.map { case (arg, idx) =>
+
+            if arg.isTerm then
+              val tpe = arg.typeOpt.widenTermRefExpr match
+                case _: MethodicType =>
+                  // Special erasure for captured function references
+                  // See `SpliceTransformer.transformCapturedApplication`
+                  defn.AnyType
+                case tpe => tpe
+              defn.QuotedExprClass.typeRef.appliedTo(tpParamMap(tpe))
+            else defn.QuotedTypeClass.typeRef.appliedTo(tpParamMap(arg.typeOpt))
+          },
+          tpParamMap(contextualResult))))
       assert(content.typeOpt =:= expectedContentType, i"expected content of the hole to be ${expectedContentType} but got ${content.typeOpt}")
 
       tree1