Check possible smart cast to enum in when subject #KT-14705 Fixed

compiler/frontend/src/org/jetbrains/kotlin/types/expressions/PatternMatchingTypingVisitor.kt

@@ -19,8 +19,8 @@ package org.jetbrains.kotlin.types.expressions
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns.isBoolean
 import org.jetbrains.kotlin.cfg.WhenChecker
-import org.jetbrains.kotlin.cfg.WhenOnSealedExhaustivenessChecker
 import org.jetbrains.kotlin.descriptors.ClassDescriptor
+import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.diagnostics.Errors
 import org.jetbrains.kotlin.diagnostics.Errors.*
@@ -100,7 +100,8 @@ class PatternMatchingTypingVisitor internal constructor(facade: ExpressionTyping
             DataFlowValueFactory.createDataFlowValue(it, subjectType, contextAfterSubject)
         } ?: DataFlowValue.nullValue(components.builtIns)
 
-        checkSmartCastsInSubjectIfRequired(expression, contextBeforeSubject, subjectType)
+        val possibleTypesForSubject = subjectTypeInfo?.dataFlowInfo?.getStableTypes(subjectDataFlowValue) ?: emptySet()
+        checkSmartCastsInSubjectIfRequired(expression, contextBeforeSubject, subjectType, possibleTypesForSubject)
 
         val dataFlowInfoForEntries = analyzeConditionsInWhenEntries(expression, contextAfterSubject, subjectDataFlowValue, subjectType)
         val whenReturnType = inferTypeForWhenExpression(expression, contextWithExpectedType, contextAfterSubject, dataFlowInfoForEntries)
