[FE] Properly support sealed interfaces in exhaustiveness checker

#KT-20423

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

@@ -146,10 +146,23 @@ internal abstract class WhenOnClassExhaustivenessChecker : WhenExhaustivenessChe
             else -> null
         }
 
-    protected val ClassDescriptor.deepSealedSubclasses: List<ClassDescriptor>
+    protected val ClassDescriptor.enumEntries: Set<ClassDescriptor>
-        get() = this.sealedSubclasses.flatMap {
+        get() = DescriptorUtils.getAllDescriptors(this.unsubstitutedInnerClassesScope)
-            if (it.modality == Modality.SEALED) it.deepSealedSubclasses
+            .filter(::isEnumEntry)
-            else setOf(it)
+            .filterIsInstance<ClassDescriptor>()
+            .toSet()
+
+
+    protected val ClassDescriptor.deepSealedSubclasses: Set<ClassDescriptor>
+        get() = this.sealedSubclasses.flatMapTo(mutableSetOf()) {
+            it.subclasses
+        }
+
+    private val ClassDescriptor.subclasses: Set<ClassDescriptor>
+        get() = when {
+            this.modality == Modality.SEALED -> this.deepSealedSubclasses
+            this.kind == ClassKind.ENUM_CLASS -> this.enumEntries
+            else -> setOf(this)
         }
 
     private val KtWhenCondition.negated
@@ -189,9 +202,7 @@ internal abstract class WhenOnClassExhaustivenessChecker : WhenExhaustivenessChe
             for (condition in whenEntry.conditions) {
                 val negated = condition.negated
                 val checkedDescriptor = condition.getCheckedDescriptor(context) ?: continue
-                val checkedDescriptorSubclasses =
+                val checkedDescriptorSubclasses = checkedDescriptor.subclasses
-                    if (checkedDescriptor.modality == Modality.SEALED) checkedDescriptor.deepSealedSubclasses
-                    else listOf(checkedDescriptor)
 
                 // Checks are important only for nested subclasses of the sealed class
                 // In additional, check without "is" is important only for objects
@@ -220,12 +231,7 @@ private object WhenOnEnumExhaustivenessChecker : WhenOnClassExhaustivenessChecke
         nullable: Boolean
     ): List<WhenMissingCase> {
         assert(isEnumClass(subjectDescriptor)) { "isWhenOnEnumExhaustive should be called with an enum class descriptor" }
-        val entryDescriptors =
+        return getMissingClassCases(expression, subjectDescriptor!!.enumEntries, context) +
-            DescriptorUtils.getAllDescriptors(subjectDescriptor!!.unsubstitutedInnerClassesScope)
-                .filter(::isEnumEntry)
-                .filterIsInstance<ClassDescriptor>()
-                .toSet()
-        return getMissingClassCases(expression, entryDescriptors, context) +
                 WhenOnNullableExhaustivenessChecker.getMissingCases(expression, context, nullable)
     }
 

compiler/testData/diagnostics/tests/sealed/interfaces/simpleSealedInterface.fir.kt

@@ -1,38 +0,0 @@
-// ISSUE: KT-20423
-// !LANGUAGE: +FreedomForSealedClasses +SealedInterfaces
-// !DIAGNOSTICS: -UNUSED_VARIABLE
-
-sealed interface Base
-
-interface A : Base
-
-sealed class B : Base {
-    class First : B()
-    class Second : B()
-}
-
-enum class C : Base {
-    SomeValue, AnotherValue
-}
-
-object D : Base
-
-fun test_1(base: Base) {
-    val x = when (base) {
-        is A -> 1
-        is B -> 2
-        is C -> 3
-        is D -> 4
-    }
-}
-
-fun test_2(base: Base) {
-    val x = when (base) {
-        is A -> 1
-        is B.First -> 2
-        is B.Second -> 3
-        C.SomeValue -> 4
-        C.AnotherValue -> 5
-        D -> 6
-    }
-}

compiler/testData/diagnostics/tests/sealed/interfaces/simpleSealedInterface.kt

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 // ISSUE: KT-20423
 // !LANGUAGE: +FreedomForSealedClasses +SealedInterfaces
 // !DIAGNOSTICS: -UNUSED_VARIABLE
@@ -27,7 +28,7 @@ fun test_1(base: Base) {
 }
 
 fun test_2(base: Base) {
-    val x = <!NO_ELSE_IN_WHEN!>when<!> (base) {
+    val x = when (base) {
         is A -> 1
         is B.First -> 2
         is B.Second -> 3