From 5a591be25f9f4fcce0d2445889ea5488160eddd5 Mon Sep 17 00:00:00 2001
From: Denis Zharkov <denis.zharkov@jetbrains.com>
Date: Tue, 13 Feb 2018 12:05:54 +0300
Subject: [PATCH] Convert SwitchCodegen to Kotlin, prettify and optimize
 codegen

---
 .../kotlin/codegen/when/SwitchCodegen.java    | 236 ------------------
 .../kotlin/codegen/when/SwitchCodegen.kt      | 181 ++++++++++++++
 .../when/exhaustiveWhenInitialization.kt      |   6 +-
 .../bytecodeText/when/exhaustiveWhenReturn.kt |   4 +-
 .../integralWhenWithNoInlinedConstants.kt     |   3 +-
 .../functionLiteralInTopLevel.kt              |   2 +-
 .../whenEnumOptimization/kt14802.kt           |   2 +-
 .../whenEnumOptimization/kt15806.kt           |   2 +-
 8 files changed, 190 insertions(+), 246 deletions(-)
 delete mode 100644 compiler/backend/src/org/jetbrains/kotlin/codegen/when/SwitchCodegen.java
 create mode 100644 compiler/backend/src/org/jetbrains/kotlin/codegen/when/SwitchCodegen.kt