@@ -230,21 +231,25 @@ class PatternMatchingTypingVisitor internal constructor(facade: ExpressionTyping
     private fun checkSmartCastsInSubjectIfRequired(
             expression: KtWhenExpression,
             contextBeforeSubject: ExpressionTypingContext,
-            subjectType: KotlinType
+            subjectType: KotlinType,
+            possibleTypesForSubject: Set<KotlinType>
     ) {
         val subjectExpression = expression.subjectExpression ?: return
-        val nullableType = TypeUtils.isNullableType(subjectType)
+        val isNullableType = TypeUtils.isNullableType(subjectType)
         val bindingContext = contextBeforeSubject.trace.bindingContext
-        if (nullableType && !WhenChecker.containsNullCase(expression, bindingContext)) {
+        if (isNullableType && !WhenChecker.containsNullCase(expression, bindingContext)) {
             val notNullableType = TypeUtils.makeNotNullable(subjectType)
-            checkSmartCastToExpectedTypeInSubject(contextBeforeSubject, subjectExpression, subjectType, notNullableType)
+            if (checkSmartCastToExpectedTypeInSubject(contextBeforeSubject, subjectExpression, subjectType,
+                                                      notNullableType)) {
+                return
+            }
         }
-        val subjectClass = subjectType.constructor.declarationDescriptor as? ClassDescriptor ?: return
-        if (subjectClass.modality == Modality.SEALED &&
-            WhenOnSealedExhaustivenessChecker.getMissingCases(expression, bindingContext, subjectClass, false).isNotEmpty()) {
-            for (descriptor in WhenOnSealedExhaustivenessChecker.getNestedSubclasses(subjectClass)) {
-                if (descriptor.modality == Modality.SEALED && DescriptorUtils.isDirectSubclass(descriptor, subjectClass)) {
-                    checkSmartCastToExpectedTypeInSubject(contextBeforeSubject, subjectExpression, subjectType, descriptor.defaultType)
+        for (possibleCastType in possibleTypesForSubject) {
+            val possibleCastClass = possibleCastType.constructor.declarationDescriptor as? ClassDescriptor ?: continue
+            if (possibleCastClass.kind == ClassKind.ENUM_CLASS || possibleCastClass.modality == Modality.SEALED) {
+                if (checkSmartCastToExpectedTypeInSubject(contextBeforeSubject, subjectExpression, subjectType,
+                                                          possibleCastType)) {
+                    return
                 }
             }
         }
@@ -255,14 +260,16 @@ class PatternMatchingTypingVisitor internal constructor(facade: ExpressionTyping
             subjectExpression: KtExpression,
             subjectType: KotlinType,
             expectedType: KotlinType
-    ) {
+    ): Boolean {
         val trace = TemporaryBindingTrace.create(contextBeforeSubject.trace, "Temporary trace for when subject nullability")
         val subjectContext = contextBeforeSubject.replaceExpectedType(expectedType).replaceBindingTrace(trace)
         val castResult = DataFlowAnalyzer.checkPossibleCast(
                 subjectType, KtPsiUtil.safeDeparenthesize(subjectExpression), subjectContext)
         if (castResult != null && castResult.isCorrect) {
             trace.commit()
+            return true
         }
+        return false
     }
 
     private fun analyzeWhenEntryConditions(

compiler/testData/diagnostics/tests/enum/enumSubjectTypeCheck.kt

@@ -0,0 +1,40 @@
+// See KT-14705
+
+enum class En { A, B, С }
+
+fun foo() {
+    // nullable variable
+    val en2: Any? = En.A
+    if (en2 is En) {
+        when (<!DEBUG_INFO_SMARTCAST!>en2<!>) {
+            En.A -> {}
+            En.B -> {}
+            En.С -> {}
+        }
+    }
+
+    // not nullable variable
+    val en1: Any = En.A
+    if (en1 is En) {
+        when (<!DEBUG_INFO_SMARTCAST!>en1<!>) {
+            En.A -> {}
+            En.B -> {}
+            En.С -> {}
+        }
+    }
+}
+
+enum class En2 { D, E, F }
+
+fun useEn(x: En) = x
+fun useEn2(x: En2) = x
+
+fun bar(x: Any) {
+    if (x is En && x is En2) {
+        when (<!DEBUG_INFO_SMARTCAST!>x<!>) {
+            En.A -> useEn(<!DEBUG_INFO_SMARTCAST!>x<!>)
+            En2.D -> useEn2(<!DEBUG_INFO_SMARTCAST!>x<!>)
+            else -> {}
+        }
+    }
+}

compiler/testData/diagnostics/tests/enum/enumSubjectTypeCheck.txt

@@ -0,0 +1,52 @@
+package
+
+public fun bar(/*0*/ x: kotlin.Any): kotlin.Unit
+public fun foo(): kotlin.Unit
+public fun useEn(/*0*/ x: En): En
+public fun useEn2(/*0*/ x: En2): En2
+
+public final enum class En : kotlin.Enum<En> {
+    enum entry A
+
+    enum entry B
+
+    enum entry С
+
+    private constructor En()
+    public final override /*1*/ /*fake_override*/ val name: kotlin.String
+    public final override /*1*/ /*fake_override*/ val ordinal: kotlin.Int
+    protected final override /*1*/ /*fake_override*/ fun clone(): kotlin.Any
+    public final override /*1*/ /*fake_override*/ fun compareTo(/*0*/ other: En): kotlin.Int
+    public final override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    protected/*protected and package*/ final override /*1*/ /*fake_override*/ /*isHiddenForResolutionEverywhereBesideSupercalls*/ fun finalize(): kotlin.Unit
+    public final override /*1*/ /*fake_override*/ /*isHiddenForResolutionEverywhereBesideSupercalls*/ fun getDeclaringClass(): java.lang.Class<En!>!
+    public final override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+
+    // Static members
+    public final /*synthesized*/ fun valueOf(/*0*/ value: kotlin.String): En
+    public final /*synthesized*/ fun values(): kotlin.Array<En>
+}
+
+public final enum class En2 : kotlin.Enum<En2> {
+    enum entry D
+
+    enum entry E
+
+    enum entry F
+
+    private constructor En2()
+    public final override /*1*/ /*fake_override*/ val name: kotlin.String
+    public final override /*1*/ /*fake_override*/ val ordinal: kotlin.Int
+    protected final override /*1*/ /*fake_override*/ fun clone(): kotlin.Any
+    public final override /*1*/ /*fake_override*/ fun compareTo(/*0*/ other: En2): kotlin.Int
+    public final override /*1*/ /*fake_override*/ fun equals(/*0*/ other: kotlin.Any?): kotlin.Boolean
+    protected/*protected and package*/ final override /*1*/ /*fake_override*/ /*isHiddenForResolutionEverywhereBesideSupercalls*/ fun finalize(): kotlin.Unit
+    public final override /*1*/ /*fake_override*/ /*isHiddenForResolutionEverywhereBesideSupercalls*/ fun getDeclaringClass(): java.lang.Class<En2!>!
+    public final override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
+    public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
+
+    // Static members
+    public final /*synthesized*/ fun valueOf(/*0*/ value: kotlin.String): En2
+    public final /*synthesized*/ fun values(): kotlin.Array<En2>
+}

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

@@ -7249,6 +7249,12 @@ public class DiagnosticsTestGenerated extends AbstractDiagnosticsTest {
                 doTest(fileName);
             }
 
+            @TestMetadata("enumSubjectTypeCheck.kt")
+            public void testEnumSubjectTypeCheck() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/tests/enum/enumSubjectTypeCheck.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("enumWithAnnotationKeyword.kt")
             public void testEnumWithAnnotationKeyword() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/diagnostics/tests/enum/enumWithAnnotationKeyword.kt");