Skip to content

Commit 004401b

Browse files
committed
Fix #854: Optimize matches on primitive constants as switches.
This does not yet unable the checks that `@switch` verifies that the compiler was indeed able to perform the optimization. This implementation also does not support guards. A match with guards will never be optimized as a switch.
1 parent 93dd1cf commit 004401b

File tree

1 file changed

+128
-2
lines changed

1 file changed

+128
-2
lines changed

src/dotty/tools/dotc/transform/PatternMatcher.scala

+128-2
Original file line numberDiff line numberDiff line change
@@ -303,8 +303,134 @@ class PatternMatcher extends MiniPhaseTransform with DenotTransformer {thisTrans
303303
def optimizeCases(prevBinder: Symbol, cases: List[List[TreeMaker]], pt: Type): (List[List[TreeMaker]], List[Tree])
304304
def analyzeCases(prevBinder: Symbol, cases: List[List[TreeMaker]], pt: Type, suppression: Suppression): Unit = {}
305305

306-
def emitSwitch(scrut: Tree, scrutSym: Symbol, cases: List[List[TreeMaker]], pt: Type, matchFailGenOverride: Option[Symbol => Tree], unchecked: Boolean): Option[Tree] =
307-
None // todo
306+
def emitSwitch(scrut: Tree, scrutSym: Symbol, cases: List[List[TreeMaker]], pt: Type, matchFailGenOverride: Option[Symbol => Tree], unchecked: Boolean): Option[Tree] = {
307+
// TODO Deal with guards?
308+
309+
def isSwitchableType(tpe: Type): Boolean = {
310+
(tpe isRef defn.IntClass) ||
311+
(tpe isRef defn.ByteClass) ||
312+
(tpe isRef defn.ShortClass) ||
313+
(tpe isRef defn.CharClass)
314+
}
315+
316+
object IntEqualityTestTreeMaker {
317+
def unapply(treeMaker: EqualityTestTreeMaker): Option[Int] = treeMaker match {
318+
case EqualityTestTreeMaker(`scrutSym`, _, Literal(const), _) =>
319+
if (const.isIntRange) Some(const.intValue)
320+
else None
321+
case _ =>
322+
None
323+
}
324+
}
325+
326+
def isSwitchCase(treeMakers: List[TreeMaker]): Boolean = treeMakers match {
327+
// case 5 =>
328+
case List(IntEqualityTestTreeMaker(_), _: BodyTreeMaker) =>
329+
true
330+
331+
// case 5 | 6 =>
332+
case List(AlternativesTreeMaker(`scrutSym`, alts, _), _: BodyTreeMaker) =>
333+
alts.forall {
334+
case List(IntEqualityTestTreeMaker(_)) => true
335+
case _ => false
336+
}
337+
338+
// case _ =>
339+
case List(_: BodyTreeMaker) =>
340+
true
341+
342+
// case x =>
343+
case List(_: SubstOnlyTreeMaker, _: BodyTreeMaker) =>
344+
true
345+
346+
case _ =>
347+
false
348+
}
349+
350+
/* (Nil, body) means that `body` is the default case
351+
* It's a bit hacky but it simplifies manipulations.
352+
*/
353+
def extractSwitchCase(treeMakers: List[TreeMaker]): (List[Int], BodyTreeMaker) = treeMakers match {
354+
// case 5 =>
355+
case List(IntEqualityTestTreeMaker(intValue), body: BodyTreeMaker) =>
356+
(List(intValue), body)
357+
358+
// case 5 | 6 =>
359+
case List(AlternativesTreeMaker(_, alts, _), body: BodyTreeMaker) =>
360+
val intValues = alts.map {
361+
case List(IntEqualityTestTreeMaker(intValue)) => intValue
362+
}
363+
(intValues, body)
364+
365+
// case _ =>
366+
case List(body: BodyTreeMaker) =>
367+
(Nil, body)
368+
369+
// case x =>
370+
case List(subst: SubstOnlyTreeMaker, body: BodyTreeMaker) =>
371+
/* Rebindings have been propagated, so `body` contains all the
372+
* necessary information. `subst` can be dropped at this point.
373+
*/
374+
(Nil, body)
375+
}
376+
377+
def doOverlap(a: List[Int], b: List[Int]): Boolean =
378+
a.exists(b.contains _)
379+
380+
def makeSwitch(valuesToCases: List[(List[Int], BodyTreeMaker)]): Tree = {
381+
def genBody(body: BodyTreeMaker): Tree = {
382+
val valDefs = body.rebindings.emitValDefs
383+
if (valDefs.isEmpty) body.body
384+
else Block(valDefs, body.body)
385+
}
386+
387+
val intScrut =
388+
if (pt isRef defn.IntClass) ref(scrutSym)
389+
else Select(ref(scrutSym), nme.toInt)
390+
391+
val (normalCases, defaultCaseAndRest) = valuesToCases.span(_._1.nonEmpty)
392+
393+
val newCases = for {
394+
(values, body) <- normalCases
395+
} yield {
396+
val literals = values.map(v => Literal(Constant(v)))
397+
val pat =
398+
if (literals.size == 1) literals.head
399+
else Alternative(literals)
400+
CaseDef(pat, EmptyTree, genBody(body))
401+
}
402+
403+
val wildcard = ctx.typeAssigner.assignType(untpd.Ident(nme.WILDCARD), defn.IntType)
404+
val catchAllDef = {
405+
if (defaultCaseAndRest.isEmpty) {
406+
matchFailGenOverride.fold[Tree](
407+
Throw(New(defn.MatchErrorType, List(ref(scrutSym)))))(
408+
_(scrutSym))
409+
} else {
410+
/* After the default case, assuming the IR even allows anything,
411+
* things are unreachable anyway and can be removed.
412+
*/
413+
genBody(defaultCaseAndRest.head._2)
414+
}
415+
}
416+
val defaultCase = CaseDef(wildcard, EmptyTree, catchAllDef)
417+
418+
Match(intScrut, newCases :+ defaultCase)
419+
}
420+
421+
if (isSwitchableType(scrut.tpe.widenDealias) && cases.forall(isSwitchCase)) {
422+
val valuesToCases = cases.map(extractSwitchCase)
423+
val values = valuesToCases.map(_._1)
424+
if (values.tails.exists { tail => tail.nonEmpty && tail.tail.exists(doOverlap(_, tail.head)) }) {
425+
// TODO Deal with overlapping cases (mostly useless without guards)
426+
None
427+
} else {
428+
Some(makeSwitch(valuesToCases))
429+
}
430+
} else {
431+
None
432+
}
433+
}
308434

309435
// for catch (no need to customize match failure)
310436
def emitTypeSwitch(bindersAndCases: List[(Symbol, List[TreeMaker])], pt: Type): Option[List[CaseDef]] =

0 commit comments

Comments
 (0)