diff --git a/compiler/backend/src/org/jetbrains/kotlin/codegen/when/SwitchCodegen.java b/compiler/backend/src/org/jetbrains/kotlin/codegen/when/SwitchCodegen.java
deleted file mode 100644
index 554b914d7a6..00000000000
--- a/compiler/backend/src/org/jetbrains/kotlin/codegen/when/SwitchCodegen.java
+++ /dev/null
@@ -1,236 +0,0 @@
-/*
- * Copyright 2010-2015 JetBrains s.r.o.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.jetbrains.kotlin.codegen.when;
-
-import org.jetbrains.annotations.NotNull;
-import org.jetbrains.annotations.Nullable;
-import org.jetbrains.kotlin.codegen.ExpressionCodegen;
-import org.jetbrains.kotlin.codegen.FrameMap;
-import org.jetbrains.kotlin.psi.KtWhenEntry;
-import org.jetbrains.kotlin.psi.KtWhenExpression;
-import org.jetbrains.kotlin.resolve.BindingContext;
-import org.jetbrains.kotlin.resolve.constants.ConstantValue;
-import org.jetbrains.kotlin.resolve.constants.NullValue;
-import org.jetbrains.kotlin.types.KotlinType;
-import org.jetbrains.kotlin.types.TypeUtils;
-import org.jetbrains.org.objectweb.asm.Label;
-import org.jetbrains.org.objectweb.asm.Type;
-import org.jetbrains.org.objectweb.asm.commons.InstructionAdapter;
-
-import java.util.*;
-
-abstract public class SwitchCodegen {
-    protected final KtWhenExpression expression;
-    protected final boolean isStatement;
-    protected final boolean isExhaustive;
-    protected final ExpressionCodegen codegen;
-    protected final BindingContext bindingContext;
-    protected final Type subjectType;
-    protected final Type resultType;
-    protected final InstructionAdapter v;
-
-    protected final NavigableMap<Integer, Label> transitionsTable = new TreeMap<>();
-    protected final List<Label> entryLabels = new ArrayList<>();
-    protected Label elseLabel = new Label();
-    protected Label endLabel = new Label();
-    protected Label defaultLabel;
-
-    protected final SwitchCodegenProvider switchCodegenProvider;
-
-    public SwitchCodegen(
-            @NotNull KtWhenExpression expression, boolean isStatement,
-            boolean isExhaustive, @NotNull ExpressionCodegen codegen,
-            @Nullable Type subjectType
-    ) {
-        this.expression = expression;
-        this.isStatement = isStatement;
-        this.isExhaustive = isExhaustive;
-        this.codegen = codegen;
-        this.bindingContext = codegen.getBindingContext();
-        this.switchCodegenProvider = new SwitchCodegenProvider(codegen);
-
-        this.subjectType = subjectType != null ? subjectType : codegen.expressionType(expression.getSubjectExpression());
-        resultType = isStatement ? Type.VOID_TYPE : codegen.expressionType(expression);
-        v = codegen.v;
-    }
-
-    /**
-     * Generates bytecode for entire when expression
-     */
-    public void generate() {
-        prepareConfiguration();
-
-        boolean hasElse = expression.getElseExpression() != null;
-
-        // if there is no else-entry and it's statement then default --- endLabel
-        defaultLabel = (hasElse || !isStatement || isExhaustive) ? elseLabel : endLabel;
-
-        generateSubject();
-
-        generateSwitchInstructionByTransitionsTable();
-
-        generateEntries();
-
-        // there is no else-entry but this is not statement, so we should return Unit
-        if (!hasElse && (!isStatement || isExhaustive)) {
-            v.visitLabel(elseLabel);
-            codegen.putUnitInstanceOntoStackForNonExhaustiveWhen(expression, isStatement);
-        }
-
-        codegen.markLineNumber(expression, isStatement);
-        v.mark(endLabel);
-    }
-
-    /**
-     * Sets up transitionsTable and maybe something else needed in a special case
-     * Behaviour may be changed by overriding processConstant
-     */
-    private void prepareConfiguration() {
-        for (KtWhenEntry entry : expression.getEntries()) {
-            Label entryLabel = new Label();
-
-            for (ConstantValue<?> constant : switchCodegenProvider.getConstantsFromEntry(entry)) {
-                if (constant instanceof NullValue) continue;
-                processConstant(constant, entryLabel);
-            }
-
-            if (entry.isElse()) {
-                elseLabel = entryLabel;
-            }
-
-            entryLabels.add(entryLabel);
-        }
-    }
-
-    abstract protected void processConstant(
-            @NotNull ConstantValue<?> constant,
-            @NotNull Label entryLabel
-    );
-
-    protected void putTransitionOnce(int value, @NotNull Label entryLabel) {
-        if (!transitionsTable.containsKey(value)) {
-            transitionsTable.put(value, entryLabel);
-        }
-    }
-
-    /**
-     * Should generate int subject on top of the stack
-     * Default implementation just run codegen for actual subject of expression
-     * May also gen nullability check if needed
-     */
-    protected void generateSubject() {
-        codegen.gen(expression.getSubjectExpression(), subjectType);
-    }
-
-    protected void generateNullCheckIfNeeded() {
-        assert expression.getSubjectExpression() != null : "subject expression can't be null";
-        KotlinType subjectJetType = bindingContext.getType(expression.getSubjectExpression());
-
-        assert subjectJetType != null : "subject type can't be null (i.e. void)";
-
-        if (TypeUtils.isNullableType(subjectJetType)) {
-            int nullEntryIndex = findNullEntryIndex(expression);
-            Label nullLabel = nullEntryIndex == -1 ? defaultLabel : entryLabels.get(nullEntryIndex);
-            Label notNullLabel = new Label();
-
-            v.dup();
-            v.ifnonnull(notNullLabel);
-
-            v.pop();
-
-            v.goTo(nullLabel);
-
-            v.visitLabel(notNullLabel);
-        }
-    }
-
-    private int findNullEntryIndex(@NotNull KtWhenExpression expression) {
-        int entryIndex = 0;
-        for (KtWhenEntry entry : expression.getEntries()) {
-            for (ConstantValue<?> constant : switchCodegenProvider.getConstantsFromEntry(entry)) {
-                if (constant instanceof NullValue) {
-                    return entryIndex;
-                }
-            }
-
-            entryIndex++;
-        }
-
-        return -1;
-    }
-
-    private void generateSwitchInstructionByTransitionsTable() {
-        int[] keys = new int[transitionsTable.size()];
-        Label[] labels = new Label[transitionsTable.size()];
-        int i = 0;
-
-        for (Map.Entry<Integer, Label> transition : transitionsTable.entrySet()) {
-            keys[i] = transition.getKey();
-            labels[i] = transition.getValue();
-
-            i++;
-        }
-
-        int nlabels = keys.length;
-        int hi = keys[nlabels - 1];
-        int lo = keys[0];
-
-        /*
-         * Heuristic estimation if it's better to use tableswitch or lookupswitch.
-         * From OpenJDK sources
-         */
-        long table_space_cost = 4 + ((long) hi - lo + 1); // words
-        long table_time_cost = 3; // comparisons
-        long lookup_space_cost = 3 + 2 * (long) nlabels;
-        //noinspection UnnecessaryLocalVariable
-        long lookup_time_cost = nlabels;
-
-        boolean useTableSwitch = nlabels > 0 &&
-                                 table_space_cost + 3 * table_time_cost <=
-                                 lookup_space_cost + 3 * lookup_time_cost;
-
-        if (!useTableSwitch) {
-            v.lookupswitch(defaultLabel, keys, labels);
-            return;
-        }
-
-        Label[] sparseLabels = new Label[hi - lo + 1];
-        Arrays.fill(sparseLabels, defaultLabel);
-
-        for (i = 0; i < keys.length; i++) {
-            sparseLabels[keys[i] - lo] = labels[i];
-        }
-
-        v.tableswitch(lo, hi, defaultLabel, sparseLabels);
-    }
-
-    protected void generateEntries() {
-        // resolving entries' entryLabels and generating entries' code
-        Iterator<Label> entryLabelsIterator = entryLabels.iterator();
-        for (KtWhenEntry entry : expression.getEntries()) {
-            v.visitLabel(entryLabelsIterator.next());
-
-            FrameMap.Mark mark = codegen.myFrameMap.mark();
-            codegen.gen(entry.getExpression(), resultType);
-            mark.dropTo();
-
-            if (!entry.isElse()) {
-                v.goTo(endLabel);
-            }
-        }
-    }
-}
diff --git a/compiler/backend/src/org/jetbrains/kotlin/codegen/when/SwitchCodegen.kt b/compiler/backend/src/org/jetbrains/kotlin/codegen/when/SwitchCodegen.kt
new file mode 100644
index 00000000000..fcd75113d36
--- /dev/null
+++ b/compiler/backend/src/org/jetbrains/kotlin/codegen/when/SwitchCodegen.kt
@@ -0,0 +1,181 @@
+/*
+ * Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.codegen.`when`
+
+import org.jetbrains.kotlin.codegen.ExpressionCodegen
+import org.jetbrains.kotlin.psi.KtWhenExpression
+import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.constants.ConstantValue
+import org.jetbrains.kotlin.resolve.constants.NullValue
+import org.jetbrains.kotlin.types.TypeUtils
+import org.jetbrains.org.objectweb.asm.Label
+import org.jetbrains.org.objectweb.asm.Type
+import org.jetbrains.org.objectweb.asm.commons.InstructionAdapter
+import java.util.*
+
+abstract class SwitchCodegen(
+    @JvmField
+    protected val expression: KtWhenExpression,
+    protected val isStatement: Boolean,
+    private val isExhaustive: Boolean,
+    @JvmField
+    protected val codegen: ExpressionCodegen,
+    subjectType: Type?
+) {
+    protected val bindingContext: BindingContext = codegen.bindingContext
+
+    @JvmField
+    protected val subjectType = subjectType ?: codegen.expressionType(expression.subjectExpression)
+
+    protected val resultType: Type = if (isStatement) Type.VOID_TYPE else codegen.expressionType(expression)
+
+    @JvmField
+    protected val v: InstructionAdapter = codegen.v
+
+    @JvmField
+    protected val transitionsTable: NavigableMap<Int, Label> = TreeMap()
+
+    private val entryLabels: MutableList<Label> = ArrayList()
+    private var elseLabel = Label()
+    private var endLabel = Label()
+    protected lateinit var defaultLabel: Label
+
+    private val switchCodegenProvider = SwitchCodegenProvider(codegen)
+
+    /**
+     * Generates bytecode for entire when expression
+     */
+    open fun generate() {
+        prepareConfiguration()
+
+        val hasElse = expression.elseExpression != null
+
+        // if there is no else-entry and it's statement then default --- endLabel
+        defaultLabel = if (hasElse || !isStatement || isExhaustive) elseLabel else endLabel
+
+        generateSubject()
+
+        generateSwitchInstructionByTransitionsTable()
+
+        generateEntries()
+
+        // there is no else-entry but this is not statement, so we should return Unit
+        if (!hasElse && (!isStatement || isExhaustive)) {
+            v.visitLabel(elseLabel)
+            codegen.putUnitInstanceOntoStackForNonExhaustiveWhen(expression, isStatement)
+        }
+
+        codegen.markLineNumber(expression, isStatement)
+        v.mark(endLabel)
+    }
+
+    /**
+     * Sets up transitionsTable and maybe something else needed in a special case
+     * Behaviour may be changed by overriding processConstant
+     */
+    private fun prepareConfiguration() {
+        for (entry in expression.entries) {
+            val entryLabel = Label()
+
+            for (constant in switchCodegenProvider.getConstantsFromEntry(entry)) {
+                if (constant is NullValue || constant == null) continue
+                processConstant(constant, entryLabel)
+            }
+
+            if (entry.isElse) {
+                elseLabel = entryLabel
+            }
+
+            entryLabels.add(entryLabel)
+        }
+    }
+
+    protected abstract fun processConstant(
+        constant: ConstantValue<*>,
+        entryLabel: Label
+    )
+
+    protected fun putTransitionOnce(value: Int, entryLabel: Label) {
+        if (!transitionsTable.containsKey(value)) {
+            transitionsTable[value] = entryLabel
+        }
+    }
+
+    /**
+     * Should generate int subject on top of the stack
+     * Default implementation just run codegen for actual subject of expression
+     * May also gen nullability check if needed
+     */
+    protected open fun generateSubject() {
+        codegen.gen(expression.subjectExpression, subjectType)
+    }
+
+    protected fun generateNullCheckIfNeeded() {
+        assert(expression.subjectExpression != null) { "subject expression can't be null" }
+        val subjectJetType = bindingContext.getType(expression.subjectExpression!!) ?: error("subject type can't be null (i.e. void)")
+
+        if (TypeUtils.isNullableType(subjectJetType)) {
+            val nullEntryIndex = findNullEntryIndex(expression)
+            val nullLabel = if (nullEntryIndex == -1) defaultLabel else entryLabels[nullEntryIndex]
+            val notNullLabel = Label()
+
+            with(v) {
+                dup()
+                ifnonnull(notNullLabel)
+                pop()
+                goTo(nullLabel)
+                visitLabel(notNullLabel)
+            }
+        }
+    }
+
+    private fun findNullEntryIndex(expression: KtWhenExpression) =
+        expression.entries.withIndex().firstOrNull { (_, entry) ->
+            switchCodegenProvider.getConstantsFromEntry(entry).any { it is NullValue }
+        }?.index ?: -1
+
+    private fun generateSwitchInstructionByTransitionsTable() {
+        val keys = transitionsTable.keys.toIntArray()
+
+        val labelsNumber = keys.size
+        val maxValue = keys.last()
+        val minValue = keys.first()
+        val rangeLength = maxValue.toLong() - minValue.toLong() + 1L
+
+        // In modern JVM implementations it shouldn't matter very much for runtime performance
+        // whether to choose lookupswitch or tableswitch.
+        // The only metric that really matters is bytecode size and here we can estimate:
+        // - lookupswitch: ~ 2 * labelsNumber
+        // - tableswitch: ~ rangeLength
+        if (rangeLength > 2L * labelsNumber || rangeLength > Int.MAX_VALUE) {
+            val labels = transitionsTable.values.toTypedArray()
+            v.lookupswitch(defaultLabel, keys, labels)
+            return
+        }
+
+        val sparseLabels = Array(rangeLength.toInt()) { index ->
+            transitionsTable[index + minValue] ?: defaultLabel
+        }
+
+        v.tableswitch(minValue, maxValue, defaultLabel, *sparseLabels)
+    }
+
+    protected open fun generateEntries() {
+        // resolving entries' entryLabels and generating entries' code
+        val entryLabelsIterator = entryLabels.iterator()
+        for (entry in expression.entries) {
+            v.visitLabel(entryLabelsIterator.next())
+
+            val mark = codegen.myFrameMap.mark()
+            codegen.gen(entry.expression, resultType)
+            mark.dropTo()
+
+            if (!entry.isElse) {
+                v.goTo(endLabel)
+            }
+        }
+    }
+}
diff --git a/compiler/testData/codegen/bytecodeText/when/exhaustiveWhenInitialization.kt b/compiler/testData/codegen/bytecodeText/when/exhaustiveWhenInitialization.kt
index 6af56c7ecc5..a00f136fced 100644
--- a/compiler/testData/codegen/bytecodeText/when/exhaustiveWhenInitialization.kt
+++ b/compiler/testData/codegen/bytecodeText/when/exhaustiveWhenInitialization.kt
@@ -9,6 +9,6 @@ fun box(): String {
   return if (b) "OK" else "FAIL"
 }
 
