JS: translate when to JsSwitch when possible

See KT-21160
2017-11-09 19:14:01 +03:00
parent 7b403211dd
commit 678b4c67c1
21 changed files with 250 additions and 24 deletions
@@ -1,4 +1,6 @@
 // WITH_RUNTIME
+// CHECK_CASES_COUNT: function=foo count=3
+// CHECK_IF_COUNT: function=foo count=0

 import kotlin.test.assertEquals

@@ -1,4 +1,6 @@
 // WITH_RUNTIME
+// CHECK_CASES_COUNT: function=foo count=4
+// CHECK_IF_COUNT: function=foo count=0

 import kotlin.test.assertEquals

@@ -1,4 +1,8 @@
 // WITH_RUNTIME
+// CHECK_CASES_COUNT: function=foo1 count=2
+// CHECK_IF_COUNT: function=foo1 count=1
+// CHECK_CASES_COUNT: function=foo2 count=4
+// CHECK_IF_COUNT: function=foo2 count=0

 import kotlin.test.assertEquals

@@ -1,4 +1,8 @@
 // WITH_RUNTIME
+// CHECK_CASES_COUNT: function=foo1 count=4
+// CHECK_IF_COUNT: function=foo1 count=0
+// CHECK_CASES_COUNT: function=foo2 count=4
+// CHECK_IF_COUNT: function=foo2 count=0

 import kotlin.test.assertEquals

@@ -1,4 +1,6 @@
 // WITH_RUNTIME
