From 8e9e34350fc1972c597865ee45a6bfe9d0118cf6 Mon Sep 17 00:00:00 2001
From: Dmitriy Novozhilov <dmitriy.novozhilov@jetbrains.com>
Date: Tue, 17 Nov 2020 10:59:30 +0300
Subject: [PATCH] [FE] Properly support sealed interfaces in exhaustiveness
 checker

#KT-20423
---
 .../org/jetbrains/kotlin/cfg/WhenChecker.kt   | 32 +++++++++-------
 .../interfaces/simpleSealedInterface.fir.kt   | 38 -------------------
 .../interfaces/simpleSealedInterface.kt       |  3 +-
 3 files changed, 21 insertions(+), 52 deletions(-)
 delete mode 100644 compiler/testData/diagnostics/tests/sealed/interfaces/simpleSealedInterface.fir.kt

diff --git a/compiler/frontend/src/org/jetbrains/kotlin/cfg/WhenChecker.kt b/compiler/frontend/src/org/jetbrains/kotlin/cfg/WhenChecker.kt
index d6bed65a6ec..be754db3ec2 100644
--- a/compiler/frontend/src/org/jetbrains/kotlin/cfg/WhenChecker.kt
+++ b/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>
-        get() = this.sealedSubclasses.flatMap {
-            if (it.modality == Modality.SEALED) it.deepSealedSubclasses
-            else setOf(it)
+    protected val ClassDescriptor.enumEntries: Set<ClassDescriptor>
+        get() = DescriptorUtils.getAllDescriptors(this.unsubstitutedInnerClassesScope)
+            .filter(::isEnumEntry)
+            .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 =
-                    if (checkedDescriptor.modality == Modality.SEALED) checkedDescriptor.deepSealedSubclasses
-                    else listOf(checkedDescriptor)
+                val checkedDescriptorSubclasses = checkedDescriptor.subclasses
 
                 // 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 =
-            DescriptorUtils.getAllDescriptors(subjectDescriptor!!.unsubstitutedInnerClassesScope)
-                .filter(::isEnumEntry)
-                .filterIsInstance<ClassDescriptor>()
-                .toSet()
-        return getMissingClassCases(expression, entryDescriptors, context) +
+        return getMissingClassCases(expression, subjectDescriptor!!.enumEntries, context) +
                 WhenOnNullableExhaustivenessChecker.getMissingCases(expression, context, nullable)
     }
 
diff --git a/compiler/testData/diagnostics/tests/sealed/interfaces/simpleSealedInterface.fir.kt b/compiler/testData/diagnostics/tests/sealed/interfaces/simpleSealedInterface.fir.kt
deleted file mode 100644
index 2d63915cbcb..00000000000
--- a/compiler/testData/diagnostics/tests/sealed/interfaces/simpleSealedInterface.fir.kt
+++ /dev/null
@@ -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
-    }
-}
diff --git a/compiler/testData/diagnostics/tests/sealed/interfaces/simpleSealedInterface.kt b/compiler/testData/diagnostics/tests/sealed/interfaces/simpleSealedInterface.kt
index 4f87a029760..284c4f84434 100644
--- a/compiler/testData/diagnostics/tests/sealed/interfaces/simpleSealedInterface.kt
+++ b/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