Skip to content

Commit d143026

Browse files
nicolasstuckinicolasstucki
committed
Add While and DoWhile extractors Tasty reflect
1 parent 90904e8 commit d143026

File tree

4 files changed

+68
-53
lines changed

4 files changed

+68
-53
lines changed

‎compiler/src/dotty/tools/dotc/tastyreflect/TastyImpl.scala

Lines changed: 47 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -374,9 +374,31 @@ class TastyImpl(val rootContext: Contexts.Context) extends scala.tasty.Tasty { s
374374
}
375375

376376
object Block extends BlockExtractor {
377-
def unapply(x: Term)(implicit ctx: Context): Option[(List[Statement], Term)] = x match {
378-
case x: tpd.Block @unchecked => Some((x.stats, x.expr))
379-
case _ => None
377+
def unapply(x: Term)(implicit ctx: Context): Option[(List[Statement], Term)] = normalizedLoops(x) match {
378+
case Trees.Block(_, expr) if expr.symbol.is(Flags.Label) => None // while or doWhile loops
379+
case Trees.Block(stats, expr) => Some((stats, expr))
380+
case _ => None
381+
}
382+
/** Normilizes non Blocks.
383+
* i) Put `while` and `doWhile` loops in thier own blocks: `{ def while$() = ...; while$() }`
384+
*/
385+
private def normalizedLoops(tree: tpd.Tree)(implicit ctx: Context): tpd.Tree = tree match {
386+
case block: tpd.Block if block.stats.size > 1 =>
387+
def normalizeInnerLoops(stats: List[tpd.Tree]): List[tpd.Tree] = stats match {
388+
case (x: tpd.DefDef) :: y :: xs if y.symbol.is(Flags.Label) =>
389+
tpd.Block(x :: Nil, y) :: normalizeInnerLoops(xs)
390+
case x :: xs => x :: normalizeInnerLoops(xs)
391+
case Nil => Nil
392+
}
393+
if (block.expr.symbol.is(Flags.Label)) {
394+
val stats1 = normalizeInnerLoops(block.stats.init)
395+
val normalLoop = tpd.Block(block.stats.last :: Nil, block.expr)
396+
tpd.Block(stats1, normalLoop)
397+
} else {
398+
val stats1 = normalizeInnerLoops(block.stats)
399+
tpd.cpy.Block(block)(stats1, block.expr)
400+
}
401+
case _ => tree
380402
}
381403
}
382404

@@ -441,6 +463,28 @@ class TastyImpl(val rootContext: Contexts.Context) extends scala.tasty.Tasty { s
441463
}
442464
}
443465

466+
object While extends WhileExtractor {
467+
def unapply(x: Term)(implicit ctx: Context): Option[(Term, Term)] = x match {
468+
case Trees.Block((ddef: tpd.DefDef) :: Nil, expr) if expr.symbol.is(Flags.Label) && expr.symbol.name == nme.WHILE_PREFIX =>
469+
val Trees.If(cond, Trees.Block(bodyStats, _), _) = ddef.rhs
470+
Some((cond, loopBody(bodyStats)))
471+
case _ => None
472+
}
473+
}
474+
475+
object DoWhile extends DoWhileExtractor {
476+
def unapply(x: Term)(implicit ctx: Context): Option[(Term, Term)] = x match {
477+
case Trees.Block((ddef: tpd.DefDef) :: Nil, expr) if expr.symbol.is(Flags.Label) && expr.symbol.name == nme.DO_WHILE_PREFIX =>
478+
val Trees.Block(bodyStats, Trees.If(cond, _, _)) = ddef.rhs
479+
Some((loopBody(bodyStats), cond))
480+
case _ => None
481+
}
482+
}
483+
484+
private def loopBody(stats: List[tpd.Tree])(implicit ctx: Context): tpd.Tree = stats match {
485+
case body :: Nil => body
486+
case stats => tpd.Block(stats.init, stats.last)
487+
}
444488
}
445489

446490
// ----- CaseDef --------------------------------------------------

‎library/src/scala/tasty/Tasty.scala

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -339,6 +339,17 @@ abstract class Tasty { tasty =>
339339
def unapply(x: Term)(implicit ctx: Context): Option[(Term, Int, Type)]
340340
}
341341

342+
val While: WhileExtractor
343+
abstract class WhileExtractor {
344+
/** Extractor for while loops. Matches `while (<cond>) <body>` and returns (<cond>, <body>) */
345+
def unapply(x: Term)(implicit ctx: Context): Option[(Term, Term)]
346+
}
347+
348+
val DoWhile: DoWhileExtractor
349+
abstract class DoWhileExtractor {
350+
/** Extractor for do while loops. Matches `do <body> while (<cond>)` and returns (<body>, <cond>) */
351+
def unapply(x: Term)(implicit ctx: Context): Option[(Term, Term)]
352+
}
342353
}
343354

344355
// ----- CaseDef --------------------------------------------------

‎library/src/scala/tasty/util/ShowSourceCode.scala

Lines changed: 8 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -204,33 +204,15 @@ class ShowSourceCode[T <: Tasty with Singleton](tasty0: T) extends Show[T](tasty
204204
this
205205
}
206206

