Skip to content

Commit a3f0700

Browse files
nicolasstuckinicolasstucki
committed
Re-architecture quote pickling
Split cross quote reference handling from pickling Fixes #8100 Fixes #12440
1 parent d2dd083 commit a3f0700

File tree

7 files changed

+666
-4
lines changed

7 files changed

+666
-4
lines changed

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

+3-1
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,9 @@ class Compiler {
5252
List(new Inlining) :: // Inline and execute macros
5353
List(new PostInlining) :: // Add mirror support for inlined code
5454
List(new Staging) :: // Check staging levels and heal staged types
55-
List(new PickleQuotes) :: // Turn quoted trees into explicit run-time data structures
55+
List(new Splicing) :: // Turn quoted trees into explicit run-time data structures
56+
List(new PickleQuotes2) :: // Turn quoted trees into explicit run-time data structures
57+
// List(new PickleQuotes) :: // Turn quoted trees into explicit run-time data structures
5658
Nil
5759

5860
/** Phases dealing with the transformation from pickled trees to backend trees */

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -75,7 +75,7 @@ object PickledQuotes {
7575
case Hole(isTerm, idx, args) =>
7676
inContext(SpliceScope.contextWithNewSpliceScope(tree.sourcePos)) {
7777
val reifiedArgs = args.map { arg =>
78-
if (arg.isTerm) (q: Quotes) ?=> new ExprImpl(arg, SpliceScope.getCurrent)
78+
if (arg.isTerm) new ExprImpl(arg, SpliceScope.getCurrent)
7979
else new TypeImpl(arg, SpliceScope.getCurrent)
8080
}
8181
if isTerm then

‎compiler/src/dotty/tools/dotc/transform/PickleQuotes2.scala

+372
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,372 @@
1+
package dotty.tools.dotc
2+
package transform
3+
4+
import core._
5+
import Decorators._
6+
import Flags._
7+
import Types._
8+
import Contexts._
9+
import Symbols._
10+
import Constants._
11+
import ast.Trees._
12+
import ast.{TreeTypeMap, untpd}
13+
import util.Spans._
14+
import tasty.TreePickler.Hole
15+
import SymUtils._
16+
import NameKinds._
17+
import dotty.tools.dotc.ast.tpd
18+
import typer.Implicits.SearchFailureType
19+
20+
import scala.collection.mutable
21+
import dotty.tools.dotc.core.Annotations._
22+
import dotty.tools.dotc.core.Names._
23+
import dotty.tools.dotc.core.StdNames._
24+
import dotty.tools.dotc.quoted._
25+
import dotty.tools.dotc.transform.TreeMapWithStages._
26+
import dotty.tools.dotc.typer.Inliner
27+
28+
import scala.annotation.constructorOnly
29+
30+
31+
/** Translates quoted terms and types to `unpickleExpr` or `unpickleType` method calls.
32+
*
33+
* Transforms top level quote
34+
* ```
35+
* '{ ...
36+
* val x1 = ???
37+
* val x2 = ???
38+
* ...
39+
* ${ ... '{ ... x1 ... x2 ...} ... }
40+
* ...
41+
* }
42+
* ```
43+
* to
44+
* ```
45+
* unpickleExpr(
46+
* pickled = [[ // PICKLED TASTY
47+
* ...
48+
* val x1 = ???
49+
* val x2 = ???
50+
* ...
51+
* Hole(<i> | x1, x2)
52+
* ...
53+
* ]],
54+
* typeHole = (idx: Int, args: List[Any]) => idx match {
55+
* case 0 => ...
56+
* },
57+
* termHole = (idx: Int, args: List[Any], quotes: Quotes) => idx match {
58+
* case 0 => ...
59+
* ...
60+
* case <i> =>
61+
* val x1$1 = args(0).asInstanceOf[Expr[T]]
62+
* val x2$1 = args(1).asInstanceOf[Expr[T]] // can be asInstanceOf[Type[T]]
63+
* ...
64+
* { ... '{ ... ${x1$1} ... ${x2$1} ...} ... }
65+
* },
66+
* )
67+
* ```
68+
* and then performs the same transformation on `'{ ... ${x1$1} ... ${x2$1} ...}`.
69+
*
70+
*/
71+
class PickleQuotes2 extends MacroTransform {
72+
import PickleQuotes2._
73+
import tpd._
74+
75+
override def phaseName: String = PickleQuotes2.name
76+
77+
override def allowsImplicitSearch: Boolean = true
78+
79+
override def checkPostCondition(tree: Tree)(using Context): Unit =
80+
tree match
81+
case tree: RefTree if !Inliner.inInlineMethod =>
82+
assert(!tree.symbol.isQuote)
83+
assert(!tree.symbol.isExprSplice)
84+
case _ : TypeDef =>
85+
assert(!tree.symbol.hasAnnotation(defn.QuotedRuntime_SplicedTypeAnnot),
86+
s"${tree.symbol} should have been removed by PickledQuotes because it has a @quoteTypeTag")
87+
case _ =>
88+
89+
override def run(using Context): Unit =
90+
if (ctx.compilationUnit.needsQuotePickling) super.run(using freshStagingContext)
91+
92+
protected def newTransformer(using Context): Transformer = new Transformer {
93+
override def transform(tree: tpd.Tree)(using Context): tpd.Tree =
94+
tree match
95+
case Apply(Select(Apply(TypeApply(fn, List(tpt)), List(code)),nme.apply), List(quotes))
96+
if fn.symbol == defn.QuotedRuntime_exprQuote =>
97+
val (splices, codeWithHoles) = makeHoles(code)
98+
val sourceRef = Inliner.inlineCallTrace(ctx.owner, tree.sourcePos)
99+
val codeWithHoles2 = Inlined(sourceRef, Nil, codeWithHoles)
100+
val pickled = ReifiedQuote(quotes, codeWithHoles2, splices, tpt.tpe, false)
101+
transform(pickled) // pickle quotes that are in the splices
102+
case Apply(TypeApply(_, List(tpt)), List(quotes)) if tree.symbol == defn.QuotedTypeModule_of =>
103+
tpt match
104+
case Select(t, _) if tpt.symbol == defn.QuotedType_splice =>
105+
// `Type.of[t.Underlying](quotes)` --> `t`
106+
ref(t.symbol)
107+
case _ =>
108+
val (splices, tptWithHoles) = makeHoles(tpt)
109+
ReifiedQuote(quotes, tptWithHoles, splices, tpt.tpe, true)
110+
case _: DefDef if tree.symbol.isInlineMethod =>
111+
tree
112+
case _ =>
113+
super.transform(tree)
114+
}
115+
116+
private def makeHoles(tree: tpd.Tree)(using Context): (List[Tree], tpd.Tree) =
117+
118+
/** Remove references to local types that will not be defined in this quote */
119+
def getTypeHoleType(using Context) = new TypeMap() {
120+
override def apply(tp: Type): Type = tp match
121+
case tp: TypeRef if tp.typeSymbol.isTypeSplice =>
122+
apply(tp.dealias)
123+
case tp @ TypeRef(pre, _) if pre == NoPrefix || pre.termSymbol.isLocal =>
124+
val hiBound = tp.typeSymbol.info match
125+
case info: ClassInfo => info.parents.reduce(_ & _)
126+
case info => info.hiBound
127+
apply(hiBound)
128+
case tp =>
129+
mapOver(tp)
130+
}
131+
132+
/** Remove references to local types that will not be defined in this quote */
133+
def getTermHoleType(using Context) = new TypeMap() {
134+
override def apply(tp: Type): Type = tp match
135+
case tp @ TypeRef(NoPrefix, _) =>
136+
// reference to term with a type defined in outer quote
137+
getTypeHoleType(tp)
138+
case tp @ TermRef(NoPrefix, _) =>
139+
// widen term refs to terms defined in outer quote
140+
apply(tp.widenTermRefExpr)
141+
case tp =>
142+
mapOver(tp)
143+
}
144+
145+
class HoleMaker extends Transformer:
146+
private var splices = List.newBuilder[Tree]
147+
private var typeHoles = mutable.Map.empty[Symbol, Hole]
148+
private var idx = -1
149+
override def transform(tree: tpd.Tree)(using Context): tpd.Tree =
150+
tree match
151+
case Apply(fn, List(splicedCode)) if fn.symbol == defn.QuotedRuntime_exprNestedSplice =>
152+
val Apply(Select(spliceFn, _), args) = splicedCode
153+
splices += spliceFn
154+
val holeArgs = args.map {
155+
case Apply(Select(Apply(_, code :: Nil), _), _) => code
156+
case Apply(TypeApply(_, List(code)), _) => code
157+
}
158+
idx += 1
159+
val holeType = getTermHoleType(tree.tpe)
160+
val hole = Hole(true, idx, holeArgs).withSpan(splicedCode.span).withType(holeType).asInstanceOf[Hole]
161+
Inlined(EmptyTree, Nil, hole).withSpan(tree.span)
162+
case Select(tp, _) if tree.symbol == defn.QuotedType_splice =>
163+
def makeTypeHole =
164+
splices += ref(tp.symbol)
165+
idx += 1
166+
val holeType = getTypeHoleType(tree.tpe)
167+
Hole(false, idx, Nil).withType(holeType).asInstanceOf[Hole]
168+
typeHoles.getOrElseUpdate(tree.symbol, makeTypeHole)
169+
case tree: DefTree =>
170+
val newAnnotations = tree.symbol.annotations.mapconserve { annot =>
171+
val newAnnotTree = transform(annot.tree)(using ctx.withOwner(tree.symbol))
172+
if (annot.tree == newAnnotTree) annot
173+
else ConcreteAnnotation(newAnnotTree)
174+
}
175+
tree.symbol.annotations = newAnnotations
176+
super.transform(tree)
177+
case _ =>
178+
super.transform(tree)
179+
def getSplices() =
180+
val res = splices.result
181+
splices.clear()
182+
res
183+
end HoleMaker
184+
185+
val holeMaker = new HoleMaker
186+
val newTree = holeMaker.transform(tree)
187+
(holeMaker.getSplices(), newTree)
188+
189+
190+
end makeHoles
191+
192+
}
193+
194+
195+
object PickleQuotes2:
196+
val name: String = "pickleQuotes2"
197+
198+
199+
object ReifiedQuote:
200+
import tpd._
201+
202+
def apply(quotes: Tree, body: Tree, splices: List[Tree], originalTp: Type, isType: Boolean)(using Context) = {
203+
/** Encode quote using Reflection.Literal
204+
*
205+
* Generate the code
206+
* ```scala
207+
* quotes => quotes.reflect.TreeMethods.asExpr(
208+
* quotes.reflect.Literal.apply(x$1.reflect.Constant.<typeName>.apply(<literalValue>))
209+
* ).asInstanceOf[scala.quoted.Expr[<body.type>]]
210+
* ```
211+
* this closure is always applied directly to the actual context and the BetaReduce phase removes it.
212+
*/
213+
def pickleAsLiteral(lit: Literal) = {
214+
val exprType = defn.QuotedExprClass.typeRef.appliedTo(body.tpe)
215+
val reflect = quotes.select("reflect".toTermName)
216+
val typeName = body.tpe.typeSymbol.name
217+
val literalValue =
218+
if lit.const.tag == Constants.NullTag || lit.const.tag == Constants.UnitTag then Nil
219+
else List(body)
220+
val constant = reflect.select(s"${typeName}Constant".toTermName).select(nme.apply).appliedToTermArgs(literalValue)
221+
val literal = reflect.select("Literal".toTermName).select(nme.apply).appliedTo(constant)
222+
reflect.select("TreeMethods".toTermName).select("asExpr".toTermName).appliedTo(literal).asInstance(exprType)
223+
}
224+
225+
/** Encode quote using Reflection.Literal
226+
*
227+
* Generate the code
228+
* ```scala
229+
* quotes => scala.quoted.ToExpr.{BooleanToExpr,ShortToExpr, ...}.apply(<literalValue>)(quotes)
230+
* ```
231+
* this closure is always applied directly to the actual context and the BetaReduce phase removes it.
232+
*/
233+
def liftedValue(lit: Literal, lifter: Symbol) =
234+
val exprType = defn.QuotedExprClass.typeRef.appliedTo(body.tpe)
235+
ref(lifter).appliedToType(originalTp).select(nme.apply).appliedTo(lit).appliedTo(quotes)
236+
237+
def pickleAsValue(lit: Literal) = {
238+
// TODO should all constants be pickled as Literals?
239+
// Should examine the generated bytecode size to decide and performance
240+
lit.const.tag match {
241+
case Constants.NullTag => pickleAsLiteral(lit)
242+
case Constants.UnitTag => pickleAsLiteral(lit)
243+
case Constants.BooleanTag => liftedValue(lit, defn.ToExprModule_BooleanToExpr)
244+
case Constants.ByteTag => liftedValue(lit, defn.ToExprModule_ByteToExpr)
245+
case Constants.ShortTag => liftedValue(lit, defn.ToExprModule_ShortToExpr)
246+
case Constants.IntTag => liftedValue(lit, defn.ToExprModule_IntToExpr)
247+
case Constants.LongTag => liftedValue(lit, defn.ToExprModule_LongToExpr)
248+
case Constants.FloatTag => liftedValue(lit, defn.ToExprModule_FloatToExpr)
249+
case Constants.DoubleTag => liftedValue(lit, defn.ToExprModule_DoubleToExpr)
250+
case Constants.CharTag => liftedValue(lit, defn.ToExprModule_CharToExpr)
251+
case Constants.StringTag => liftedValue(lit, defn.ToExprModule_StringToExpr)
252+
}
253+
}
254+
255+
/** Encode quote using QuoteUnpickler.{unpickleExpr, unpickleType}
256+
*
257+
* Generate the code
258+
* ```scala
259+
* quotes => quotes.asInstanceOf[QuoteUnpickler].<unpickleExpr|unpickleType>[<type>](
260+
* <pickledQuote>,
261+
* <typeHole>,
262+
* <termHole>,
263+
* )
264+
* ```
265+
* this closure is always applied directly to the actual context and the BetaReduce phase removes it.
266+
*/
267+
def pickleAsTasty() = {
268+
def liftList(list: List[Tree], tpe: Type)(using Context): Tree =
269+
list.foldRight[Tree](ref(defn.NilModule)) { (x, acc) =>
270+
acc.select("::".toTermName).appliedToType(tpe).appliedTo(x)
271+
}
272+
273+
val pickleQuote = PickledQuotes.pickleQuote(body)
274+
val pickledQuoteStrings = pickleQuote match
275+
case x :: Nil => Literal(Constant(x))
276+
case xs => liftList(xs.map(x => Literal(Constant(x))), defn.StringType)
277+
278+
// TODO split holes earlier into types and terms. This all holes in each category can have consecutive indices
279+
val (typeSplices, termSplices) = splices.zipWithIndex.partition {
280+
case (splice, _) => splice.tpe.derivesFrom(defn.QuotedTypeClass)
281+
}
282+
283+
// This and all closures in typeSplices are removed by the BetaReduce phase
284+
val typeHoles =
285+
if typeSplices.isEmpty then Literal(Constant(null)) // keep pickled quote without splices as small as possible
286+
else
287+
Lambda(
288+
MethodType(
289+
List("idx", "splices").map(name => UniqueName.fresh(name.toTermName).toTermName),
290+
List(defn.IntType, defn.SeqType.appliedTo(defn.AnyType)),
291+
defn.QuotedTypeClass.typeRef.appliedTo(WildcardType)),
292+
args => {
293+
val cases = typeSplices.map { case (splice, idx) =>
294+
CaseDef(Literal(Constant(idx)), EmptyTree, splice)
295+
}
296+
cases match
297+
case CaseDef(_, _, rhs) :: Nil => rhs
298+
case _ => Match(args(0).annotated(New(ref(defn.UncheckedAnnot.typeRef))), cases)
299+
}
300+
)
301+
302+
// This and all closures in termSplices are removed by the BetaReduce phase
303+
val termHoles =
304+
if termSplices.isEmpty then Literal(Constant(null)) // keep pickled quote without splices as small as possible
305+
else
306+
Lambda(
307+
MethodType(
308+
List("idx", "splices", "quotes").map(name => UniqueName.fresh(name.toTermName).toTermName),
309+
List(defn.IntType, defn.SeqType.appliedTo(defn.AnyType), defn.QuotesClass.typeRef),
310+
defn.QuotedExprClass.typeRef.appliedTo(defn.AnyType)),
311+
args => {
312+
val cases = termSplices.map { case (splice, idx) =>
313+
val defn.FunctionOf(argTypes, defn.FunctionOf(quotesType :: _, _, _, _), _, _) = splice.tpe
314+
val rhs = {
315+
val spliceArgs = argTypes.zipWithIndex.map { (argType, i) =>
316+
args(1).select(nme.apply).appliedTo(Literal(Constant(i))).select(defn.Any_asInstanceOf).appliedToType(argType)
317+
}
318+
val Block(List(ddef: DefDef), _) = splice
319+
// TODO: beta reduce inner closure? Or wait until BetaReduce phase?
320+
BetaReduce(ddef, spliceArgs).select(nme.apply).appliedTo(args(2).asInstance(quotesType))
321+
}
322+
CaseDef(Literal(Constant(idx)), EmptyTree, rhs)
323+
}
324+
cases match
325+
case CaseDef(_, _, rhs) :: Nil => rhs
326+
case _ => Match(args(0).annotated(New(ref(defn.UncheckedAnnot.typeRef))), cases)
327+
}
328+
)
329+
330+
val quoteClass = if isType then defn.QuotedTypeClass else defn.QuotedExprClass
331+
val quotedType = quoteClass.typeRef.appliedTo(originalTp)
332+
val lambdaTpe = MethodType(defn.QuotesClass.typeRef :: Nil, quotedType)
333+
val unpickleMeth = if isType then defn.QuoteUnpickler_unpickleType else defn.QuoteUnpickler_unpickleExpr
334+
quotes
335+
.asInstance(defn.QuoteUnpicklerClass.typeRef)
336+
.select(unpickleMeth).appliedToType(originalTp)
337+
.appliedTo(pickledQuoteStrings, typeHoles, termHoles).withSpan(body.span)
338+
}
339+
340+
/** Encode quote using Reflection.TypeRepr.typeConstructorOf
341+
*
342+
* Generate the code
343+
* ```scala
344+
* quotes.reflect.TypeReprMethods.asType(
345+
* quotes.reflect.TypeRepr.typeConstructorOf(classOf[<type>]])
346+
* ).asInstanceOf[scala.quoted.Type[<type>]]
347+
* ```
348+
* this closure is always applied directly to the actual context and the BetaReduce phase removes it.
349+
*/
350+
def taggedType() =
351+
val typeType = defn.QuotedTypeClass.typeRef.appliedTo(body.tpe)
352+
val classTree = TypeApply(ref(defn.Predef_classOf.termRef), body :: Nil)
353+
val reflect = quotes.select("reflect".toTermName)
354+
val typeRepr = reflect.select("TypeRepr".toTermName).select("typeConstructorOf".toTermName).appliedTo(classTree)
355+
reflect.select("TypeReprMethods".toTermName).select("asType".toTermName).appliedTo(typeRepr).asInstance(typeType)
356+
357+
def getLiteral(tree: tpd.Tree): Option[Literal] = tree match
358+
case tree: Literal => Some(tree)
359+
case Block(Nil, e) => getLiteral(e)
360+
case Inlined(_, Nil, e) => getLiteral(e)
361+
case _ => None
362+
363+
if (isType) then
364+
if splices.isEmpty && body.symbol.isPrimitiveValueClass then taggedType()
365+
else pickleAsTasty()
366+
else
367+
getLiteral(body) match
368+
case Some(lit) => pickleAsValue(lit)
369+
case _ => pickleAsTasty()
370+
}
371+
372+
end ReifiedQuote

0 commit comments

Comments
 (0)