-// 0 TABLESWITCH
-// 1 LOOKUPSWITCH
-// 1 ATHROW
\ No newline at end of file
+// 1 TABLESWITCH
+// 0 LOOKUPSWITCH
+// 1 ATHROW
diff --git a/compiler/testData/codegen/bytecodeText/when/exhaustiveWhenReturn.kt b/compiler/testData/codegen/bytecodeText/when/exhaustiveWhenReturn.kt
index 5d58827502f..0cd44207502 100644
--- a/compiler/testData/codegen/bytecodeText/when/exhaustiveWhenReturn.kt
+++ b/compiler/testData/codegen/bytecodeText/when/exhaustiveWhenReturn.kt
@@ -7,6 +7,6 @@ fun box(): String {
     }
 }
 
-// 0 TABLESWITCH
-// 1 LOOKUPSWITCH
+// 1 TABLESWITCH
+// 0 LOOKUPSWITCH
 // 1 ATHROW
diff --git a/compiler/testData/codegen/bytecodeText/when/integralWhenWithNoInlinedConstants.kt b/compiler/testData/codegen/bytecodeText/when/integralWhenWithNoInlinedConstants.kt
index f01c101af8c..044521366a1 100644
--- a/compiler/testData/codegen/bytecodeText/when/integralWhenWithNoInlinedConstants.kt
+++ b/compiler/testData/codegen/bytecodeText/when/integralWhenWithNoInlinedConstants.kt
@@ -12,5 +12,4 @@ fun foo2(x: Int): Boolean {
     }
 }
 