207-
case While(cond, stats) =>
207+
case Term.While(cond, body) =>
208208
this += "while ("
209209
printTree(cond)
210210
this += ") "
211-
stats match {
212-
case stat :: Nil =>
213-
printTree(stat)
214-
case stats =>
215-
this += "{"
216-
indented {
217-
printStats(stats.init, stats.last)
218-
}
219-
this += lineBreak() += "}"
220-
}
211+
printTree(body)
221212

222-
case DoWhile(stats, cond) =>
213+
case Term.DoWhile(body, cond) =>
223214
this += "do "
224-
stats match {
225-
case stat :: Nil =>
226-
printTree(stat)
227-
case stats =>
228-
this += "{"
229-
indented {
230-
printStats(stats.init, stats.last)
231-
}
232-
this += lineBreak() += "}"
233-
}
215+
printTree(body)
234216
this += " while ("
235217
printTree(cond)
236218
this += ")"
@@ -349,14 +331,11 @@ class ShowSourceCode[T <: Tasty with Singleton](tasty0: T) extends Show[T](tasty
349331
printTree(rhs)
350332

351333
case Term.Block(stats0, expr) =>
352-
def shouldNotBePrinted(tree: Tree): Boolean = tree match {
353-
case Term.Apply(Term.Ident("while$" | "doWhile$"), _) => true
354-
case tree @ ValDef(_, _, _) => tree.flags.isObject
355-
case _ => false
334+
val stats = stats0.filter {
335+
case tree @ ValDef(_, _, _) => !tree.flags.isObject
336+
case _ => true
356337
}
357338

358-
val stats = stats0.filterNot(shouldNotBePrinted)
359-
360339
expr match {
361340
case Term.Lambda(_, _) =>
362341
// Decompile lambda from { def annon$(...) = ...; closure(annon$, ...)}
@@ -368,11 +347,8 @@ class ShowSourceCode[T <: Tasty with Singleton](tasty0: T) extends Show[T](tasty
368347
this += ")"
369348
case _ =>
370349
this += "{"
371-
val (stats1, expr1) =
372-
if (shouldNotBePrinted(expr)) (stats.init, stats.last)
373-
else (stats, expr)
374350
indented {
375-
printStats(stats1, expr1)
351+
printStats(stats, expr)
376352
}
377353
this += lineBreak() += "}"
378354
}
@@ -441,7 +417,6 @@ class ShowSourceCode[T <: Tasty with Singleton](tasty0: T) extends Show[T](tasty
441417
// Avoid accidental application of opening `{` on next line with a double break
442418
val next = if (nextStats.isEmpty) expr else nextStats.head
443419
next match {
444-
case Term.Block(DefDef("while$" | "doWhile$", _, _, _, _) :: Nil, _) => this += lineBreak()
445420
case Term.Block(_, _) => this += doubleLineBreak()
446421
case Term.Inlined(_, _, _) => this += doubleLineBreak()
447422
case _ => this += lineBreak()
@@ -1073,7 +1048,6 @@ class ShowSourceCode[T <: Tasty with Singleton](tasty0: T) extends Show[T](tasty
10731048
private def escapedString(str: String): String = str flatMap escapedChar
10741049
}
10751050

1076-
10771051
private object SpecialOp {
10781052
def unapply(arg: Term)(implicit ctx: Context): Option[(String, List[Term])] = arg match {
10791053
case arg@Term.Apply(fn, args) =>
@@ -1086,22 +1060,6 @@ class ShowSourceCode[T <: Tasty with Singleton](tasty0: T) extends Show[T](tasty
10861060
}
10871061
}
10881062

1089-
private object While {
1090-
def unapply(arg: Tree)(implicit ctx: Context): Option[(Term, List[Statement])] = arg match {
1091-
case DefDef("while$", _, _, _, Some(Term.If(cond, Term.Block(bodyStats, _), _))) => Some((cond, bodyStats))
1092-
case Term.Block(List(tree), _) => unapply(tree)
1093-
case _ => None
1094-
}
1095-
}
1096-
1097-
private object DoWhile {
1098-
def unapply(arg: Tree)(implicit ctx: Context): Option[(List[Statement], Term)] = arg match {
1099-
case DefDef("doWhile$", _, _, _, Some(Term.Block(body, Term.If(cond, _, _)))) => Some((body, cond))
1100-
case Term.Block(List(tree), _) => unapply(tree)
1101-
case _ => None
1102-
}
1103-
}
1104-
11051063
private object Annotation {
11061064
def unapply(arg: Tree)(implicit ctx: Context): Option[(TypeTree, List[Term])] = arg match {
11071065
case Term.New(annot) => Some((annot, Nil))

‎tests/pos/tasty/definitions.scala

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,8 @@ object definitions {
8585
case Return(expr: Term)
8686
case Repeated(args: List[Term])
8787
case SelectOuter(from: Term, levels: Int, target: Type) // can be generated by inlining
88+
case While(cond: Term, body: Term)
89+
case DoWhile(body: Term, cond: Term)
8890
}
8991

9092
/** Trees denoting types */

0 commit comments

Comments
 (0)