+// CHECK_CASES_COUNT: function=dense count=10
+// CHECK_IF_COUNT: function=dense count=0

 fun dense(x: Int): Int {
    return when (x) {
@@ -1,4 +1,6 @@
 // WITH_RUNTIME
+// CHECK_CASES_COUNT: function=foo count=9
+// CHECK_IF_COUNT: function=foo count=0

 fun foo(x: Int): Int {
    return when (x) {
@@ -1,4 +1,6 @@
 // WITH_RUNTIME
+// CHECK_CASES_COUNT: function=sparse count=3
+// CHECK_IF_COUNT: function=sparse count=0

 fun sparse(x: Int): Int {
    return when ((x % 4) * 100) {
@@ -1,4 +1,8 @@
 // WITH_RUNTIME
+// CHECK_CASES_COUNT: function=exhaustive count=3
+// CHECK_IF_COUNT: function=exhaustive count=0
+// CHECK_CASES_COUNT: function=nonExhaustive count=3
+// CHECK_IF_COUNT: function=nonExhaustive count=0

 fun exhaustive(x: Int): Int {
    var r: Int
@@ -1,4 +1,12 @@
 // WITH_RUNTIME
+// CHECK_CASES_COUNT: function=intFoo count=3
+// CHECK_IF_COUNT: function=intFoo count=0
+// CHECK_CASES_COUNT: function=shortFoo count=3
+// CHECK_IF_COUNT: function=shortFoo count=0
+// CHECK_CASES_COUNT: function=byteFoo count=3
+// CHECK_IF_COUNT: function=byteFoo count=0
+// CHECK_CASES_COUNT: function=charFoo count=3
+// CHECK_IF_COUNT: function=charFoo count=0

 fun intFoo(x: Int): Int {
    return when (x) {
@@ -1,4 +1,6 @@
 // WITH_RUNTIME
+// CHECK_CASES_COUNT: function=foo count=3
+// CHECK_IF_COUNT: function=foo count=0

 fun foo(x: Int): Int {
    return when (x) {
@@ -0,0 +1,31 @@
+// WITH_RUNTIME
+// CHECK_CASES_COUNT: function=foo count=4
+// CHECK_IF_COUNT: function=foo count=1
+
+var log = ""
+
+fun foo(x: Int): String {
+    log += "foo($x);"
+    return when (x) {
+        1 -> "one"
+        2 -> "two"
+        three(x) -> "three"
+        4 -> "four"
+        5 -> "five"
+        else -> "many"
+    }
+}
+
+fun three(x: Int): Int {
+    log += "three($x);"
+    return 3
+}
+
+fun box(): String {
+    var result = (1..7).map(::foo).joinToString()
+
+    if (result != "one, two, three, four, five, many, many") return "fail1: $result"
+    if (log != "foo(1);foo(2);foo(3);three(3);foo(4);three(4);foo(5);three(5);foo(6);three(6);foo(7);three(7);") return "fail2: $log"
+
+    return "OK"
+}
@@ -20747,6 +20747,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
            doTest(fileName);
        }

+        @TestMetadata("switchOptimizationWithGap.kt")
+        public void testSwitchOptimizationWithGap() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/when/switchOptimizationWithGap.kt");
+            doTest(fileName);
+        }
+
        @TestMetadata("typeDisjunction.kt")
        public void testTypeDisjunction() throws Exception {
            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/when/typeDisjunction.kt");
@@ -20747,6 +20747,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
            doTest(fileName);
        }

+        @TestMetadata("switchOptimizationWithGap.kt")
+        public void testSwitchOptimizationWithGap() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/when/switchOptimizationWithGap.kt");
+            doTest(fileName);
+        }
+
        @TestMetadata("typeDisjunction.kt")
        public void testTypeDisjunction() throws Exception {
            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/when/typeDisjunction.kt");
@@ -20747,6 +20747,12 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
            doTest(fileName);
        }

+        @TestMetadata("switchOptimizationWithGap.kt")
+        public void testSwitchOptimizationWithGap() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/when/switchOptimizationWithGap.kt");
+            doTest(fileName);
+        }
+
        @TestMetadata("typeDisjunction.kt")
        public void testTypeDisjunction() throws Exception {
            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/when/typeDisjunction.kt");
@@ -306,6 +306,12 @@ public class JsLineNumberTestGenerated extends AbstractJsLineNumberTest {
        doTest(fileName);
    }

+    @TestMetadata("whenEntryWithMultipleConditionsNonOptimized.kt")
+    public void testWhenEntryWithMultipleConditionsNonOptimized() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("js/js.translator/testData/lineNumbers/whenEntryWithMultipleConditionsNonOptimized.kt");
+        doTest(fileName);
+    }
+
    @TestMetadata("whenIn.kt")
    public void testWhenIn() throws Exception {
        String fileName = KotlinTestUtils.navigationMetadata("js/js.translator/testData/lineNumbers/whenIn.kt");
@@ -24503,6 +24503,12 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
            doTest(fileName);
        }

+        @TestMetadata("switchOptimizationWithGap.kt")
+        public void testSwitchOptimizationWithGap() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/when/switchOptimizationWithGap.kt");
+            doTest(fileName);
+        }
+
        @TestMetadata("typeDisjunction.kt")
        public void testTypeDisjunction() throws Exception {
            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/when/typeDisjunction.kt");
@@ -205,6 +205,10 @@ public class DirectiveTestUtils {

    private static final DirectiveHandler COUNT_NULLS = new CountNodesDirective<>("CHECK_NULLS_COUNT", JsNullLiteral.class);

+    private static final DirectiveHandler COUNT_CASES = new CountNodesDirective<>("CHECK_CASES_COUNT", JsCase.class);
+
+    private static final DirectiveHandler COUNT_IF = new CountNodesDirective<>("CHECK_IF_COUNT", JsIf.class);
+
    private static final DirectiveHandler COUNT_DEBUGGER = new CountNodesDirective<>("CHECK_DEBUGGER_COUNT", JsDebugger.class);

    private static final DirectiveHandler NOT_REFERENCED = new DirectiveHandler("CHECK_NOT_REFERENCED") {
@@ -339,6 +343,8 @@ public class DirectiveTestUtils {
            COUNT_VARS,
            COUNT_BREAKS,
            COUNT_NULLS,
+            COUNT_CASES,
+            COUNT_IF,
            COUNT_DEBUGGER,
            NOT_REFERENCED,
            HAS_INLINE_METADATA,
@@ -31,12 +31,14 @@ import org.jetbrains.kotlin.js.translate.utils.mutator.LastExpressionMutator
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getTextWithLocation
+import org.jetbrains.kotlin.resolve.constants.evaluate.ConstantExpressionEvaluator
 import org.jetbrains.kotlin.types.KotlinType

 class WhenTranslator
 private constructor(private val whenExpression: KtWhenExpression, context: TranslationContext) : AbstractTranslator(context) {
    private val expressionToMatch: JsExpression?
    private val type: KotlinType?
+    private val uniqueConstants = mutableSetOf<Any>()

    private val isExhaustive: Boolean
        get() {
@@ -53,9 +55,27 @@ private constructor(private val whenExpression: KtWhenExpression, context: Trans
    }

    private fun translate(): JsNode {
-        var currentIf: JsIf? = null
-        var resultIf: JsIf? = null
-        for (entry in whenExpression.entries) {
+        var resultIf: JsNode? = null
+        var setWhenStatement: (JsStatement) -> Unit = { resultIf = it }
+
+        var i = 0
+        var hasElse = false
+        while (i < whenExpression.entries.size) {
+            val asSwitch = translateAsSwitch(i)
+            if (asSwitch != null) {
+                val (jsSwitch, next) = asSwitch
+                setWhenStatement(jsSwitch)
+                setWhenStatement = { whenStatement ->
+                    jsSwitch.cases += JsDefault().apply {
+                        statements += whenStatement
+                        statements += JsBreak().apply { source = whenExpression }
+                    }
+                }
+                i = next
+                continue
+            }
+
+            val entry = whenExpression.entries[i++]
            val statementBlock = JsBlock()
            var statement = translateEntryExpression(entry, context(), statementBlock)

@@ -65,32 +85,90 @@ private constructor(private val whenExpression: KtWhenExpression, context: Trans
            }
            statement = JsAstUtils.mergeStatementInBlockIfNeeded(statement, statementBlock)

-            if (resultIf == null) {
-                currentIf = JsAstUtils.newJsIf(translateConditions(entry, context()), statement)
-                currentIf.source = entry
-                resultIf = currentIf
-            }
-            else {
-                currentIf!!
-                if (entry.isElse) {
-                    currentIf.elseStatement = statement
-                    return resultIf
-                }
-                val conditionsBlock = JsBlock()
-                val nextIf = JsAstUtils.newJsIf(translateConditions(entry, context().innerBlock(conditionsBlock)), statement)
-                nextIf.source = entry
-                val statementToAdd = JsAstUtils.mergeStatementInBlockIfNeeded(nextIf, conditionsBlock)
-                currentIf.elseStatement = statementToAdd
-                currentIf = nextIf
+            val conditionsBlock = JsBlock()
+            if (entry.isElse) {
+                hasElse = true
+                setWhenStatement(statement)
+                break
            }
+            val jsIf = JsAstUtils.newJsIf(translateConditions(entry, context().innerBlock(conditionsBlock)), statement)
+            jsIf.source = entry
+
+            val statementToAdd = JsAstUtils.mergeStatementInBlockIfNeeded(jsIf, conditionsBlock)
+            setWhenStatement(statementToAdd)
+            setWhenStatement = { jsIf.elseStatement = it }
        }

-        if (currentIf != null && currentIf.elseStatement == null && isExhaustive) {
+        if (isExhaustive && !hasElse) {
            val noWhenMatchedInvocation = JsInvocation(JsAstUtils.pureFqn("noWhenBranchMatched", Namer.kotlinObject()))
-            currentIf.elseStatement = JsAstUtils.asSyntheticStatement(noWhenMatchedInvocation)
+            setWhenStatement(JsAstUtils.asSyntheticStatement(noWhenMatchedInvocation))
        }

-        return if (resultIf != null) resultIf else JsNullLiteral()
+        return if (resultIf != null) resultIf!! else JsNullLiteral()
+    }
+
+    private fun translateAsSwitch(fromIndex: Int): Pair<JsSwitch, Int>? {
+        val expectedType = type ?: return null
+        val subject = expressionToMatch ?: return null
+
+        val entries = whenExpression.entries
+        val entriesForSwitch = mutableListOf<Pair<List<JsExpression>, KtWhenEntry>>()
+        var i = fromIndex
+        while (i < entries.size) {
+            val entry = entries[i]
+            if (entry.isElse) break
+
+            var hasImproperConstants = false
+            val constantValues = entry.conditions.mapNotNull { condition ->
+                val expression = (condition as? KtWhenConditionWithExpression)?.expression
+                expression?.let { ConstantExpressionEvaluator.getConstant(it, bindingContext())?.getValue(expectedType) } ?: run {
+                    hasImproperConstants = true
+                    null
+                }
+            }
+            if (hasImproperConstants) break
+
+            val constants = constantValues.filter { uniqueConstants.add(it) }.mapNotNull { value ->
+                when (value) {
+                    is String -> JsStringLiteral(value)
+                    is Int -> JsIntLiteral(value)
+                    is Short -> JsIntLiteral(value.toInt())
+                    is Byte -> JsIntLiteral(value.toInt())
+                    else -> {
+                        hasImproperConstants = true
+                        null
+                    }
+                }
+            }
+            if (hasImproperConstants) break
+
+            if (constants.isNotEmpty()) {
+                entriesForSwitch += Pair(constants, entry)
+            }
+            i++
+        }
+
+        return if (entriesForSwitch.asSequence().map { it.first.size }.sum() > 1) {
+            val switchEntries = mutableListOf<JsSwitchMember>()
+            entriesForSwitch.flatMapTo(switchEntries) { (conditions, entry) ->
+                val members = conditions.map {
+                    JsCase().apply {
+                        caseExpression = it.source(entry)
+                    }
+                }
+                val block = JsBlock()
+                val statement = translateEntryExpression(entry, context(), block)
+                val lastEntry = members.last()
+                lastEntry.statements += block.statements
+                lastEntry.statements += statement
+                lastEntry.statements += JsBreak().apply { source = entry }
+                members
+            }
+            Pair(JsSwitch(subject, switchEntries).apply { source = expression }, i)
+        }
+        else {
+            null
+        }
    }

    private fun translateEntryExpression(
@@ -51,6 +51,9 @@ public final class LastExpressionMutator {
        if (node instanceof JsExpressionStatement) {
            return applyToStatement((JsExpressionStatement) node);
        }
+        if (node instanceof JsSwitch) {
+            return applyToSwitch((JsSwitch) node);
+        }
        return mutator.mutate(node);
    }

@@ -88,4 +91,19 @@ public final class LastExpressionMutator {
        statements.set(size - 1, convertToStatement(apply(statements.get(size - 1))));
        return node;
    }
+
+    @NotNull
+    private JsNode applyToSwitch(@NotNull JsSwitch node) {
+        for (JsSwitchMember member : node.getCases()) {
+            int size = member.getStatements().size();
+            if (size < 2) continue;
+
+            JsNode lastStatement = apply(member.getStatements().get(size - 1));
+            if (!(lastStatement instanceof JsBreak)) continue;
+
+            member.getStatements().set(size - 2, convertToStatement(apply(member.getStatements().get(size - 2))));
+        }
+
+        return node;
+    }
 }
@@ -18,4 +18,4 @@ fun box(x: Int) {
    )
 }

-// LINES: 19 4 4 3 6 4 6 4 7 4 8 9 9 11 4 11 4 12 13 13 16 16 2
+// LINES: 19 4 4 3 4 6 9 9 6 11 13 13 11 16 16 3 2
@@ -0,0 +1,31 @@
+fun box(x: Int) {
+    println(
+            when (
+                x
+            ) {
+                one(),
+                two(),
+                three() ->
+                    55
+
+                four(),
+                five() ->
+                    66
+
+                else ->
+                    77
+            }
+    )
+}
+
+fun one() = 1
+
+fun two() = 2
+
+fun three() = 3
+
+fun four() = 4
+
+fun five() = 5
+
+// LINES: 19 4 4 3 6 4 6 4 7 4 8 9 9 11 4 11 4 12 13 13 16 16 2 21 21 21 23 23 23 25 25 25 27 27 27 29 29 29