Support 'when' with subject variable in psi2ir

compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/BranchingExpressionGenerator.kt

@@ -84,9 +84,7 @@ class BranchingExpressionGenerator(statementGenerator: StatementGenerator) : Sta
     }
 
     fun generateWhenExpression(expression: KtWhenExpression): IrExpression {
-        val irSubject = expression.subjectExpression?.let {
-            scope.createTemporaryVariable(it.genExpr(), "subject")
-        }
+        val irSubject = generateWhenSubject(expression)
 
         val inferredType = getInferredTypeWithImplicitCastsOrFail(expression)
 
@@ -127,6 +125,16 @@ class BranchingExpressionGenerator(statementGenerator: StatementGenerator) : Sta
         return generateWhenBody(expression, irSubject, irWhen)
     }
 
+    private fun generateWhenSubject(expression: KtWhenExpression): IrVariable? {
+        val subjectVariable = expression.subjectVariable
+        val subjectExpression = expression.subjectExpression
+        return when {
+            subjectVariable != null -> statementGenerator.visitProperty(subjectVariable, null) as IrVariable
+            subjectExpression != null -> scope.createTemporaryVariable(subjectExpression.genExpr(), "subject")
+            else -> null
+        }
+    }
+
     private fun addElseBranchForExhaustiveWhenIfNeeded(irWhen: IrWhen, whenExpression: KtWhenExpression) {
         if (irWhen.branches.filterIsInstance<IrElseBranch>().isEmpty()) {
             //TODO: check condition: seems it's safe to always generate exception
@@ -144,8 +152,8 @@ class BranchingExpressionGenerator(statementGenerator: StatementGenerator) : Sta
                 || true == get(BindingContext.EXHAUSTIVE_WHEN, this)
                 || true == get(BindingContext.IMPLICIT_EXHAUSTIVE_WHEN, this)
 
-    private fun generateWhenBody(expression: KtWhenExpression, irSubject: IrVariable?, irWhen: IrWhen): IrExpression {
-        return if (irSubject == null) {
+    private fun generateWhenBody(expression: KtWhenExpression, irSubject: IrVariable?, irWhen: IrWhen): IrExpression =
+        if (irSubject == null) {
             if (irWhen.branches.isEmpty())
                 IrBlockImpl(expression.startOffset, expression.endOffset, context.builtIns.unitType, IrStatementOrigin.WHEN)
             else
@@ -162,7 +170,6 @@ class BranchingExpressionGenerator(statementGenerator: StatementGenerator) : Sta
                 irBlock
             }
         }
-    }
 
     private fun generateWhenConditionNoSubject(ktCondition: KtWhenCondition): IrExpression =
         (ktCondition as KtWhenConditionWithExpression).expression!!.genExpr()

compiler/testData/ir/irText/expressions/whenWithSubjectVariable.kt

@@ -0,0 +1,13 @@
+// !LANGUAGE: +VariableDeclarationInWhenSubject
+
+fun foo(): Any = 1
+
+fun test() =
+    when (val y = foo()) {
+        42 -> 1
+        is String -> y.length
+        !is Int -> 2
+        in 0..10 -> 3
+        !in 10..20 -> 4
+        else -> -1
+    }

compiler/testData/ir/irText/expressions/whenWithSubjectVariable.txt

@@ -0,0 +1,52 @@
+FILE fqName:<root> fileName:/whenWithSubjectVariable.kt
+  FUN name:foo visibility:public modality:FINAL <> () returnType:Any flags:
+    BLOCK_BODY
+      RETURN type=kotlin.Nothing from='foo(): Any'
+        CONST Int type=kotlin.Int value=1
+  FUN name:test visibility:public modality:FINAL <> () returnType:Int flags:
+    BLOCK_BODY
+      RETURN type=kotlin.Nothing from='test(): Int'
+        BLOCK type=kotlin.Int origin=WHEN
+          VAR name:y type:kotlin.Any flags:val
+            CALL 'foo(): Any' type=kotlin.Any origin=null
+          WHEN type=kotlin.Int origin=WHEN
+            BRANCH
+              if: CALL 'EQEQ(Any?, Any?): Boolean' type=kotlin.Boolean origin=EQEQ
+                arg0: GET_VAR 'y: Any' type=kotlin.Any origin=null
+                arg1: CONST Int type=kotlin.Int value=42
+              then: CONST Int type=kotlin.Int value=1
+            BRANCH
+              if: TYPE_OP type=kotlin.Boolean origin=INSTANCEOF typeOperand=kotlin.String
+                typeOperand: CLASS IR_EXTERNAL_DECLARATION_STUB CLASS name:String modality:FINAL visibility:public flags:
+                GET_VAR 'y: Any' type=kotlin.Any origin=null
+              then: CALL '<get-length>(): Int' type=kotlin.Int origin=GET_PROPERTY
+                $this: TYPE_OP type=kotlin.String origin=IMPLICIT_CAST typeOperand=kotlin.String
+                  typeOperand: CLASS IR_EXTERNAL_DECLARATION_STUB CLASS name:String modality:FINAL visibility:public flags:
+                  GET_VAR 'y: Any' type=kotlin.Any origin=null
+            BRANCH
+              if: TYPE_OP type=kotlin.Boolean origin=INSTANCEOF typeOperand=kotlin.Int
+                typeOperand: CLASS IR_EXTERNAL_DECLARATION_STUB CLASS name:Int modality:FINAL visibility:public flags:
+                GET_VAR 'y: Any' type=kotlin.Any origin=null
+              then: CONST Int type=kotlin.Int value=2
+            BRANCH
+              if: CALL 'contains(Int): Boolean' type=kotlin.Boolean origin=IN
+                $this: CALL 'rangeTo(Int): IntRange' type=kotlin.ranges.IntRange origin=RANGE
+                  $this: CONST Int type=kotlin.Int value=0
+                  other: CONST Int type=kotlin.Int value=10
+                value: TYPE_OP type=kotlin.Int origin=IMPLICIT_CAST typeOperand=kotlin.Int
+                  typeOperand: CLASS IR_EXTERNAL_DECLARATION_STUB CLASS name:Int modality:FINAL visibility:public flags:
+                  GET_VAR 'y: Any' type=kotlin.Any origin=null
+              then: CONST Int type=kotlin.Int value=3
+            BRANCH
+              if: CALL 'NOT(Boolean): Boolean' type=kotlin.Boolean origin=EXCL
+                arg0: CALL 'contains(Int): Boolean' type=kotlin.Boolean origin=NOT_IN
+                  $this: CALL 'rangeTo(Int): IntRange' type=kotlin.ranges.IntRange origin=RANGE
+                    $this: CONST Int type=kotlin.Int value=10
+                    other: CONST Int type=kotlin.Int value=20
+                  value: TYPE_OP type=kotlin.Int origin=IMPLICIT_CAST typeOperand=kotlin.Int
+                    typeOperand: CLASS IR_EXTERNAL_DECLARATION_STUB CLASS name:Int modality:FINAL visibility:public flags:
+                    GET_VAR 'y: Any' type=kotlin.Any origin=null
+              then: CONST Int type=kotlin.Int value=4
+            BRANCH
+              if: CONST Boolean type=kotlin.Boolean value=true
+              then: CONST Int type=kotlin.Int value=-1

compiler/tests/org/jetbrains/kotlin/ir/IrTextTestCaseGenerated.java

@@ -1127,6 +1127,11 @@ public class IrTextTestCaseGenerated extends AbstractIrTextTestCase {
             runTest("compiler/testData/ir/irText/expressions/whenReturn.kt");
         }
 
+        @TestMetadata("whenWithSubjectVariable.kt")
+        public void testWhenWithSubjectVariable() throws Exception {
+            runTest("compiler/testData/ir/irText/expressions/whenWithSubjectVariable.kt");
+        }
+
         @TestMetadata("whileDoWhile.kt")
         public void testWhileDoWhile() throws Exception {
             runTest("compiler/testData/ir/irText/expressions/whileDoWhile.kt");