From c73533d214dbd40f655b9c8e4390ca794a005792 Mon Sep 17 00:00:00 2001 From: Denis Zharkov Date: Fri, 4 Apr 2014 17:28:30 +0400 Subject: [PATCH] Switch Optimization --- .../jet/codegen/ExpressionCodegen.java | 132 ++++++++++++++++++ .../codegen/box/when/switchOptimization.kt | 75 ++++++++++ .../BlackBoxCodegenTestGenerated.java | 5 + 3 files changed, 212 insertions(+) create mode 100644 compiler/testData/codegen/box/when/switchOptimization.kt diff --git a/compiler/backend/src/org/jetbrains/jet/codegen/ExpressionCodegen.java b/compiler/backend/src/org/jetbrains/jet/codegen/ExpressionCodegen.java index 2713f8ec425..dd57fb6c3ea 100644 --- a/compiler/backend/src/org/jetbrains/jet/codegen/ExpressionCodegen.java +++ b/compiler/backend/src/org/jetbrains/jet/codegen/ExpressionCodegen.java @@ -20,6 +20,7 @@ import com.google.common.collect.Lists; import com.google.common.collect.Maps; import com.intellij.openapi.editor.Document; import com.intellij.openapi.progress.ProcessCanceledException; +import com.intellij.openapi.util.Pair; import com.intellij.psi.PsiElement; import com.intellij.psi.tree.IElementType; import com.intellij.util.ArrayUtil; @@ -46,6 +47,7 @@ import org.jetbrains.jet.lang.resolve.DescriptorUtils; import org.jetbrains.jet.lang.resolve.calls.model.*; import org.jetbrains.jet.lang.resolve.calls.util.CallMaker; import org.jetbrains.jet.lang.resolve.constants.CompileTimeConstant; +import org.jetbrains.jet.lang.resolve.constants.IntValue; import org.jetbrains.jet.lang.resolve.constants.IntegerValueTypeConstant; import org.jetbrains.jet.lang.resolve.java.AsmTypeConstants; import org.jetbrains.jet.lang.resolve.java.JvmAbi; @@ -3641,7 +3643,13 @@ The "returned" value of try expression with no finally is either the last expres JetExpression expr = expression.getSubjectExpression(); JetType subjectJetType = bindingContext.get(BindingContext.EXPRESSION_TYPE, expr); Type subjectType = asmTypeOrVoid(subjectJetType); + Type resultType = isStatement ? Type.VOID_TYPE : expressionType(expression); + + if (canSwitchBeUsedIn(expression, subjectType)) { + return generateSwitch(expression, subjectType, resultType, isStatement); + } + int subjectLocal = expr != null ? myFrameMap.enterTemp(subjectType) : -1; if (subjectLocal != -1) { gen(expr, subjectType); @@ -3724,6 +3732,130 @@ The "returned" value of try expression with no finally is either the last expres } } + private StackValue generateSwitch(JetWhenExpression expression, Type subjectType, Type resultType, boolean isStatement) { + JetType subjectJetType = bindingContext.get(BindingContext.EXPRESSION_TYPE, expression.getSubjectExpression()); + assert subjectJetType != null; + + Map transitions = Maps.newTreeMap(); + + Label[] entryLabels = new Label[expression.getEntries().size()]; + int entryLabelsCounter = 0; + + Label elseLabel = new Label(); + Label endLabel = new Label(); + boolean hasElse = expression.getElseExpression() != null; + + for (JetWhenEntry entry : expression.getEntries()) { + Label entryLabel = new Label(); + + for (JetWhenCondition condition : entry.getConditions()) { + assert condition instanceof JetWhenConditionWithExpression; + + JetExpression conditionExpression = ((JetWhenConditionWithExpression) condition).getExpression(); + assert conditionExpression != null; + + CompileTimeConstant constant = getCompileTimeConstant(conditionExpression, bindingContext); + assert constant != null && constant instanceof IntValue; + + int value = ((IntValue) constant).getValue(); + + if (!transitions.containsKey(value)) { + transitions.put(value, entryLabel); + } + } + + if (entry.isElse()) { + elseLabel = entryLabel; + } + + entryLabels[entryLabelsCounter++] = entryLabel; + } + + gen(expression.getSubjectExpression(), subjectType); + generateSwitchLookupCall( + transitions, + //if there is no else-entry and it's statement then default --- endLabel + (hasElse || !isStatement) ? elseLabel : endLabel + ); + + //resolving entries' labels + int i = 0; + for (JetWhenEntry entry : expression.getEntries()) { + v.visitLabel(entryLabels[i++]); + + FrameMap.Mark mark = myFrameMap.mark(); + gen(entry.getExpression(), resultType); + mark.dropTo(); + + if (!entry.isElse()) { + v.goTo(endLabel); + } + } + + //there is no else-entry but this is not statement, so we should return Unit + if (!hasElse && !isStatement) { + v.visitLabel(elseLabel); + // a result is expected + if (Boolean.TRUE.equals(bindingContext.get(BindingContext.EXHAUSTIVE_WHEN, expression))) { + // when() is supposed to be exhaustive + throwNewException("kotlin/NoWhenBranchMatchedException"); + } + else { + // non-exhaustive when() with no else -> Unit must be expected + StackValue.putUnitInstance(v); + } + } + + markLineNumber(expression); + v.mark(endLabel); + + return StackValue.onStack(resultType); + } + + private void generateSwitchLookupCall(Map transitions, Label defaultLabel) { + int[] keys = new int[transitions.size()]; + Label[] labels = new Label[transitions.size()]; + int i = 0; + + for (Map.Entry transition : transitions.entrySet()) { + keys[i] = transition.getKey(); + labels[i] = transition.getValue(); + + i++; + } + + //switch code + v.lookupswitch(defaultLabel, keys, labels); + } + + private boolean canSwitchBeUsedIn(JetWhenExpression expression, Type subjectType) { + if (subjectType.getSort() != Type.INT) { + return false; + } + + for (JetWhenEntry entry : expression.getEntries()) { + for (JetWhenCondition condition : entry.getConditions()) { + if (!(condition instanceof JetWhenConditionWithExpression)) { + return false; + } + + //ensure that expression is constant + JetExpression patternExpression = ((JetWhenConditionWithExpression) condition).getExpression(); + + if (patternExpression == null) { + return false; + } + + CompileTimeConstant constant = getCompileTimeConstant(patternExpression, bindingContext); + if (constant == null || !(constant instanceof IntValue)) { + return false; + } + } + } + + return true; + } + private void invokeFunctionByReference(JetSimpleNameExpression operationReference) { ResolvedCall resolvedCall = bindingContext.get(RESOLVED_CALL, operationReference); Call call = bindingContext.get(CALL, operationReference); diff --git a/compiler/testData/codegen/box/when/switchOptimization.kt b/compiler/testData/codegen/box/when/switchOptimization.kt new file mode 100644 index 00000000000..26aba979fd6 --- /dev/null +++ b/compiler/testData/codegen/box/when/switchOptimization.kt @@ -0,0 +1,75 @@ +fun foo1(x : Int) : Int { + return when (x % 10) { + 1,2,3 -> 1 + 4,5,6 -> 2 + 7,8,9 -> 3 + else -> 4 + } +} + +fun foo2(x : Int) : Int { + when (x % 10) { + 1,2,3 -> return 1 + 4,5,6 -> return 2 + 7,8,9 -> return 3 + } + + return 4 +} + +fun foo3(x : Int) : Int { + when (x % 10) { + 1 -> return 1 + 2 -> return 1 + 3 -> return 1 + 4 -> return 2 + 5 -> return 2 + 6 -> return 2 + 7 -> return 3 + 8 -> return 3 + 9 -> return 3 + } + + return 4 +} + +fun foo4(x : Int) : Int { + return when (x % 10) { + 1 -> 1 + 2 -> 1 + 3 -> 1 + 4 -> 2 + 5 -> 2 + 6 -> 2 + 7 -> 3 + 8 -> 3 + 9 -> 3 + else -> 4 + } +} + +fun test(foo : (Int) -> Int, name : String) : String { + val foo0 = foo(10) + + var result = if (foo0 == 4) "" else "$name[10/$foo0/4]" + + for (i in 11..19) { + val shouldBe = (i-11) / 3 + 1 + val fooI = foo(i) + if (fooI != shouldBe) { + result += "$name[$i/$fooI/$shouldBe]" + } + } + + return result +} + +fun box(): String { + val testResult = test({x -> foo1(x)}, "foo1") + + test({x -> foo2(x)}, "foo2") + + test({x -> foo3(x)}, "foo3") + + test({x -> foo4(x)}, "foo4") + if (testResult != "") return testResult + return "OK" +} + diff --git a/compiler/tests/org/jetbrains/jet/codegen/generated/BlackBoxCodegenTestGenerated.java b/compiler/tests/org/jetbrains/jet/codegen/generated/BlackBoxCodegenTestGenerated.java index 339aab09911..d2949f2c81b 100644 --- a/compiler/tests/org/jetbrains/jet/codegen/generated/BlackBoxCodegenTestGenerated.java +++ b/compiler/tests/org/jetbrains/jet/codegen/generated/BlackBoxCodegenTestGenerated.java @@ -5583,6 +5583,11 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest { doTest("compiler/testData/codegen/box/when/range.kt"); } + @TestMetadata("switchOptimization.kt") + public void testSwitchOptimization() throws Exception { + doTest("compiler/testData/codegen/box/when/switchOptimization.kt"); + } + @TestMetadata("whenArgumentIsEvaluatedOnlyOnce.kt") public void testWhenArgumentIsEvaluatedOnlyOnce() throws Exception { doTest("compiler/testData/codegen/box/when/whenArgumentIsEvaluatedOnlyOnce.kt");