Skip to content

Commit e474800

Browse files
nicolasstuckinicolasstucki
committed
Re-architecture quote pickling
Split cross quote reference handling from pickling Fixes scala#8100 Fixes scala#12443
1 parent 4bf2f04 commit e474800

File tree

6 files changed

+623
-3
lines changed

6 files changed

+623
-3
lines changed

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

Lines changed: 3 additions & 1 deletion
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/transform/PickleQuotes2.scala

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

0 commit comments

Comments
 (0)