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/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

@@ -319,17 +319,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)
               }
@@ -338,18 +349,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

@@ -24,6 +24,7 @@ import dotty.tools.dotc.config.ScalaRelease.*
 import dotty.tools.dotc.staging.QuoteContext.*
 import dotty.tools.dotc.staging.StagingLevel.*
 import dotty.tools.dotc.staging.QuoteTypeTags
+import dotty.tools.dotc.staging.DirectTypeOf
 
 import scala.annotation.constructorOnly
 
@@ -132,7 +133,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)
@@ -154,7 +155,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)
       }
@@ -198,10 +199,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))
 
@@ -255,6 +302,9 @@ class Splicing extends MacroTransform:
         if tree.symbol == defn.QuotedTypeModule_of && containsCapturedType(tpt.tpe) =>
           val newContent = capturedPartTypes(tpt)
           newContent match
+            case DirectTypeOf.Healed(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)
@@ -354,7 +404,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

@@ -670,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
@@ -687,7 +690,29 @@ 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"unexpected content of hole\nexpected: ${expectedContentType}\nwas: ${content.typeOpt}")
 
       tree1

tests/pos-macros/captured-type.scala

@@ -0,0 +1,6 @@
+import scala.quoted.*
+
+object Foo:
+  def baz(using Quotes): Unit = '{
+    def f[T](x: T): T = ${ identity('{ x: T }) }
+  }