Skip to content

Commit 1039b5b

Browse files
nicolasstuckinicolasstucki
committed
Re-architecture quote pickling
Split cross quote reference handling from pickling Fixes #8100 Fixes #12440
1 parent 73f099f commit 1039b5b

File tree

7 files changed

+670
-4
lines changed

7 files changed

+670
-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

+376
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,376 @@
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 tree: DefDef if tree.symbol.is(Macro) =>
111+
// Shrink size of the tree. The methods have already been inlined.
112+
// TODO move to FirstTransform to trigger even without quotes
113+
cpy.DefDef(tree)(rhs = defaultValue(tree.rhs.tpe))
114+
case _: DefDef if tree.symbol.isInlineMethod =>
115+
tree
116+
case _ =>
117+
super.transform(tree)
118+
}
119+
120+
private def makeHoles(tree: tpd.Tree)(using Context): (List[Tree], tpd.Tree) =
121+
122+
/** Remove references to local types that will not be defined in this quote */
123+
def getTypeHoleType(using Context) = new TypeMap() {
124+
override def apply(tp: Type): Type = tp match
125+
case tp: TypeRef if tp.typeSymbol.isTypeSplice =>
126+
apply(tp.dealias)
127+
case tp @ TypeRef(pre, _) if pre == NoPrefix || pre.termSymbol.isLocal =>
128+
val hiBound = tp.typeSymbol.info match
129+
case info: ClassInfo => info.parents.reduce(_ & _)
130+
case info => info.hiBound
131+
apply(hiBound)
132+
case tp =>
133+
mapOver(tp)
134+
}
135+
136+
/** Remove references to local types that will not be defined in this quote */
137+
def getTermHoleType(using Context) = new TypeMap() {
138+
override def apply(tp: Type): Type = tp match
139+
case tp @ TypeRef(NoPrefix, _) =>
140+
// reference to term with a type defined in outer quote
141+
getTypeHoleType(tp)
142+
case tp @ TermRef(NoPrefix, _) =>
143+
// widen term refs to terms defined in outer quote
144+
apply(tp.widenTermRefExpr)
145+
case tp =>
146+
mapOver(tp)
147+
}
148+
149+
class HoleMaker extends Transformer:
150+
private var splices = List.newBuilder[Tree]
151+
private var typeHoles = mutable.Map.empty[Symbol, Hole]
152+
private var idx = -1
153+
override def transform(tree: tpd.Tree)(using Context): tpd.Tree =
154+
tree match
155+
case Apply(fn, List(splicedCode)) if fn.symbol == defn.QuotedRuntime_exprNestedSplice =>
156+
val Apply(Select(spliceFn, _), args) = splicedCode
157+
splices += spliceFn
158+
val holeArgs = args.map {
159+
case Apply(Select(Apply(_, code :: Nil), _), _) => code
160+
case Apply(TypeApply(_, List(code)), _) => code
161+
}
162+
idx += 1
163+
val holeType = getTermHoleType(tree.tpe)
164+
val hole = Hole(true, idx, holeArgs).withSpan(splicedCode.span).withType(holeType).asInstanceOf[Hole]
165+
Inlined(EmptyTree, Nil, hole).withSpan(tree.span)
166+
case Select(tp, _) if tree.symbol == defn.QuotedType_splice =>
167+
def makeTypeHole =
168+
splices += ref(tp.symbol)
169+
idx += 1
170+
val holeType = getTypeHoleType(tree.tpe)
171+
Hole(false, idx, Nil).withType(holeType).asInstanceOf[Hole]
172+
typeHoles.getOrElseUpdate(tree.symbol, makeTypeHole)
173+
case tree: DefTree =>
174+
val newAnnotations = tree.symbol.annotations.mapconserve { annot =>
175+
val newAnnotTree = transform(annot.tree)(using ctx.withOwner(tree.symbol))
176+
if (annot.tree == newAnnotTree) annot
177+
else ConcreteAnnotation(newAnnotTree)
178+
}
179+
tree.symbol.annotations = newAnnotations
180+
super.transform(tree)
181+
case _ =>
182+
super.transform(tree)
183+
def getSplices() =
184+
val res = splices.result
185+
splices.clear()
186+
res
187+
end HoleMaker
188+
189+
val holeMaker = new HoleMaker
190+
val newTree = holeMaker.transform(tree)
191+
(holeMaker.getSplices(), newTree)
192+
193+
194+
end makeHoles
195+
196+
}
197+
198+
199+
object PickleQuotes2:
200+
val name: String = "pickleQuotes2"
201+
202+
203+
object ReifiedQuote:
204+
import tpd._
205+
206+
def apply(quotes: Tree, body: Tree, splices: List[Tree], originalTp: Type, isType: Boolean)(using Context) = {
207+
/** Encode quote using Reflection.Literal
208+
*
209+
* Generate the code
210+
* ```scala
211+
* quotes => quotes.reflect.TreeMethods.asExpr(
212+
* quotes.reflect.Literal.apply(x$1.reflect.Constant.<typeName>.apply(<literalValue>))
213+
* ).asInstanceOf[scala.quoted.Expr[<body.type>]]
214+
* ```
215+
* this closure is always applied directly to the actual context and the BetaReduce phase removes it.
216+
*/
217+
def pickleAsLiteral(lit: Literal) = {
218+
val exprType = defn.QuotedExprClass.typeRef.appliedTo(body.tpe)
219+
val reflect = quotes.select("reflect".toTermName)
220+
val typeName = body.tpe.typeSymbol.name
221+
val literalValue =
222+
if lit.const.tag == Constants.NullTag || lit.const.tag == Constants.UnitTag then Nil
223+
else List(body)
224+
val constant = reflect.select(s"${typeName}Constant".toTermName).select(nme.apply).appliedToTermArgs(literalValue)
225+
val literal = reflect.select("Literal".toTermName).select(nme.apply).appliedTo(constant)
226+
reflect.select("TreeMethods".toTermName).select("asExpr".toTermName).appliedTo(literal).asInstance(exprType)
227+
}
228+
229+
/** Encode quote using Reflection.Literal
230+
*
231+
* Generate the code
232+
* ```scala
233+
* quotes => scala.quoted.ToExpr.{BooleanToExpr,ShortToExpr, ...}.apply(<literalValue>)(quotes)
234+
* ```
235+
* this closure is always applied directly to the actual context and the BetaReduce phase removes it.
236+
*/
237+
def liftedValue(lit: Literal, lifter: Symbol) =
238+
val exprType = defn.QuotedExprClass.typeRef.appliedTo(body.tpe)
239+
ref(lifter).appliedToType(originalTp).select(nme.apply).appliedTo(lit).appliedTo(quotes)
240+
241+
def pickleAsValue(lit: Literal) = {
242+
// TODO should all constants be pickled as Literals?
243+
// Should examine the generated bytecode size to decide and performance
244+
lit.const.tag match {
245+
case Constants.NullTag => pickleAsLiteral(lit)
246+
case Constants.UnitTag => pickleAsLiteral(lit)
247+
case Constants.BooleanTag => liftedValue(lit, defn.ToExprModule_BooleanToExpr)
248+
case Constants.ByteTag => liftedValue(lit, defn.ToExprModule_ByteToExpr)
249+
case Constants.ShortTag => liftedValue(lit, defn.ToExprModule_ShortToExpr)
250+
case Constants.IntTag => liftedValue(lit, defn.ToExprModule_IntToExpr)
251+
case Constants.LongTag => liftedValue(lit, defn.ToExprModule_LongToExpr)
252+
case Constants.FloatTag => liftedValue(lit, defn.ToExprModule_FloatToExpr)
253+
case Constants.DoubleTag => liftedValue(lit, defn.ToExprModule_DoubleToExpr)
254+
case Constants.CharTag => liftedValue(lit, defn.ToExprModule_CharToExpr)
255+
case Constants.StringTag => liftedValue(lit, defn.ToExprModule_StringToExpr)
256+
}
257+
}
258+
259+
/** Encode quote using QuoteUnpickler.{unpickleExpr, unpickleType}
260+
*
261+
* Generate the code
262+
* ```scala
263+
* quotes => quotes.asInstanceOf[QuoteUnpickler].<unpickleExpr|unpickleType>[<type>](
264+
* <pickledQuote>,
265+
* <typeHole>,
266+
* <termHole>,
267+
* )
268+
* ```
269+
* this closure is always applied directly to the actual context and the BetaReduce phase removes it.
270+
*/
271+
def pickleAsTasty() = {
272+
def liftList(list: List[Tree], tpe: Type)(using Context): Tree =
273+
list.foldRight[Tree](ref(defn.NilModule)) { (x, acc) =>
274+
acc.select("::".toTermName).appliedToType(tpe).appliedTo(x)
275+
}
276+
277+
val pickleQuote = PickledQuotes.pickleQuote(body)
278+
val pickledQuoteStrings = pickleQuote match
279+
case x :: Nil => Literal(Constant(x))
280+
case xs => liftList(xs.map(x => Literal(Constant(x))), defn.StringType)
281+
282+
// TODO split holes earlier into types and terms. This all holes in each category can have consecutive indices
283+
val (typeSplices, termSplices) = splices.zipWithIndex.partition {
284+
case (splice, _) => splice.tpe.derivesFrom(defn.QuotedTypeClass)
285+
}
286+
287+
// This and all closures in typeSplices are removed by the BetaReduce phase
288+
val typeHoles =
289+
if typeSplices.isEmpty then Literal(Constant(null)) // keep pickled quote without splices as small as possible
290+
else
291+
Lambda(
292+
MethodType(
293+
List("idx", "splices").map(name => UniqueName.fresh(name.toTermName).toTermName),
294+
List(defn.IntType, defn.SeqType.appliedTo(defn.AnyType)),
295+
defn.QuotedTypeClass.typeRef.appliedTo(WildcardType)),
296+
args => {
297+
val cases = typeSplices.map { case (splice, idx) =>
298+
CaseDef(Literal(Constant(idx)), EmptyTree, splice)
299+
}
300+
cases match
301+
case CaseDef(_, _, rhs) :: Nil => rhs
302+
case _ => Match(args(0).annotated(New(ref(defn.UncheckedAnnot.typeRef))), cases)
303+
}
304+
)
305+
306+
// This and all closures in termSplices are removed by the BetaReduce phase
307+
val termHoles =
308+
if termSplices.isEmpty then Literal(Constant(null)) // keep pickled quote without splices as small as possible
309+
else
310+
Lambda(
311+
MethodType(
312+
List("idx", "splices", "quotes").map(name => UniqueName.fresh(name.toTermName).toTermName),
313+
List(defn.IntType, defn.SeqType.appliedTo(defn.AnyType), defn.QuotesClass.typeRef),
314+
defn.QuotedExprClass.typeRef.appliedTo(defn.AnyType)),
315+
args => {
316+
val cases = termSplices.map { case (splice, idx) =>
317+
val defn.FunctionOf(argTypes, defn.FunctionOf(quotesType :: _, _, _, _), _, _) = splice.tpe
318+
val rhs = {
319+
val spliceArgs = argTypes.zipWithIndex.map { (argType, i) =>
320+
args(1).select(nme.apply).appliedTo(Literal(Constant(i))).select(defn.Any_asInstanceOf).appliedToType(argType)
321+
}
322+
val Block(List(ddef: DefDef), _) = splice
323+
// TODO: beta reduce inner closure? Or wait until BetaReduce phase?
324+
BetaReduce(ddef, spliceArgs).select(nme.apply).appliedTo(args(2).asInstance(quotesType))
325+
}
326+
CaseDef(Literal(Constant(idx)), EmptyTree, rhs)
327+
}
328+
cases match
329+
case CaseDef(_, _, rhs) :: Nil => rhs
330+
case _ => Match(args(0).annotated(New(ref(defn.UncheckedAnnot.typeRef))), cases)
331+
}
332+
)
333+
334+
val quoteClass = if isType then defn.QuotedTypeClass else defn.QuotedExprClass
335+
val quotedType = quoteClass.typeRef.appliedTo(originalTp)
336+
val lambdaTpe = MethodType(defn.QuotesClass.typeRef :: Nil, quotedType)
337+
val unpickleMeth = if isType then defn.QuoteUnpickler_unpickleType else defn.QuoteUnpickler_unpickleExpr
338+
quotes
339+
.asInstance(defn.QuoteUnpicklerClass.typeRef)
340+
.select(unpickleMeth).appliedToType(originalTp)
341+
.appliedTo(pickledQuoteStrings, typeHoles, termHoles).withSpan(body.span)
342+
}
343+
344+
/** Encode quote using Reflection.TypeRepr.typeConstructorOf
345+
*
346+
* Generate the code
347+
* ```scala
348+
* quotes.reflect.TypeReprMethods.asType(
349+
* quotes.reflect.TypeRepr.typeConstructorOf(classOf[<type>]])
350+
* ).asInstanceOf[scala.quoted.Type[<type>]]
351+
* ```
352+
* this closure is always applied directly to the actual context and the BetaReduce phase removes it.
353+
*/
354+
def taggedType() =
355+
val typeType = defn.QuotedTypeClass.typeRef.appliedTo(body.tpe)
356+
val classTree = TypeApply(ref(defn.Predef_classOf.termRef), body :: Nil)
357+
val reflect = quotes.select("reflect".toTermName)
358+
val typeRepr = reflect.select("TypeRepr".toTermName).select("typeConstructorOf".toTermName).appliedTo(classTree)
359+
reflect.select("TypeReprMethods".toTermName).select("asType".toTermName).appliedTo(typeRepr).asInstance(typeType)
360+
361+
def getLiteral(tree: tpd.Tree): Option[Literal] = tree match
362+
case tree: Literal => Some(tree)
363+
case Block(Nil, e) => getLiteral(e)
364+
case Inlined(_, Nil, e) => getLiteral(e)
365+
case _ => None
366+
367+
if (isType) then
368+
if splices.isEmpty && body.symbol.isPrimitiveValueClass then taggedType()
369+
else pickleAsTasty()
370+
else
371+
getLiteral(body) match
372+
case Some(lit) => pickleAsValue(lit)
373+
case _ => pickleAsTasty()
374+
}
375+
376+
end ReifiedQuote

0 commit comments

Comments
 (0)