Fix exhaustiveness check for when with subject variable

compiler/frontend/src/org/jetbrains/kotlin/cfg/WhenChecker.kt

@@ -27,6 +27,7 @@ import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.checkReservedPrefixWord
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.BindingContext.SMARTCAST
+import org.jetbrains.kotlin.resolve.BindingContext.VARIABLE
 import org.jetbrains.kotlin.resolve.BindingTrace
 import org.jetbrains.kotlin.resolve.CompileTimeConstantUtils
 import org.jetbrains.kotlin.resolve.DescriptorUtils
@@ -282,8 +283,15 @@ object WhenChecker {
 
 
     @JvmStatic
-    fun whenSubjectType(expression: KtWhenExpression, context: BindingContext) =
-        expression.subjectExpression?.let { context.get(SMARTCAST, it)?.defaultType ?: context.getType(it) }
+    fun whenSubjectType(expression: KtWhenExpression, context: BindingContext): KotlinType? {
+        val subjectVariable = expression.subjectVariable
+        val subjectExpression = expression.subjectExpression
+        return when {
+            subjectVariable != null -> context.get(VARIABLE, subjectVariable)?.type
+            subjectExpression != null -> context.get(SMARTCAST, subjectExpression)?.defaultType ?: context.getType(subjectExpression)
+            else -> null
+        }
+    }
 
     @JvmStatic
     fun getEnumMissingCases(

compiler/testData/diagnostics/tests/when/withSubjectVariable/smartcastToEnum.kt

@@ -3,21 +3,16 @@
 
 enum class E { FIRST, SECOND }
 
-fun testSmartcastToEnumInSubjectInitializer(e: E?) {
-    val x = <!NO_ELSE_IN_WHEN!>when<!> (val ne = e!!) { // TODO fix
+fun testSmartcastToEnumInSubjectInitializer1(e: E?) {
+    val x1 = when (val ne = e!!) {
         E.FIRST -> "f"
         E.SECOND -> "s"
     }
 }
 
-
-sealed class Either
-class Left : Either()
-class Right : Either()
-
-fun testSmartcastToSealedInSubjectInitializer(x: Any?) {
-    val y = <!NO_ELSE_IN_WHEN!>when<!> (val either = x as Either) { // TODO fix
-        is Left -> "L"
-        is Right -> "R"
+fun testSmartcastToEnumInSubjectInitializer2(e: E?) {
+    val x2 = <!NO_ELSE_IN_WHEN!>when<!> (val ne: Any = e!!) { // NB explicit type annotation
+        E.FIRST -> "f"
+        E.SECOND -> "s"
     }
-}
+}

compiler/testData/diagnostics/tests/when/withSubjectVariable/smartcastToEnum.txt

@@ -1,7 +1,7 @@
 package
 
-public fun testSmartcastToEnumInSubjectInitializer(/*0*/ e: E?): kotlin.Unit
-public fun testSmartcastToSealedInSubjectInitializer(/*0*/ x: kotlin.Any?): kotlin.Unit
+public fun testSmartcastToEnumInSubjectInitializer1(/*0*/ e: E?): kotlin.Unit
+public fun testSmartcastToEnumInSubjectInitializer2(/*0*/ e: E?): kotlin.Unit
 
 public final enum class E : kotlin.Enum<E> {
     enum entry FIRST
@@ -23,24 +23,3 @@ public final enum class E : kotlin.Enum<E> {
     public final /*synthesized*/ fun valueOf(/*0*/ value: kotlin.String): E
     public final /*synthesized*/ fun values(): kotlin.Array<E>
 }
-
-public sealed class Either {
-    private constructor Either()
-    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
-    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
-    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
-}
-
-public final class Left : Either {
-    public constructor Left()
-    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
-    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
-    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
-}
-
-public final class Right : Either {
-    public constructor Right()
-    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
-    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
-    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
-}

compiler/testData/diagnostics/tests/when/withSubjectVariable/smartcastToSealed.kt

@@ -0,0 +1,20 @@
+// !LANGUAGE: +VariableDeclarationInWhenSubject
+// !DIAGNOSTICS: -UNUSED_VARIABLE
+
+sealed class Either
+class Left : Either()
+class Right : Either()
+
+fun testSmartcastToSealedInSubjectInitializer1(x: Any?) {
+    val y1 = when (val either = x as Either) {
+        is Left -> "L"
+        is Right -> "R"
+    }
+}
+
+fun testSmartcastToSealedInSubjectInitializer2(x: Any?) {
+    val y2 = <!NO_ELSE_IN_WHEN!>when<!> (val either: Any = x as Either) { // NB explicit type annotation
+        is Left -> "L"
+        is Right -> "R"
+    }
+}

compiler/testData/diagnostics/tests/when/withSubjectVariable/smartcastToSealed.txt

@@ -0,0 +1,25 @@
+package
+
+public fun testSmartcastToSealedInSubjectInitializer1(/*0*/ x: kotlin.Any?): kotlin.Unit
+public fun testSmartcastToSealedInSubjectInitializer2(/*0*/ x: kotlin.Any?): kotlin.Unit
+
+public sealed class Either {
+    private constructor Either()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}
+
+public final class Left : Either {
+    public constructor Left()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}
+
+public final class Right : Either {
+    public constructor Right()
+    public open override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+}

compiler/tests/org/jetbrains/kotlin/checkers/DiagnosticsTestGenerated.java

@@ -22520,6 +22520,11 @@ public class DiagnosticsTestGenerated extends AbstractDiagnosticsTest {
                     runTest("compiler/testData/diagnostics/tests/when/withSubjectVariable/smartcastToEnum.kt");
                 }
 
+                @TestMetadata("smartcastToSealed.kt")
+                public void testSmartcastToSealed() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/when/withSubjectVariable/smartcastToSealed.kt");
+                }
+
                 @TestMetadata("softModifierName.kt")
                 public void testSoftModifierName() throws Exception {
                     runTest("compiler/testData/diagnostics/tests/when/withSubjectVariable/softModifierName.kt");