KT-36047 Support when-with-subject in optimized 'when' generators

compiler/backend/src/org/jetbrains/kotlin/codegen/when/StringSwitchCodegen.java

@@ -62,7 +62,6 @@ public class StringSwitchCodegen extends SwitchCodegen {
     @Override
     public void generate() {
         super.generate();
-        codegen.myFrameMap.leaveTemp(subjectType);
     }
 
     @Override

compiler/backend/src/org/jetbrains/kotlin/codegen/when/SwitchCodegen.kt

@@ -20,7 +20,9 @@ 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.*
+import kotlin.collections.ArrayList
 
+@Suppress("MemberVisibilityCanBePrivate")
 abstract class SwitchCodegen(
     @JvmField
     protected val expression: KtWhenExpression,
@@ -36,7 +38,7 @@ abstract class SwitchCodegen(
     protected val subjectExpression = expression.subjectExpression ?: throw AssertionError("No subject expression: ${expression.text}")
 
     protected val subjectKotlinType = WhenChecker.whenSubjectTypeWithoutSmartCasts(expression, bindingContext)
-            ?: throw AssertionError("No subject type: ${expression}")
+        ?: throw AssertionError("No subject type: $expression")
 
     @JvmField
     protected val subjectType = subjectType ?: codegen.asmType(subjectKotlinType)
@@ -64,6 +66,8 @@ abstract class SwitchCodegen(
      * Generates bytecode for entire when expression
      */
     open fun generate() {
+        val frameMapAtStart = codegen.frameMap.mark()
+
         prepareConfiguration()
 
         val hasElse = expression.elseExpression != null
@@ -90,8 +94,9 @@ abstract class SwitchCodegen(
         codegen.markLineNumber(expression, isStatement)
         v.mark(endLabel)
 
+        frameMapAtStart.dropTo()
+
         subjectVariableDescriptor?.let {
-            codegen.frameMap.leave(it)
             v.visitLocalVariable(
                 it.name.asString(), subjectType.descriptor, null,
                 beginLabel, endLabel, subjectLocal
@@ -137,7 +142,7 @@ abstract class SwitchCodegen(
     private fun generateSubjectValue() {
         if (subjectVariable != null) {
             val mySubjectVariable = bindingContext[BindingContext.VARIABLE, subjectVariable]
-                    ?: throw AssertionError("Unresolved subject variable: $expression")
+                ?: throw AssertionError("Unresolved subject variable: $expression")
             subjectLocal = codegen.frameMap.enter(mySubjectVariable, subjectType)
             codegen.visitProperty(subjectVariable, null)
             StackValue.local(subjectLocal, subjectType, subjectKotlinType).put(subjectType, subjectKotlinType, codegen.v)

compiler/backend/src/org/jetbrains/kotlin/codegen/when/SwitchCodegenProvider.kt

@@ -20,12 +20,13 @@ import org.jetbrains.kotlin.codegen.AsmUtil
 import org.jetbrains.kotlin.codegen.ExpressionCodegen
 import org.jetbrains.kotlin.codegen.binding.CodegenBinding
 import org.jetbrains.kotlin.codegen.state.GenerationState
-import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.KtWhenConditionWithExpression
+import org.jetbrains.kotlin.psi.KtWhenEntry
+import org.jetbrains.kotlin.psi.KtWhenExpression
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.constants.*
 import org.jetbrains.org.objectweb.asm.Type
-
+import java.util.*
-import java.util.ArrayList
 
 class SwitchCodegenProvider
 private constructor(
@@ -78,7 +79,9 @@ private constructor(
             return null
         }
 
-        val subjectType = codegen.expressionType(expression.subjectExpression)
+        val subjectType =
+            expression.subjectVariable?.let { codegen.expressionType(it.initializer) }
+                ?: codegen.expressionType(expression.subjectExpression)
 
         val mapping = codegen.bindingContext.get(CodegenBinding.MAPPING_FOR_WHEN_BY_ENUM, expression)
 

compiler/testData/codegen/bytecodeText/when/lookupSwitchWithSubjectVal.kt

@@ -0,0 +1,11 @@
+fun foo(x: Int): String {
+    return when (val y = x) {
+        100 -> "1"
+        200 -> "2"
+        300 -> "3"
+        else -> "else"
+    }
+}
+
+// 1 LOOKUPSWITCH
+// 0 TABLESWITCH

compiler/testData/codegen/bytecodeText/when/stringSwitchWithSubjectVal.kt

@@ -0,0 +1,12 @@
+fun foo(x: String): String {
+    return when (val y = x) {
+        "1" -> "one"
+        "2" -> "two"
+        "3" -> "two and one"
+        "4" -> "two and two"
+        else -> "many"
+    }
+}
+
+// 0 LOOKUPSWITCH
+// 1 TABLESWITCH

compiler/testData/codegen/bytecodeText/when/tableSwitchWithSubjectVal.kt

@@ -0,0 +1,11 @@
+fun foo(x: Int): String {
+    return when (val y = x) {
+        101 -> "1"
+        102 -> "2"
+        103 -> "3"
+        else -> "else"
+    }
+}
+
+// 0 LOOKUPSWITCH
+// 1 TABLESWITCH

compiler/tests/org/jetbrains/kotlin/codegen/BytecodeTextTestGenerated.java

@@ -4387,6 +4387,11 @@ public class BytecodeTextTestGenerated extends AbstractBytecodeTextTest {
             runTest("compiler/testData/codegen/bytecodeText/when/lookupSwitch.kt");
         }
 
+        @TestMetadata("lookupSwitchWithSubjectVal.kt")
+        public void testLookupSwitchWithSubjectVal() throws Exception {
+            runTest("compiler/testData/codegen/bytecodeText/when/lookupSwitchWithSubjectVal.kt");
+        }
+
         @TestMetadata("noBoxingInDefaultWhenWithSpecialCases.kt")
         public void testNoBoxingInDefaultWhenWithSpecialCases() throws Exception {
             runTest("compiler/testData/codegen/bytecodeText/when/noBoxingInDefaultWhenWithSpecialCases.kt");
@@ -4407,6 +4412,11 @@ public class BytecodeTextTestGenerated extends AbstractBytecodeTextTest {
             runTest("compiler/testData/codegen/bytecodeText/when/simpleConstValsInsideWhen.kt");
         }
 
+        @TestMetadata("stringSwitchWithSubjectVal.kt")
+        public void testStringSwitchWithSubjectVal() throws Exception {
+            runTest("compiler/testData/codegen/bytecodeText/when/stringSwitchWithSubjectVal.kt");
+        }
+
         @TestMetadata("subjectValHasLocalVariableSlot.kt")
         public void testSubjectValHasLocalVariableSlot() throws Exception {
             runTest("compiler/testData/codegen/bytecodeText/when/subjectValHasLocalVariableSlot.kt");
@@ -4437,6 +4447,11 @@ public class BytecodeTextTestGenerated extends AbstractBytecodeTextTest {
             runTest("compiler/testData/codegen/bytecodeText/when/tableSwitch.kt");
         }
 
+        @TestMetadata("tableSwitchWithSubjectVal.kt")
+        public void testTableSwitchWithSubjectVal() throws Exception {
+            runTest("compiler/testData/codegen/bytecodeText/when/tableSwitchWithSubjectVal.kt");
+        }
+
         @TestMetadata("whenNull.kt")
         public void testWhenNull() throws Exception {
             runTest("compiler/testData/codegen/bytecodeText/when/whenNull.kt");

compiler/tests/org/jetbrains/kotlin/codegen/ir/IrBytecodeTextTestGenerated.java

@@ -4305,6 +4305,11 @@ public class IrBytecodeTextTestGenerated extends AbstractIrBytecodeTextTest {
             runTest("compiler/testData/codegen/bytecodeText/when/lookupSwitch.kt");
         }
 
+        @TestMetadata("lookupSwitchWithSubjectVal.kt")
+        public void testLookupSwitchWithSubjectVal() throws Exception {
+            runTest("compiler/testData/codegen/bytecodeText/when/lookupSwitchWithSubjectVal.kt");
+        }
+
         @TestMetadata("noBoxingInDefaultWhenWithSpecialCases.kt")
         public void testNoBoxingInDefaultWhenWithSpecialCases() throws Exception {
             runTest("compiler/testData/codegen/bytecodeText/when/noBoxingInDefaultWhenWithSpecialCases.kt");
@@ -4325,6 +4330,11 @@ public class IrBytecodeTextTestGenerated extends AbstractIrBytecodeTextTest {
             runTest("compiler/testData/codegen/bytecodeText/when/simpleConstValsInsideWhen.kt");
         }
 
+        @TestMetadata("stringSwitchWithSubjectVal.kt")
+        public void testStringSwitchWithSubjectVal() throws Exception {
+            runTest("compiler/testData/codegen/bytecodeText/when/stringSwitchWithSubjectVal.kt");
+        }
+
         @TestMetadata("subjectValHasLocalVariableSlot.kt")
         public void testSubjectValHasLocalVariableSlot() throws Exception {
             runTest("compiler/testData/codegen/bytecodeText/when/subjectValHasLocalVariableSlot.kt");
@@ -4355,6 +4365,11 @@ public class IrBytecodeTextTestGenerated extends AbstractIrBytecodeTextTest {
             runTest("compiler/testData/codegen/bytecodeText/when/tableSwitch.kt");
         }
 
+        @TestMetadata("tableSwitchWithSubjectVal.kt")
+        public void testTableSwitchWithSubjectVal() throws Exception {
+            runTest("compiler/testData/codegen/bytecodeText/when/tableSwitchWithSubjectVal.kt");
+        }
+
         @TestMetadata("whenNull.kt")
         public void testWhenNull() throws Exception {
             runTest("compiler/testData/codegen/bytecodeText/when/whenNull.kt");