-// 0 TABLESWITCH
-// 2 LOOKUPSWITCH
+// 2 TABLESWITCH
diff --git a/compiler/testData/codegen/bytecodeText/whenEnumOptimization/functionLiteralInTopLevel.kt b/compiler/testData/codegen/bytecodeText/whenEnumOptimization/functionLiteralInTopLevel.kt
index a59001bb961..b026893fff0 100644
--- a/compiler/testData/codegen/bytecodeText/whenEnumOptimization/functionLiteralInTopLevel.kt
+++ b/compiler/testData/codegen/bytecodeText/whenEnumOptimization/functionLiteralInTopLevel.kt
@@ -18,4 +18,4 @@ fun box() : String {
     }
 }
 
-// 1 LOOKUPSWITCH
+// 1 TABLESWITCH
diff --git a/compiler/testData/codegen/bytecodeText/whenEnumOptimization/kt14802.kt b/compiler/testData/codegen/bytecodeText/whenEnumOptimization/kt14802.kt
index 5ab98eb9c4b..edd99cad712 100644
--- a/compiler/testData/codegen/bytecodeText/whenEnumOptimization/kt14802.kt
+++ b/compiler/testData/codegen/bytecodeText/whenEnumOptimization/kt14802.kt
@@ -24,4 +24,4 @@ fun box(): String {
     return res
 }
 
-// 1 LOOKUPSWITCH
\ No newline at end of file
+// 1 TABLESWITCH
diff --git a/compiler/testData/codegen/bytecodeText/whenEnumOptimization/kt15806.kt b/compiler/testData/codegen/bytecodeText/whenEnumOptimization/kt15806.kt
index 64b88c22caa..048800172cc 100644
--- a/compiler/testData/codegen/bytecodeText/whenEnumOptimization/kt15806.kt
+++ b/compiler/testData/codegen/bytecodeText/whenEnumOptimization/kt15806.kt
@@ -24,4 +24,4 @@ fun box(): String {
     return "O".doTheThing() + Level.K.doTheThing()
 }
 
-// 1 LOOKUPSWITCH
\ No newline at end of file
+// 1 TABLESWITCH